mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor ai serving modules and crates
This commit is contained in:
fawney19
@@ -8,7 +8,8 @@ description = "Rust ingress gateway for Aether phase 3a transparent proxy"
|
||||
|
||||
[dependencies]
|
||||
aether-admin.workspace = true
|
||||
aether-ai-pipeline.workspace = true
|
||||
aether-ai-serving.workspace = true
|
||||
aether-ai-surfaces.workspace = true
|
||||
aether-billing.workspace = true
|
||||
aether-cache.workspace = true
|
||||
aether-contracts.workspace = true
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
|
||||
const MAX_BUFFER_SIZE: usize = MAX_MESSAGE_SIZE;
|
||||
const MAX_ERRORS: usize = 5;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct KiroToClaudeCliStreamState {
|
||||
decoder: EventStreamDecoder,
|
||||
state: KiroClaudeStreamState,
|
||||
started: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct KiroClaudeStreamState {
|
||||
model: String,
|
||||
thinking_enabled: bool,
|
||||
estimated_input_tokens: usize,
|
||||
message_id: String,
|
||||
output_tokens: usize,
|
||||
context_input_tokens: Option<usize>,
|
||||
next_block_index: usize,
|
||||
open_blocks: BTreeMap<usize, String>,
|
||||
text_block_index: Option<usize>,
|
||||
thinking_block_index: Option<usize>,
|
||||
tool_block_indices: BTreeMap<String, usize>,
|
||||
thinking_buffer: String,
|
||||
in_thinking_block: bool,
|
||||
thinking_extracted: bool,
|
||||
strip_thinking_leading_newline: bool,
|
||||
has_tool_use: bool,
|
||||
stop_reason_override: Option<String>,
|
||||
had_error: bool,
|
||||
last_content: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct EventStreamDecoder {
|
||||
buffer: Vec<u8>,
|
||||
error_count: usize,
|
||||
stopped: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct AwsHeaders {
|
||||
values: BTreeMap<String, AwsHeaderValue>,
|
||||
}
|
||||
|
||||
enum AwsHeaderValue {
|
||||
Ignored,
|
||||
String(String),
|
||||
}
|
||||
|
||||
struct AwsEventFrame {
|
||||
headers: AwsHeaders,
|
||||
payload: Vec<u8>,
|
||||
}
|
||||
|
||||
enum FrameParseError {
|
||||
Incomplete,
|
||||
Invalid(String),
|
||||
}
|
||||
|
||||
#[path = "stream/decoder.rs"]
|
||||
mod decoder;
|
||||
#[path = "stream/state.rs"]
|
||||
mod state;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "stream/tests.rs"]
|
||||
mod tests;
|
||||
@@ -1,222 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
AwsEventFrame, AwsHeaderValue, AwsHeaders, EventStreamDecoder, FrameParseError,
|
||||
MAX_BUFFER_SIZE, MAX_ERRORS, MAX_MESSAGE_SIZE,
|
||||
};
|
||||
use crate::ai_pipeline::kiro_crc32 as crc32;
|
||||
|
||||
impl EventStreamDecoder {
|
||||
pub(super) fn feed(&mut self, data: &[u8]) -> Result<(), String> {
|
||||
if self.stopped || data.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let new_size = self.buffer.len() + data.len();
|
||||
if new_size > MAX_BUFFER_SIZE {
|
||||
self.stopped = true;
|
||||
return Err(format!(
|
||||
"buffer overflow: size={new_size} max={MAX_BUFFER_SIZE}"
|
||||
));
|
||||
}
|
||||
self.buffer.extend_from_slice(data);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn decode_available(&mut self) -> Result<Vec<AwsEventFrame>, String> {
|
||||
let mut out = Vec::new();
|
||||
if self.stopped {
|
||||
return Ok(out);
|
||||
}
|
||||
|
||||
loop {
|
||||
match parse_frame(&self.buffer) {
|
||||
Ok(Some((frame, consumed))) => {
|
||||
if consumed == 0 {
|
||||
break;
|
||||
}
|
||||
out.push(frame);
|
||||
self.buffer.drain(..consumed);
|
||||
self.error_count = 0;
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(FrameParseError::Incomplete) => break,
|
||||
Err(FrameParseError::Invalid(message)) => {
|
||||
self.error_count += 1;
|
||||
if self.error_count >= MAX_ERRORS {
|
||||
self.stopped = true;
|
||||
return Err(message);
|
||||
}
|
||||
if self.buffer.is_empty() {
|
||||
break;
|
||||
}
|
||||
self.buffer.drain(..1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
impl AwsHeaders {
|
||||
fn get_string(&self, name: &str) -> Option<&str> {
|
||||
match self.values.get(name) {
|
||||
Some(AwsHeaderValue::String(value)) => Some(value.as_str()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn message_type(&self) -> Option<&str> {
|
||||
self.get_string(":message-type")
|
||||
}
|
||||
|
||||
pub(super) fn event_type(&self) -> Option<&str> {
|
||||
self.get_string(":event-type")
|
||||
}
|
||||
|
||||
pub(super) fn exception_type(&self) -> Option<&str> {
|
||||
self.get_string(":exception-type")
|
||||
}
|
||||
|
||||
pub(super) fn error_code(&self) -> Option<&str> {
|
||||
self.get_string(":error-code")
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_frame(buffer: &[u8]) -> Result<Option<(AwsEventFrame, usize)>, FrameParseError> {
|
||||
if buffer.len() < 12 {
|
||||
return Ok(None);
|
||||
}
|
||||
let total_length = u32::from_be_bytes(buffer[0..4].try_into().expect("slice size")) as usize;
|
||||
let header_length = u32::from_be_bytes(buffer[4..8].try_into().expect("slice size")) as usize;
|
||||
let prelude_crc = u32::from_be_bytes(buffer[8..12].try_into().expect("slice size"));
|
||||
|
||||
if total_length < 16 {
|
||||
return Err(FrameParseError::Invalid(format!(
|
||||
"message too small: length={total_length}"
|
||||
)));
|
||||
}
|
||||
if total_length > MAX_MESSAGE_SIZE {
|
||||
return Err(FrameParseError::Invalid(format!(
|
||||
"message too large: length={total_length}"
|
||||
)));
|
||||
}
|
||||
if buffer.len() < total_length {
|
||||
return Ok(None);
|
||||
}
|
||||
if crc32(&buffer[0..8]) != prelude_crc {
|
||||
return Err(FrameParseError::Invalid("prelude crc mismatch".to_string()));
|
||||
}
|
||||
let message_crc = u32::from_be_bytes(
|
||||
buffer[total_length - 4..total_length]
|
||||
.try_into()
|
||||
.expect("slice size"),
|
||||
);
|
||||
if crc32(&buffer[..total_length - 4]) != message_crc {
|
||||
return Err(FrameParseError::Invalid("message crc mismatch".to_string()));
|
||||
}
|
||||
|
||||
let headers_start = 12;
|
||||
let headers_end = headers_start + header_length;
|
||||
if headers_end > total_length - 4 {
|
||||
return Err(FrameParseError::Invalid(
|
||||
"header length exceeds frame boundary".to_string(),
|
||||
));
|
||||
}
|
||||
let headers = parse_headers(&buffer[headers_start..headers_end], header_length)?;
|
||||
let payload = buffer[headers_end..total_length - 4].to_vec();
|
||||
Ok(Some((AwsEventFrame { headers, payload }, total_length)))
|
||||
}
|
||||
|
||||
fn parse_headers(data: &[u8], header_length: usize) -> Result<AwsHeaders, FrameParseError> {
|
||||
if data.len() < header_length {
|
||||
return Err(FrameParseError::Incomplete);
|
||||
}
|
||||
let mut values = BTreeMap::new();
|
||||
let mut offset = 0usize;
|
||||
while offset < header_length {
|
||||
ensure_header_bytes(data, offset, 1)?;
|
||||
let name_len = data[offset] as usize;
|
||||
offset += 1;
|
||||
if name_len == 0 {
|
||||
return Err(FrameParseError::Invalid(
|
||||
"header name length cannot be 0".to_string(),
|
||||
));
|
||||
}
|
||||
ensure_header_bytes(data, offset, name_len)?;
|
||||
let name = String::from_utf8_lossy(&data[offset..offset + name_len]).to_string();
|
||||
offset += name_len;
|
||||
|
||||
ensure_header_bytes(data, offset, 1)?;
|
||||
let value_type = data[offset];
|
||||
offset += 1;
|
||||
|
||||
let value = match value_type {
|
||||
0 => AwsHeaderValue::Ignored,
|
||||
1 => AwsHeaderValue::Ignored,
|
||||
2 => {
|
||||
ensure_header_bytes(data, offset, 1)?;
|
||||
let _ = i8::from_be_bytes([data[offset]]);
|
||||
offset += 1;
|
||||
AwsHeaderValue::Ignored
|
||||
}
|
||||
3 => {
|
||||
ensure_header_bytes(data, offset, 2)?;
|
||||
let _ = i16::from_be_bytes(data[offset..offset + 2].try_into().expect("slice"));
|
||||
offset += 2;
|
||||
AwsHeaderValue::Ignored
|
||||
}
|
||||
4 => {
|
||||
ensure_header_bytes(data, offset, 4)?;
|
||||
let _ = i32::from_be_bytes(data[offset..offset + 4].try_into().expect("slice"));
|
||||
offset += 4;
|
||||
AwsHeaderValue::Ignored
|
||||
}
|
||||
5 | 8 => {
|
||||
ensure_header_bytes(data, offset, 8)?;
|
||||
let _ = i64::from_be_bytes(data[offset..offset + 8].try_into().expect("slice"));
|
||||
offset += 8;
|
||||
AwsHeaderValue::Ignored
|
||||
}
|
||||
6 => {
|
||||
ensure_header_bytes(data, offset, 2)?;
|
||||
let length = u16::from_be_bytes(data[offset..offset + 2].try_into().expect("slice"))
|
||||
as usize;
|
||||
offset += 2;
|
||||
ensure_header_bytes(data, offset, length)?;
|
||||
offset += length;
|
||||
AwsHeaderValue::Ignored
|
||||
}
|
||||
7 => {
|
||||
ensure_header_bytes(data, offset, 2)?;
|
||||
let length = u16::from_be_bytes(data[offset..offset + 2].try_into().expect("slice"))
|
||||
as usize;
|
||||
offset += 2;
|
||||
ensure_header_bytes(data, offset, length)?;
|
||||
let out = String::from_utf8_lossy(&data[offset..offset + length]).to_string();
|
||||
offset += length;
|
||||
AwsHeaderValue::String(out)
|
||||
}
|
||||
9 => {
|
||||
ensure_header_bytes(data, offset, 16)?;
|
||||
offset += 16;
|
||||
AwsHeaderValue::Ignored
|
||||
}
|
||||
other => {
|
||||
return Err(FrameParseError::Invalid(format!(
|
||||
"invalid header type: {other}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
values.insert(name, value);
|
||||
}
|
||||
Ok(AwsHeaders { values })
|
||||
}
|
||||
|
||||
fn ensure_header_bytes(data: &[u8], offset: usize, needed: usize) -> Result<(), FrameParseError> {
|
||||
let available = data.len().saturating_sub(offset);
|
||||
if available < needed {
|
||||
return Err(FrameParseError::Incomplete);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,8 +0,0 @@
|
||||
#[path = "state/blocks.rs"]
|
||||
mod blocks;
|
||||
#[path = "state/events.rs"]
|
||||
mod events;
|
||||
#[path = "state/finalize.rs"]
|
||||
mod finalize;
|
||||
#[path = "state/lifecycle.rs"]
|
||||
mod lifecycle;
|
||||
@@ -1,97 +0,0 @@
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::super::KiroClaudeStreamState;
|
||||
|
||||
impl KiroClaudeStreamState {
|
||||
pub(super) fn ensure_text_block_open(&mut self) -> Vec<Value> {
|
||||
if let Some(idx) = self.text_block_index {
|
||||
if self
|
||||
.open_blocks
|
||||
.get(&idx)
|
||||
.map(|value| value == "text")
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Vec::new();
|
||||
}
|
||||
}
|
||||
let idx = self.next_block_index;
|
||||
self.next_block_index += 1;
|
||||
self.text_block_index = Some(idx);
|
||||
self.open_blocks.insert(idx, "text".to_string());
|
||||
vec![json!({
|
||||
"type": "content_block_start",
|
||||
"index": idx,
|
||||
"content_block": {"type": "text", "text": ""}
|
||||
})]
|
||||
}
|
||||
|
||||
pub(super) fn ensure_thinking_block_open(&mut self) -> Vec<Value> {
|
||||
if let Some(idx) = self.thinking_block_index {
|
||||
if self
|
||||
.open_blocks
|
||||
.get(&idx)
|
||||
.map(|value| value == "thinking")
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Vec::new();
|
||||
}
|
||||
}
|
||||
let idx = self.next_block_index;
|
||||
self.next_block_index += 1;
|
||||
self.thinking_block_index = Some(idx);
|
||||
self.open_blocks.insert(idx, "thinking".to_string());
|
||||
vec![json!({
|
||||
"type": "content_block_start",
|
||||
"index": idx,
|
||||
"content_block": {"type": "thinking", "thinking": ""}
|
||||
})]
|
||||
}
|
||||
|
||||
pub(super) fn close_block(&mut self, idx: usize) -> Vec<Value> {
|
||||
if self.open_blocks.remove(&idx).is_none() {
|
||||
return Vec::new();
|
||||
}
|
||||
vec![json!({"type": "content_block_stop", "index": idx})]
|
||||
}
|
||||
|
||||
pub(super) fn emit_text_delta(&mut self, text: &str) -> Vec<Value> {
|
||||
if text.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let mut events = self.ensure_text_block_open();
|
||||
let idx = self.text_block_index.unwrap_or_default();
|
||||
events.push(json!({
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "text_delta", "text": text}
|
||||
}));
|
||||
events
|
||||
}
|
||||
|
||||
pub(super) fn emit_thinking_delta(&mut self, thinking: &str) -> Vec<Value> {
|
||||
if thinking.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let mut events = self.ensure_thinking_block_open();
|
||||
let idx = self.thinking_block_index.unwrap_or_default();
|
||||
events.push(json!({
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "thinking_delta", "thinking": thinking}
|
||||
}));
|
||||
events
|
||||
}
|
||||
|
||||
pub(super) fn close_thinking_block(&mut self) -> Vec<Value> {
|
||||
let Some(idx) = self.thinking_block_index else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut events = vec![json!({
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "thinking_delta", "thinking": ""}
|
||||
})];
|
||||
events.extend(self.close_block(idx));
|
||||
events
|
||||
}
|
||||
}
|
||||
@@ -1,327 +0,0 @@
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
calculate_kiro_context_input_tokens, encode_kiro_sse_events, estimate_kiro_tokens,
|
||||
find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
|
||||
find_kiro_real_thinking_start_tag, KIRO_MAX_THINKING_BUFFER,
|
||||
};
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
use super::super::AwsEventFrame;
|
||||
use super::super::KiroClaudeStreamState;
|
||||
|
||||
fn floor_char_boundary(text: &str, index: usize) -> usize {
|
||||
let mut boundary = index.min(text.len());
|
||||
while boundary > 0 && !text.is_char_boundary(boundary) {
|
||||
boundary -= 1;
|
||||
}
|
||||
boundary
|
||||
}
|
||||
|
||||
fn split_preserving_trailing_bytes(
|
||||
buffer: &str,
|
||||
trailing_bytes: usize,
|
||||
) -> Option<(String, String)> {
|
||||
if buffer.len() <= trailing_bytes {
|
||||
return None;
|
||||
}
|
||||
|
||||
let split = floor_char_boundary(buffer, buffer.len() - trailing_bytes);
|
||||
if split == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some((buffer[..split].to_string(), buffer[split..].to_string()))
|
||||
}
|
||||
|
||||
impl KiroClaudeStreamState {
|
||||
pub(super) fn process_frame(&mut self, frame: AwsEventFrame) -> Result<Vec<u8>, GatewayError> {
|
||||
let message_type = frame.headers.message_type().unwrap_or("event");
|
||||
match message_type {
|
||||
"event" => self.process_event_frame(frame),
|
||||
"exception" => self.process_exception_frame(frame),
|
||||
"error" => self.process_error_frame(frame),
|
||||
_ => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn process_event_frame(
|
||||
&mut self,
|
||||
frame: AwsEventFrame,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let event_type = frame.headers.event_type().unwrap_or_default();
|
||||
let payload: Value = if frame.payload.is_empty() {
|
||||
json!({})
|
||||
} else {
|
||||
serde_json::from_slice(&frame.payload).unwrap_or_else(|_| json!({}))
|
||||
};
|
||||
let payload_object = payload.as_object();
|
||||
let mut events = Vec::new();
|
||||
match event_type {
|
||||
"assistantResponseEvent" => {
|
||||
if let Some(content) = payload_object
|
||||
.and_then(|value| value.get("content"))
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
events.extend(self.process_assistant_response(content));
|
||||
}
|
||||
}
|
||||
"toolUseEvent" => {
|
||||
if let Some(payload_object) = payload_object {
|
||||
let name = payload_object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let tool_use_id = payload_object
|
||||
.get("toolUseId")
|
||||
.or_else(|| payload_object.get("tool_use_id"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let input_json = match payload_object.get("input") {
|
||||
None | Some(Value::Null) => String::new(),
|
||||
Some(Value::String(text)) => text.clone(),
|
||||
Some(other) => serde_json::to_string(other)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
};
|
||||
let stop = payload_object
|
||||
.get("stop")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
events.extend(self.process_tool_use(name, tool_use_id, &input_json, stop));
|
||||
}
|
||||
}
|
||||
"contextUsageEvent" => {
|
||||
if let Some(percentage) = payload_object
|
||||
.and_then(|value| value.get("contextUsagePercentage"))
|
||||
.and_then(Value::as_f64)
|
||||
{
|
||||
self.context_input_tokens =
|
||||
Some(calculate_kiro_context_input_tokens(percentage));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
encode_kiro_sse_events(events).map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(super) fn process_exception_frame(
|
||||
&mut self,
|
||||
frame: AwsEventFrame,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let exception_type = frame
|
||||
.headers
|
||||
.exception_type()
|
||||
.unwrap_or("UnknownException")
|
||||
.to_string();
|
||||
if exception_type == "ContentLengthExceededException" {
|
||||
self.stop_reason_override = Some("max_tokens".to_string());
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.emit_stream_error("upstream_exception", &exception_type)
|
||||
}
|
||||
|
||||
pub(super) fn process_error_frame(
|
||||
&mut self,
|
||||
frame: AwsEventFrame,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let error_code = frame
|
||||
.headers
|
||||
.error_code()
|
||||
.unwrap_or("UnknownError")
|
||||
.to_string();
|
||||
self.emit_stream_error("upstream_error", &error_code)
|
||||
}
|
||||
|
||||
pub(super) fn process_assistant_response(&mut self, content: &str) -> Vec<Value> {
|
||||
if content.is_empty() || content == self.last_content {
|
||||
return Vec::new();
|
||||
}
|
||||
self.last_content = content.to_string();
|
||||
self.output_tokens += estimate_kiro_tokens(content);
|
||||
|
||||
if !self.thinking_enabled {
|
||||
return self.emit_text_delta(content);
|
||||
}
|
||||
|
||||
self.thinking_buffer.push_str(content);
|
||||
if self.thinking_buffer.len() > KIRO_MAX_THINKING_BUFFER {
|
||||
let overflow = std::mem::take(&mut self.thinking_buffer);
|
||||
if self.in_thinking_block {
|
||||
let mut events = self.emit_thinking_delta(&overflow);
|
||||
events.extend(self.close_thinking_block());
|
||||
self.in_thinking_block = false;
|
||||
self.thinking_extracted = true;
|
||||
return events;
|
||||
}
|
||||
return self.emit_text_delta(&overflow);
|
||||
}
|
||||
|
||||
let mut events = Vec::new();
|
||||
loop {
|
||||
if !self.in_thinking_block && !self.thinking_extracted {
|
||||
if let Some(start_pos) = find_kiro_real_thinking_start_tag(&self.thinking_buffer) {
|
||||
let before = self.thinking_buffer[..start_pos].to_string();
|
||||
if !before.trim().is_empty() {
|
||||
events.extend(self.emit_text_delta(&before));
|
||||
}
|
||||
self.in_thinking_block = true;
|
||||
self.strip_thinking_leading_newline = true;
|
||||
self.thinking_buffer =
|
||||
self.thinking_buffer[start_pos + "<thinking>".len()..].to_string();
|
||||
events.extend(self.ensure_thinking_block_open());
|
||||
continue;
|
||||
}
|
||||
|
||||
let keep = "<thinking>".len();
|
||||
if let Some((safe, remaining)) =
|
||||
split_preserving_trailing_bytes(&self.thinking_buffer, keep)
|
||||
{
|
||||
if !safe.trim().is_empty() {
|
||||
events.extend(self.emit_text_delta(&safe));
|
||||
self.thinking_buffer = remaining;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
if self.in_thinking_block {
|
||||
if self.strip_thinking_leading_newline {
|
||||
if self.thinking_buffer.starts_with('\n') {
|
||||
self.thinking_buffer.remove(0);
|
||||
self.strip_thinking_leading_newline = false;
|
||||
} else if !self.thinking_buffer.is_empty() {
|
||||
self.strip_thinking_leading_newline = false;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(end_pos) = find_kiro_real_thinking_end_tag(&self.thinking_buffer) {
|
||||
let thinking_text = self.thinking_buffer[..end_pos].to_string();
|
||||
if !thinking_text.is_empty() {
|
||||
events.extend(self.emit_thinking_delta(&thinking_text));
|
||||
}
|
||||
events.extend(self.close_thinking_block());
|
||||
self.in_thinking_block = false;
|
||||
self.thinking_extracted = true;
|
||||
self.thinking_buffer =
|
||||
self.thinking_buffer[end_pos + "</thinking>".len()..].to_string();
|
||||
continue;
|
||||
}
|
||||
|
||||
let keep = "</thinking>".len();
|
||||
if let Some((safe, remaining)) =
|
||||
split_preserving_trailing_bytes(&self.thinking_buffer, keep)
|
||||
{
|
||||
if !safe.is_empty() {
|
||||
events.extend(self.emit_thinking_delta(&safe));
|
||||
self.thinking_buffer = remaining;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
if !self.thinking_buffer.is_empty() {
|
||||
let remaining = std::mem::take(&mut self.thinking_buffer);
|
||||
events.extend(self.emit_text_delta(&remaining));
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
events
|
||||
}
|
||||
|
||||
pub(super) fn process_tool_use(
|
||||
&mut self,
|
||||
name: &str,
|
||||
tool_use_id: &str,
|
||||
input_json: &str,
|
||||
stop: bool,
|
||||
) -> Vec<Value> {
|
||||
if tool_use_id.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
self.has_tool_use = true;
|
||||
let mut events = Vec::new();
|
||||
|
||||
if self.thinking_enabled && self.in_thinking_block && !self.thinking_buffer.is_empty() {
|
||||
if let Some(end_pos) =
|
||||
find_kiro_real_thinking_end_tag_at_buffer_end(&self.thinking_buffer)
|
||||
{
|
||||
let thinking_text = self.thinking_buffer[..end_pos].to_string();
|
||||
if !thinking_text.is_empty() {
|
||||
events.extend(self.emit_thinking_delta(&thinking_text));
|
||||
}
|
||||
events.extend(self.close_thinking_block());
|
||||
let remaining = self.thinking_buffer[end_pos + "</thinking>".len()..].to_string();
|
||||
self.thinking_buffer.clear();
|
||||
self.in_thinking_block = false;
|
||||
self.thinking_extracted = true;
|
||||
if !remaining.is_empty() {
|
||||
events.extend(self.emit_text_delta(&remaining));
|
||||
}
|
||||
} else {
|
||||
let thinking = std::mem::take(&mut self.thinking_buffer);
|
||||
events.extend(self.emit_thinking_delta(&thinking));
|
||||
events.extend(self.close_thinking_block());
|
||||
self.in_thinking_block = false;
|
||||
self.thinking_extracted = true;
|
||||
}
|
||||
}
|
||||
|
||||
if self.thinking_enabled
|
||||
&& !self.in_thinking_block
|
||||
&& !self.thinking_extracted
|
||||
&& !self.thinking_buffer.is_empty()
|
||||
{
|
||||
let buffered = std::mem::take(&mut self.thinking_buffer);
|
||||
events.extend(self.emit_text_delta(&buffered));
|
||||
}
|
||||
|
||||
if let Some(idx) = self.text_block_index.take() {
|
||||
events.extend(self.close_block(idx));
|
||||
}
|
||||
|
||||
let block_index = if let Some(block_index) = self.tool_block_indices.get(tool_use_id) {
|
||||
*block_index
|
||||
} else {
|
||||
let block_index = self.next_block_index;
|
||||
self.next_block_index += 1;
|
||||
self.tool_block_indices
|
||||
.insert(tool_use_id.to_string(), block_index);
|
||||
block_index
|
||||
};
|
||||
|
||||
if let std::collections::btree_map::Entry::Vacant(e) = self.open_blocks.entry(block_index) {
|
||||
e.insert("tool_use".to_string());
|
||||
events.push(json!({
|
||||
"type": "content_block_start",
|
||||
"index": block_index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": tool_use_id,
|
||||
"name": name,
|
||||
"input": {},
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
if !input_json.is_empty() {
|
||||
self.output_tokens += estimate_kiro_tokens(input_json);
|
||||
events.push(json!({
|
||||
"type": "content_block_delta",
|
||||
"index": block_index,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": input_json,
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
if stop {
|
||||
events.extend(self.close_block(block_index));
|
||||
}
|
||||
|
||||
events
|
||||
}
|
||||
}
|
||||
@@ -1,96 +0,0 @@
|
||||
use crate::ai_pipeline::{
|
||||
build_kiro_final_message_sse_events, encode_kiro_sse_events,
|
||||
find_kiro_real_thinking_end_tag_at_buffer_end,
|
||||
};
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
use super::super::KiroClaudeStreamState;
|
||||
|
||||
impl KiroClaudeStreamState {
|
||||
pub(super) fn finalize(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
if self.thinking_enabled && !self.thinking_buffer.is_empty() {
|
||||
let flush_events = if self.in_thinking_block {
|
||||
if let Some(end_pos) =
|
||||
find_kiro_real_thinking_end_tag_at_buffer_end(&self.thinking_buffer)
|
||||
{
|
||||
let thinking_text = self.thinking_buffer[..end_pos].to_string();
|
||||
let mut events = Vec::new();
|
||||
if !thinking_text.is_empty() {
|
||||
events.extend(self.emit_thinking_delta(&thinking_text));
|
||||
}
|
||||
events.extend(self.close_thinking_block());
|
||||
let remaining =
|
||||
self.thinking_buffer[end_pos + "</thinking>".len()..].to_string();
|
||||
if !remaining.is_empty() {
|
||||
events.extend(self.emit_text_delta(&remaining));
|
||||
}
|
||||
events
|
||||
} else {
|
||||
let mut events = self.emit_thinking_delta(&self.thinking_buffer.clone());
|
||||
events.extend(self.close_thinking_block());
|
||||
events
|
||||
}
|
||||
} else {
|
||||
self.emit_text_delta(&self.thinking_buffer.clone())
|
||||
};
|
||||
self.thinking_buffer.clear();
|
||||
self.in_thinking_block = false;
|
||||
self.thinking_extracted = true;
|
||||
let mut output = encode_kiro_sse_events(flush_events)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
for idx in self
|
||||
.open_blocks
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
{
|
||||
output.extend(
|
||||
encode_kiro_sse_events(self.close_block(idx))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(self.final_message_bytes()?);
|
||||
return Ok(output);
|
||||
}
|
||||
|
||||
let mut output = Vec::new();
|
||||
for idx in self
|
||||
.open_blocks
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
{
|
||||
output.extend(
|
||||
encode_kiro_sse_events(self.close_block(idx))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(self.final_message_bytes()?);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub(super) fn final_message_bytes(&self) -> Result<Vec<u8>, GatewayError> {
|
||||
let stop_reason = self.stop_reason_override.clone().unwrap_or_else(|| {
|
||||
if self.has_tool_use {
|
||||
"tool_use"
|
||||
} else {
|
||||
"end_turn"
|
||||
}
|
||||
.to_string()
|
||||
});
|
||||
let input_tokens = self
|
||||
.context_input_tokens
|
||||
.unwrap_or(self.estimated_input_tokens) as u64;
|
||||
encode_kiro_sse_events(build_kiro_final_message_sse_events(
|
||||
&stop_reason,
|
||||
input_tokens as usize,
|
||||
self.output_tokens,
|
||||
))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
}
|
||||
@@ -1,129 +0,0 @@
|
||||
use serde_json::Value;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
build_kiro_initial_sse_events, build_kiro_stream_error_sse_events, encode_kiro_sse_events,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
use super::super::{EventStreamDecoder, KiroClaudeStreamState, KiroToClaudeCliStreamState};
|
||||
|
||||
impl KiroToClaudeCliStreamState {
|
||||
pub(crate) fn new(report_context: &Value) -> Self {
|
||||
Self {
|
||||
decoder: EventStreamDecoder::default(),
|
||||
state: KiroClaudeStreamState::new(report_context),
|
||||
started: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn push_chunk(
|
||||
&mut self,
|
||||
_report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let mut output = Vec::new();
|
||||
if !self.started {
|
||||
self.started = true;
|
||||
output.extend(self.state.generate_initial_bytes()?);
|
||||
}
|
||||
|
||||
if let Err(err) = self.decoder.feed(chunk) {
|
||||
output.extend(
|
||||
self.state
|
||||
.emit_stream_error("upstream_stream_error", &err)?,
|
||||
);
|
||||
return Ok(output);
|
||||
}
|
||||
|
||||
match self.decoder.decode_available() {
|
||||
Ok(frames) => {
|
||||
for frame in frames {
|
||||
output.extend(self.state.process_frame(frame)?);
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
output.extend(
|
||||
self.state
|
||||
.emit_stream_error("upstream_stream_error", &err)?,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub(crate) fn finish(&mut self, _report_context: &Value) -> Result<Vec<u8>, GatewayError> {
|
||||
if !self.started || self.state.had_error {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.state.finalize()
|
||||
}
|
||||
}
|
||||
|
||||
impl KiroClaudeStreamState {
|
||||
pub(super) fn new(report_context: &Value) -> Self {
|
||||
let model = report_context
|
||||
.get("mapped_model")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| {
|
||||
report_context
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
})
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
let thinking_enabled = report_context
|
||||
.get("original_request_body")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|body| body.get("thinking"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|thinking| thinking.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.map(|value| {
|
||||
value.trim().eq_ignore_ascii_case("enabled")
|
||||
|| value.trim().eq_ignore_ascii_case("adaptive")
|
||||
})
|
||||
.unwrap_or(false);
|
||||
let estimated_input_tokens = report_context
|
||||
.get("input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|value| value as usize)
|
||||
.unwrap_or(0);
|
||||
Self {
|
||||
model,
|
||||
thinking_enabled,
|
||||
estimated_input_tokens,
|
||||
message_id: format!("msg_{}", Uuid::new_v4().simple()),
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn generate_initial_bytes(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
let events = build_kiro_initial_sse_events(
|
||||
&self.message_id,
|
||||
&self.model,
|
||||
self.estimated_input_tokens,
|
||||
);
|
||||
let mut events = events;
|
||||
if !self.thinking_enabled {
|
||||
events.extend(self.ensure_text_block_open());
|
||||
}
|
||||
encode_kiro_sse_events(events).map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(super) fn emit_stream_error(
|
||||
&mut self,
|
||||
error_type: &str,
|
||||
message: &str,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
if self.had_error {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.had_error = true;
|
||||
encode_kiro_sse_events(build_kiro_stream_error_sse_events(error_type, message))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
}
|
||||
@@ -1,166 +0,0 @@
|
||||
use aether_ai_pipeline::adaptation::kiro_stream::kiro_crc32 as crc32;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::KiroToClaudeCliStreamState;
|
||||
|
||||
fn encode_string_header(name: &str, value: &str) -> Vec<u8> {
|
||||
let mut out = Vec::new();
|
||||
out.push(name.len() as u8);
|
||||
out.extend_from_slice(name.as_bytes());
|
||||
out.push(7);
|
||||
out.extend_from_slice(&(value.len() as u16).to_be_bytes());
|
||||
out.extend_from_slice(value.as_bytes());
|
||||
out
|
||||
}
|
||||
|
||||
fn encode_event_frame(message_type: &str, event_type: Option<&str>, payload: &Value) -> Vec<u8> {
|
||||
let mut headers = encode_string_header(":message-type", message_type);
|
||||
if let Some(event_type) = event_type {
|
||||
headers.extend_from_slice(&encode_string_header(":event-type", event_type));
|
||||
}
|
||||
let payload_bytes = serde_json::to_vec(payload).expect("payload should encode");
|
||||
encode_frame(headers, payload_bytes)
|
||||
}
|
||||
|
||||
fn encode_frame(headers: Vec<u8>, payload: Vec<u8>) -> Vec<u8> {
|
||||
let total_len = 12 + headers.len() + payload.len() + 4;
|
||||
let header_len = headers.len();
|
||||
let mut out = Vec::with_capacity(total_len);
|
||||
out.extend_from_slice(&(total_len as u32).to_be_bytes());
|
||||
out.extend_from_slice(&(header_len as u32).to_be_bytes());
|
||||
let prelude_crc = crc32(&out[..8]);
|
||||
out.extend_from_slice(&prelude_crc.to_be_bytes());
|
||||
out.extend_from_slice(&headers);
|
||||
out.extend_from_slice(&payload);
|
||||
let message_crc = crc32(&out);
|
||||
out.extend_from_slice(&message_crc.to_be_bytes());
|
||||
out
|
||||
}
|
||||
|
||||
fn kiro_report_context(thinking_enabled: bool) -> Value {
|
||||
let mut context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"envelope_name": "kiro:generateAssistantResponse",
|
||||
"mapped_model": "claude-sonnet-4.5"
|
||||
});
|
||||
if thinking_enabled {
|
||||
context["original_request_body"] = json!({
|
||||
"thinking": {
|
||||
"type": "enabled"
|
||||
}
|
||||
});
|
||||
}
|
||||
context
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_rewriter_converts_text_events_to_claude_sse() {
|
||||
let report_context = kiro_report_context(false);
|
||||
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
|
||||
let chunk = [
|
||||
encode_event_frame(
|
||||
"event",
|
||||
Some("assistantResponseEvent"),
|
||||
&json!({"content": "Hello from Kiro"}),
|
||||
),
|
||||
encode_event_frame(
|
||||
"event",
|
||||
Some("contextUsageEvent"),
|
||||
&json!({"contextUsagePercentage": 1.0}),
|
||||
),
|
||||
]
|
||||
.concat();
|
||||
|
||||
let first = rewriter
|
||||
.push_chunk(&report_context, &chunk)
|
||||
.expect("rewrite should succeed");
|
||||
let rest = rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed");
|
||||
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
|
||||
assert!(text.contains("event: message_start"));
|
||||
assert!(text.contains("\"type\":\"content_block_delta\""));
|
||||
assert!(text.contains("Hello from Kiro"));
|
||||
assert!(text.contains("\"stop_reason\":\"end_turn\""));
|
||||
assert!(text.contains("\"input_tokens\":2000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_rewriter_converts_tool_use_to_claude_events() {
|
||||
let report_context = kiro_report_context(false);
|
||||
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
|
||||
let chunk = [
|
||||
encode_event_frame(
|
||||
"event",
|
||||
Some("assistantResponseEvent"),
|
||||
&json!({"content": "Need a tool."}),
|
||||
),
|
||||
encode_event_frame(
|
||||
"event",
|
||||
Some("toolUseEvent"),
|
||||
&json!({
|
||||
"name": "get_weather",
|
||||
"toolUseId": "tool_123",
|
||||
"input": {"city": "SF"},
|
||||
"stop": true
|
||||
}),
|
||||
),
|
||||
]
|
||||
.concat();
|
||||
|
||||
let first = rewriter
|
||||
.push_chunk(&report_context, &chunk)
|
||||
.expect("rewrite should succeed");
|
||||
let rest = rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed");
|
||||
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
|
||||
assert!(text.contains("\"type\":\"tool_use\""));
|
||||
assert!(text.contains("\"id\":\"tool_123\""));
|
||||
assert!(text.contains("\"name\":\"get_weather\""));
|
||||
assert!(text.contains("\"partial_json\":\"{\\\"city\\\":\\\"SF\\\"}\""));
|
||||
assert!(text.contains("\"stop_reason\":\"tool_use\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_rewriter_handles_multibyte_text_without_thinking_tag() {
|
||||
let report_context = kiro_report_context(true);
|
||||
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
|
||||
let chunk = encode_event_frame(
|
||||
"event",
|
||||
Some("assistantResponseEvent"),
|
||||
&json!({"content": "\n\n你好！有"}),
|
||||
);
|
||||
|
||||
let first = rewriter
|
||||
.push_chunk(&report_context, &chunk)
|
||||
.expect("rewrite should succeed");
|
||||
let rest = rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed");
|
||||
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
|
||||
assert!(text.contains("\"type\":\"text_delta\""));
|
||||
assert!(text.contains("你好！有"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_rewriter_handles_multibyte_text_inside_thinking_block() {
|
||||
let report_context = kiro_report_context(true);
|
||||
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
|
||||
let chunk = encode_event_frame(
|
||||
"event",
|
||||
Some("assistantResponseEvent"),
|
||||
&json!({"content": "<thinking>\n\n你好！有"}),
|
||||
);
|
||||
|
||||
let first = rewriter
|
||||
.push_chunk(&report_context, &chunk)
|
||||
.expect("rewrite should succeed");
|
||||
let rest = rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed");
|
||||
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
|
||||
assert!(text.contains("\"type\":\"thinking_delta\""));
|
||||
assert!(text.contains("你好！有"));
|
||||
}
|
||||
@@ -1,17 +0,0 @@
|
||||
#[path = "private_envelope/stream.rs"]
|
||||
mod stream;
|
||||
#[path = "private_envelope/sync.rs"]
|
||||
mod sync;
|
||||
#[cfg(test)]
|
||||
#[path = "private_envelope/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
pub(crate) use self::stream::{
|
||||
maybe_build_provider_private_stream_normalizer, ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
pub(crate) use self::sync::maybe_normalize_provider_private_sync_report_payload;
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
normalize_provider_private_report_context, normalize_provider_private_response_value,
|
||||
provider_private_response_allows_sync_finalize, stream_body_contains_error_event,
|
||||
transform_provider_private_stream_line,
|
||||
};
|
||||
@@ -1,96 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::adaptation::KiroToClaudeCliStreamState;
|
||||
use crate::ai_pipeline::{provider_adaptation_descriptor_for_envelope, KIRO_ENVELOPE_NAME};
|
||||
use crate::GatewayError;
|
||||
|
||||
use super::transform_provider_private_stream_line;
|
||||
|
||||
enum ProviderPrivateStreamNormalizeMode {
|
||||
EnvelopeUnwrap,
|
||||
KiroToClaudeCli(KiroToClaudeCliStreamState),
|
||||
}
|
||||
|
||||
pub(crate) struct ProviderPrivateStreamNormalizer<'a> {
|
||||
report_context: &'a Value,
|
||||
buffered: Vec<u8>,
|
||||
mode: ProviderPrivateStreamNormalizeMode,
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_provider_private_stream_normalizer<'a>(
|
||||
report_context: Option<&'a Value>,
|
||||
) -> Option<ProviderPrivateStreamNormalizer<'a>> {
|
||||
let report_context = report_context?;
|
||||
if !report_context
|
||||
.get("has_envelope")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let envelope_name = report_context
|
||||
.get("envelope_name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let provider_api_format = report_context
|
||||
.get("provider_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let descriptor =
|
||||
provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format)?;
|
||||
let mode = if descriptor
|
||||
.envelope_name
|
||||
.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
|
||||
{
|
||||
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(KiroToClaudeCliStreamState::new(
|
||||
report_context,
|
||||
))
|
||||
} else if descriptor.unwraps_response_envelope {
|
||||
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
Some(ProviderPrivateStreamNormalizer {
|
||||
report_context,
|
||||
buffered: Vec::new(),
|
||||
mode,
|
||||
})
|
||||
}
|
||||
|
||||
impl ProviderPrivateStreamNormalizer<'_> {
|
||||
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
|
||||
match &mut self.mode {
|
||||
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
|
||||
let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
|
||||
output.extend(
|
||||
transform_provider_private_stream_line(self.report_context, line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
match &mut self.mode {
|
||||
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => {
|
||||
state.finish(self.report_context)
|
||||
}
|
||||
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => {
|
||||
if self.buffered.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let line = std::mem::take(&mut self.buffered);
|
||||
transform_provider_private_stream_line(self.report_context, line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,39 +0,0 @@
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn normalizes_supported_private_report_context() {
|
||||
let report_context = json!({
|
||||
"has_envelope": true,
|
||||
"envelope_name": "antigravity:v1internal",
|
||||
"provider_api_format": "gemini:generate_content",
|
||||
});
|
||||
let normalized = normalize_provider_private_report_context(Some(&report_context))
|
||||
.expect("context should normalize");
|
||||
assert_eq!(normalized["has_envelope"], json!(false));
|
||||
assert!(normalized.get("envelope_name").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn private_stream_normalizer_unwraps_antigravity_stream() {
|
||||
let report_context = json!({
|
||||
"has_envelope": true,
|
||||
"provider_api_format": "gemini:generate_content",
|
||||
"client_api_format": "gemini:generate_content",
|
||||
"envelope_name": "antigravity:v1internal",
|
||||
"mapped_model": "claude-sonnet-4-5",
|
||||
});
|
||||
let mut normalizer = maybe_build_provider_private_stream_normalizer(Some(&report_context))
|
||||
.expect("normalizer should exist");
|
||||
let output = normalizer
|
||||
.push_chunk(
|
||||
b"data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{\"city\":\"SF\"}}}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"claude-sonnet-4-5\"},\"responseId\":\"resp_123\"}\n\n",
|
||||
)
|
||||
.expect("unwrap should succeed");
|
||||
let output_text = String::from_utf8(output).expect("text should decode");
|
||||
assert!(output_text.contains("\"_v1internal_response_id\":\"resp_123\""));
|
||||
assert!(output_text.contains("\"id\":\"call_get_weather_0\""));
|
||||
}
|
||||
@@ -1,22 +0,0 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
generic_decision_missing_exact_provider_request as generic_decision_missing_exact_provider_request_impl,
|
||||
GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
|
||||
pub(crate) fn generic_decision_missing_exact_provider_request(
|
||||
payload: &GatewayControlSyncDecisionResponse,
|
||||
) -> bool {
|
||||
if !generic_decision_missing_exact_provider_request_impl(payload) {
|
||||
return false;
|
||||
}
|
||||
|
||||
warn!(
|
||||
decision_kind = payload.decision_kind.as_deref().unwrap_or_default(),
|
||||
provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default(),
|
||||
client_api_format = payload.client_api_format.as_deref().unwrap_or_default(),
|
||||
"gateway generic decision missing exact provider request; falling back to plan"
|
||||
);
|
||||
true
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
pub(crate) mod control_payloads;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
core_error_background_report_kind, core_error_default_client_api_format,
|
||||
core_success_background_report_kind, implicit_sync_finalize_report_kind,
|
||||
ExecutionRuntimeAuthContext, GatewayControlPlanRequest, GatewayControlPlanResponse,
|
||||
GatewayControlSyncDecisionResponse, CLAUDE_CHAT_STREAM_PLAN_KIND,
|
||||
CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,
|
||||
CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND,
|
||||
CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND,
|
||||
CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
|
||||
GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND,
|
||||
GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
|
||||
GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
|
||||
GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND,
|
||||
GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
|
||||
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
|
||||
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
|
||||
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND,
|
||||
OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND,
|
||||
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND,
|
||||
OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
pub(crate) use control_payloads::generic_decision_missing_exact_provider_request;
|
||||
@@ -1,156 +0,0 @@
|
||||
#[cfg(test)]
|
||||
pub(crate) use crate::ai_pipeline::core_success_background_report_kind;
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
build_core_error_body_for_client_format, core_error_background_report_kind,
|
||||
core_error_default_client_api_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
request_candidate_api_format_preference, request_candidate_api_formats,
|
||||
request_conversion_direct_auth, request_conversion_enabled_for_transport,
|
||||
request_conversion_kind, request_conversion_requires_enable_flag,
|
||||
request_conversion_transport_supported, request_conversion_transport_unsupported_reason,
|
||||
request_pair_allowed_for_transport, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
|
||||
SyncCliResponseConversionKind,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
request_candidate_api_format_preference, request_candidate_api_formats,
|
||||
request_conversion_kind, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
|
||||
SyncCliResponseConversionKind,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn request_conversion_registry_supports_bidirectional_standard_matrix() {
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:chat", "openai:responses"),
|
||||
Some(RequestConversionKind::ToOpenAiResponses)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:chat", "claude:messages"),
|
||||
Some(RequestConversionKind::ToClaudeStandard)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:responses", "openai:chat"),
|
||||
Some(RequestConversionKind::ToOpenAIChat)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:generate_content", "claude:messages"),
|
||||
Some(RequestConversionKind::ToClaudeStandard)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:responses:compact", "gemini:generate_content"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:generate_content", "openai:responses:compact"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:chat", "openai:responses:compact"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("claude:messages", "claude:messages"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_response_conversion_registry_supports_bidirectional_standard_matrix() {
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("openai:chat", "claude:messages"),
|
||||
Some(SyncChatResponseConversionKind::ToClaudeChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("claude:messages", "gemini:generate_content"),
|
||||
Some(SyncChatResponseConversionKind::ToGeminiChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("gemini:generate_content", "openai:chat"),
|
||||
Some(SyncChatResponseConversionKind::ToOpenAIChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("openai:responses", "gemini:generate_content"),
|
||||
Some(SyncCliResponseConversionKind::ToGeminiCli)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("claude:messages", "openai:responses"),
|
||||
Some(SyncCliResponseConversionKind::ToOpenAiResponses)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("claude:messages", "openai:responses:compact"),
|
||||
Some(SyncCliResponseConversionKind::ToOpenAiResponses)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("openai:responses:compact", "claude:messages"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("gemini:generate_content", "claude:messages"),
|
||||
Some(SyncCliResponseConversionKind::ToClaudeCli)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_candidate_registry_excludes_compact_as_cross_format_target() {
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:chat", false),
|
||||
vec![
|
||||
"openai:chat",
|
||||
"openai:responses",
|
||||
"claude:messages",
|
||||
"gemini:generate_content",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:responses", false),
|
||||
vec![
|
||||
"openai:responses",
|
||||
"openai:chat",
|
||||
"claude:messages",
|
||||
"gemini:generate_content",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("claude:messages", false),
|
||||
vec![
|
||||
"claude:messages",
|
||||
"openai:chat",
|
||||
"openai:responses",
|
||||
"gemini:generate_content",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:cli", false),
|
||||
Vec::<&'static str>::new()
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("claude:cli", false),
|
||||
Vec::<&'static str>::new()
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:compact", false),
|
||||
Vec::<&'static str>::new()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_candidate_registry_prefers_same_kind_before_same_family_fallbacks() {
|
||||
assert_eq!(
|
||||
request_candidate_api_format_preference("claude:cli", "openai:responses"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_format_preference("claude:cli", "claude:chat"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_format_preference("claude:cli", "openai:chat"),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,501 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::adaptation::private_envelope::transform_provider_private_stream_line as transform_envelope_line;
|
||||
use crate::ai_pipeline::adaptation::KiroToClaudeCliStreamState;
|
||||
use crate::ai_pipeline::finalize::sse::encode_json_sse;
|
||||
use crate::ai_pipeline::finalize::standard::StreamingStandardConversionState;
|
||||
use crate::ai_pipeline::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
|
||||
use crate::GatewayError;
|
||||
|
||||
enum RewriteMode {
|
||||
EnvelopeUnwrap,
|
||||
OpenAiImage(OpenAiImageStreamState),
|
||||
Standard(StreamingStandardConversionState),
|
||||
KiroToClaudeCli(KiroToClaudeCliStreamState),
|
||||
KiroToClaudeCliThenStandard {
|
||||
kiro: KiroToClaudeCliStreamState,
|
||||
standard: StreamingStandardConversionState,
|
||||
},
|
||||
}
|
||||
|
||||
pub(crate) struct LocalStreamRewriter<'a> {
|
||||
report_context: &'a Value,
|
||||
buffered: Vec<u8>,
|
||||
mode: RewriteMode,
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_stream_rewriter<'a>(
|
||||
report_context: Option<&'a Value>,
|
||||
) -> Option<LocalStreamRewriter<'a>> {
|
||||
let report_context = report_context?;
|
||||
let mode = match resolve_finalize_stream_rewrite_mode(report_context)? {
|
||||
FinalizeStreamRewriteMode::EnvelopeUnwrap => RewriteMode::EnvelopeUnwrap,
|
||||
FinalizeStreamRewriteMode::OpenAiImage => {
|
||||
RewriteMode::OpenAiImage(OpenAiImageStreamState::default())
|
||||
}
|
||||
FinalizeStreamRewriteMode::Standard => {
|
||||
RewriteMode::Standard(StreamingStandardConversionState::default())
|
||||
}
|
||||
FinalizeStreamRewriteMode::KiroToClaudeCli => {
|
||||
RewriteMode::KiroToClaudeCli(KiroToClaudeCliStreamState::new(report_context))
|
||||
}
|
||||
FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard => {
|
||||
RewriteMode::KiroToClaudeCliThenStandard {
|
||||
kiro: KiroToClaudeCliStreamState::new(report_context),
|
||||
standard: StreamingStandardConversionState::default(),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalStreamRewriter {
|
||||
report_context,
|
||||
buffered: Vec::new(),
|
||||
mode,
|
||||
})
|
||||
}
|
||||
|
||||
impl LocalStreamRewriter<'_> {
|
||||
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
|
||||
if let RewriteMode::OpenAiImage(state) = &mut self.mode {
|
||||
return state.push_chunk(self.report_context, chunk);
|
||||
}
|
||||
if let RewriteMode::KiroToClaudeCli(state) = &mut self.mode {
|
||||
return state.push_chunk(self.report_context, chunk);
|
||||
}
|
||||
if let RewriteMode::KiroToClaudeCliThenStandard { kiro, standard } = &mut self.mode {
|
||||
let claude_bytes = kiro.push_chunk(self.report_context, chunk)?;
|
||||
return transform_standard_bytes(standard, self.report_context, claude_bytes);
|
||||
}
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
|
||||
let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
|
||||
output.extend(self.transform_line(line)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
if let RewriteMode::OpenAiImage(state) = &mut self.mode {
|
||||
return state.finish(self.report_context);
|
||||
}
|
||||
if let RewriteMode::KiroToClaudeCli(state) = &mut self.mode {
|
||||
return state.finish(self.report_context);
|
||||
}
|
||||
if let RewriteMode::KiroToClaudeCliThenStandard { kiro, standard } = &mut self.mode {
|
||||
let mut output = transform_standard_bytes(
|
||||
standard,
|
||||
self.report_context,
|
||||
kiro.finish(self.report_context)?,
|
||||
)?;
|
||||
output.extend(standard.finish(self.report_context)?);
|
||||
return Ok(output);
|
||||
}
|
||||
if self.buffered.is_empty() {
|
||||
match &mut self.mode {
|
||||
RewriteMode::Standard(state) => return state.finish(self.report_context),
|
||||
RewriteMode::OpenAiImage(_) => {}
|
||||
RewriteMode::KiroToClaudeCli(_) => {}
|
||||
RewriteMode::KiroToClaudeCliThenStandard { .. } => {}
|
||||
RewriteMode::EnvelopeUnwrap => {}
|
||||
}
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let line = std::mem::take(&mut self.buffered);
|
||||
let mut output = self.transform_line(line)?;
|
||||
match &mut self.mode {
|
||||
RewriteMode::Standard(state) => {
|
||||
output.extend(state.finish(self.report_context)?);
|
||||
}
|
||||
RewriteMode::OpenAiImage(_) => {}
|
||||
RewriteMode::KiroToClaudeCli(_) => {}
|
||||
RewriteMode::KiroToClaudeCliThenStandard { .. } => {}
|
||||
RewriteMode::EnvelopeUnwrap => {}
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn transform_line(&mut self, line: Vec<u8>) -> Result<Vec<u8>, GatewayError> {
|
||||
match &mut self.mode {
|
||||
RewriteMode::EnvelopeUnwrap => transform_envelope_line(self.report_context, line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string())),
|
||||
RewriteMode::OpenAiImage(_) => Ok(Vec::new()),
|
||||
RewriteMode::Standard(state) => state.transform_line(self.report_context, line),
|
||||
RewriteMode::KiroToClaudeCli(_) => Ok(Vec::new()),
|
||||
RewriteMode::KiroToClaudeCliThenStandard { .. } => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn transform_standard_bytes(
|
||||
standard: &mut StreamingStandardConversionState,
|
||||
report_context: &Value,
|
||||
bytes: Vec<u8>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
if bytes.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut output = Vec::new();
|
||||
for line in bytes.split_inclusive(|byte| *byte == b'\n') {
|
||||
output.extend(standard.transform_line(report_context, line.to_vec())?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct OpenAiImageStreamState {
|
||||
buffered: Vec<u8>,
|
||||
latest_image: Option<OpenAiImageFrame>,
|
||||
emitted_partial_count: u64,
|
||||
saw_upstream_partial: bool,
|
||||
emitted_failure: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct OpenAiImageFrame {
|
||||
b64_json: String,
|
||||
}
|
||||
|
||||
impl OpenAiImageStreamState {
|
||||
fn push_chunk(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(block_end) = find_sse_block_end(&self.buffered) {
|
||||
let block = self.buffered.drain(..block_end).collect::<Vec<_>>();
|
||||
output.extend(self.transform_block(report_context, &block)?);
|
||||
drain_sse_separator(&mut self.buffered);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, GatewayError> {
|
||||
if self.buffered.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let block = std::mem::take(&mut self.buffered);
|
||||
self.transform_block(report_context, &block)
|
||||
}
|
||||
|
||||
fn transform_block(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
block: &[u8],
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let text =
|
||||
std::str::from_utf8(block).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let mut event_name = None::<String>;
|
||||
let mut data_lines = Vec::new();
|
||||
for raw_line in text.lines() {
|
||||
let line = raw_line.trim_end_matches('\r');
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_name = Some(value.trim().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim().to_string());
|
||||
}
|
||||
}
|
||||
let data = data_lines.join("\n");
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let event: Value =
|
||||
serde_json::from_str(&data).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let event_type = event
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.or(event_name.as_deref())
|
||||
.unwrap_or_default();
|
||||
match event_type {
|
||||
"error" | "response.failed" => self.handle_failed(report_context, &event),
|
||||
"response.image_generation_call.partial_image" => {
|
||||
self.handle_image_generation_partial(report_context, &event)
|
||||
}
|
||||
"response.output_item.done" => self.handle_output_item_done(report_context, &event),
|
||||
"response.completed" => self.handle_completed(report_context, &event),
|
||||
_ => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_image_generation_partial(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if requested_partial_images(report_context) == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(result) = event
|
||||
.get("partial_image_b64")
|
||||
.or_else(|| event.get("b64_json"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let partial_image_index = event
|
||||
.get("partial_image_index")
|
||||
.or_else(|| event.get("output_index"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(self.emitted_partial_count);
|
||||
self.emitted_partial_count = self
|
||||
.emitted_partial_count
|
||||
.max(partial_image_index.saturating_add(1));
|
||||
self.saw_upstream_partial = true;
|
||||
self.latest_image = Some(OpenAiImageFrame {
|
||||
b64_json: result.to_string(),
|
||||
});
|
||||
|
||||
encode_json_sse(
|
||||
Some(image_partial_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_partial_event_name(report_context),
|
||||
"b64_json": result,
|
||||
"partial_image_index": partial_image_index,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn handle_output_item_done(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(item) = event.get("item").and_then(Value::as_object) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(result) = item.get("result").and_then(Value::as_str).map(str::trim) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if result.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.latest_image = Some(OpenAiImageFrame {
|
||||
b64_json: result.to_string(),
|
||||
});
|
||||
|
||||
if requested_partial_images(report_context) == 0 || self.saw_upstream_partial {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let partial_image_index = event
|
||||
.get("output_index")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(self.emitted_partial_count);
|
||||
self.emitted_partial_count = partial_image_index.saturating_add(1);
|
||||
|
||||
encode_json_sse(
|
||||
Some(image_partial_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_partial_event_name(report_context),
|
||||
"b64_json": result,
|
||||
"partial_image_index": partial_image_index,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn handle_completed(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if self.latest_image.is_none() {
|
||||
if let Some(result) = completed_response_image_result(event) {
|
||||
self.latest_image = Some(OpenAiImageFrame {
|
||||
b64_json: result.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
let Some(latest_image) = self.latest_image.clone() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let usage = event
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| {
|
||||
response
|
||||
.get("tool_usage")
|
||||
.and_then(|value| value.get("image_gen"))
|
||||
.cloned()
|
||||
.or_else(|| response.get("usage").cloned())
|
||||
})
|
||||
.unwrap_or(Value::Null);
|
||||
|
||||
encode_json_sse(
|
||||
Some(image_completed_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_completed_event_name(report_context),
|
||||
"b64_json": latest_image.b64_json,
|
||||
"usage": usage,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn handle_failed(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.emitted_failure = true;
|
||||
let error = image_failure_error(event);
|
||||
encode_json_sse(
|
||||
Some(image_failed_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_failed_event_name(report_context),
|
||||
"error": error,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn image_failure_error(event: &Value) -> Value {
|
||||
let mut error = event
|
||||
.get("error")
|
||||
.or_else(|| event.get("response").and_then(|value| value.get("error")))
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
if !error.contains_key("message") {
|
||||
if let Some(message) = event
|
||||
.get("message")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| {
|
||||
event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(|value| value.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
error.insert("message".to_string(), Value::String(message.to_string()));
|
||||
}
|
||||
}
|
||||
if !error.contains_key("code") {
|
||||
if let Some(code) = event
|
||||
.get("code")
|
||||
.or_else(|| {
|
||||
event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(|value| value.get("code"))
|
||||
})
|
||||
.cloned()
|
||||
{
|
||||
error.insert("code".to_string(), code);
|
||||
}
|
||||
}
|
||||
if !error.contains_key("type") {
|
||||
let inferred_type = error
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("upstream_error");
|
||||
error.insert("type".to_string(), Value::String(inferred_type.to_string()));
|
||||
}
|
||||
if !error.contains_key("message") {
|
||||
error.insert(
|
||||
"message".to_string(),
|
||||
Value::String("Image generation failed".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
Value::Object(error)
|
||||
}
|
||||
|
||||
fn completed_response_image_result(event: &Value) -> Option<&str> {
|
||||
event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("output"))
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
|
||||
.filter_map(|item| item.get("result").and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.find(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn requested_partial_images(report_context: &Value) -> u64 {
|
||||
report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("partial_images"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn image_partial_event_name(report_context: &Value) -> &'static str {
|
||||
if image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.partial_image"
|
||||
} else {
|
||||
"image_generation.partial_image"
|
||||
}
|
||||
}
|
||||
|
||||
fn image_completed_event_name(report_context: &Value) -> &'static str {
|
||||
if image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.completed"
|
||||
} else {
|
||||
"image_generation.completed"
|
||||
}
|
||||
}
|
||||
|
||||
fn image_failed_event_name(report_context: &Value) -> &'static str {
|
||||
if image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.failed"
|
||||
} else {
|
||||
"image_generation.failed"
|
||||
}
|
||||
}
|
||||
|
||||
fn image_request_operation(report_context: &Value) -> Option<&str> {
|
||||
report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("operation"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn find_sse_block_end(buffer: &[u8]) -> Option<usize> {
|
||||
buffer
|
||||
.windows(2)
|
||||
.position(|window| window == b"\n\n")
|
||||
.map(|index| index + 2)
|
||||
.or_else(|| {
|
||||
buffer
|
||||
.windows(4)
|
||||
.position(|window| window == b"\r\n\r\n")
|
||||
.map(|index| index + 4)
|
||||
})
|
||||
}
|
||||
|
||||
fn drain_sse_separator(buffer: &mut Vec<u8>) {
|
||||
while matches!(buffer.first(), Some(b'\n' | b'\r')) {
|
||||
buffer.remove(0);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_stream.rs"]
|
||||
mod tests;
|
||||
@@ -1,246 +0,0 @@
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
|
||||
use crate::ai_pipeline::{build_generated_tool_call_id, canonicalize_tool_arguments};
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
use base64::Engine as _;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::finalize::common::{
|
||||
build_local_success_outcome, build_local_success_outcome_with_conversion_report,
|
||||
local_finalize_allows_envelope, unwrap_local_finalize_response_value,
|
||||
LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::finalize::standard::{
|
||||
maybe_build_standard_sync_finalize_product_from_normalized_payload,
|
||||
StandardSyncFinalizeNormalizedProduct,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
|
||||
aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses,
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_build_local_core_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
|
||||
if let Some(outcome) =
|
||||
maybe_build_local_openai_image_sync_finalize_response(trace_id, decision, payload)?
|
||||
{
|
||||
return Ok(Some(outcome));
|
||||
}
|
||||
|
||||
let Some(normalized_payload) =
|
||||
crate::ai_pipeline::adaptation::private_envelope::maybe_normalize_provider_private_sync_report_payload(payload)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let payload = &normalized_payload;
|
||||
let Some(report_context) = payload.report_context.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !local_finalize_allows_envelope(report_context) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(product) = maybe_build_standard_sync_finalize_product_from_normalized_payload(
|
||||
payload.report_kind.as_str(),
|
||||
payload.status_code,
|
||||
Some(report_context),
|
||||
payload.body_json.as_ref(),
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.map_err(GatewayError::from)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match product {
|
||||
StandardSyncFinalizeNormalizedProduct::SuccessBody(body_json) => {
|
||||
let Some(body_json) = unwrap_local_finalize_response_value(body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(build_local_success_outcome(
|
||||
trace_id, decision, payload, body_json,
|
||||
)?))
|
||||
}
|
||||
StandardSyncFinalizeNormalizedProduct::CrossFormat(product) => {
|
||||
let Some(provider_body_json) =
|
||||
unwrap_local_finalize_response_value(product.provider_body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(build_local_success_outcome_with_conversion_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload,
|
||||
product.client_body_json,
|
||||
provider_body_json,
|
||||
)?))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
|
||||
if payload.report_kind != "openai_image_sync_finalize" || payload.status_code >= 400 {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(report_context) = payload.report_context.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if report_context
|
||||
.get("client_api_format")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
!= Some("openai:image")
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(body_base64) = payload.body_base64.as_deref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let default_output_format = report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("output_format"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT);
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let text =
|
||||
std::str::from_utf8(&body_bytes).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
|
||||
let mut created = None;
|
||||
let mut completed_response = None;
|
||||
let mut images = Vec::new();
|
||||
|
||||
for raw_block in text.split("\n\n") {
|
||||
let block = raw_block.trim();
|
||||
if block.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let data_line = block
|
||||
.lines()
|
||||
.find_map(|line| line.trim().strip_prefix("data:").map(str::trim));
|
||||
let Some(data_line) = data_line else {
|
||||
continue;
|
||||
};
|
||||
if data_line.is_empty() || data_line == "[DONE]" {
|
||||
continue;
|
||||
}
|
||||
let event: serde_json::Value = serde_json::from_str(data_line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
match event
|
||||
.get("type")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"response.created" => {
|
||||
created = event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("created_at"))
|
||||
.and_then(serde_json::Value::as_i64)
|
||||
.or(created);
|
||||
}
|
||||
"response.output_item.done" => {
|
||||
let Some(item) = event.get("item").and_then(serde_json::Value::as_object) else {
|
||||
continue;
|
||||
};
|
||||
if item.get("type").and_then(serde_json::Value::as_str)
|
||||
!= Some("image_generation_call")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(result) = item.get("result").and_then(serde_json::Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
images.push(serde_json::json!({
|
||||
"b64_json": result,
|
||||
"output_format": item.get("output_format").cloned().unwrap_or(serde_json::Value::String(default_output_format.to_string())),
|
||||
"revised_prompt": item.get("revised_prompt").cloned().unwrap_or(serde_json::Value::Null),
|
||||
}));
|
||||
}
|
||||
"response.completed" => {
|
||||
completed_response = event
|
||||
.get("response")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.cloned();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if images.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let completed_response = completed_response.unwrap_or_default();
|
||||
let provider_usage = completed_response
|
||||
.get("tool_usage")
|
||||
.and_then(|value| value.get("image_gen"))
|
||||
.cloned()
|
||||
.or_else(|| completed_response.get("usage").cloned());
|
||||
let provider_body_json = serde_json::json!({
|
||||
"id": completed_response.get("id").cloned().unwrap_or(serde_json::Value::Null),
|
||||
"object": "response",
|
||||
"model": completed_response.get("model").cloned().unwrap_or(serde_json::Value::Null),
|
||||
"status": completed_response.get("status").cloned().unwrap_or(serde_json::Value::String("completed".to_string())),
|
||||
"usage": provider_usage,
|
||||
"tool_usage": completed_response.get("tool_usage").cloned().unwrap_or(serde_json::Value::Null),
|
||||
"output": images
|
||||
.iter()
|
||||
.map(|image| serde_json::json!({
|
||||
"type": "image_generation_call",
|
||||
"output_format": image.get("output_format").cloned().unwrap_or(serde_json::Value::Null),
|
||||
"revised_prompt": image.get("revised_prompt").cloned().unwrap_or(serde_json::Value::Null),
|
||||
}))
|
||||
.collect::<Vec<_>>(),
|
||||
});
|
||||
let client_images = images
|
||||
.iter()
|
||||
.map(|image| {
|
||||
let revised_prompt = image
|
||||
.get("revised_prompt")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null);
|
||||
let b64_json = image
|
||||
.get("b64_json")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let output_format = image
|
||||
.get("output_format")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or(default_output_format);
|
||||
serde_json::json!({
|
||||
"b64_json": b64_json,
|
||||
"revised_prompt": revised_prompt,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let client_body_json = serde_json::json!({
|
||||
"created": created.unwrap_or_default(),
|
||||
"data": client_images,
|
||||
"usage": provider_body_json.get("usage").cloned().unwrap_or(serde_json::Value::Null),
|
||||
});
|
||||
|
||||
Ok(Some(build_local_success_outcome_with_conversion_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload,
|
||||
client_body_json,
|
||||
provider_body_json,
|
||||
)?))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_sync.rs"]
|
||||
mod tests;
|
||||
@@ -1,564 +0,0 @@
|
||||
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::ai_pipeline::finalize::sse::encode_json_sse;
|
||||
use crate::ai_pipeline::{
|
||||
convert_claude_response_to_openai_responses, convert_gemini_response_to_openai_responses,
|
||||
convert_openai_chat_response_to_openai_responses, ClaudeClientEmitter, GeminiClientEmitter,
|
||||
OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
|
||||
};
|
||||
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 provider_api_format == "openai:image" && client_api_format == "openai:image" {
|
||||
return maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context);
|
||||
}
|
||||
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_responses_response) = convert_provider_sync_response_to_openai_responses(
|
||||
provider_body_json,
|
||||
provider_api_format.as_str(),
|
||||
&bridge_context,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let terminal_summary =
|
||||
build_terminal_summary_from_openai_responses_response(&openai_responses_response);
|
||||
let canonical_frames = build_canonical_frames_from_openai_responses_response(
|
||||
&openai_responses_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 maybe_bridge_openai_image_sync_json_to_stream(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, GatewayError> {
|
||||
let Some(response) = provider_body_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(image) = response
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.find_map(extract_openai_image_sync_b64_json)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let usage = response.get("usage").cloned().unwrap_or(Value::Null);
|
||||
let event_name = openai_image_completed_event_name(report_context);
|
||||
let sse_body = encode_json_sse(
|
||||
Some(event_name),
|
||||
&json!({
|
||||
"type": event_name,
|
||||
"b64_json": image,
|
||||
"usage": usage,
|
||||
}),
|
||||
)?;
|
||||
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary: Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: response
|
||||
.get("usage")
|
||||
.and_then(standardized_usage_from_openai_usage),
|
||||
finish_reason: Some("stop".to_string()),
|
||||
response_id: response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned),
|
||||
model: response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| image_bridge_model(report_context)),
|
||||
observed_finish: true,
|
||||
unknown_event_count: 0,
|
||||
parser_error: None,
|
||||
}),
|
||||
}))
|
||||
}
|
||||
|
||||
fn normalize_api_format(value: &str) -> String {
|
||||
crate::ai_pipeline::normalize_api_format_alias(value)
|
||||
}
|
||||
|
||||
fn is_standard_api_format(value: &str) -> bool {
|
||||
matches!(
|
||||
value,
|
||||
"openai:chat"
|
||||
| "openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "claude:messages"
|
||||
| "gemini:generate_content"
|
||||
)
|
||||
}
|
||||
|
||||
fn extract_openai_image_sync_b64_json(item: &serde_json::Map<String, Value>) -> Option<String> {
|
||||
item.get("b64_json")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| {
|
||||
item.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(extract_base64_from_data_url)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_base64_from_data_url(value: &str) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
let (metadata, payload) = trimmed.split_once(',')?;
|
||||
if !metadata.starts_with("data:") || !metadata.ends_with(";base64") {
|
||||
return None;
|
||||
}
|
||||
(!payload.trim().is_empty()).then(|| payload.trim().to_string())
|
||||
}
|
||||
|
||||
fn openai_image_completed_event_name(report_context: Option<&Value>) -> &'static str {
|
||||
if openai_image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.completed"
|
||||
} else {
|
||||
"image_generation.completed"
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_image_request_operation(report_context: Option<&Value>) -> Option<&str> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("operation"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn image_bridge_model(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context.and_then(|context| {
|
||||
context
|
||||
.get("mapped_model")
|
||||
.or_else(|| context.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
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_responses(
|
||||
provider_body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
report_context: &Value,
|
||||
) -> Option<Value> {
|
||||
match provider_api_format {
|
||||
"openai:responses" | "openai:responses:compact" => Some(provider_body_json.clone()),
|
||||
"openai:chat" => convert_openai_chat_response_to_openai_responses(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
false,
|
||||
),
|
||||
"claude:messages" => {
|
||||
convert_claude_response_to_openai_responses(provider_body_json, report_context)
|
||||
}
|
||||
"gemini:generate_content" => {
|
||||
convert_gemini_response_to_openai_responses(provider_body_json, report_context)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_canonical_frames_from_openai_responses_response(
|
||||
openai_responses_response: &Value,
|
||||
report_context: &Value,
|
||||
) -> Result<Vec<crate::ai_pipeline::CanonicalStreamFrame>, GatewayError> {
|
||||
let mut state = OpenAIResponsesProviderState::default();
|
||||
let line = format!(
|
||||
"data: {}\n",
|
||||
serde_json::to_string(&json!({
|
||||
"type": "response.completed",
|
||||
"response": openai_responses_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:responses" | "openai:responses:compact" => {
|
||||
let mut emitter = OpenAIResponsesClientEmitter::default();
|
||||
emit_with_openai_responses_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"claude:messages" => {
|
||||
let mut emitter = ClaudeClientEmitter::default();
|
||||
emit_with_claude_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"gemini:generate_content" => {
|
||||
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_responses_emitter(
|
||||
emitter: &mut OpenAIResponsesClientEmitter,
|
||||
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_responses_response(
|
||||
openai_responses_response: &Value,
|
||||
) -> Option<ExecutionStreamTerminalSummary> {
|
||||
let response = openai_responses_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_responses_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,
|
||||
unknown_event_count: 0,
|
||||
parser_error: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_openai_responses_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 mut 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),
|
||||
);
|
||||
if input_tokens == 0 && total_tokens > output_tokens {
|
||||
input_tokens = total_tokens.saturating_sub(output_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())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{maybe_bridge_standard_sync_json_to_stream, standardized_usage_from_openai_usage};
|
||||
|
||||
fn utf8(bytes: Vec<u8>) -> String {
|
||||
String::from_utf8(bytes).expect("utf8 should decode")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_sync_usage_derives_missing_input_tokens_from_total() {
|
||||
let usage = standardized_usage_from_openai_usage(&json!({
|
||||
"output_tokens": 177,
|
||||
"total_tokens": 20_612,
|
||||
"input_tokens_details": {
|
||||
"cached_tokens": 19_840,
|
||||
},
|
||||
}))
|
||||
.expect("usage should parse");
|
||||
|
||||
assert_eq!(usage.input_tokens, 20_435);
|
||||
assert_eq!(usage.output_tokens, 177);
|
||||
assert_eq!(usage.cache_read_tokens, 19_840);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridges_openai_image_sync_json_to_generation_completed_sse() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"mapped_model": "gpt-image-1",
|
||||
"image_request": {
|
||||
"operation": "generate"
|
||||
}
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"created": 1776971267,
|
||||
"data": [{
|
||||
"b64_json": "aGVsbG8="
|
||||
}],
|
||||
"usage": {
|
||||
"total_tokens": 100,
|
||||
"input_tokens": 50,
|
||||
"output_tokens": 50,
|
||||
"input_tokens_details": {
|
||||
"text_tokens": 10,
|
||||
"image_tokens": 40
|
||||
}
|
||||
}
|
||||
}),
|
||||
"openai:image",
|
||||
"openai:image",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("bridge should produce sse");
|
||||
|
||||
let output = utf8(outcome.sse_body);
|
||||
assert!(output.contains("event: image_generation.completed"));
|
||||
assert!(output.contains("\"type\":\"image_generation.completed\""));
|
||||
assert!(output.contains("\"b64_json\":\"aGVsbG8=\""));
|
||||
assert!(output.contains("\"total_tokens\":100"));
|
||||
|
||||
let summary = outcome
|
||||
.terminal_summary
|
||||
.expect("terminal summary should exist");
|
||||
assert_eq!(summary.model.as_deref(), Some("gpt-image-1"));
|
||||
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!(100))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridges_openai_image_sync_data_url_to_edit_completed_sse() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"image_request": {
|
||||
"operation": "edit"
|
||||
}
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"created": 1776971267,
|
||||
"data": [{
|
||||
"url": "data:image/webp;base64,d29ybGQ="
|
||||
}],
|
||||
"usage": {
|
||||
"total_tokens": 9,
|
||||
"input_tokens": 4,
|
||||
"output_tokens": 5
|
||||
}
|
||||
}),
|
||||
"openai:image",
|
||||
"openai:image",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("bridge should produce sse");
|
||||
|
||||
let output = utf8(outcome.sse_body);
|
||||
assert!(output.contains("event: image_edit.completed"));
|
||||
assert!(output.contains("\"type\":\"image_edit.completed\""));
|
||||
assert!(output.contains("\"b64_json\":\"d29ybGQ=\""));
|
||||
assert!(output.contains("\"total_tokens\":9"));
|
||||
}
|
||||
}
|
||||
@@ -1,6 +0,0 @@
|
||||
//! Standard finalize streaming conversion helpers.
|
||||
pub(crate) use crate::ai_pipeline::CanonicalStreamFrame;
|
||||
|
||||
mod orchestrator;
|
||||
|
||||
pub(crate) use orchestrator::StreamingStandardConversionState;
|
||||
@@ -1,49 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::adaptation::private_envelope::transform_provider_private_stream_line as transform_envelope_line;
|
||||
use crate::ai_pipeline::{
|
||||
provider_adaptation_should_unwrap_stream_envelope, StreamingStandardFormatMatrix,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct StreamingStandardConversionState {
|
||||
matrix: StreamingStandardFormatMatrix,
|
||||
}
|
||||
|
||||
impl StreamingStandardConversionState {
|
||||
pub(crate) fn transform_line(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let line = if should_unwrap_envelope(report_context) {
|
||||
transform_envelope_line(report_context, line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
} else {
|
||||
line
|
||||
};
|
||||
if line.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.matrix
|
||||
.transform_line(report_context, line)
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
pub(crate) fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, GatewayError> {
|
||||
self.matrix.finish(report_context).map_err(Into::into)
|
||||
}
|
||||
}
|
||||
|
||||
fn should_unwrap_envelope(report_context: &Value) -> bool {
|
||||
let envelope_name = report_context
|
||||
.get("envelope_name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let provider_api_format = report_context
|
||||
.get("provider_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format)
|
||||
}
|
||||
@@ -1,622 +0,0 @@
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_pipeline::planner::passthrough::resolve_same_format_provider_transport_unsupported_reason_for_trace;
|
||||
use crate::ai_pipeline::transport::{
|
||||
body_rules_are_locally_supported, header_rules_are_locally_supported,
|
||||
local_gemini_transport_unsupported_reason_with_network,
|
||||
local_openai_chat_transport_unsupported_reason,
|
||||
local_standard_transport_unsupported_reason_with_network, resolve_transport_tls_profile,
|
||||
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
request_conversion_enabled_for_transport, request_conversion_kind,
|
||||
request_conversion_requires_enable_flag, request_conversion_transport_unsupported_reason,
|
||||
request_pair_allowed_for_transport, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
|
||||
use crate::append_execution_contract_fields_to_value;
|
||||
|
||||
pub(crate) struct LocalExecutionCandidateMetadataParts<'a> {
|
||||
pub(crate) eligible: &'a EligibleLocalExecutionCandidate,
|
||||
pub(crate) provider_api_format: &'a str,
|
||||
pub(crate) client_api_format: &'a str,
|
||||
pub(crate) extra_fields: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub(crate) fn append_ranking_metadata_to_object(
|
||||
object: &mut Map<String, Value>,
|
||||
ranking: &SchedulerRankingOutcome,
|
||||
) {
|
||||
object.insert(
|
||||
"ranking_mode".to_string(),
|
||||
Value::String(format!("{:?}", ranking.ranking_mode)),
|
||||
);
|
||||
object.insert(
|
||||
"priority_mode".to_string(),
|
||||
Value::String(format!("{:?}", ranking.priority_mode)),
|
||||
);
|
||||
object.insert(
|
||||
"ranking_index".to_string(),
|
||||
Value::Number(serde_json::Number::from(ranking.ranking_index as u64)),
|
||||
);
|
||||
object.insert(
|
||||
"priority_slot".to_string(),
|
||||
Value::Number(serde_json::Number::from(i64::from(ranking.priority_slot))),
|
||||
);
|
||||
if let Some(promoted_by) = ranking.promoted_by {
|
||||
object.insert(
|
||||
"promoted_by".to_string(),
|
||||
Value::String(promoted_by.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(demoted_by) = ranking.demoted_by {
|
||||
object.insert(
|
||||
"demoted_by".to_string(),
|
||||
Value::String(demoted_by.to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_request_trace_proxy_value(
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
resolved_proxy: Option<&ProxySnapshot>,
|
||||
) -> Option<Value> {
|
||||
let resolved_proxy = resolved_proxy?;
|
||||
let mut object = Map::new();
|
||||
|
||||
if let Some(node_id) = trimmed_non_empty(resolved_proxy.node_id.as_deref()) {
|
||||
object.insert("node_id".to_string(), Value::String(node_id));
|
||||
}
|
||||
if let Some(node_name) = trimmed_non_empty(resolved_proxy.label.as_deref()) {
|
||||
object.insert("node_name".to_string(), Value::String(node_name));
|
||||
}
|
||||
if let Some(url) = sanitize_trace_proxy_url(resolved_proxy.url.as_deref()) {
|
||||
object.insert("url".to_string(), Value::String(url));
|
||||
}
|
||||
if let Some(source) = resolve_request_trace_proxy_source(transport, true) {
|
||||
object.insert("source".to_string(), Value::String(source.to_string()));
|
||||
}
|
||||
|
||||
(!object.is_empty()).then_some(Value::Object(object))
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_metadata(
|
||||
parts: LocalExecutionCandidateMetadataParts<'_>,
|
||||
) -> Value {
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
&parts.eligible.candidate,
|
||||
Some(parts.eligible.transport.as_ref()),
|
||||
parts.provider_api_format,
|
||||
parts.client_api_format,
|
||||
parts.extra_fields,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_metadata_for_candidate(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
extra_fields: Map<String, Value>,
|
||||
) -> Value {
|
||||
let mut object = Map::new();
|
||||
object.insert(
|
||||
"provider_api_format".to_string(),
|
||||
Value::String(provider_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"client_api_format".to_string(),
|
||||
Value::String(client_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"global_model_id".to_string(),
|
||||
Value::String(candidate.global_model_id.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"global_model_name".to_string(),
|
||||
Value::String(candidate.global_model_name.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"model_id".to_string(),
|
||||
Value::String(candidate.model_id.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"selected_provider_model_name".to_string(),
|
||||
Value::String(candidate.selected_provider_model_name.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"mapping_matched_model".to_string(),
|
||||
candidate
|
||||
.mapping_matched_model
|
||||
.clone()
|
||||
.map(Value::String)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
object.insert(
|
||||
"provider_name".to_string(),
|
||||
Value::String(candidate.provider_name.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"key_name".to_string(),
|
||||
Value::String(candidate.key_name.clone()),
|
||||
);
|
||||
object.extend(extra_fields);
|
||||
append_transport_diagnostics_to_value(
|
||||
Value::Object(object),
|
||||
transport,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_contract_metadata(
|
||||
parts: LocalExecutionCandidateMetadataParts<'_>,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
&parts.eligible.candidate,
|
||||
Some(parts.eligible.transport.as_ref()),
|
||||
parts.provider_api_format,
|
||||
parts.client_api_format,
|
||||
parts.extra_fields,
|
||||
),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
parts.client_api_format,
|
||||
provider_contract,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
extra_fields: Map<String, Value>,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
candidate,
|
||||
transport,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
extra_fields,
|
||||
),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
client_api_format,
|
||||
provider_contract,
|
||||
)
|
||||
}
|
||||
|
||||
fn append_transport_diagnostics_to_value(
|
||||
value: Value,
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> Value {
|
||||
let Value::Object(mut object) = value else {
|
||||
return value;
|
||||
};
|
||||
object.insert(
|
||||
"transport_diagnostics".to_string(),
|
||||
transport
|
||||
.map(|transport| {
|
||||
build_transport_diagnostics(transport, client_api_format, provider_api_format)
|
||||
})
|
||||
.unwrap_or_else(|| json!({ "transport_snapshot_available": false })),
|
||||
);
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
fn build_transport_diagnostics(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> Value {
|
||||
let resolved_tls_profile = resolve_transport_tls_profile(transport);
|
||||
let configured_tls_profile = transport
|
||||
.key
|
||||
.fingerprint
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|value| value.get("tls_profile"))
|
||||
.cloned()
|
||||
.unwrap_or(Value::Null);
|
||||
let has_oauth_config = transport.key.decrypted_auth_config.is_some();
|
||||
let oauth_resolution_supported =
|
||||
!has_oauth_config || supports_local_oauth_request_auth_resolution(transport);
|
||||
let request_transport_unsupported_reason = resolve_request_transport_unsupported_reason(
|
||||
transport,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
);
|
||||
|
||||
json!({
|
||||
"transport_snapshot_available": true,
|
||||
"provider_type": transport.provider.provider_type,
|
||||
"provider_is_active": transport.provider.is_active,
|
||||
"endpoint_is_active": transport.endpoint.is_active,
|
||||
"key_is_active": transport.key.is_active,
|
||||
"provider_enable_format_conversion": transport.provider.enable_format_conversion,
|
||||
"provider_keep_priority_on_conversion": transport.provider.keep_priority_on_conversion,
|
||||
"endpoint_format_acceptance_config": transport.endpoint.format_acceptance_config,
|
||||
"endpoint_custom_path": transport.endpoint.custom_path,
|
||||
"header_rules": transport.endpoint.header_rules,
|
||||
"header_rules_supported": header_rules_are_locally_supported(transport.endpoint.header_rules.as_ref()),
|
||||
"body_rules": transport.endpoint.body_rules,
|
||||
"body_rules_supported": body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref()),
|
||||
"proxy": {
|
||||
"locally_supported": transport_proxy_is_locally_supported(transport),
|
||||
"provider": summarize_proxy_config(transport.provider.proxy.as_ref()),
|
||||
"endpoint": summarize_proxy_config(transport.endpoint.proxy.as_ref()),
|
||||
"key": summarize_proxy_config(transport.key.proxy.as_ref()),
|
||||
},
|
||||
"auth": {
|
||||
"key_auth_type": transport.key.auth_type,
|
||||
"has_oauth_config": has_oauth_config,
|
||||
"oauth_request_auth_resolution_supported": oauth_resolution_supported,
|
||||
},
|
||||
"fingerprint": transport.key.fingerprint,
|
||||
"configured_tls_profile": configured_tls_profile,
|
||||
"resolved_tls_profile": resolved_tls_profile,
|
||||
"request_pair": {
|
||||
"client_api_format": client_api_format,
|
||||
"provider_api_format": provider_api_format,
|
||||
"requires_conversion_enable_flag": request_conversion_requires_enable_flag(
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
),
|
||||
"conversion_enabled": request_conversion_enabled_for_transport(
|
||||
transport,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
),
|
||||
"pair_allowed": request_pair_allowed_for_transport(
|
||||
transport,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
),
|
||||
"transport_unsupported_reason": request_transport_unsupported_reason,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn summarize_proxy_config(proxy: Option<&Value>) -> Value {
|
||||
let Some(object) = proxy.and_then(Value::as_object) else {
|
||||
return Value::Null;
|
||||
};
|
||||
let has_url = object
|
||||
.get("url")
|
||||
.or_else(|| object.get("proxy_url"))
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| !value.trim().is_empty());
|
||||
json!({
|
||||
"enabled": object.get("enabled").cloned().unwrap_or(Value::Null),
|
||||
"mode": object.get("mode").cloned().unwrap_or(Value::Null),
|
||||
"node_id": object.get("node_id").cloned().unwrap_or(Value::Null),
|
||||
"label": object.get("label").cloned().unwrap_or(Value::Null),
|
||||
"has_url": has_url,
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_request_trace_proxy_source(
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
has_resolved_proxy: bool,
|
||||
) -> Option<&'static str> {
|
||||
let transport = transport?;
|
||||
if transport_has_explicit_proxy(transport.key.proxy.as_ref()) {
|
||||
return Some("key");
|
||||
}
|
||||
if transport_has_explicit_proxy(transport.endpoint.proxy.as_ref()) {
|
||||
return Some("endpoint");
|
||||
}
|
||||
if transport_has_explicit_proxy(transport.provider.proxy.as_ref()) {
|
||||
return Some("provider");
|
||||
}
|
||||
has_resolved_proxy.then_some("system")
|
||||
}
|
||||
|
||||
fn transport_has_explicit_proxy(proxy: Option<&Value>) -> bool {
|
||||
let Some(object) = proxy.and_then(Value::as_object) else {
|
||||
return false;
|
||||
};
|
||||
let enabled = object
|
||||
.get("enabled")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(true);
|
||||
if !enabled {
|
||||
return false;
|
||||
}
|
||||
|
||||
object
|
||||
.get("node_id")
|
||||
.or_else(|| object.get("url"))
|
||||
.or_else(|| object.get("proxy_url"))
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| !value.trim().is_empty())
|
||||
}
|
||||
|
||||
fn sanitize_trace_proxy_url(url: Option<&str>) -> Option<String> {
|
||||
let raw = url.map(str::trim).filter(|value| !value.is_empty())?;
|
||||
let parsed = url::Url::parse(raw).ok()?;
|
||||
let scheme = parsed.scheme().trim();
|
||||
let host = parsed.host_str()?.trim();
|
||||
if scheme.is_empty() || host.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut safe = format!("{scheme}://{host}");
|
||||
if let Some(port) = parsed.port() {
|
||||
safe.push(':');
|
||||
safe.push_str(port.to_string().as_str());
|
||||
}
|
||||
Some(safe)
|
||||
}
|
||||
|
||||
fn trimmed_non_empty(value: Option<&str>) -> Option<String> {
|
||||
value
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn resolve_request_transport_unsupported_reason(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> Option<&'static str> {
|
||||
let client_api_format = client_api_format.trim().to_ascii_lowercase();
|
||||
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
|
||||
if client_api_format == provider_api_format {
|
||||
if let Some(skip_reason) =
|
||||
resolve_same_format_provider_transport_unsupported_reason_for_trace(
|
||||
transport,
|
||||
provider_api_format.as_str(),
|
||||
)
|
||||
{
|
||||
return Some(skip_reason);
|
||||
}
|
||||
return match provider_api_format.as_str() {
|
||||
"openai:chat" => local_openai_chat_transport_unsupported_reason(transport),
|
||||
"gemini:generate_content" => local_gemini_transport_unsupported_reason_with_network(
|
||||
transport,
|
||||
provider_api_format.as_str(),
|
||||
),
|
||||
_ => local_standard_transport_unsupported_reason_with_network(
|
||||
transport,
|
||||
provider_api_format.as_str(),
|
||||
),
|
||||
};
|
||||
}
|
||||
match request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str()) {
|
||||
Some(kind) => request_conversion_transport_unsupported_reason(transport, kind),
|
||||
None => Some("transport_api_format_unsupported"),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
build_local_execution_candidate_metadata_for_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider,
|
||||
};
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, GatewayProviderTransportSnapshot};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn sample_candidate() -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
SchedulerMinimalCandidateSelectionCandidate {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_name: "RightCode".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 22,
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
endpoint_api_format: "openai:responses".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
key_name: "codex".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_internal_priority: 10,
|
||||
key_global_priority_for_format: None,
|
||||
key_capabilities: None,
|
||||
model_id: "model-1".to_string(),
|
||||
global_model_id: "global-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
selected_provider_model_name: "gpt-5.4".to_string(),
|
||||
mapping_matched_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "RightCode".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: Some(json!({"enabled": true, "mode": "node", "node_id": "proxy-node-1"})),
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "openai:responses".to_string(),
|
||||
api_family: None,
|
||||
endpoint_kind: None,
|
||||
is_active: true,
|
||||
base_url: "https://example.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: Some("/v1/responses".to_string()),
|
||||
config: None,
|
||||
format_acceptance_config: Some(json!({
|
||||
"enabled": true,
|
||||
"accept_formats": ["claude:messages"]
|
||||
})),
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "codex".to_string(),
|
||||
auth_type: "oauth".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: Some(json!({
|
||||
"tls_profile": "chrome_136",
|
||||
"user_agent": "Mozilla/5.0"
|
||||
})),
|
||||
decrypted_api_key: "sk-test".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_claude_code_transport_without_auth() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-cc-1".to_string(),
|
||||
name: "NekoCode".to_string(),
|
||||
provider_type: "claude_code".to_string(),
|
||||
website: Some("https://nekocode.ai".to_string()),
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-cc-1".to_string(),
|
||||
provider_id: "provider-cc-1".to_string(),
|
||||
api_format: "claude:messages".to_string(),
|
||||
api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://api.anthropic.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-cc-1".to_string(),
|
||||
provider_id: "provider-cc-1".to_string(),
|
||||
name: "CC-特价-0.4".to_string(),
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:messages".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_contract_metadata_includes_transport_diagnostics() {
|
||||
let metadata = build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
Some(&sample_transport()),
|
||||
"openai:responses",
|
||||
"claude:messages",
|
||||
serde_json::Map::new(),
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
"openai:responses",
|
||||
);
|
||||
|
||||
assert_eq!(metadata["transport_diagnostics"]["provider_type"], "codex");
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["fingerprint"]["tls_profile"],
|
||||
"chrome_136"
|
||||
);
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["resolved_tls_profile"],
|
||||
"chrome_136"
|
||||
);
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["conversion_enabled"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
assert!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["transport_unsupported_reason"]
|
||||
.is_null()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_metadata_marks_missing_transport_snapshot() {
|
||||
let metadata = build_local_execution_candidate_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
None,
|
||||
"openai:responses",
|
||||
"openai:responses",
|
||||
serde_json::Map::new(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["transport_snapshot_available"],
|
||||
Value::Bool(false)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_metadata_uses_same_format_provider_specific_transport_reason() {
|
||||
let metadata = build_local_execution_candidate_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
Some(&sample_claude_code_transport_without_auth()),
|
||||
"claude:messages",
|
||||
"claude:messages",
|
||||
serde_json::Map::new(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["transport_unsupported_reason"],
|
||||
Value::String("transport_auth_unavailable".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,417 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_provider_transport::provider_types::provider_type_is_fixed;
|
||||
use tracing::warn;
|
||||
|
||||
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, PlannerAppState,
|
||||
};
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
|
||||
use super::candidate_ranking::rank_eligible_local_execution_candidates;
|
||||
use super::pool_scheduler::apply_local_execution_pool_scheduler;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct EligibleLocalExecutionCandidate {
|
||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(crate) provider_api_format: String,
|
||||
pub(crate) orchestration: LocalExecutionCandidateMetadata,
|
||||
pub(crate) ranking: Option<SchedulerRankingOutcome>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct SkippedLocalExecutionCandidate {
|
||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) skip_reason: &'static str,
|
||||
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
|
||||
pub(crate) ranking: Option<SchedulerRankingOutcome>,
|
||||
pub(crate) extra_data: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl SkippedLocalExecutionCandidate {
|
||||
pub(crate) fn transport_ref(&self) -> Option<&GatewayProviderTransportSnapshot> {
|
||||
self.transport.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_and_rank_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let requested_model = requested_model.trim();
|
||||
resolve_and_rank_local_execution_candidates_with_gate(
|
||||
state,
|
||||
candidates,
|
||||
client_api_format,
|
||||
Some(requested_model),
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
|candidate, transport, normalized_client_api_format| {
|
||||
current_local_execution_candidate_skip_reason_with_transport(
|
||||
candidate,
|
||||
transport,
|
||||
normalized_client_api_format,
|
||||
requested_model,
|
||||
)
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transport_pair_gate(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let requested_model = requested_model.map(str::trim);
|
||||
resolve_and_rank_local_execution_candidates_with_gate(
|
||||
state,
|
||||
candidates,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
|candidate, transport, _normalized_client_api_format| {
|
||||
current_local_execution_candidate_common_skip_reason_with_transport(
|
||||
candidate,
|
||||
transport,
|
||||
requested_model,
|
||||
)
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_and_rank_local_execution_candidates_with_gate<F>(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
runtime_skip_reason: F,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
)
|
||||
where
|
||||
F: Fn(
|
||||
&SchedulerMinimalCandidateSelectionCandidate,
|
||||
&GatewayProviderTransportSnapshot,
|
||||
&str,
|
||||
) -> Option<&'static str>,
|
||||
{
|
||||
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
|
||||
let mut selectable = Vec::with_capacity(candidates.len());
|
||||
let mut skipped = Vec::with_capacity(candidates.len());
|
||||
|
||||
for candidate in candidates {
|
||||
let Some(transport) = read_candidate_transport_snapshot(state, &candidate).await else {
|
||||
skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason: "transport_snapshot_missing",
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
});
|
||||
continue;
|
||||
};
|
||||
let transport = Arc::new(transport);
|
||||
match runtime_skip_reason(
|
||||
&candidate,
|
||||
transport.as_ref(),
|
||||
normalized_client_api_format.as_str(),
|
||||
) {
|
||||
Some(skip_reason) => skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason,
|
||||
transport: Some(transport),
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
}),
|
||||
None => selectable.push(EligibleLocalExecutionCandidate {
|
||||
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
|
||||
candidate,
|
||||
transport,
|
||||
orchestration: LocalExecutionCandidateMetadata::default(),
|
||||
ranking: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
let ranked = rank_eligible_local_execution_candidates(
|
||||
state,
|
||||
selectable,
|
||||
normalized_client_api_format.as_str(),
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
)
|
||||
.await;
|
||||
let (ranked, pool_skipped) =
|
||||
apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
|
||||
skipped.extend(pool_skipped);
|
||||
|
||||
(ranked, skipped)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
|
||||
fn non_empty_str(value: Option<&serde_json::Value>) -> Option<&str> {
|
||||
value
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
let object = body_json.as_object()?;
|
||||
|
||||
non_empty_str(object.get("prompt_cache_key"))
|
||||
.or_else(|| non_empty_str(object.get("conversation_id")))
|
||||
.or_else(|| non_empty_str(object.get("conversationId")))
|
||||
.or_else(|| non_empty_str(object.get("session_id")))
|
||||
.or_else(|| non_empty_str(object.get("sessionId")))
|
||||
.or_else(|| {
|
||||
object
|
||||
.get("metadata")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|metadata| {
|
||||
non_empty_str(metadata.get("session_id"))
|
||||
.or_else(|| non_empty_str(metadata.get("conversation_id")))
|
||||
})
|
||||
})
|
||||
.or_else(|| {
|
||||
object
|
||||
.get("conversationState")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|state| {
|
||||
non_empty_str(state.get("conversationId"))
|
||||
.or_else(|| non_empty_str(state.get("sessionId")))
|
||||
})
|
||||
})
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn current_local_execution_candidate_common_skip_reason_with_transport(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<&'static str> {
|
||||
let requested_model = requested_model.unwrap_or_default();
|
||||
|
||||
if !transport.provider.is_active {
|
||||
return Some("provider_inactive");
|
||||
}
|
||||
if !transport.endpoint.is_active {
|
||||
return Some("endpoint_inactive");
|
||||
}
|
||||
if !transport.key.is_active {
|
||||
return Some("key_inactive");
|
||||
}
|
||||
|
||||
let endpoint_api_format = transport.endpoint.api_format.trim();
|
||||
if !candidate
|
||||
.endpoint_api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(endpoint_api_format)
|
||||
&& !api_format_matches(&candidate.endpoint_api_format, endpoint_api_format)
|
||||
{
|
||||
return Some("endpoint_api_format_changed");
|
||||
}
|
||||
|
||||
if !transport_key_supports_api_format(transport, endpoint_api_format) {
|
||||
return Some("key_api_format_disabled");
|
||||
}
|
||||
if !transport_key_allows_candidate_model(transport, requested_model, candidate) {
|
||||
return Some("key_model_disabled");
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn disabled_format_conversion_skip_reason(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
normalized_client_api_format: &str,
|
||||
) -> Option<&'static str> {
|
||||
let endpoint_api_format = transport.endpoint.api_format.trim();
|
||||
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
|
||||
return None;
|
||||
}
|
||||
|
||||
if crate::ai_pipeline::conversion::request_conversion_kind(
|
||||
normalized_client_api_format,
|
||||
endpoint_api_format,
|
||||
)
|
||||
.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
if crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
|
||||
normalized_client_api_format,
|
||||
endpoint_api_format,
|
||||
) && !crate::ai_pipeline::conversion::request_conversion_enabled_for_transport(
|
||||
transport,
|
||||
normalized_client_api_format,
|
||||
endpoint_api_format,
|
||||
) {
|
||||
return Some("format_conversion_disabled");
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn current_local_execution_candidate_skip_reason_with_transport(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
normalized_client_api_format: &str,
|
||||
requested_model: &str,
|
||||
) -> Option<&'static str> {
|
||||
if let Some(skip_reason) = current_local_execution_candidate_common_skip_reason_with_transport(
|
||||
candidate,
|
||||
transport,
|
||||
Some(requested_model),
|
||||
) {
|
||||
return Some(skip_reason);
|
||||
}
|
||||
|
||||
let endpoint_api_format = transport.endpoint.api_format.trim();
|
||||
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(skip_reason) =
|
||||
disabled_format_conversion_skip_reason(transport, normalized_client_api_format)
|
||||
{
|
||||
return Some(skip_reason);
|
||||
}
|
||||
|
||||
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
|
||||
transport,
|
||||
normalized_client_api_format,
|
||||
endpoint_api_format,
|
||||
) {
|
||||
return Some("transport_unsupported");
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn transport_key_supports_api_format(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
endpoint_api_format: &str,
|
||||
) -> bool {
|
||||
let provider_type = transport.provider.provider_type.trim();
|
||||
let auth_type = transport.key.auth_type.trim();
|
||||
let inherits_provider_api_formats = provider_type_is_fixed(provider_type)
|
||||
&& (auth_type.eq_ignore_ascii_case("oauth")
|
||||
|| (provider_type.eq_ignore_ascii_case("kiro")
|
||||
&& auth_type.eq_ignore_ascii_case("bearer")
|
||||
&& transport
|
||||
.key
|
||||
.decrypted_auth_config
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())));
|
||||
if inherits_provider_api_formats {
|
||||
return true;
|
||||
}
|
||||
|
||||
match transport.key.api_formats.as_deref() {
|
||||
None => true,
|
||||
Some(formats) => formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, endpoint_api_format)),
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_api_format_alias(value: &str) -> String {
|
||||
crate::ai_pipeline::normalize_api_format_alias(value)
|
||||
}
|
||||
|
||||
fn api_format_matches(left: &str, right: &str) -> bool {
|
||||
normalize_api_format_alias(left) == normalize_api_format_alias(right)
|
||||
}
|
||||
|
||||
fn transport_key_allows_candidate_model(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
requested_model: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> bool {
|
||||
let Some(allowed_models) = transport.key.allowed_models.as_deref() else {
|
||||
return true;
|
||||
};
|
||||
|
||||
let requested_model = requested_model.trim();
|
||||
let global_model_name = candidate.global_model_name.trim();
|
||||
let selected_provider_model_name = candidate.selected_provider_model_name.trim();
|
||||
let mapping_matched_model = candidate
|
||||
.mapping_matched_model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
for allowed_model in allowed_models.iter().map(String::as_str).map(str::trim) {
|
||||
if allowed_model.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if allowed_model == requested_model
|
||||
|| allowed_model == global_model_name
|
||||
|| allowed_model == selected_provider_model_name
|
||||
|| mapping_matched_model.is_some_and(|value| value == allowed_model)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
pub(crate) async fn read_candidate_transport_snapshot(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> Option<GatewayProviderTransportSnapshot> {
|
||||
match state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(transport)) => Some(transport),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "candidate_resolution_transport_load_failed",
|
||||
log_type = "event",
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
error = ?error,
|
||||
"failed to load provider transport while evaluating local candidate eligibility"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,34 +0,0 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
|
||||
|
||||
pub(crate) fn auth_snapshot_allows_cross_format_candidate(
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
requested_model: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> bool {
|
||||
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
|
||||
let provider_allowed = allowed_providers.iter().any(|value| {
|
||||
aether_scheduler_core::provider_matches_allowed_value(
|
||||
value,
|
||||
&candidate.provider_id,
|
||||
&candidate.provider_name,
|
||||
&candidate.provider_type,
|
||||
)
|
||||
});
|
||||
if !provider_allowed {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(allowed_models) = auth_snapshot.effective_allowed_models() {
|
||||
let model_allowed = allowed_models
|
||||
.iter()
|
||||
.any(|value| value == requested_model || value == &candidate.global_model_name);
|
||||
if !model_allowed {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
@@ -1,260 +0,0 @@
|
||||
use axum::body::Bytes;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::contracts::{
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
|
||||
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
|
||||
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
|
||||
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
extract_gemini_model_from_path as extract_gemini_model_from_path_impl,
|
||||
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl,
|
||||
is_json_request, parse_direct_request_body as parse_direct_request_body_impl,
|
||||
};
|
||||
use crate::LocalExecutionRuntimeMissDiagnostic;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum RequestedModelFamily {
|
||||
Standard,
|
||||
Gemini,
|
||||
}
|
||||
|
||||
pub(crate) fn parse_direct_request_body(
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &Bytes,
|
||||
) -> Option<(serde_json::Value, Option<String>)> {
|
||||
parse_direct_request_body_impl(is_json_request(&parts.headers), body_bytes.as_ref())
|
||||
}
|
||||
|
||||
pub(crate) fn force_upstream_streaming_for_provider(
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
) -> bool {
|
||||
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> {
|
||||
body_json
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_requested_model_from_request(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
family: RequestedModelFamily,
|
||||
) -> Option<String> {
|
||||
match family {
|
||||
RequestedModelFamily::Standard => extract_standard_requested_model(body_json),
|
||||
RequestedModelFamily::Gemini => extract_gemini_model_from_path_impl(parts.uri.path()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_runtime_miss_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: reason.to_string(),
|
||||
route_family: decision.route_family.clone(),
|
||||
route_kind: decision.route_kind.clone(),
|
||||
public_path: Some(decision.public_path.clone()),
|
||||
plan_kind: Some(plan_kind.to_string()),
|
||||
requested_model: requested_model.map(ToOwned::to_owned),
|
||||
candidate_count: None,
|
||||
skipped_candidate_count: None,
|
||||
skip_reasons: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_candidate_evaluation_progress(
|
||||
diagnostic: &mut LocalExecutionRuntimeMissDiagnostic,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_candidate_terminal_plan_reason(
|
||||
diagnostic: &mut LocalExecutionRuntimeMissDiagnostic,
|
||||
no_plan_reason: &'static str,
|
||||
) {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count
|
||||
&& diagnostic.skip_reasons.len() == 1
|
||||
&& diagnostic
|
||||
.skip_reasons
|
||||
.get("api_key_concurrency_limit_reached")
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
> 0
|
||||
{
|
||||
"api_key_concurrency_limit_reached".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
no_plan_reason.to_string()
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
apply_local_candidate_evaluation_progress, apply_local_candidate_terminal_plan_reason,
|
||||
build_local_runtime_miss_diagnostic, extract_requested_model_from_request,
|
||||
extract_standard_requested_model, force_upstream_streaming_for_provider,
|
||||
RequestedModelFamily,
|
||||
};
|
||||
use axum::http::Request;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn forces_streaming_for_codex_openai_responses() {
|
||||
assert!(force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_force_streaming_for_compact_or_other_provider_types() {
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"openai",
|
||||
"openai:responses"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_standard_requested_model_from_request_body() {
|
||||
let requested_model =
|
||||
extract_standard_requested_model(&json!({ "model": " claude-sonnet-4 " }));
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_family_helper_delegates_standard_model_extraction() {
|
||||
let request = Request::builder()
|
||||
.uri("https://example.test/v1/chat/completions")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
&parts,
|
||||
&json!({ "model": " claude-sonnet-4 " }),
|
||||
RequestedModelFamily::Standard,
|
||||
);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_gemini_requested_model_from_request_path() {
|
||||
let request = Request::builder()
|
||||
.uri("https://example.test/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let requested_model =
|
||||
extract_requested_model_from_request(&parts, &json!({}), RequestedModelFamily::Gemini);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("gemini-2.5-pro"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_evaluation_progress_sets_candidate_count_and_reason() {
|
||||
let mut diagnostic = build_local_runtime_miss_diagnostic(
|
||||
&crate::ai_pipeline::GatewayControlDecision::synthetic(
|
||||
"/v1/test",
|
||||
Some("passthrough".to_string()),
|
||||
Some("ai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("test#sync".to_string()),
|
||||
),
|
||||
"test_plan",
|
||||
Some("test-model"),
|
||||
"seed",
|
||||
);
|
||||
|
||||
apply_local_candidate_evaluation_progress(&mut diagnostic, 0);
|
||||
assert_eq!(diagnostic.candidate_count, Some(0));
|
||||
assert_eq!(diagnostic.reason, "candidate_list_empty");
|
||||
|
||||
apply_local_candidate_evaluation_progress(&mut diagnostic, 3);
|
||||
assert_eq!(diagnostic.candidate_count, Some(3));
|
||||
assert_eq!(diagnostic.reason, "candidate_evaluation_incomplete");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_terminal_reason_prefers_empty_then_skipped_then_fallback() {
|
||||
let mut diagnostic = build_local_runtime_miss_diagnostic(
|
||||
&crate::ai_pipeline::GatewayControlDecision::synthetic(
|
||||
"/v1/test",
|
||||
Some("passthrough".to_string()),
|
||||
Some("ai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("test#sync".to_string()),
|
||||
),
|
||||
"test_plan",
|
||||
Some("test-model"),
|
||||
"seed",
|
||||
);
|
||||
|
||||
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
|
||||
assert_eq!(diagnostic.reason, "candidate_list_empty");
|
||||
|
||||
diagnostic.candidate_count = Some(2);
|
||||
diagnostic.skipped_candidate_count = Some(2);
|
||||
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
|
||||
assert_eq!(diagnostic.reason, "all_candidates_skipped");
|
||||
|
||||
diagnostic.skip_reasons = std::collections::BTreeMap::from([(
|
||||
"api_key_concurrency_limit_reached".to_string(),
|
||||
2,
|
||||
)]);
|
||||
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
|
||||
assert_eq!(diagnostic.reason, "api_key_concurrency_limit_reached");
|
||||
|
||||
diagnostic.skipped_candidate_count = Some(1);
|
||||
diagnostic.skip_reasons.clear();
|
||||
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
|
||||
assert_eq!(diagnostic.reason, "no_local_sync_plans");
|
||||
}
|
||||
}
|
||||
@@ -1,174 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::planner::route::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
resolve_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if !is_matching_stream_request(plan_kind, parts, body_json, body_base64) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Some(payload) = maybe_build_local_video_task_content_stream_decision_payload(
|
||||
state, parts, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_image_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_openai_responses_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_standard_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_same_format_provider_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_gemini_files_decision_payload(
|
||||
state, parts, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_task_content_stream_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
if plan_kind != OPENAI_VIDEO_CONTENT_PLAN_KIND
|
||||
|| decision.route_family.as_deref() != Some("openai")
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
|
||||
.await?;
|
||||
|
||||
let Some(action) = state.video_tasks.prepare_openai_content_stream_action(
|
||||
parts.uri.path(),
|
||||
parts.uri.query(),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) = action else {
|
||||
return Ok(None);
|
||||
};
|
||||
let plan = *plan;
|
||||
let provider_contract = plan.provider_api_format.clone();
|
||||
let client_contract = plan.client_api_format.clone();
|
||||
let execution_strategy = if plan.provider_api_format == plan.client_api_format {
|
||||
ExecutionStrategy::LocalSameFormat
|
||||
} else {
|
||||
ExecutionStrategy::LocalCrossFormat
|
||||
};
|
||||
let conversion_mode = if plan.provider_api_format == plan.client_api_format {
|
||||
ConversionMode::None
|
||||
} else {
|
||||
ConversionMode::Bidirectional
|
||||
};
|
||||
|
||||
Ok(Some(GatewayControlSyncDecisionResponse {
|
||||
action: EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(plan_kind.to_string()),
|
||||
execution_strategy: Some(execution_strategy.as_str().to_string()),
|
||||
conversion_mode: Some(conversion_mode.as_str().to_string()),
|
||||
request_id: Some(plan.request_id),
|
||||
candidate_id: plan.candidate_id,
|
||||
provider_name: plan.provider_name,
|
||||
provider_id: Some(plan.provider_id),
|
||||
endpoint_id: Some(plan.endpoint_id),
|
||||
key_id: Some(plan.key_id),
|
||||
upstream_base_url: None,
|
||||
upstream_url: Some(plan.url),
|
||||
provider_request_method: Some(plan.method),
|
||||
auth_header: None,
|
||||
auth_value: None,
|
||||
provider_api_format: Some(plan.provider_api_format),
|
||||
client_api_format: Some(plan.client_api_format),
|
||||
provider_contract: Some(provider_contract),
|
||||
client_contract: Some(client_contract),
|
||||
model_name: plan.model_name,
|
||||
mapped_model: None,
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: plan.headers,
|
||||
provider_request_body: None,
|
||||
provider_request_body_base64: None,
|
||||
content_type: plan.content_type,
|
||||
proxy: plan.proxy,
|
||||
tls_profile: plan.tls_profile,
|
||||
timeouts: plan.timeouts,
|
||||
upstream_is_stream: true,
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
auth_context: resolve_decision_execution_runtime_auth_context(decision),
|
||||
}))
|
||||
}
|
||||
@@ -1,366 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use tracing::debug;
|
||||
use url::Url;
|
||||
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DELETE_PLAN_KIND,
|
||||
GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::planner::route::resolve_execution_runtime_sync_plan_kind;
|
||||
use crate::ai_pipeline::{
|
||||
build_execution_runtime_auth_context, resolve_execution_runtime_auth_context, ConversionMode,
|
||||
ExecutionStrategy, GatewayControlDecision,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_sync_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if let Some(payload) = maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
state, parts, body_json, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_sync_local_video_decision_payload(
|
||||
state, parts, body_json, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_sync_local_image_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_sync_local_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_sync_local_openai_responses_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_sync_local_standard_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_sync_local_same_format_provider_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if matches!(
|
||||
plan_kind,
|
||||
GEMINI_FILES_LIST_PLAN_KIND | GEMINI_FILES_GET_PLAN_KIND | GEMINI_FILES_DELETE_PLAN_KIND
|
||||
) {
|
||||
if let Some(payload) = super::maybe_build_sync_local_gemini_files_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
if !matches!(
|
||||
plan_kind,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND
|
||||
| GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
|
||||
.await?;
|
||||
|
||||
let auth_context = resolve_execution_runtime_auth_context(
|
||||
state,
|
||||
decision,
|
||||
&parts.headers,
|
||||
&parts.uri,
|
||||
trace_id,
|
||||
)
|
||||
.await?;
|
||||
let Some(auth_context) = auth_context else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(follow_up) = state.video_tasks.prepare_follow_up_sync_plan(
|
||||
plan_kind,
|
||||
parts.uri.path(),
|
||||
Some(body_json),
|
||||
Some(&auth_context),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let aether_video_tasks_core::LocalVideoTaskFollowUpPlan {
|
||||
plan,
|
||||
report_kind,
|
||||
report_context,
|
||||
} = follow_up;
|
||||
let aether_contracts::ExecutionPlan {
|
||||
request_id: _request_id,
|
||||
candidate_id,
|
||||
provider_name,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method,
|
||||
url,
|
||||
headers,
|
||||
content_type,
|
||||
content_encoding: _content_encoding,
|
||||
body,
|
||||
stream: _stream,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name,
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts,
|
||||
} = plan;
|
||||
let auth_pair = extract_auth_header_pair(&headers);
|
||||
let execution_strategy = if provider_api_format == client_api_format {
|
||||
ExecutionStrategy::LocalSameFormat
|
||||
} else {
|
||||
ExecutionStrategy::LocalCrossFormat
|
||||
};
|
||||
let conversion_mode = if provider_api_format == client_api_format {
|
||||
ConversionMode::None
|
||||
} else {
|
||||
ConversionMode::Bidirectional
|
||||
};
|
||||
let upstream_base_url = infer_upstream_base_url(&url);
|
||||
let provider_contract = provider_api_format.clone();
|
||||
let client_contract = client_api_format.clone();
|
||||
let auth_header = auth_pair.map(|(name, _)| name.to_string());
|
||||
let auth_value = auth_pair.map(|(_, value)| value.to_string());
|
||||
let aether_contracts::RequestBody {
|
||||
json_body,
|
||||
body_bytes_b64,
|
||||
body_ref: _body_ref,
|
||||
} = body;
|
||||
|
||||
debug!(
|
||||
event_name = "local_video_follow_up_sync_decision_payload_built",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
request_id = %trace_id,
|
||||
candidate_id = ?candidate_id,
|
||||
provider_id = %provider_id,
|
||||
endpoint_id = %endpoint_id,
|
||||
key_id = %key_id,
|
||||
plan_kind,
|
||||
downstream_path = %parts.uri.path(),
|
||||
provider_api_format = %provider_api_format,
|
||||
client_api_format = %client_api_format,
|
||||
upstream_base_url = ?upstream_base_url,
|
||||
upstream_url = %url,
|
||||
"gateway built local video follow-up sync decision payload"
|
||||
);
|
||||
|
||||
Ok(Some(GatewayControlSyncDecisionResponse {
|
||||
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(plan_kind.to_string()),
|
||||
execution_strategy: Some(execution_strategy.as_str().to_string()),
|
||||
conversion_mode: Some(conversion_mode.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id,
|
||||
provider_name,
|
||||
provider_id: Some(provider_id),
|
||||
endpoint_id: Some(endpoint_id),
|
||||
key_id: Some(key_id),
|
||||
upstream_base_url,
|
||||
upstream_url: Some(url),
|
||||
provider_request_method: Some(method),
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_api_format: Some(provider_api_format),
|
||||
client_api_format: Some(client_api_format),
|
||||
provider_contract: Some(provider_contract),
|
||||
client_contract: Some(client_contract),
|
||||
model_name,
|
||||
mapped_model: None,
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: headers,
|
||||
provider_request_body: json_body,
|
||||
provider_request_body_base64: body_bytes_b64,
|
||||
content_type,
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts,
|
||||
upstream_is_stream: false,
|
||||
report_kind,
|
||||
report_context,
|
||||
auth_context: Some(build_execution_runtime_auth_context(&auth_context)),
|
||||
}))
|
||||
}
|
||||
|
||||
fn extract_auth_header_pair<'a>(
|
||||
headers: &'a BTreeMap<String, String>,
|
||||
) -> Option<(&'a str, &'a str)> {
|
||||
[
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"api-key",
|
||||
"x-goog-api-key",
|
||||
"proxy-authorization",
|
||||
]
|
||||
.into_iter()
|
||||
.find_map(|name| {
|
||||
headers
|
||||
.iter()
|
||||
.find(|(header_name, _)| header_name.eq_ignore_ascii_case(name))
|
||||
.map(|(header_name, value)| (header_name.as_str(), value.as_str()))
|
||||
})
|
||||
}
|
||||
|
||||
fn infer_upstream_base_url(upstream_url: &str) -> Option<String> {
|
||||
let parsed = Url::parse(upstream_url).ok()?;
|
||||
let host = parsed.host_str()?;
|
||||
let mut base = format!("{}://{}", parsed.scheme(), host);
|
||||
if let Some(port) = parsed.port() {
|
||||
base.push(':');
|
||||
base.push_str(port.to_string().as_str());
|
||||
}
|
||||
let base_path = infer_upstream_base_path(parsed.path());
|
||||
if !base_path.is_empty() {
|
||||
base.push_str(base_path);
|
||||
}
|
||||
Some(base)
|
||||
}
|
||||
|
||||
fn infer_upstream_base_path(path: &str) -> &str {
|
||||
let trimmed = path.trim_end_matches('/');
|
||||
if trimmed.is_empty() || trimmed == "/" {
|
||||
return "";
|
||||
}
|
||||
|
||||
for suffix in [
|
||||
"/responses/compact",
|
||||
"/responses",
|
||||
"/chat/completions",
|
||||
"/messages",
|
||||
] {
|
||||
if let Some(prefix) = trimmed.strip_suffix(suffix) {
|
||||
return normalize_inferred_base_path(prefix);
|
||||
}
|
||||
}
|
||||
|
||||
for marker in ["/v1/videos", "/v1beta/"] {
|
||||
if let Some((prefix, _)) = trimmed.split_once(marker) {
|
||||
return normalize_inferred_base_path(prefix);
|
||||
}
|
||||
}
|
||||
|
||||
normalize_inferred_base_path(trimmed)
|
||||
}
|
||||
|
||||
fn normalize_inferred_base_path(path: &str) -> &str {
|
||||
let trimmed = path.trim_end_matches('/');
|
||||
if trimmed.is_empty() || trimmed == "/" {
|
||||
""
|
||||
} else {
|
||||
trimmed
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::infer_upstream_base_url;
|
||||
|
||||
#[test]
|
||||
fn infer_upstream_base_url_preserves_codex_base_path() {
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://tiger.bookapi.cc/codex/responses").as_deref(),
|
||||
Some("https://tiger.bookapi.cc/codex")
|
||||
);
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://chatgpt.com/backend-api/codex/responses").as_deref(),
|
||||
Some("https://chatgpt.com/backend-api/codex")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_upstream_base_url_preserves_nested_v1_prefix() {
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://api.openai.example/custom/v1/chat/completions?mode=1")
|
||||
.as_deref(),
|
||||
Some("https://api.openai.example/custom/v1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_upstream_base_url_strips_video_operation_path() {
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://video.example/nested/v1/videos/task-123/content")
|
||||
.as_deref(),
|
||||
Some("https://video.example/nested")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,83 +0,0 @@
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct ResolvedLocalDecisionAuthInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalRequestedModelDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalAuthenticatedDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_requested_model_decision_input(
|
||||
resolved_input: ResolvedLocalDecisionAuthInput,
|
||||
requested_model: String,
|
||||
) -> LocalRequestedModelDecisionInput {
|
||||
LocalRequestedModelDecisionInput {
|
||||
auth_context: resolved_input.auth_context,
|
||||
requested_model,
|
||||
auth_snapshot: resolved_input.auth_snapshot,
|
||||
required_capabilities: resolved_input.required_capabilities,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_authenticated_decision_input(
|
||||
resolved_input: ResolvedLocalDecisionAuthInput,
|
||||
) -> LocalAuthenticatedDecisionInput {
|
||||
LocalAuthenticatedDecisionInput {
|
||||
auth_context: resolved_input.auth_context,
|
||||
auth_snapshot: resolved_input.auth_snapshot,
|
||||
required_capabilities: resolved_input.required_capabilities,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_local_authenticated_decision_input(
|
||||
state: &AppState,
|
||||
auth_context: ExecutionRuntimeAuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let auth_snapshot = match planner_state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(snapshot) => snapshot,
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let required_capabilities = planner_state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
requested_model,
|
||||
explicit_required_capabilities,
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(Some(ResolvedLocalDecisionAuthInput {
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
}))
|
||||
}
|
||||
@@ -1,190 +0,0 @@
|
||||
use serde_json::{json, Value};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum CandidateFailureDiagnosticKind {
|
||||
RequestBodyBuild,
|
||||
RequestConversion,
|
||||
BodyRules,
|
||||
HeaderRules,
|
||||
UrlBuild,
|
||||
TransportAuth,
|
||||
EnvelopeBuild,
|
||||
}
|
||||
|
||||
impl CandidateFailureDiagnosticKind {
|
||||
pub(crate) fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::RequestBodyBuild => "request_body_build",
|
||||
Self::RequestConversion => "request_conversion",
|
||||
Self::BodyRules => "body_rules",
|
||||
Self::HeaderRules => "header_rules",
|
||||
Self::UrlBuild => "url_build",
|
||||
Self::TransportAuth => "transport_auth",
|
||||
Self::EnvelopeBuild => "envelope_build",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct CandidateFailureDiagnostic {
|
||||
pub(crate) kind: CandidateFailureDiagnosticKind,
|
||||
pub(crate) path: String,
|
||||
pub(crate) message: String,
|
||||
pub(crate) source: Option<String>,
|
||||
pub(crate) client_api_format: Option<String>,
|
||||
pub(crate) provider_api_format: Option<String>,
|
||||
pub(crate) safe_to_show: bool,
|
||||
}
|
||||
|
||||
impl CandidateFailureDiagnostic {
|
||||
pub(crate) fn new(
|
||||
kind: CandidateFailureDiagnosticKind,
|
||||
path: impl Into<String>,
|
||||
message: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
path: path.into(),
|
||||
message: message.into(),
|
||||
source: None,
|
||||
client_api_format: None,
|
||||
provider_api_format: None,
|
||||
safe_to_show: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn source(mut self, source: impl Into<String>) -> Self {
|
||||
self.source = Some(source.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn formats(
|
||||
mut self,
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
) -> Self {
|
||||
self.client_api_format = Some(client_api_format.into());
|
||||
self.provider_api_format = Some(provider_api_format.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn to_extra_data(&self) -> Value {
|
||||
let diagnostic = self.to_value();
|
||||
let mut extra_data = json!({
|
||||
"failure_diagnostic": diagnostic,
|
||||
});
|
||||
|
||||
// Compatibility for current usage UI and already persisted trace readers.
|
||||
if self.kind == CandidateFailureDiagnosticKind::RequestBodyBuild {
|
||||
if let Some(object) = extra_data.as_object_mut() {
|
||||
object.insert(
|
||||
"request_body_build_error".to_string(),
|
||||
json!({
|
||||
"path": self.path,
|
||||
"message": self.message,
|
||||
"client_api_format": self.client_api_format,
|
||||
"provider_api_format": self.provider_api_format,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
extra_data
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_url_missing(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::UrlBuild,
|
||||
"$.endpoint",
|
||||
"无法构建上游请求地址；请检查 base_url、custom_path、API 格式和模型映射",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub(crate) fn header_rules_apply_failed(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::HeaderRules,
|
||||
"$.endpoint.header_rules",
|
||||
"Header 规则应用失败；请检查规则格式、条件配置，或是否试图覆盖受保护认证头",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub(crate) fn body_rules_apply_failed(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::BodyRules,
|
||||
"$.endpoint.body_rules",
|
||||
"Body 规则应用失败；请检查规则格式、条件配置，或规则输出是否仍是当前上游支持的请求体",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub(crate) fn body_rules_unsupported_for_binary_upload(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::BodyRules,
|
||||
"$.endpoint.body_rules",
|
||||
"二进制上传暂不支持本地应用 Body 规则；请移除该 Endpoint 的 Body 规则或改用 JSON 请求体",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_request_body_missing(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::RequestBodyBuild,
|
||||
"$",
|
||||
"无法构建上游请求体；请检查请求体是否为支持的 JSON object，以及该任务类型必需字段是否存在且取值受支持",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
pub(crate) fn envelope_build_failed(
|
||||
client_api_format: impl Into<String>,
|
||||
provider_api_format: impl Into<String>,
|
||||
source: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
CandidateFailureDiagnosticKind::EnvelopeBuild,
|
||||
"$",
|
||||
"无法构建上游请求封装；请检查该 Provider 的认证配置、模型映射、Endpoint Body 规则和当前请求体是否兼容",
|
||||
)
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(source)
|
||||
}
|
||||
|
||||
fn to_value(&self) -> Value {
|
||||
json!({
|
||||
"kind": self.kind.as_str(),
|
||||
"path": self.path,
|
||||
"message": self.message,
|
||||
"source": self.source,
|
||||
"client_api_format": self.client_api_format,
|
||||
"provider_api_format": self.provider_api_format,
|
||||
"safe_to_show": self.safe_to_show,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,85 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||
LocalAvailableCandidatePersistenceContext, LocalSkippedCandidatePersistenceContext,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[allow(clippy::enum_variant_names)]
|
||||
pub(crate) enum LocalCandidatePersistencePolicyKind {
|
||||
StandardDecision,
|
||||
SameFormatProviderDecision,
|
||||
OpenAiChatDecision,
|
||||
OpenAiResponsesDecision,
|
||||
ImageDecision,
|
||||
GeminiFilesDecision,
|
||||
VideoDecision,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct LocalCandidatePersistencePolicy<'a> {
|
||||
pub(crate) available: LocalAvailableCandidatePersistenceContext<'a>,
|
||||
pub(crate) skipped: LocalSkippedCandidatePersistenceContext<'a>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_candidate_persistence_policy<'a>(
|
||||
auth_context: &'a ExecutionRuntimeAuthContext,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
kind: LocalCandidatePersistencePolicyKind,
|
||||
) -> LocalCandidatePersistencePolicy<'a> {
|
||||
let (available_error_context, skipped_error_context, record_runtime_miss_diagnostic) =
|
||||
match kind {
|
||||
LocalCandidatePersistencePolicyKind::StandardDecision => (
|
||||
"gateway local standard decision request candidate upsert failed",
|
||||
"gateway local standard decision failed to persist skipped candidate",
|
||||
true,
|
||||
),
|
||||
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision => (
|
||||
"gateway local same-format decision request candidate upsert failed",
|
||||
"gateway local same-format decision failed to persist skipped candidate",
|
||||
true,
|
||||
),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiChatDecision => (
|
||||
"gateway local openai chat decision request candidate upsert failed",
|
||||
"gateway local openai chat decision failed to persist skipped candidate",
|
||||
true,
|
||||
),
|
||||
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision => (
|
||||
"gateway local openai responses decision request candidate upsert failed",
|
||||
"gateway local openai responses decision failed to persist skipped candidate",
|
||||
true,
|
||||
),
|
||||
LocalCandidatePersistencePolicyKind::ImageDecision => (
|
||||
"gateway local openai image decision request candidate upsert failed",
|
||||
"gateway local openai image decision failed to persist skipped candidate",
|
||||
false,
|
||||
),
|
||||
LocalCandidatePersistencePolicyKind::GeminiFilesDecision => (
|
||||
"gateway local gemini files request candidate upsert failed",
|
||||
"gateway local gemini files failed to persist skipped candidate",
|
||||
false,
|
||||
),
|
||||
LocalCandidatePersistencePolicyKind::VideoDecision => (
|
||||
"gateway local video decision request candidate upsert failed",
|
||||
"gateway local video decision failed to persist skipped candidate",
|
||||
false,
|
||||
),
|
||||
};
|
||||
|
||||
LocalCandidatePersistencePolicy {
|
||||
available: LocalAvailableCandidatePersistenceContext {
|
||||
user_id: &auth_context.user_id,
|
||||
api_key_id: &auth_context.api_key_id,
|
||||
required_capabilities,
|
||||
error_context: available_error_context,
|
||||
},
|
||||
skipped: LocalSkippedCandidatePersistenceContext {
|
||||
user_id: &auth_context.user_id,
|
||||
api_key_id: &auth_context.api_key_id,
|
||||
required_capabilities,
|
||||
error_context: skipped_error_context,
|
||||
record_runtime_miss_diagnostic,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1,177 +0,0 @@
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, take_non_empty_string, LocalStreamPlanAndReport,
|
||||
LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::{GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) fn build_passthrough_sync_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<LocalSyncPlanAndReport>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(request_id) = take_non_empty_string(&mut payload.request_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_id) = take_non_empty_string(&mut payload.provider_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(endpoint_id) = take_non_empty_string(&mut payload.endpoint_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(client_api_format) = take_non_empty_string(&mut payload.client_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let ignored_provider_request_body = serde_json::Value::Null;
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&ignored_provider_request_body,
|
||||
)?;
|
||||
let request_body = resolve_passthrough_sync_request_body(
|
||||
payload.provider_request_body.take(),
|
||||
payload.provider_request_body_base64.take(),
|
||||
);
|
||||
let provider_request_method = take_non_empty_string(&mut payload.provider_request_method);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.take(),
|
||||
provider_name: payload.provider_name.take(),
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method: provider_request_method.unwrap_or_else(|| parts.method.to_string()),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body: request_body,
|
||||
stream: false,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name: payload.model_name.take(),
|
||||
proxy: payload.proxy.take(),
|
||||
tls_profile: payload.tls_profile.take(),
|
||||
timeouts: payload.timeouts.take(),
|
||||
};
|
||||
|
||||
Ok(Some(LocalSyncPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_passthrough_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<LocalStreamPlanAndReport>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(request_id) = take_non_empty_string(&mut payload.request_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_id) = take_non_empty_string(&mut payload.provider_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(endpoint_id) = take_non_empty_string(&mut payload.endpoint_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(client_api_format) = take_non_empty_string(&mut payload.client_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.take(),
|
||||
provider_name: payload.provider_name.take(),
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method: parts.method.to_string(),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: true,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name: payload.model_name.take(),
|
||||
proxy: payload.proxy.take(),
|
||||
tls_profile: payload.tls_profile.take(),
|
||||
timeouts: payload.timeouts.take(),
|
||||
};
|
||||
|
||||
Ok(Some(LocalStreamPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context: payload.report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
fn resolve_passthrough_sync_request_body(
|
||||
provider_request_body: Option<serde_json::Value>,
|
||||
provider_request_body_base64: Option<String>,
|
||||
) -> RequestBody {
|
||||
if let Some(body_bytes_b64) = provider_request_body_base64.and_then(trim_owned_non_empty_string)
|
||||
{
|
||||
return RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: Some(body_bytes_b64),
|
||||
body_ref: None,
|
||||
};
|
||||
}
|
||||
|
||||
match provider_request_body.unwrap_or(serde_json::Value::Null) {
|
||||
serde_json::Value::Null => RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
other => RequestBody::from_json(other),
|
||||
}
|
||||
}
|
||||
|
||||
fn trim_owned_non_empty_string(value: String) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if trimmed.len() == value.len() {
|
||||
return Some(value);
|
||||
}
|
||||
Some(trimmed.to_owned())
|
||||
}
|
||||
@@ -1,134 +0,0 @@
|
||||
use crate::ai_pipeline::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata;
|
||||
use crate::ai_pipeline::transport::auth::{resolve_local_gemini_auth, resolve_local_standard_auth};
|
||||
use crate::ai_pipeline::transport::claude_code::local_claude_code_transport_unsupported_reason_with_network;
|
||||
use crate::ai_pipeline::transport::kiro::local_kiro_request_transport_unsupported_reason_with_network;
|
||||
use crate::ai_pipeline::transport::policy::{
|
||||
local_gemini_transport_unsupported_reason_with_network,
|
||||
local_standard_transport_unsupported_reason_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::transport::vertex::{
|
||||
is_vertex_api_key_transport_context,
|
||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
|
||||
|
||||
use super::super::LocalSameFormatProviderFamily;
|
||||
|
||||
pub(super) struct SameFormatProviderRequestBehavior {
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) is_claude_code: bool,
|
||||
pub(super) is_vertex: bool,
|
||||
pub(super) is_kiro: bool,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
pub(super) report_kind: &'static str,
|
||||
}
|
||||
|
||||
pub(super) fn classify_same_format_provider_request_behavior(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
spec_metadata: LocalExecutionSurfaceSpecMetadata,
|
||||
) -> SameFormatProviderRequestBehavior {
|
||||
let is_antigravity = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("antigravity");
|
||||
let is_claude_code = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("claude_code");
|
||||
let is_vertex = is_vertex_api_key_transport_context(transport);
|
||||
let is_kiro = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("kiro");
|
||||
let default_report_kind = spec_metadata
|
||||
.report_kind
|
||||
.expect("same-format provider specs should declare report kind");
|
||||
let upstream_is_stream = is_kiro || is_antigravity || spec_metadata.require_streaming;
|
||||
let report_kind = if is_kiro && !spec_metadata.require_streaming {
|
||||
"claude_cli_sync_finalize"
|
||||
} else if is_antigravity && !spec_metadata.require_streaming {
|
||||
match default_report_kind {
|
||||
"gemini_chat_sync_success" => "gemini_chat_sync_finalize",
|
||||
"gemini_cli_sync_success" => "gemini_cli_sync_finalize",
|
||||
_ => default_report_kind,
|
||||
}
|
||||
} else {
|
||||
default_report_kind
|
||||
};
|
||||
|
||||
SameFormatProviderRequestBehavior {
|
||||
is_antigravity,
|
||||
is_claude_code,
|
||||
is_vertex,
|
||||
is_kiro,
|
||||
upstream_is_stream,
|
||||
report_kind,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_supported(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
same_format_provider_transport_unsupported_reason(behavior, transport, family, api_format)
|
||||
.is_none()
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_unsupported_reason(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> Option<&'static str> {
|
||||
if behavior.is_kiro {
|
||||
local_kiro_request_transport_unsupported_reason_with_network(transport)
|
||||
} else if behavior.is_antigravity {
|
||||
None
|
||||
} else if behavior.is_claude_code {
|
||||
local_claude_code_transport_unsupported_reason_with_network(transport, api_format)
|
||||
} else if behavior.is_vertex {
|
||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport)
|
||||
} else {
|
||||
match family {
|
||||
LocalSameFormatProviderFamily::Standard => {
|
||||
local_standard_transport_unsupported_reason_with_network(transport, api_format)
|
||||
}
|
||||
LocalSameFormatProviderFamily::Gemini => {
|
||||
local_gemini_transport_unsupported_reason_with_network(transport, api_format)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn should_try_same_format_provider_oauth_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> bool {
|
||||
behavior.is_kiro
|
||||
|| matches!(family, LocalSameFormatProviderFamily::Standard)
|
||||
&& resolve_local_standard_auth(transport).is_none()
|
||||
|| matches!(family, LocalSameFormatProviderFamily::Gemini)
|
||||
&& !behavior.is_vertex
|
||||
&& resolve_local_gemini_auth(transport).is_none()
|
||||
}
|
||||
|
||||
pub(super) fn resolve_same_format_provider_direct_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> Option<(String, String)> {
|
||||
if behavior.is_vertex {
|
||||
None
|
||||
} else {
|
||||
match family {
|
||||
LocalSameFormatProviderFamily::Standard => resolve_local_standard_auth(transport),
|
||||
LocalSameFormatProviderFamily::Gemini => resolve_local_gemini_auth(transport),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::{
|
||||
apply_local_body_rules, build_kiro_provider_request_body, sanitize_claude_code_request_body,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) fn build_same_format_provider_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
body_rules: Option<&Value>,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
|
||||
is_claude_code: bool,
|
||||
) -> Option<Value> {
|
||||
if let Some(kiro_auth) = kiro_auth {
|
||||
return build_kiro_provider_request_body(
|
||||
body_json,
|
||||
mapped_model,
|
||||
&kiro_auth.auth_config,
|
||||
body_rules,
|
||||
);
|
||||
}
|
||||
|
||||
let request_body_object = body_json.as_object()?;
|
||||
let mut provider_request_body = serde_json::Map::from_iter(
|
||||
request_body_object
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), value.clone())),
|
||||
);
|
||||
match spec.family {
|
||||
LocalSameFormatProviderFamily::Standard => {
|
||||
provider_request_body
|
||||
.insert("model".to_string(), Value::String(mapped_model.to_string()));
|
||||
if upstream_is_stream {
|
||||
provider_request_body.insert("stream".to_string(), Value::Bool(true));
|
||||
}
|
||||
}
|
||||
LocalSameFormatProviderFamily::Gemini => {
|
||||
provider_request_body.remove("model");
|
||||
}
|
||||
}
|
||||
let mut provider_request_body = Value::Object(provider_request_body);
|
||||
if is_claude_code {
|
||||
sanitize_claude_code_request_body(&mut provider_request_body);
|
||||
}
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
Some(provider_request_body)
|
||||
}
|
||||
@@ -1,25 +0,0 @@
|
||||
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
|
||||
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
|
||||
pub(crate) fn build_same_format_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
|
||||
) -> Option<String> {
|
||||
maybe_add_gemini_stream_alt_sse(crate::ai_pipeline::build_provider_transport_request_url(
|
||||
transport,
|
||||
spec.api_format,
|
||||
Some(mapped_model),
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
|
||||
))
|
||||
}
|
||||
|
||||
fn maybe_add_gemini_stream_alt_sse(url: Option<String>) -> Option<String> {
|
||||
url
|
||||
}
|
||||
@@ -1,90 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::{ExecutionTimeouts, ProxySnapshot};
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, GatewayControlSyncDecisionResponse};
|
||||
use crate::{EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION};
|
||||
|
||||
pub(crate) struct LocalExecutionDecisionResponseParts {
|
||||
pub(crate) decision_is_stream: bool,
|
||||
pub(crate) decision_kind: String,
|
||||
pub(crate) execution_strategy: ExecutionStrategy,
|
||||
pub(crate) conversion_mode: ConversionMode,
|
||||
pub(crate) request_id: String,
|
||||
pub(crate) candidate_id: String,
|
||||
pub(crate) provider_name: String,
|
||||
pub(crate) provider_id: String,
|
||||
pub(crate) endpoint_id: String,
|
||||
pub(crate) key_id: String,
|
||||
pub(crate) upstream_base_url: String,
|
||||
pub(crate) upstream_url: String,
|
||||
pub(crate) provider_request_method: Option<String>,
|
||||
pub(crate) auth_header: Option<String>,
|
||||
pub(crate) auth_value: Option<String>,
|
||||
pub(crate) provider_api_format: String,
|
||||
pub(crate) client_api_format: String,
|
||||
pub(crate) model_name: String,
|
||||
pub(crate) mapped_model: String,
|
||||
pub(crate) prompt_cache_key: Option<String>,
|
||||
pub(crate) provider_request_headers: BTreeMap<String, String>,
|
||||
pub(crate) provider_request_body: Option<serde_json::Value>,
|
||||
pub(crate) provider_request_body_base64: Option<String>,
|
||||
pub(crate) content_type: Option<String>,
|
||||
pub(crate) proxy: Option<ProxySnapshot>,
|
||||
pub(crate) tls_profile: Option<String>,
|
||||
pub(crate) timeouts: Option<ExecutionTimeouts>,
|
||||
pub(crate) upstream_is_stream: bool,
|
||||
pub(crate) report_kind: Option<String>,
|
||||
pub(crate) report_context: Option<serde_json::Value>,
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_decision_response(
|
||||
parts: LocalExecutionDecisionResponseParts,
|
||||
) -> GatewayControlSyncDecisionResponse {
|
||||
GatewayControlSyncDecisionResponse {
|
||||
action: local_execution_decision_action(parts.decision_is_stream).to_string(),
|
||||
decision_kind: Some(parts.decision_kind),
|
||||
execution_strategy: Some(parts.execution_strategy.as_str().to_string()),
|
||||
conversion_mode: Some(parts.conversion_mode.as_str().to_string()),
|
||||
request_id: Some(parts.request_id),
|
||||
candidate_id: Some(parts.candidate_id),
|
||||
provider_name: Some(parts.provider_name),
|
||||
provider_id: Some(parts.provider_id),
|
||||
endpoint_id: Some(parts.endpoint_id),
|
||||
key_id: Some(parts.key_id),
|
||||
upstream_base_url: Some(parts.upstream_base_url),
|
||||
upstream_url: Some(parts.upstream_url),
|
||||
provider_request_method: parts.provider_request_method,
|
||||
auth_header: parts.auth_header,
|
||||
auth_value: parts.auth_value,
|
||||
provider_api_format: Some(parts.provider_api_format.clone()),
|
||||
client_api_format: Some(parts.client_api_format.clone()),
|
||||
provider_contract: Some(parts.provider_api_format),
|
||||
client_contract: Some(parts.client_api_format),
|
||||
model_name: Some(parts.model_name),
|
||||
mapped_model: Some(parts.mapped_model),
|
||||
prompt_cache_key: parts.prompt_cache_key,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: parts.provider_request_headers,
|
||||
provider_request_body: parts.provider_request_body,
|
||||
provider_request_body_base64: parts.provider_request_body_base64,
|
||||
content_type: parts.content_type,
|
||||
proxy: parts.proxy,
|
||||
tls_profile: parts.tls_profile,
|
||||
timeouts: parts.timeouts,
|
||||
upstream_is_stream: parts.upstream_is_stream,
|
||||
report_kind: parts.report_kind,
|
||||
report_context: parts.report_context,
|
||||
auth_context: Some(parts.auth_context),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_execution_decision_action(decision_is_stream: bool) -> &'static str {
|
||||
if decision_is_stream {
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION
|
||||
} else {
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION
|
||||
}
|
||||
}
|
||||
@@ -1,352 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_scheduler_core::SchedulerRankingOutcome;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_metadata::append_ranking_metadata_to_object;
|
||||
use crate::ai_pipeline::{request_origin_from_headers, RequestOrigin};
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
|
||||
pub(crate) request_id: &'a str,
|
||||
pub(crate) candidate_id: &'a str,
|
||||
pub(crate) attempt_identity: ExecutionAttemptIdentity,
|
||||
pub(crate) model: &'a str,
|
||||
pub(crate) provider_name: &'a str,
|
||||
pub(crate) provider_id: &'a str,
|
||||
pub(crate) endpoint_id: &'a str,
|
||||
pub(crate) key_id: &'a str,
|
||||
pub(crate) key_name: Option<&'a str>,
|
||||
pub(crate) model_id: Option<&'a str>,
|
||||
pub(crate) global_model_id: Option<&'a str>,
|
||||
pub(crate) global_model_name: Option<&'a str>,
|
||||
pub(crate) provider_api_format: &'a str,
|
||||
pub(crate) client_api_format: &'a str,
|
||||
pub(crate) mapped_model: Option<&'a str>,
|
||||
pub(crate) candidate_group_id: Option<&'a str>,
|
||||
pub(crate) ranking: Option<&'a SchedulerRankingOutcome>,
|
||||
pub(crate) upstream_url: Option<&'a str>,
|
||||
pub(crate) header_rules: Option<&'a Value>,
|
||||
pub(crate) body_rules: Option<&'a Value>,
|
||||
pub(crate) provider_request_method: Option<Value>,
|
||||
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
||||
pub(crate) original_headers: &'a http::HeaderMap,
|
||||
pub(crate) request_origin: Option<RequestOrigin>,
|
||||
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>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_report_context(
|
||||
parts: LocalExecutionReportContextParts<'_>,
|
||||
) -> Value {
|
||||
let mut object = Map::new();
|
||||
object.insert(
|
||||
"user_id".to_string(),
|
||||
Value::String(parts.auth_context.user_id.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"api_key_id".to_string(),
|
||||
Value::String(parts.auth_context.api_key_id.clone()),
|
||||
);
|
||||
object.insert(
|
||||
"api_key_is_standalone".to_string(),
|
||||
Value::Bool(parts.auth_context.api_key_is_standalone),
|
||||
);
|
||||
object.insert(
|
||||
"username".to_string(),
|
||||
parts
|
||||
.auth_context
|
||||
.username
|
||||
.clone()
|
||||
.map(Value::String)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
object.insert(
|
||||
"api_key_name".to_string(),
|
||||
parts
|
||||
.auth_context
|
||||
.api_key_name
|
||||
.clone()
|
||||
.map(Value::String)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
object.insert(
|
||||
"request_id".to_string(),
|
||||
Value::String(parts.request_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"candidate_id".to_string(),
|
||||
Value::String(parts.candidate_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"candidate_index".to_string(),
|
||||
Value::Number(parts.attempt_identity.candidate_index.into()),
|
||||
);
|
||||
object.insert(
|
||||
"retry_index".to_string(),
|
||||
Value::Number(parts.attempt_identity.retry_index.into()),
|
||||
);
|
||||
object.insert("model".to_string(), Value::String(parts.model.to_string()));
|
||||
object.insert(
|
||||
"provider_name".to_string(),
|
||||
Value::String(parts.provider_name.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"provider_id".to_string(),
|
||||
Value::String(parts.provider_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"endpoint_id".to_string(),
|
||||
Value::String(parts.endpoint_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"key_id".to_string(),
|
||||
Value::String(parts.key_id.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"provider_api_format".to_string(),
|
||||
Value::String(parts.provider_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"client_api_format".to_string(),
|
||||
Value::String(parts.client_api_format.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"original_headers".to_string(),
|
||||
serde_json::to_value(crate::ai_pipeline::collect_control_headers(
|
||||
parts.original_headers,
|
||||
))
|
||||
.expect("control headers should serialize"),
|
||||
);
|
||||
object.insert(
|
||||
"original_request_body".to_string(),
|
||||
crate::ai_pipeline::build_report_context_original_request_echo(
|
||||
parts.original_request_body_json,
|
||||
parts.original_request_body_base64,
|
||||
)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
let RequestOrigin {
|
||||
client_ip,
|
||||
user_agent,
|
||||
} = parts
|
||||
.request_origin
|
||||
.unwrap_or_else(|| request_origin_from_headers(parts.original_headers));
|
||||
if let Some(client_ip) = client_ip {
|
||||
object.insert("client_ip".to_string(), Value::String(client_ip));
|
||||
}
|
||||
if let Some(user_agent) = user_agent {
|
||||
object.insert("user_agent".to_string(), Value::String(user_agent));
|
||||
}
|
||||
object.insert(
|
||||
"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(),
|
||||
Value::Bool(parts.needs_conversion),
|
||||
);
|
||||
|
||||
if let Some(key_name) = parts.key_name {
|
||||
object.insert("key_name".to_string(), Value::String(key_name.to_string()));
|
||||
}
|
||||
if let Some(model_id) = parts.model_id {
|
||||
object.insert("model_id".to_string(), Value::String(model_id.to_string()));
|
||||
}
|
||||
if let Some(global_model_id) = parts.global_model_id {
|
||||
object.insert(
|
||||
"global_model_id".to_string(),
|
||||
Value::String(global_model_id.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(global_model_name) = parts.global_model_name {
|
||||
object.insert(
|
||||
"global_model_name".to_string(),
|
||||
Value::String(global_model_name.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(mapped_model) = parts.mapped_model {
|
||||
object.insert(
|
||||
"mapped_model".to_string(),
|
||||
Value::String(mapped_model.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(candidate_group_id) = parts.candidate_group_id {
|
||||
object.insert(
|
||||
"candidate_group_id".to_string(),
|
||||
Value::String(candidate_group_id.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(ranking) = parts.ranking {
|
||||
append_ranking_metadata_to_object(&mut object, ranking);
|
||||
}
|
||||
if let Some(upstream_url) = parts.upstream_url {
|
||||
object.insert(
|
||||
"upstream_url".to_string(),
|
||||
Value::String(upstream_url.to_string()),
|
||||
);
|
||||
}
|
||||
if let Some(header_rules) = parts.header_rules {
|
||||
object.insert("header_rules".to_string(), header_rules.clone());
|
||||
}
|
||||
if let Some(body_rules) = parts.body_rules {
|
||||
object.insert("body_rules".to_string(), body_rules.clone());
|
||||
}
|
||||
if let Some(provider_request_method) = parts.provider_request_method {
|
||||
object.insert(
|
||||
"provider_request_method".to_string(),
|
||||
provider_request_method,
|
||||
);
|
||||
}
|
||||
if let Some(provider_request_headers) = parts.provider_request_headers {
|
||||
object.insert(
|
||||
"provider_request_headers".to_string(),
|
||||
serde_json::to_value(provider_request_headers)
|
||||
.expect("provider request headers should serialize"),
|
||||
);
|
||||
}
|
||||
if let Some(pool_key_index) = parts.attempt_identity.pool_key_index {
|
||||
object.insert(
|
||||
"pool_key_index".to_string(),
|
||||
Value::Number(pool_key_index.into()),
|
||||
);
|
||||
}
|
||||
|
||||
object.extend(parts.extra_fields);
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_stream_event_api_format_for_provider_type(
|
||||
provider_type: &str,
|
||||
) -> Option<&'static str> {
|
||||
match provider_type.trim().to_ascii_lowercase().as_str() {
|
||||
"codex" => Some("openai:responses"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn insert_provider_stream_event_api_format(
|
||||
extra_fields: &mut Map<String, Value>,
|
||||
provider_type: &str,
|
||||
) {
|
||||
if let Some(api_format) = provider_stream_event_api_format_for_provider_type(provider_type) {
|
||||
extra_fields.insert(
|
||||
"provider_stream_event_api_format".to_string(),
|
||||
Value::String(api_format.to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use super::{
|
||||
build_local_execution_report_context, provider_stream_event_api_format_for_provider_type,
|
||||
LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::RequestOrigin;
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
#[test]
|
||||
fn codex_provider_uses_openai_responses_stream_event_format() {
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("codex"),
|
||||
Some("openai:responses")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("CODEX"),
|
||||
Some("openai:responses")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ordinary_providers_do_not_override_stream_event_format() {
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("openai"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("anthropic"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_execution_report_context_records_request_origin() {
|
||||
let auth_context = ExecutionRuntimeAuthContext {
|
||||
user_id: "user-1".to_string(),
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
balance_remaining: None,
|
||||
access_allowed: true,
|
||||
api_key_is_standalone: false,
|
||||
};
|
||||
let original_headers = http::HeaderMap::new();
|
||||
let provider_request_headers = BTreeMap::new();
|
||||
|
||||
let report_context =
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &auth_context,
|
||||
request_id: "trace-1",
|
||||
candidate_id: "candidate-1",
|
||||
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
|
||||
model: "gpt-5",
|
||||
provider_name: "OpenAI",
|
||||
provider_id: "provider-1",
|
||||
endpoint_id: "endpoint-1",
|
||||
key_id: "key-1",
|
||||
key_name: None,
|
||||
model_id: None,
|
||||
global_model_id: None,
|
||||
global_model_name: None,
|
||||
provider_api_format: "openai:chat",
|
||||
client_api_format: "openai:chat",
|
||||
mapped_model: None,
|
||||
candidate_group_id: None,
|
||||
ranking: None,
|
||||
upstream_url: None,
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
provider_request_method: None,
|
||||
provider_request_headers: Some(&provider_request_headers),
|
||||
original_headers: &original_headers,
|
||||
request_origin: Some(RequestOrigin {
|
||||
client_ip: Some("203.0.113.8".to_string()),
|
||||
user_agent: Some("Claude-Code/1.0".to_string()),
|
||||
}),
|
||||
original_request_body_json: Some(&json!({"model": "gpt-5"})),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields: Map::new(),
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
report_context["client_ip"],
|
||||
Value::String("203.0.113.8".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
report_context["user_agent"],
|
||||
Value::String("Claude-Code/1.0".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,137 +0,0 @@
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
apply_local_candidate_evaluation_progress, apply_local_candidate_terminal_plan_reason,
|
||||
build_local_runtime_miss_diagnostic,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
|
||||
|
||||
pub(crate) fn set_local_runtime_miss_diagnostic_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_runtime_miss_diagnostic(decision, plan_kind, requested_model, reason),
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_runtime_execution_exhausted_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
let mut diagnostic = build_local_runtime_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
"execution_runtime_candidates_exhausted",
|
||||
);
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_execution_exhausted_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_runtime_execution_exhausted_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_runtime_candidate_evaluation_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
let mut diagnostic = build_local_runtime_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
apply_local_candidate_evaluation_progress(&mut diagnostic, candidate_count);
|
||||
diagnostic
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_candidate_evaluation_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_runtime_candidate_evaluation_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_evaluation_progress(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
apply_local_candidate_evaluation_progress(diagnostic, candidate_count);
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let preserve_existing_candidate_signal = candidate_count == 0
|
||||
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
|
||||
if preserve_existing_candidate_signal {
|
||||
return;
|
||||
}
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_terminal_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
no_plan_reason: &'static str,
|
||||
) {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
apply_local_candidate_terminal_plan_reason(diagnostic, no_plan_reason);
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn record_local_runtime_candidate_skip_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
*diagnostic
|
||||
.skip_reasons
|
||||
.entry(skip_reason.to_string())
|
||||
.or_insert(0) += 1;
|
||||
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
|
||||
});
|
||||
}
|
||||
@@ -1,152 +0,0 @@
|
||||
use crate::ai_pipeline::planner::common::RequestedModelFamily;
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
GatewayControlSyncDecisionResponse, LocalGeminiFilesSpec, LocalOpenAiImageSpec,
|
||||
LocalOpenAiResponsesSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
LocalStandardSourceFamily, LocalStandardSpec, LocalVideoCreateFamily, LocalVideoCreateSpec,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) struct LocalExecutionSurfaceSpecMetadata {
|
||||
pub(crate) api_format: &'static str,
|
||||
pub(crate) decision_kind: &'static str,
|
||||
pub(crate) report_kind: Option<&'static str>,
|
||||
pub(crate) require_streaming: bool,
|
||||
pub(crate) requested_model_family: Option<RequestedModelFamily>,
|
||||
}
|
||||
|
||||
pub(crate) fn requested_model_family_for_standard_source(
|
||||
family: LocalStandardSourceFamily,
|
||||
) -> RequestedModelFamily {
|
||||
match family {
|
||||
LocalStandardSourceFamily::Standard => RequestedModelFamily::Standard,
|
||||
LocalStandardSourceFamily::Gemini => RequestedModelFamily::Gemini,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_standard_spec_metadata(
|
||||
spec: LocalStandardSpec,
|
||||
) -> LocalExecutionSurfaceSpecMetadata {
|
||||
LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: Some(requested_model_family_for_standard_source(spec.family)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_same_format_provider_spec_metadata(
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> LocalExecutionSurfaceSpecMetadata {
|
||||
LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: Some(requested_model_family_for_same_format_provider(spec.family)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_openai_responses_spec_metadata(
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
) -> LocalExecutionSurfaceSpecMetadata {
|
||||
LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_gemini_files_spec_metadata(
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> LocalExecutionSurfaceSpecMetadata {
|
||||
LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: "gemini:files",
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: spec.report_kind,
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_openai_image_spec_metadata(
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> LocalExecutionSurfaceSpecMetadata {
|
||||
LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: Some(RequestedModelFamily::Standard),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_video_create_spec_metadata(
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> LocalExecutionSurfaceSpecMetadata {
|
||||
LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: false,
|
||||
requested_model_family: Some(requested_model_family_for_video_create(spec.family)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn requested_model_family_for_same_format_provider(
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> RequestedModelFamily {
|
||||
match family {
|
||||
LocalSameFormatProviderFamily::Standard => RequestedModelFamily::Standard,
|
||||
LocalSameFormatProviderFamily::Gemini => RequestedModelFamily::Gemini,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn requested_model_family_for_video_create(
|
||||
family: LocalVideoCreateFamily,
|
||||
) -> RequestedModelFamily {
|
||||
match family {
|
||||
LocalVideoCreateFamily::OpenAi => RequestedModelFamily::Standard,
|
||||
LocalVideoCreateFamily::Gemini => RequestedModelFamily::Gemini,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_sync_plan_from_requested_model_family(
|
||||
family: RequestedModelFamily,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<LocalSyncPlanAndReport>, GatewayError> {
|
||||
match family {
|
||||
RequestedModelFamily::Standard => {
|
||||
build_standard_sync_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
RequestedModelFamily::Gemini => {
|
||||
build_gemini_sync_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_stream_plan_from_requested_model_family(
|
||||
family: RequestedModelFamily,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<LocalStreamPlanAndReport>, GatewayError> {
|
||||
match family {
|
||||
RequestedModelFamily::Standard => {
|
||||
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
|
||||
}
|
||||
RequestedModelFamily::Gemini => {
|
||||
build_gemini_stream_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,285 +0,0 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
|
||||
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||
persist_available_local_execution_candidates_with_context,
|
||||
persist_skipped_local_execution_candidates_with_context,
|
||||
remember_first_local_candidate_affinity,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_contract_metadata,
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_resolution::{
|
||||
extract_pool_sticky_session_token, resolve_and_rank_local_execution_candidates,
|
||||
SkippedLocalExecutionCandidate,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
|
||||
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_pipeline::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_pipeline::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_standard_spec_metadata;
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
|
||||
|
||||
pub(super) async fn resolve_local_standard_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<LocalStandardDecisionInput> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("standard specs should declare requested-model family"),
|
||||
)?;
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
auth_context,
|
||||
Some(requested_model.as_str()),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local standard decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalStandardDecisionInput,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let sticky_session_token = extract_pool_sticky_session_token(body_json);
|
||||
let persistence_policy = build_local_candidate_persistence_policy(
|
||||
&input.auth_context,
|
||||
input.required_capabilities.as_ref(),
|
||||
LocalCandidatePersistencePolicyKind::StandardDecision,
|
||||
);
|
||||
let mut seen_candidates = BTreeSet::new();
|
||||
let mut seen_skipped_candidates = BTreeSet::new();
|
||||
let mut candidates = Vec::new();
|
||||
let mut preselection_skipped = Vec::new();
|
||||
for candidate_api_format in
|
||||
request_candidate_api_formats(spec_metadata.api_format, spec_metadata.require_streaming)
|
||||
{
|
||||
let auth_snapshot = if candidate_api_format == spec_metadata.api_format {
|
||||
Some(&input.auth_snapshot)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let (mut selected_candidates, skipped_candidates) = planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
candidate_api_format,
|
||||
&input.requested_model,
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
if auth_snapshot.is_none() {
|
||||
selected_candidates.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
candidate,
|
||||
)
|
||||
});
|
||||
}
|
||||
for skipped_candidate in skipped_candidates {
|
||||
if auth_snapshot.is_none()
|
||||
&& !auth_snapshot_allows_cross_format_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let candidate_key = format!(
|
||||
"{}:{}:{}:{}:{}:{}",
|
||||
skipped_candidate.candidate.provider_id,
|
||||
skipped_candidate.candidate.endpoint_id,
|
||||
skipped_candidate.candidate.key_id,
|
||||
skipped_candidate.candidate.model_id,
|
||||
skipped_candidate.candidate.selected_provider_model_name,
|
||||
skipped_candidate.candidate.endpoint_api_format,
|
||||
);
|
||||
if seen_skipped_candidates.insert(candidate_key) {
|
||||
preselection_skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate: skipped_candidate.candidate,
|
||||
skip_reason: skipped_candidate.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
for candidate in selected_candidates {
|
||||
let candidate_key = format!(
|
||||
"{}:{}:{}:{}:{}:{}",
|
||||
candidate.provider_id,
|
||||
candidate.endpoint_id,
|
||||
candidate.key_id,
|
||||
candidate.model_id,
|
||||
candidate.selected_provider_model_name,
|
||||
candidate.endpoint_api_format,
|
||||
);
|
||||
if seen_candidates.insert(candidate_key) {
|
||||
candidates.push(candidate);
|
||||
}
|
||||
}
|
||||
}
|
||||
let (candidates, skipped_candidates) = resolve_and_rank_local_execution_candidates(
|
||||
planner_state,
|
||||
candidates,
|
||||
spec_metadata.api_format,
|
||||
&input.requested_model,
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
)
|
||||
.await;
|
||||
let skipped_candidates = preselection_skipped
|
||||
.into_iter()
|
||||
.chain(skipped_candidates)
|
||||
.map(|mut skipped_candidate| {
|
||||
let provider_api_format = skipped_candidate
|
||||
.transport
|
||||
.as_ref()
|
||||
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| {
|
||||
skipped_candidate
|
||||
.candidate
|
||||
.endpoint_api_format
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
});
|
||||
let execution_strategy = if provider_api_format == spec_metadata.api_format {
|
||||
ExecutionStrategy::LocalSameFormat
|
||||
} else {
|
||||
ExecutionStrategy::LocalCrossFormat
|
||||
};
|
||||
let conversion_mode =
|
||||
if request_conversion_kind(spec_metadata.api_format, provider_api_format.as_str())
|
||||
.is_some()
|
||||
{
|
||||
ConversionMode::Bidirectional
|
||||
} else {
|
||||
ConversionMode::None
|
||||
};
|
||||
skipped_candidate.extra_data = Some(
|
||||
build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
provider_api_format.as_str(),
|
||||
spec_metadata.api_format,
|
||||
serde_json::Map::new(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
provider_api_format.as_str(),
|
||||
),
|
||||
);
|
||||
skipped_candidate
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let candidate_count = candidates.len() + skipped_candidates.len();
|
||||
|
||||
remember_first_local_candidate_affinity(
|
||||
planner_state,
|
||||
Some(&input.auth_snapshot),
|
||||
spec_metadata.api_format,
|
||||
Some(&input.requested_model),
|
||||
&candidates,
|
||||
);
|
||||
let available_candidate_count = candidates.len() as u32;
|
||||
let attempts = persist_available_local_execution_candidates_with_context(
|
||||
planner_state,
|
||||
trace_id,
|
||||
persistence_policy.available,
|
||||
candidates,
|
||||
|eligible| {
|
||||
let provider_api_format = eligible.provider_api_format.clone();
|
||||
let execution_strategy = if provider_api_format == spec_metadata.api_format {
|
||||
ExecutionStrategy::LocalSameFormat
|
||||
} else {
|
||||
ExecutionStrategy::LocalCrossFormat
|
||||
};
|
||||
let conversion_mode =
|
||||
if request_conversion_kind(spec_metadata.api_format, provider_api_format.as_str())
|
||||
.is_some()
|
||||
{
|
||||
ConversionMode::Bidirectional
|
||||
} else {
|
||||
ConversionMode::None
|
||||
};
|
||||
Some(build_local_execution_candidate_contract_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
provider_api_format: provider_api_format.as_str(),
|
||||
client_api_format: spec_metadata.api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
eligible.candidate.endpoint_api_format.as_str(),
|
||||
))
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
persist_skipped_local_execution_candidates_with_context(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
available_candidate_count,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok((attempts, candidate_count))
|
||||
}
|
||||
@@ -1,176 +0,0 @@
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, generic_decision_missing_exact_provider_request,
|
||||
take_non_empty_string, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::transport::ensure_upstream_auth_header;
|
||||
use crate::{GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) fn build_gemini_sync_plan_from_decision(
|
||||
_parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<LocalSyncPlanAndReport>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(request_id) = take_non_empty_string(&mut payload.request_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_id) = take_non_empty_string(&mut payload.provider_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(endpoint_id) = take_non_empty_string(&mut payload.endpoint_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let auth_header = take_non_empty_string(&mut payload.auth_header);
|
||||
let auth_value = take_non_empty_string(&mut payload.auth_value);
|
||||
if auth_header.is_some() != auth_value.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(client_api_format) = take_non_empty_string(&mut payload.client_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
if let (Some(auth_header), Some(auth_value)) = (auth_header.as_deref(), auth_value.as_deref()) {
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
|
||||
}
|
||||
if payload.upstream_is_stream {
|
||||
provider_request_headers
|
||||
.entry("accept".to_string())
|
||||
.or_insert_with(|| "text/event-stream".to_string());
|
||||
}
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.take(),
|
||||
provider_name: payload.provider_name.take(),
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream: payload.upstream_is_stream,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name: payload.model_name.take(),
|
||||
proxy: payload.proxy.take(),
|
||||
tls_profile: payload.tls_profile.take(),
|
||||
timeouts: payload.timeouts.take(),
|
||||
};
|
||||
|
||||
Ok(Some(LocalSyncPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_gemini_stream_plan_from_decision(
|
||||
_parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<LocalStreamPlanAndReport>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(request_id) = take_non_empty_string(&mut payload.request_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_id) = take_non_empty_string(&mut payload.provider_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(endpoint_id) = take_non_empty_string(&mut payload.endpoint_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let auth_header = take_non_empty_string(&mut payload.auth_header);
|
||||
let auth_value = take_non_empty_string(&mut payload.auth_value);
|
||||
if auth_header.is_some() != auth_value.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(client_api_format) = take_non_empty_string(&mut payload.client_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
if let (Some(auth_header), Some(auth_value)) = (auth_header.as_deref(), auth_value.as_deref()) {
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
|
||||
}
|
||||
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.take(),
|
||||
provider_name: payload.provider_name.take(),
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream: true,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name: payload.model_name.take(),
|
||||
proxy: payload.proxy.take(),
|
||||
tls_profile: payload.tls_profile.take(),
|
||||
timeouts: payload.timeouts.take(),
|
||||
};
|
||||
|
||||
Ok(Some(LocalStreamPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
@@ -1,127 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::conversion::{request_conversion_kind, RequestConversionKind};
|
||||
use crate::ai_pipeline::transport::apply_local_body_rules;
|
||||
use crate::ai_pipeline::transport::url::{
|
||||
build_claude_messages_url, build_openai_chat_url, build_openai_responses_url,
|
||||
build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
apply_codex_openai_responses_special_body_edits,
|
||||
apply_openai_responses_compact_special_body_edits,
|
||||
build_cross_format_openai_chat_request_body as pipeline_build_cross_format_openai_chat_request_body,
|
||||
build_local_openai_chat_request_body as pipeline_build_local_openai_chat_request_body,
|
||||
GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
pub(crate) fn build_local_openai_chat_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
upstream_is_stream: bool,
|
||||
body_rules: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let mut provider_request_body =
|
||||
pipeline_build_local_openai_chat_request_body(body_json, mapped_model, upstream_is_stream)?;
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_chat_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<String> {
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
match custom_path {
|
||||
Some(path) => {
|
||||
build_passthrough_path_url(&transport.endpoint.base_url, path, parts.uri.query(), &[])
|
||||
}
|
||||
None => Some(build_openai_chat_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_chat_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let mut provider_request_body = pipeline_build_cross_format_openai_chat_request_body(
|
||||
body_json,
|
||||
mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
)?;
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
apply_openai_responses_compact_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_chat_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<String> {
|
||||
let conversion_kind = request_conversion_kind("openai:chat", provider_api_format)?;
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
match custom_path {
|
||||
Some(path) => {
|
||||
build_passthrough_path_url(&transport.endpoint.base_url, path, parts.uri.query(), &[])
|
||||
}
|
||||
None => match conversion_kind {
|
||||
RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
RequestConversionKind::ToGeminiStandard => {
|
||||
crate::ai_pipeline::build_provider_transport_request_url(
|
||||
transport,
|
||||
provider_api_format,
|
||||
Some(mapped_model),
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
RequestConversionKind::ToOpenAiResponses => Some(build_openai_responses_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1,180 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
use url::form_urlencoded;
|
||||
|
||||
use crate::ai_pipeline::conversion::{request_conversion_kind, RequestConversionKind};
|
||||
use crate::ai_pipeline::transport::antigravity::{
|
||||
build_antigravity_v1internal_url, AntigravityRequestUrlAction,
|
||||
};
|
||||
use crate::ai_pipeline::transport::apply_local_body_rules;
|
||||
use crate::ai_pipeline::transport::url::{
|
||||
build_claude_messages_url, build_openai_chat_url, build_openai_responses_url,
|
||||
build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
apply_codex_openai_responses_special_body_edits,
|
||||
apply_openai_responses_compact_special_body_edits,
|
||||
build_cross_format_openai_responses_request_body as pipeline_build_cross_format_openai_responses_request_body,
|
||||
build_local_openai_responses_request_body as pipeline_build_local_openai_responses_request_body,
|
||||
GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
pub(crate) fn build_local_openai_responses_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
require_streaming: bool,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let mut provider_request_body = pipeline_build_local_openai_responses_request_body(
|
||||
body_json,
|
||||
mapped_model,
|
||||
require_streaming,
|
||||
)?;
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
apply_openai_responses_compact_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_responses_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
provider_type: &str,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let mut provider_request_body = pipeline_build_cross_format_openai_responses_request_body(
|
||||
body_json,
|
||||
mapped_model,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
)?;
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
apply_openai_responses_compact_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_api_format,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_responses_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
compact: bool,
|
||||
) -> Option<String> {
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
match custom_path {
|
||||
Some(path) => {
|
||||
build_passthrough_path_url(&transport.endpoint.base_url, path, parts.uri.query(), &[])
|
||||
}
|
||||
None => Some(build_openai_responses_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
compact,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_responses_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<String> {
|
||||
let conversion_kind = request_conversion_kind(client_api_format, provider_api_format)?;
|
||||
if transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("antigravity")
|
||||
{
|
||||
let query = parts.uri.query().map(|query| {
|
||||
form_urlencoded::parse(query.as_bytes())
|
||||
.into_owned()
|
||||
.collect::<BTreeMap<String, String>>()
|
||||
});
|
||||
return build_antigravity_v1internal_url(
|
||||
&transport.endpoint.base_url,
|
||||
if upstream_is_stream {
|
||||
AntigravityRequestUrlAction::StreamGenerateContent
|
||||
} else {
|
||||
AntigravityRequestUrlAction::GenerateContent
|
||||
},
|
||||
query.as_ref(),
|
||||
);
|
||||
}
|
||||
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
match custom_path {
|
||||
Some(path) => {
|
||||
build_passthrough_path_url(&transport.endpoint.base_url, path, parts.uri.query(), &[])
|
||||
}
|
||||
None => match conversion_kind {
|
||||
RequestConversionKind::ToOpenAIChat => Some(build_openai_chat_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
RequestConversionKind::ToOpenAiResponses => Some(build_openai_responses_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
RequestConversionKind::ToGeminiStandard => {
|
||||
crate::ai_pipeline::build_provider_transport_request_url(
|
||||
transport,
|
||||
provider_api_format,
|
||||
Some(mapped_model),
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1,102 +0,0 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
|
||||
use crate::ai_pipeline::conversion::request_candidate_api_formats;
|
||||
use crate::ai_pipeline::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) async fn list_local_openai_chat_candidates(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
require_streaming: bool,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let mut combined = Vec::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
let mut skipped = Vec::new();
|
||||
let mut seen_skipped = BTreeSet::new();
|
||||
|
||||
let api_formats = request_candidate_api_formats("openai:chat", require_streaming);
|
||||
|
||||
for api_format in api_formats {
|
||||
let auth_snapshot = if api_format == "openai:chat" {
|
||||
Some(&input.auth_snapshot)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let (mut candidates, skipped_candidates) = planner_state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await?;
|
||||
if api_format != "openai:chat" {
|
||||
candidates.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
candidate,
|
||||
)
|
||||
});
|
||||
}
|
||||
for skipped_candidate in skipped_candidates {
|
||||
if api_format != "openai:chat"
|
||||
&& !auth_snapshot_allows_cross_format_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let candidate_key = format!(
|
||||
"{}:{}:{}:{}:{}",
|
||||
skipped_candidate.candidate.provider_id,
|
||||
skipped_candidate.candidate.endpoint_id,
|
||||
skipped_candidate.candidate.key_id,
|
||||
skipped_candidate.candidate.model_id,
|
||||
skipped_candidate.candidate.selected_provider_model_name,
|
||||
);
|
||||
if seen_skipped.insert(candidate_key) {
|
||||
skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate: skipped_candidate.candidate,
|
||||
skip_reason: skipped_candidate.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
for candidate in candidates {
|
||||
let candidate_key = format!(
|
||||
"{}:{}:{}:{}:{}",
|
||||
candidate.provider_id,
|
||||
candidate.endpoint_id,
|
||||
candidate.key_id,
|
||||
candidate.model_id,
|
||||
candidate.selected_provider_model_name,
|
||||
);
|
||||
if seen.insert(candidate_key) {
|
||||
combined.push(candidate);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((combined, skipped))
|
||||
}
|
||||
@@ -1,190 +0,0 @@
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, take_non_empty_string, LocalStreamPlanAndReport,
|
||||
LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::contracts::generic_decision_missing_exact_provider_request;
|
||||
use crate::ai_pipeline::provider_adaptation_requires_eventstream_accept;
|
||||
use crate::ai_pipeline::transport::ensure_upstream_auth_header;
|
||||
use crate::{GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) fn build_standard_sync_plan_from_decision(
|
||||
_parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<LocalSyncPlanAndReport>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(request_id) = take_non_empty_string(&mut payload.request_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_id) = take_non_empty_string(&mut payload.provider_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(endpoint_id) = take_non_empty_string(&mut payload.endpoint_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let auth_header = take_non_empty_string(&mut payload.auth_header);
|
||||
let auth_value = take_non_empty_string(&mut payload.auth_value);
|
||||
if auth_header.is_some() != auth_value.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(client_api_format) = take_non_empty_string(&mut payload.client_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
if let (Some(auth_header), Some(auth_value)) = (auth_header.as_deref(), auth_value.as_deref()) {
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
|
||||
}
|
||||
if payload.upstream_is_stream {
|
||||
provider_request_headers
|
||||
.entry("accept".to_string())
|
||||
.or_insert_with(|| "text/event-stream".to_string());
|
||||
}
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.take(),
|
||||
provider_name: payload.provider_name.take(),
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream: payload.upstream_is_stream,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name: payload.model_name.take(),
|
||||
proxy: payload.proxy.take(),
|
||||
tls_profile: payload.tls_profile.take(),
|
||||
timeouts: payload.timeouts.take(),
|
||||
};
|
||||
|
||||
Ok(Some(LocalSyncPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_standard_stream_plan_from_decision(
|
||||
_parts: &http::request::Parts,
|
||||
_body_json: &serde_json::Value,
|
||||
payload: GatewayControlSyncDecisionResponse,
|
||||
_inject_stream_flag: bool,
|
||||
) -> Result<Option<LocalStreamPlanAndReport>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
if generic_decision_missing_exact_provider_request(&payload) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(request_id) = take_non_empty_string(&mut payload.request_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_id) = take_non_empty_string(&mut payload.provider_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(endpoint_id) = take_non_empty_string(&mut payload.endpoint_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let auth_header = take_non_empty_string(&mut payload.auth_header);
|
||||
let auth_value = take_non_empty_string(&mut payload.auth_value);
|
||||
if auth_header.is_some() != auth_value.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(client_api_format) = take_non_empty_string(&mut payload.client_api_format) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let envelope_name = payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("envelope_name"))
|
||||
.and_then(serde_json::Value::as_str);
|
||||
let mut provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
if let (Some(auth_header), Some(auth_value)) = (auth_header.as_deref(), auth_value.as_deref()) {
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
|
||||
}
|
||||
if provider_adaptation_requires_eventstream_accept(envelope_name, provider_api_format.as_str())
|
||||
{
|
||||
provider_request_headers
|
||||
.entry("accept".to_string())
|
||||
.or_insert_with(|| "application/vnd.amazon.eventstream".to_string());
|
||||
} else {
|
||||
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
|
||||
}
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| Some("application/json".to_string()));
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.take(),
|
||||
provider_name: payload.provider_name.take(),
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: std::mem::take(&mut provider_request_headers),
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(provider_request_body_value),
|
||||
stream: true,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name: payload.model_name.take(),
|
||||
proxy: payload.proxy.take(),
|
||||
tls_profile: payload.tls_profile.take(),
|
||||
timeouts: payload.timeouts.take(),
|
||||
};
|
||||
|
||||
Ok(Some(LocalStreamPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
@@ -1,749 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::planner::{CandidateFailureDiagnostic, CandidateFailureDiagnosticKind};
|
||||
|
||||
pub(crate) fn request_body_build_failure_extra_data(
|
||||
body_json: &Value,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> Option<Value> {
|
||||
let diagnostic =
|
||||
diagnose_request_body_build_failure(body_json, client_api_format, provider_api_format)?;
|
||||
Some(
|
||||
diagnostic
|
||||
.formats(client_api_format, provider_api_format)
|
||||
.source(request_body_build_source(
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
))
|
||||
.to_extra_data(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn same_format_provider_request_body_failure_extra_data(
|
||||
body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
body_rules: Option<&Value>,
|
||||
context: &str,
|
||||
) -> Option<Value> {
|
||||
let diagnostic =
|
||||
diagnose_same_format_provider_request_body_failure(body_json, body_rules, context)?;
|
||||
Some(
|
||||
diagnostic
|
||||
.formats(provider_api_format, provider_api_format)
|
||||
.source(context)
|
||||
.to_extra_data(),
|
||||
)
|
||||
}
|
||||
|
||||
type RequestBodyBuildDiagnostic = CandidateFailureDiagnostic;
|
||||
|
||||
fn diagnose_request_body_build_failure(
|
||||
body_json: &Value,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
if !body_json.is_object() {
|
||||
return Some(diagnostic("$", "请求体必须是 JSON object"));
|
||||
}
|
||||
|
||||
if is_openai_responses_client_format(client_api_format) {
|
||||
if let Some(diagnostic) = diagnose_openai_responses_request(body_json) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
return Some(diagnostic(
|
||||
"$",
|
||||
"OpenAI Responses 请求体初步结构检查通过；失败可能发生在后续跨格式转换或 Body 规则应用",
|
||||
));
|
||||
}
|
||||
|
||||
if client_api_format == "openai:chat"
|
||||
&& (provider_api_format.starts_with("claude:")
|
||||
|| provider_api_format.starts_with("gemini:"))
|
||||
{
|
||||
return diagnose_openai_chat_cross_format_request(body_json, provider_api_format);
|
||||
}
|
||||
|
||||
Some(diagnostic(
|
||||
"$",
|
||||
"请求体转换失败；当前转换器未返回更细的字段路径",
|
||||
))
|
||||
}
|
||||
|
||||
fn is_openai_responses_client_format(client_api_format: &str) -> bool {
|
||||
crate::ai_pipeline::is_openai_responses_family_format(client_api_format)
|
||||
}
|
||||
|
||||
fn diagnose_same_format_provider_request_body_failure(
|
||||
body_json: &Value,
|
||||
body_rules: Option<&Value>,
|
||||
context: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
if !body_json.is_object() {
|
||||
return Some(diagnostic("$", "反代请求体必须是 JSON object"));
|
||||
}
|
||||
if body_rules.is_some_and(|rules| !rules.is_array()) {
|
||||
return Some(diagnostic(
|
||||
"$.endpoint.body_rules",
|
||||
"Endpoint Body 规则必须是数组，本地反代无法应用该配置",
|
||||
));
|
||||
}
|
||||
match context {
|
||||
"kiro_envelope" => Some(diagnostic(
|
||||
"$",
|
||||
"Kiro 反代请求体包装失败；请检查 Kiro auth_config 与 Endpoint Body 规则",
|
||||
)),
|
||||
"antigravity_envelope" => Some(diagnostic(
|
||||
"$",
|
||||
"Antigravity 反代请求体包装失败；请检查请求体是否满足该传输封装要求",
|
||||
)),
|
||||
_ => Some(diagnostic(
|
||||
"$",
|
||||
"反代请求体构建失败；当前路径未返回更细的字段信息",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_chat_cross_format_request(
|
||||
body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let request = body_json.as_object()?;
|
||||
|
||||
if let Some(messages) = request.get("messages") {
|
||||
let Some(messages) = messages.as_array() else {
|
||||
return Some(diagnostic(
|
||||
"$.messages",
|
||||
"OpenAI Chat 的 messages 必须是数组",
|
||||
));
|
||||
};
|
||||
for (message_index, message) in messages.iter().enumerate() {
|
||||
let Some(message_object) = message.as_object() else {
|
||||
return Some(diagnostic(
|
||||
format!("$.messages[{message_index}]"),
|
||||
"message 必须是 object",
|
||||
));
|
||||
};
|
||||
let role = message_object
|
||||
.get("role")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
match role.as_str() {
|
||||
"system" | "developer" => {
|
||||
if let Some(diagnostic) = diagnose_openai_text_content(
|
||||
message_object.get("content"),
|
||||
format!("$.messages[{message_index}].content"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
}
|
||||
"user" | "assistant" => {
|
||||
if let Some(diagnostic) = diagnose_openai_content_blocks(
|
||||
message_object.get("content"),
|
||||
format!("$.messages[{message_index}].content"),
|
||||
role.as_str(),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
if role == "assistant" {
|
||||
if let Some(diagnostic) = diagnose_openai_assistant_tool_calls(
|
||||
message_object.get("tool_calls"),
|
||||
format!("$.messages[{message_index}].tool_calls"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
}
|
||||
}
|
||||
"tool" => {
|
||||
let valid_tool_call_id = message_object
|
||||
.get("tool_call_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_tool_call_id {
|
||||
return Some(diagnostic(
|
||||
format!("$.messages[{message_index}].tool_call_id"),
|
||||
"tool 消息必须包含非空 tool_call_id",
|
||||
));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(diagnostic) = diagnose_openai_tools(request.get("tools"), provider_api_format) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
diagnose_openai_tool_choice(request.get("tool_choice"))
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_request(body_json: &Value) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let request = body_json.as_object()?;
|
||||
|
||||
if let Some(diagnostic) = diagnose_openai_responses_text_content(
|
||||
request.get("instructions"),
|
||||
"$.instructions".to_string(),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
|
||||
if let Some(diagnostic) = diagnose_openai_responses_input(request.get("input")) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
if let Some(diagnostic) = diagnose_openai_responses_tools(request.get("tools")) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
diagnose_openai_responses_tool_choice(request.get("tool_choice"))
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_input(input: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(input) = input else {
|
||||
return None;
|
||||
};
|
||||
match input {
|
||||
Value::Null | Value::String(_) => None,
|
||||
Value::Array(items) => {
|
||||
for (item_index, item) in items.iter().enumerate() {
|
||||
if item.is_string() {
|
||||
continue;
|
||||
}
|
||||
let item_path = format!("$.input[{item_index}]");
|
||||
let Some(item_object) = item.as_object() else {
|
||||
return Some(diagnostic(
|
||||
item_path,
|
||||
"OpenAI Responses input 数组项必须是 string 或 object",
|
||||
));
|
||||
};
|
||||
let item_type = item_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("message")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
match item_type.as_str() {
|
||||
"message" => {
|
||||
let role = item_object
|
||||
.get("role")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("user")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
if role == "system" || role == "developer" {
|
||||
if let Some(diagnostic) = diagnose_openai_responses_text_content(
|
||||
item_object.get("content"),
|
||||
format!("{item_path}.content"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
} else if let Some(diagnostic) = diagnose_openai_responses_message_content(
|
||||
item_object.get("content"),
|
||||
format!("{item_path}.content"),
|
||||
) {
|
||||
return Some(diagnostic);
|
||||
}
|
||||
}
|
||||
"function_call" => {
|
||||
let valid_name = item_object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{item_path}.name"),
|
||||
"function_call 必须包含非空 name",
|
||||
));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
_ => Some(diagnostic(
|
||||
"$.input",
|
||||
"OpenAI Responses input 必须是 string、array 或 null",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_text_content(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
if !part.is_object() {
|
||||
return Some(diagnostic(
|
||||
format!("{path}[{part_index}]"),
|
||||
"文本 content 数组项必须是 object",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => Some(diagnostic(
|
||||
path,
|
||||
"文本 content 必须是 string、array 或 null",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_message_content(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
let part_path = format!("{path}[{part_index}]");
|
||||
let Some(part_object) = part.as_object() else {
|
||||
return Some(diagnostic(part_path, "message content 数组项必须是 object"));
|
||||
};
|
||||
let part_type = part_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
if matches!(
|
||||
part_type.as_str(),
|
||||
"input_image" | "output_image" | "image_url"
|
||||
) && image_part_url(part_object).is_none()
|
||||
{
|
||||
return Some(diagnostic(
|
||||
part_path,
|
||||
"图片 content 缺少 image_url/url，无法规范化为 OpenAI Chat 图片内容",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_text_content(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
if !part.is_object() {
|
||||
return Some(diagnostic(
|
||||
format!("{path}[{part_index}]"),
|
||||
"content 数组项必须是 object",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => Some(diagnostic(path, "content 必须是 string、array 或 null")),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_content_blocks(
|
||||
content: Option<&Value>,
|
||||
path: String,
|
||||
role: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
match content {
|
||||
None | Some(Value::Null) | Some(Value::String(_)) => None,
|
||||
Some(Value::Array(parts)) => {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
let part_path = format!("{path}[{part_index}]");
|
||||
let Some(part_object) = part.as_object() else {
|
||||
return Some(diagnostic(part_path, "content 数组项必须是 object"));
|
||||
};
|
||||
let part_type = part_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
if matches!(part_type, "image_url" | "input_image" | "output_image")
|
||||
&& role == "user"
|
||||
&& image_part_url(part_object).is_none()
|
||||
{
|
||||
return Some(diagnostic(
|
||||
part_path,
|
||||
"图片 content 缺少 image_url/url，无法转换为 Claude image block",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
Some(_) => Some(diagnostic(path, "content 必须是 string、array 或 null")),
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_assistant_tool_calls(
|
||||
tool_calls: Option<&Value>,
|
||||
path: String,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(tool_calls) = tool_calls else {
|
||||
return None;
|
||||
};
|
||||
let Some(tool_calls) = tool_calls.as_array() else {
|
||||
return Some(diagnostic(path, "assistant.tool_calls 必须是数组"));
|
||||
};
|
||||
for (tool_call_index, tool_call) in tool_calls.iter().enumerate() {
|
||||
let tool_call_path = format!("{path}[{tool_call_index}]");
|
||||
let Some(tool_call_object) = tool_call.as_object() else {
|
||||
return Some(diagnostic(tool_call_path, "tool_call 必须是 object"));
|
||||
};
|
||||
let Some(function) = tool_call_object.get("function").and_then(Value::as_object) else {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_call_path}.function"),
|
||||
"tool_call 必须包含 function object",
|
||||
));
|
||||
};
|
||||
let valid_name = function
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_call_path}.function.name"),
|
||||
"tool_call.function.name 必须是非空字符串",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn diagnose_openai_tools(
|
||||
tools: Option<&Value>,
|
||||
provider_api_format: &str,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(tools) = tools else {
|
||||
return None;
|
||||
};
|
||||
let Some(tools) = tools.as_array() else {
|
||||
return Some(diagnostic("$.tools", "OpenAI Chat 的 tools 必须是数组"));
|
||||
};
|
||||
for (tool_index, tool) in tools.iter().enumerate() {
|
||||
let tool_path = format!("$.tools[{tool_index}]");
|
||||
let Some(tool_object) = tool.as_object() else {
|
||||
return Some(diagnostic(tool_path, "tool 必须是 object"));
|
||||
};
|
||||
if tool_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value != "function")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(function) = tool_object.get("function").and_then(Value::as_object) else {
|
||||
let native_tool_hint = if provider_api_format.starts_with("claude:") {
|
||||
"；如果这是 Claude 原生 tool，请改为 OpenAI function tool 格式"
|
||||
} else if provider_api_format.starts_with("gemini:") {
|
||||
"；如果这是 Gemini 原生 tool，请改为 OpenAI function tool 格式"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
return Some(diagnostic(
|
||||
format!("{tool_path}.function"),
|
||||
format!("OpenAI tool 必须包含 function object{native_tool_hint}"),
|
||||
));
|
||||
};
|
||||
let valid_name = function
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_path}.function.name"),
|
||||
"OpenAI tool 的 function.name 必须是非空字符串",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_tools(tools: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(tools) = tools else {
|
||||
return None;
|
||||
};
|
||||
let Some(tool_values) = tools.as_array() else {
|
||||
return None;
|
||||
};
|
||||
for (tool_index, tool) in tool_values.iter().enumerate() {
|
||||
let tool_path = format!("$.tools[{tool_index}]");
|
||||
let Some(tool_object) = tool.as_object() else {
|
||||
return Some(diagnostic(tool_path, "OpenAI Responses tool 必须是 object"));
|
||||
};
|
||||
let tool_type = tool_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("function")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
if tool_type.starts_with("web_search")
|
||||
|| tool_object.get("function").is_some()
|
||||
|| tool_type != "function"
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let valid_name = tool_object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if !valid_name {
|
||||
return Some(diagnostic(
|
||||
format!("{tool_path}.name"),
|
||||
"OpenAI Responses function tool 必须包含非空 name",
|
||||
));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn diagnose_openai_responses_tool_choice(
|
||||
tool_choice: Option<&Value>,
|
||||
) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(Value::Object(object)) = tool_choice else {
|
||||
return None;
|
||||
};
|
||||
let is_cli_function_choice = object.get("function").is_none()
|
||||
&& object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("function"));
|
||||
if !is_cli_function_choice {
|
||||
return None;
|
||||
}
|
||||
let valid_name = object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if valid_name {
|
||||
None
|
||||
} else {
|
||||
Some(diagnostic(
|
||||
"$.tool_choice.name",
|
||||
"OpenAI Responses tool_choice 指定 function 时必须包含非空 name",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnose_openai_tool_choice(tool_choice: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
|
||||
let Some(Value::Object(object)) = tool_choice else {
|
||||
return None;
|
||||
};
|
||||
let valid_name = object
|
||||
.get("function")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|function| function.get("name"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if valid_name {
|
||||
None
|
||||
} else {
|
||||
Some(diagnostic(
|
||||
"$.tool_choice.function.name",
|
||||
"tool_choice 指定具体工具时必须包含非空 function.name",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn image_part_url(part_object: &serde_json::Map<String, Value>) -> Option<&str> {
|
||||
part_object
|
||||
.get("image_url")
|
||||
.and_then(|value| {
|
||||
value.as_str().or_else(|| {
|
||||
value
|
||||
.as_object()
|
||||
.and_then(|object| object.get("url"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
})
|
||||
.or_else(|| part_object.get("url").and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn diagnostic(path: impl Into<String>, message: impl Into<String>) -> RequestBodyBuildDiagnostic {
|
||||
CandidateFailureDiagnostic::new(
|
||||
CandidateFailureDiagnosticKind::RequestBodyBuild,
|
||||
path,
|
||||
message,
|
||||
)
|
||||
}
|
||||
|
||||
fn request_body_build_source(client_api_format: &str, provider_api_format: &str) -> String {
|
||||
format!("{client_api_format}_to_{provider_api_format}")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::request_body_build_failure_extra_data;
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_claude_reports_claude_native_tool_shape() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{ "role": "user", "content": "hello" }],
|
||||
"tools": [{
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"input_schema": { "type": "object" }
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:chat", "claude:messages")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.tools[0].function"
|
||||
);
|
||||
assert_eq!(
|
||||
diagnostic["failure_diagnostic"]["kind"],
|
||||
"request_body_build"
|
||||
);
|
||||
assert_eq!(
|
||||
diagnostic["failure_diagnostic"]["source"],
|
||||
"openai:chat_to_claude:messages"
|
||||
);
|
||||
assert!(diagnostic["request_body_build_error"]["message"]
|
||||
.as_str()
|
||||
.expect("message")
|
||||
.contains("Claude 原生 tool"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_claude_reports_invalid_message_content_part() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": ["not-an-object"]
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:chat", "claude:messages")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.messages[0].content[0]"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_chat_to_gemini_reports_gemini_native_tool_shape() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{ "role": "user", "content": "hello" }],
|
||||
"tools": [{
|
||||
"functionDeclarations": [{
|
||||
"name": "search",
|
||||
"parameters": { "type": "object" }
|
||||
}]
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:chat", "gemini:generate_content")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.tools[0].function"
|
||||
);
|
||||
assert!(diagnostic["request_body_build_error"]["message"]
|
||||
.as_str()
|
||||
.expect("message")
|
||||
.contains("Gemini 原生 tool"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_responses_reports_invalid_function_call_name() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": [{
|
||||
"type": "function_call",
|
||||
"arguments": "{}"
|
||||
}]
|
||||
});
|
||||
|
||||
let diagnostic =
|
||||
request_body_build_failure_extra_data(&body, "openai:responses", "claude:messages")
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.input[0].name"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_responses_reports_invalid_tool_choice_name() {
|
||||
let body = json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": "hello",
|
||||
"tool_choice": { "type": "function" }
|
||||
});
|
||||
|
||||
let diagnostic = request_body_build_failure_extra_data(
|
||||
&body,
|
||||
"openai:responses",
|
||||
"gemini:generate_content",
|
||||
)
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.tool_choice.name"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_provider_reports_non_object_body() {
|
||||
let diagnostic = super::same_format_provider_request_body_failure_extra_data(
|
||||
&json!("raw"),
|
||||
"openai:chat",
|
||||
None,
|
||||
"same_format",
|
||||
)
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(diagnostic["request_body_build_error"]["path"], "$");
|
||||
assert!(diagnostic["request_body_build_error"]["message"]
|
||||
.as_str()
|
||||
.expect("message")
|
||||
.contains("反代请求体必须是 JSON object"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_provider_reports_invalid_body_rules_shape() {
|
||||
let diagnostic = super::same_format_provider_request_body_failure_extra_data(
|
||||
&json!({ "model": "gpt-5.4" }),
|
||||
"openai:chat",
|
||||
Some(&json!({ "action": "set" })),
|
||||
"same_format",
|
||||
)
|
||||
.expect("diagnostic");
|
||||
|
||||
assert_eq!(
|
||||
diagnostic["request_body_build_error"]["path"],
|
||||
"$.endpoint.body_rules"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,57 +0,0 @@
|
||||
pub(crate) mod antigravity {
|
||||
pub(crate) use aether_ai_pipeline::transport::antigravity::*;
|
||||
}
|
||||
|
||||
pub(crate) mod auth {
|
||||
pub(crate) use aether_ai_pipeline::transport::auth::*;
|
||||
}
|
||||
|
||||
pub(crate) mod claude_code {
|
||||
pub(crate) use aether_ai_pipeline::transport::claude_code::*;
|
||||
}
|
||||
|
||||
pub(crate) mod kiro {
|
||||
pub(crate) use aether_ai_pipeline::transport::kiro::*;
|
||||
}
|
||||
|
||||
pub(crate) mod oauth_refresh {
|
||||
pub(crate) use aether_ai_pipeline::transport::oauth_refresh::*;
|
||||
}
|
||||
|
||||
pub(crate) mod policy {
|
||||
pub(crate) use aether_ai_pipeline::transport::policy::*;
|
||||
}
|
||||
|
||||
pub(crate) mod provider_types {
|
||||
pub(crate) use aether_ai_pipeline::transport::provider_types::*;
|
||||
}
|
||||
|
||||
pub(crate) mod rules {
|
||||
pub(crate) use aether_ai_pipeline::transport::rules::*;
|
||||
}
|
||||
|
||||
pub(crate) mod snapshot {
|
||||
pub(crate) use aether_ai_pipeline::transport::snapshot::*;
|
||||
}
|
||||
|
||||
pub(crate) mod url {
|
||||
pub(crate) use aether_ai_pipeline::transport::url::*;
|
||||
}
|
||||
|
||||
pub(crate) mod vertex {
|
||||
pub(crate) use aether_ai_pipeline::transport::vertex::*;
|
||||
}
|
||||
|
||||
pub(crate) use aether_ai_pipeline::transport::{
|
||||
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
|
||||
body_rules_handle_path, build_passthrough_headers, ensure_upstream_auth_header,
|
||||
header_rules_are_locally_supported, local_gemini_transport_unsupported_reason_with_network,
|
||||
local_openai_chat_transport_unsupported_reason,
|
||||
local_standard_transport_unsupported_reason_with_network, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot, resolve_transport_proxy_snapshot_with_tunnel_affinity,
|
||||
resolve_transport_tls_profile, should_skip_upstream_passthrough_header,
|
||||
supports_local_gemini_transport_with_network,
|
||||
supports_local_generic_oauth_request_auth_resolution,
|
||||
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
@@ -1,7 +1,10 @@
|
||||
pub(crate) mod kiro;
|
||||
pub(crate) mod private_envelope;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
pub(crate) mod kiro {
|
||||
pub(crate) use crate::ai_serving::pure::KiroToClaudeCliStreamState;
|
||||
}
|
||||
|
||||
pub(crate) use crate::ai_serving::{
|
||||
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format,
|
||||
provider_adaptation_descriptor_for_envelope, provider_adaptation_descriptor_for_provider_type,
|
||||
provider_adaptation_requires_eventstream_accept,
|
||||
@@ -0,0 +1,10 @@
|
||||
#[path = "private_envelope/sync.rs"]
|
||||
mod sync;
|
||||
|
||||
pub(crate) use self::sync::maybe_normalize_provider_private_sync_report_payload;
|
||||
pub(crate) use crate::ai_serving::{
|
||||
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
||||
normalize_provider_private_response_value, provider_private_response_allows_sync_finalize,
|
||||
stream_body_contains_error_event, transform_provider_private_stream_line,
|
||||
ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
@@ -3,11 +3,10 @@ use serde_json::Value;
|
||||
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
use super::stream::maybe_build_provider_private_stream_normalizer;
|
||||
use super::{
|
||||
normalize_provider_private_report_context, normalize_provider_private_response_value,
|
||||
provider_private_response_allows_sync_finalize, stream_body_contains_error_event,
|
||||
ProviderPrivateStreamNormalizer,
|
||||
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
||||
normalize_provider_private_response_value, provider_private_response_allows_sync_finalize,
|
||||
stream_body_contains_error_event, ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_normalize_provider_private_sync_report_payload(
|
||||
@@ -65,7 +64,7 @@ fn normalize_provider_private_stream_bytes(
|
||||
else {
|
||||
return Ok(Some(body.to_vec()));
|
||||
};
|
||||
let mut normalized = normalizer.push_chunk(body)?;
|
||||
normalized.extend(normalizer.finish()?);
|
||||
let mut normalized = normalizer.push_chunk(body).map_err(GatewayError::from)?;
|
||||
normalized.extend(normalizer.finish().map_err(GatewayError::from)?);
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
use crate::ai_pipeline::{is_json_request, GatewayControlDecision};
|
||||
use crate::ai_serving::{is_json_request, GatewayControlDecision};
|
||||
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
pub(crate) use crate::ai_serving::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
@@ -20,12 +20,15 @@ pub(crate) use crate::ai_pipeline::{
|
||||
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
pub(crate) use crate::ai_serving::{
|
||||
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::{
|
||||
pub(crate) use crate::ai_serving::{
|
||||
AiExecutionDecision, AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
pub(crate) use aether_ai_surfaces::api::{
|
||||
build_core_error_body_for_client_format, core_error_background_report_kind,
|
||||
core_error_default_client_api_format, core_success_background_report_kind,
|
||||
encode_kiro_sse_events, implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
|
||||
@@ -34,10 +37,9 @@ pub(crate) use aether_ai_pipeline::api::{
|
||||
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
|
||||
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
|
||||
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
|
||||
ExecutionRuntimeAuthContext, GatewayControlPlanRequest, GatewayControlPlanResponse,
|
||||
GatewayControlSyncDecisionResponse, LocalCoreSyncErrorKind, LocalOpenAiImageSpec,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
|
||||
LocalStandardSourceMode, LocalStandardSpec, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
ExecutionRuntimeAuthContext, GatewayControlPlanRequest, LocalCoreSyncErrorKind,
|
||||
LocalOpenAiImageSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
|
||||
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
@@ -50,7 +52,7 @@ pub(crate) fn parse_direct_request_body(
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &axum::body::Bytes,
|
||||
) -> Option<(serde_json::Value, Option<String>)> {
|
||||
aether_ai_pipeline::api::parse_direct_request_body(
|
||||
aether_ai_surfaces::api::parse_direct_request_body(
|
||||
is_json_request(&parts.headers),
|
||||
body_bytes.as_ref(),
|
||||
)
|
||||
@@ -60,7 +62,7 @@ pub(crate) fn resolve_execution_runtime_stream_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
aether_ai_pipeline::api::resolve_execution_runtime_stream_plan_kind(
|
||||
aether_ai_surfaces::api::resolve_execution_runtime_stream_plan_kind(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
@@ -73,7 +75,7 @@ pub(crate) fn resolve_execution_runtime_sync_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
aether_ai_pipeline::api::resolve_execution_runtime_sync_plan_kind(
|
||||
aether_ai_surfaces::api::resolve_execution_runtime_sync_plan_kind(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
@@ -88,31 +90,31 @@ pub(crate) fn is_matching_stream_request(
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> bool {
|
||||
crate::ai_pipeline::planner_is_matching_stream_request(plan_kind, parts, body_json, body_base64)
|
||||
crate::ai_serving::planner_is_matching_stream_request(plan_kind, parts, body_json, body_base64)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
aether_ai_pipeline::api::supports_sync_scheduler_decision_kind(plan_kind)
|
||||
aether_ai_surfaces::api::supports_sync_scheduler_decision_kind(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
aether_ai_pipeline::api::supports_stream_scheduler_decision_kind(plan_kind)
|
||||
aether_ai_surfaces::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)
|
||||
aether_ai_surfaces::api::aggregate_openai_chat_stream_sync_response(body)
|
||||
}
|
||||
|
||||
pub(crate) fn aggregate_openai_responses_stream_sync_response(
|
||||
body: &[u8],
|
||||
) -> Option<serde_json::Value> {
|
||||
aether_ai_pipeline::api::aggregate_openai_responses_stream_sync_response(body)
|
||||
aether_ai_surfaces::api::aggregate_openai_responses_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)
|
||||
aether_ai_surfaces::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)
|
||||
aether_ai_surfaces::api::aggregate_gemini_stream_sync_response(body)
|
||||
}
|
||||
@@ -4,17 +4,17 @@ use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
pub(crate) use crate::ai_serving::api::{
|
||||
normalize_provider_private_response_value as unwrap_local_finalize_response_value,
|
||||
provider_private_response_allows_sync_finalize as local_finalize_allows_envelope,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
build_generated_tool_call_id,
|
||||
build_local_success_background_report as build_local_success_background_report_impl,
|
||||
build_local_success_conversion_background_report as build_local_success_conversion_background_report_impl,
|
||||
canonicalize_tool_arguments,
|
||||
prepare_local_success_response_parts as prepare_local_success_response_parts_impl,
|
||||
GatewayControlDecision,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline_api::{
|
||||
normalize_provider_private_response_value as unwrap_local_finalize_response_value,
|
||||
provider_private_response_allows_sync_finalize as local_finalize_allows_envelope,
|
||||
GatewayControlDecision, LocalSyncReportParts,
|
||||
};
|
||||
use crate::api::response::build_client_response_from_parts;
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
@@ -40,6 +40,36 @@ fn build_local_success_response(
|
||||
)
|
||||
}
|
||||
|
||||
fn surface_report_parts_from_gateway(payload: &GatewaySyncReportRequest) -> LocalSyncReportParts {
|
||||
LocalSyncReportParts {
|
||||
trace_id: payload.trace_id.clone(),
|
||||
report_kind: payload.report_kind.clone(),
|
||||
report_context: payload.report_context.clone(),
|
||||
status_code: payload.status_code,
|
||||
headers: payload.headers.clone(),
|
||||
body_json: payload.body_json.clone(),
|
||||
client_body_json: payload.client_body_json.clone(),
|
||||
body_base64: payload.body_base64.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn gateway_report_from_surface(
|
||||
source: &GatewaySyncReportRequest,
|
||||
report: LocalSyncReportParts,
|
||||
) -> GatewaySyncReportRequest {
|
||||
GatewaySyncReportRequest {
|
||||
trace_id: report.trace_id,
|
||||
report_kind: report.report_kind,
|
||||
report_context: report.report_context,
|
||||
status_code: report.status_code,
|
||||
headers: report.headers,
|
||||
body_json: report.body_json,
|
||||
client_body_json: report.client_body_json,
|
||||
body_base64: report.body_base64,
|
||||
telemetry: source.telemetry.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_success_outcome(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
@@ -50,8 +80,10 @@ pub(crate) fn build_local_success_outcome(
|
||||
let (body_bytes, response_headers) =
|
||||
prepare_local_success_response_parts_impl(&payload.headers, &body_json)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let surface_payload = surface_report_parts_from_gateway(payload);
|
||||
let background_report =
|
||||
build_local_success_background_report_impl(payload, body_json, report_headers);
|
||||
build_local_success_background_report_impl(&surface_payload, body_json, report_headers)
|
||||
.map(|report| gateway_report_from_surface(payload, report));
|
||||
build_local_success_outcome_with_report(
|
||||
trace_id,
|
||||
decision,
|
||||
@@ -88,11 +120,13 @@ pub(crate) fn build_local_success_outcome_with_conversion_report(
|
||||
let (body_bytes, response_headers) =
|
||||
prepare_local_success_response_parts_impl(&payload.headers, &client_body_json)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let surface_payload = surface_report_parts_from_gateway(payload);
|
||||
let report_payload = build_local_success_conversion_background_report_impl(
|
||||
payload,
|
||||
&surface_payload,
|
||||
client_body_json,
|
||||
provider_body_json,
|
||||
);
|
||||
)
|
||||
.map(|report| gateway_report_from_surface(payload, report));
|
||||
|
||||
build_local_success_outcome_with_report(
|
||||
trace_id,
|
||||
@@ -2,7 +2,7 @@ use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
#[path = "stream_rewrite.rs"]
|
||||
@@ -0,0 +1,35 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::{
|
||||
maybe_build_ai_surface_stream_rewriter, AiSurfaceFinalizeError, AiSurfaceStreamRewriter,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) struct LocalStreamRewriter<'a> {
|
||||
inner: AiSurfaceStreamRewriter<'a>,
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_stream_rewriter<'a>(
|
||||
report_context: Option<&'a Value>,
|
||||
) -> Option<LocalStreamRewriter<'a>> {
|
||||
maybe_build_ai_surface_stream_rewriter(report_context)
|
||||
.map(|inner| LocalStreamRewriter { inner })
|
||||
}
|
||||
|
||||
impl LocalStreamRewriter<'_> {
|
||||
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
|
||||
self.inner.push_chunk(chunk).map_err(map_surface_error)
|
||||
}
|
||||
|
||||
pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
self.inner.finish().map_err(map_surface_error)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_surface_error(error: AiSurfaceFinalizeError) -> GatewayError {
|
||||
error.into()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_stream.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,113 @@
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{build_generated_tool_call_id, canonicalize_tool_arguments};
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
pub(crate) use crate::ai_serving::finalize::common::{
|
||||
build_local_success_outcome, build_local_success_outcome_with_conversion_report,
|
||||
local_finalize_allows_envelope, unwrap_local_finalize_response_value,
|
||||
LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::finalize::standard::{
|
||||
maybe_build_standard_sync_finalize_product_from_normalized_payload,
|
||||
StandardSyncFinalizeNormalizedProduct,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
|
||||
aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response,
|
||||
maybe_build_openai_image_sync_finalize_product,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses,
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_build_local_core_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
|
||||
if let Some(outcome) =
|
||||
maybe_build_local_openai_image_sync_finalize_response(trace_id, decision, payload)?
|
||||
{
|
||||
return Ok(Some(outcome));
|
||||
}
|
||||
|
||||
let Some(normalized_payload) =
|
||||
crate::ai_serving::adaptation::private_envelope::maybe_normalize_provider_private_sync_report_payload(payload)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let payload = &normalized_payload;
|
||||
let Some(report_context) = payload.report_context.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !local_finalize_allows_envelope(report_context) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(product) = maybe_build_standard_sync_finalize_product_from_normalized_payload(
|
||||
payload.report_kind.as_str(),
|
||||
payload.status_code,
|
||||
Some(report_context),
|
||||
payload.body_json.as_ref(),
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.map_err(GatewayError::from)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match product {
|
||||
StandardSyncFinalizeNormalizedProduct::SuccessBody(body_json) => {
|
||||
let Some(body_json) = unwrap_local_finalize_response_value(body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(build_local_success_outcome(
|
||||
trace_id, decision, payload, body_json,
|
||||
)?))
|
||||
}
|
||||
StandardSyncFinalizeNormalizedProduct::CrossFormat(product) => {
|
||||
let Some(provider_body_json) =
|
||||
unwrap_local_finalize_response_value(product.provider_body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(build_local_success_outcome_with_conversion_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload,
|
||||
product.client_body_json,
|
||||
provider_body_json,
|
||||
)?))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
|
||||
let Some(product) = maybe_build_openai_image_sync_finalize_product(
|
||||
payload.report_kind.as_str(),
|
||||
payload.status_code,
|
||||
payload.report_context.as_ref(),
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.map_err(GatewayError::from)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
Ok(Some(build_local_success_outcome_with_conversion_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload,
|
||||
product.client_body_json,
|
||||
product.provider_body_json,
|
||||
)?))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_sync.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,20 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) use crate::ai_serving::pure::SyncToStreamBridgeOutcome;
|
||||
|
||||
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> {
|
||||
crate::ai_serving::pure::maybe_bridge_standard_sync_json_to_stream(
|
||||
provider_body_json,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
report_context,
|
||||
)
|
||||
.map_err(GatewayError::from)
|
||||
}
|
||||
@@ -1,12 +1,12 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
use crate::ai_serving::{
|
||||
encode_done_sse, encode_json_sse as encode_json_sse_impl, map_claude_stop_reason,
|
||||
PipelineFinalizeError,
|
||||
AiSurfaceFinalizeError,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
fn map_error(err: PipelineFinalizeError) -> GatewayError {
|
||||
fn map_error(err: AiSurfaceFinalizeError) -> GatewayError {
|
||||
err.into()
|
||||
}
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
//! Standard finalize surface for standard contract sync/stream compilation.
|
||||
|
||||
#[path = "stream_core/mod.rs"]
|
||||
mod stream;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
pub(crate) use crate::ai_serving::{
|
||||
aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response,
|
||||
build_openai_responses_response, convert_claude_chat_response_to_openai_chat,
|
||||
convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat,
|
||||
@@ -20,4 +17,3 @@ pub(crate) use crate::ai_pipeline::{
|
||||
maybe_build_standard_sync_finalize_product_from_normalized_payload,
|
||||
StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
|
||||
};
|
||||
pub(crate) use stream::*;
|
||||
@@ -1,6 +1,6 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_pipeline::maybe_bridge_standard_sync_json_to_stream;
|
||||
use crate::ai_serving::maybe_bridge_standard_sync_json_to_stream;
|
||||
|
||||
use super::maybe_build_local_stream_rewriter;
|
||||
|
||||
@@ -11,8 +11,8 @@ use super::{
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
|
||||
maybe_build_local_core_sync_finalize_response,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
convert_openai_chat_response_to_openai_responses,
|
||||
convert_openai_responses_response_to_openai_chat,
|
||||
};
|
||||
@@ -1,6 +1,5 @@
|
||||
mod adaptation;
|
||||
mod contracts;
|
||||
mod conversion;
|
||||
pub(crate) mod api;
|
||||
mod finalize;
|
||||
mod planner;
|
||||
mod pure;
|
||||
@@ -8,12 +7,9 @@ pub(crate) mod transport;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Response, Uri};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
||||
|
||||
use self::contracts::ExecutionRuntimeAuthContext;
|
||||
|
||||
pub(crate) use self::adaptation::{
|
||||
maybe_build_provider_private_stream_normalizer, ProviderPrivateStreamNormalizer,
|
||||
};
|
||||
@@ -47,9 +43,23 @@ pub(crate) use self::planner::{
|
||||
LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
pub(crate) use self::pure::*;
|
||||
pub(crate) use self::transport::{
|
||||
append_transport_diagnostics_to_value, build_request_trace_proxy_value,
|
||||
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
|
||||
request_conversion_direct_auth, request_conversion_enabled_for_transport,
|
||||
request_conversion_transport_supported, request_conversion_transport_unsupported_reason,
|
||||
request_pair_allowed_for_transport, CandidateTransportPolicyFacts,
|
||||
};
|
||||
pub(crate) use crate::control::GatewayControlDecision;
|
||||
pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
|
||||
pub(crate) use crate::headers::RequestOrigin;
|
||||
pub(crate) use aether_ai_serving::{
|
||||
ai_local_execution_contract_for_formats, augment_sync_report_context,
|
||||
build_ai_report_context_original_request_echo as build_report_context_original_request_echo,
|
||||
extract_ai_gemini_model_from_path as extract_gemini_model_from_path,
|
||||
generic_decision_missing_exact_provider_request as generic_decision_missing_exact_provider_request_impl,
|
||||
AiExecutionDecision, AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
|
||||
pub(crate) fn build_provider_transport_request_url(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
@@ -59,9 +69,9 @@ pub(crate) fn build_provider_transport_request_url(
|
||||
request_query: Option<&str>,
|
||||
kiro_api_region: Option<&str>,
|
||||
) -> Option<String> {
|
||||
crate::provider_transport::build_transport_request_url(
|
||||
self::transport::build_transport_request_url(
|
||||
transport,
|
||||
crate::provider_transport::TransportRequestUrlParams {
|
||||
self::transport::TransportRequestUrlParams {
|
||||
provider_api_format,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
@@ -96,38 +106,10 @@ pub(crate) fn request_origin_from_parts(parts: &http::request::Parts) -> Request
|
||||
crate::headers::request_origin_from_parts(parts)
|
||||
}
|
||||
|
||||
pub(crate) fn build_report_context_original_request_echo(
|
||||
body_json: Option<&Value>,
|
||||
body_bytes_b64: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
if let Some(body_bytes_b64) = body_bytes_b64
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return Some(json!({ "body_bytes_b64": body_bytes_b64 }));
|
||||
}
|
||||
|
||||
body_json.filter(|body| !body.is_null()).cloned()
|
||||
}
|
||||
|
||||
pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
|
||||
crate::headers::is_json_request(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_gemini_model_from_path(path: &str) -> Option<String> {
|
||||
let (_, suffix) = path.split_once("/models/")?;
|
||||
let model = suffix
|
||||
.split_once(':')
|
||||
.map(|(value, _)| value)
|
||||
.unwrap_or(suffix);
|
||||
let model = model.trim();
|
||||
if model.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(model.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_execution_runtime_auth_context(
|
||||
auth_context: &crate::control::GatewayControlAuthContext,
|
||||
) -> ExecutionRuntimeAuthContext {
|
||||
@@ -159,6 +141,22 @@ pub(crate) fn resolve_local_decision_execution_runtime_auth_context(
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn generic_decision_missing_exact_provider_request(
|
||||
payload: &AiExecutionDecision,
|
||||
) -> bool {
|
||||
if !generic_decision_missing_exact_provider_request_impl(payload) {
|
||||
return false;
|
||||
}
|
||||
|
||||
tracing::warn!(
|
||||
decision_kind = payload.decision_kind.as_deref().unwrap_or_default(),
|
||||
provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default(),
|
||||
client_api_format = payload.client_api_format.as_deref().unwrap_or_default(),
|
||||
"gateway generic decision missing exact provider request; falling back to plan"
|
||||
);
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
@@ -166,45 +164,3 @@ pub(crate) fn maybe_build_local_sync_finalize_response(
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::execution_runtime::maybe_build_local_sync_finalize_response(trace_id, decision, payload)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{build_report_context_original_request_echo, extract_gemini_model_from_path};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn build_report_context_original_request_echo_preserves_full_request_body() {
|
||||
let body = json!({
|
||||
"messages": [{"role": "user", "content": "large payload should be omitted"}],
|
||||
"service_tier": "default",
|
||||
"instructions": "Be concise.",
|
||||
"thinking": {"type": "enabled", "budget_tokens": 512},
|
||||
"metadata": {"trace": "keep"},
|
||||
"body_bytes_b64": "aGVsbG8=",
|
||||
});
|
||||
|
||||
let echo = build_report_context_original_request_echo(Some(&body), None)
|
||||
.expect("echo should be produced");
|
||||
|
||||
assert_eq!(echo, body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_report_context_original_request_echo_prefers_binary_body_bytes() {
|
||||
let echo = build_report_context_original_request_echo(
|
||||
Some(&json!({"ignored": true})),
|
||||
Some("aGVsbG8="),
|
||||
)
|
||||
.expect("echo should be produced");
|
||||
|
||||
assert_eq!(echo, json!({"body_bytes_b64": "aGVsbG8="}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_gemini_model_from_path_trims_method_suffix() {
|
||||
let model =
|
||||
extract_gemini_model_from_path("/v1beta/models/gemini-2.5-pro:streamGenerateContent");
|
||||
|
||||
assert_eq!(model.as_deref(), Some("gemini-2.5-pro"));
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,7 @@ use aether_scheduler_core::{
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
};
|
||||
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
|
||||
|
||||
const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
|
||||
@@ -1,16 +1,27 @@
|
||||
use aether_ai_serving::{
|
||||
ai_candidate_extra_data_with_ranking, ai_should_persist_available_candidate_for_pool_key,
|
||||
ai_should_persist_skipped_candidate_for_pool_membership,
|
||||
run_ai_available_candidate_persistence, run_ai_candidate_materialization,
|
||||
run_ai_skipped_candidate_persistence, AiAvailableCandidatePersistencePort,
|
||||
AiCandidateMaterializationOutcome, AiCandidateMaterializationPort, AiCandidateResolutionMode,
|
||||
AiSkippedCandidatePersistencePort,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use aether_scheduler_core::SchedulerRankingOutcome;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use std::convert::Infallible;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
|
||||
use crate::ai_pipeline::planner::candidate_metadata::append_ranking_metadata_to_object;
|
||||
use crate::ai_pipeline::planner::candidate_resolution::{
|
||||
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
|
||||
use crate::ai_serving::planner::candidate_resolution::{
|
||||
resolve_and_rank_local_execution_candidates,
|
||||
resolve_and_rank_local_execution_candidates_without_transport_pair_gate,
|
||||
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
|
||||
};
|
||||
use crate::ai_pipeline::planner::failure_diagnostic::CandidateFailureDiagnostic;
|
||||
use crate::ai_pipeline::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::ai_serving::planner::materialization_policy::LocalCandidatePersistencePolicy;
|
||||
use crate::ai_serving::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_ms;
|
||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
||||
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
|
||||
@@ -48,6 +59,292 @@ pub(crate) struct LocalSkippedCandidatePersistenceContext<'a> {
|
||||
pub(crate) record_runtime_miss_diagnostic: bool,
|
||||
}
|
||||
|
||||
pub(crate) use aether_ai_serving::AiCandidateResolutionMode as LocalCandidateResolutionMode;
|
||||
|
||||
struct GatewayLocalCandidateMaterializationPort<'a, F, G> {
|
||||
state: PlannerAppState<'a>,
|
||||
trace_id: &'a str,
|
||||
client_api_format: &'a str,
|
||||
requested_model: Option<&'a str>,
|
||||
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
sticky_session_token: Option<&'a str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'a>,
|
||||
resolution_mode: LocalCandidateResolutionMode,
|
||||
build_available_extra_data: F,
|
||||
decorate_skipped_candidate: G,
|
||||
}
|
||||
|
||||
struct GatewayAvailableCandidatePersistencePort<'a, F> {
|
||||
state: PlannerAppState<'a>,
|
||||
trace_id: &'a str,
|
||||
user_id: &'a str,
|
||||
api_key_id: &'a str,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
error_context: &'static str,
|
||||
created_at_unix_ms: u64,
|
||||
build_extra_data: F,
|
||||
}
|
||||
|
||||
struct GatewaySkippedCandidatePersistencePort<'a> {
|
||||
state: &'a AppState,
|
||||
trace_id: &'a str,
|
||||
user_id: &'a str,
|
||||
api_key_id: &'a str,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
error_context: &'static str,
|
||||
record_runtime_miss_diagnostic: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<F, G> AiCandidateMaterializationPort for GatewayLocalCandidateMaterializationPort<'_, F, G>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
|
||||
{
|
||||
type Candidate = SchedulerMinimalCandidateSelectionCandidate;
|
||||
type Eligible = EligibleLocalExecutionCandidate;
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type Attempt = LocalExecutionCandidateAttempt;
|
||||
type Error = Infallible;
|
||||
|
||||
async fn resolve_and_rank_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Candidate>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
|
||||
let requested_model = self.requested_model.map(str::to_string);
|
||||
let resolved = match self.resolution_mode {
|
||||
AiCandidateResolutionMode::Standard => {
|
||||
resolve_and_rank_local_execution_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
self.client_api_format,
|
||||
requested_model.as_deref().unwrap_or_default(),
|
||||
self.auth_snapshot,
|
||||
self.required_capabilities,
|
||||
self.sticky_session_token,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiCandidateResolutionMode::WithoutTransportPairGate => {
|
||||
resolve_and_rank_local_execution_candidates_without_transport_pair_gate(
|
||||
self.state,
|
||||
candidates,
|
||||
self.client_api_format,
|
||||
requested_model.as_deref(),
|
||||
self.auth_snapshot,
|
||||
self.required_capabilities,
|
||||
self.sticky_session_token,
|
||||
)
|
||||
.await
|
||||
}
|
||||
};
|
||||
Ok(resolved)
|
||||
}
|
||||
|
||||
fn decorate_skipped_candidate(&self, skipped: Self::Skipped) -> Self::Skipped {
|
||||
(self.decorate_skipped_candidate)(skipped)
|
||||
}
|
||||
|
||||
fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]) {
|
||||
remember_first_local_candidate_affinity(
|
||||
self.state,
|
||||
self.auth_snapshot,
|
||||
self.client_api_format,
|
||||
self.requested_model,
|
||||
candidates,
|
||||
);
|
||||
}
|
||||
|
||||
async fn persist_available_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<Vec<Self::Attempt>, Self::Error> {
|
||||
Ok(persist_available_local_execution_candidates_with_context(
|
||||
self.state,
|
||||
self.trace_id,
|
||||
self.persistence_policy.available,
|
||||
candidates,
|
||||
&self.build_available_extra_data,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
async fn persist_skipped_candidates(
|
||||
&self,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<Self::Skipped>,
|
||||
) -> Result<(), Self::Error> {
|
||||
persist_skipped_local_execution_candidates_with_context(
|
||||
self.state.app(),
|
||||
self.trace_id,
|
||||
self.persistence_policy.skipped,
|
||||
starting_candidate_index,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<F> AiAvailableCandidatePersistencePort for GatewayAvailableCandidatePersistencePort<'_, F>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
{
|
||||
type Candidate = EligibleLocalExecutionCandidate;
|
||||
type Attempt = LocalExecutionCandidateAttempt;
|
||||
type ExtraData = Value;
|
||||
type Error = Infallible;
|
||||
|
||||
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
|
||||
local_attempt_slot_count(&candidate.transport)
|
||||
}
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
|
||||
ai_candidate_extra_data_with_ranking(
|
||||
(self.build_extra_data)(candidate),
|
||||
candidate.ranking.as_ref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn generate_candidate_id(&self) -> String {
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
|
||||
fn should_persist_available_candidate(&self, candidate: &Self::Candidate) -> bool {
|
||||
should_persist_available_local_candidate(candidate)
|
||||
}
|
||||
|
||||
async fn persist_available_candidate(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<String, Self::Error> {
|
||||
Ok(self
|
||||
.state
|
||||
.persist_available_local_candidate(
|
||||
self.trace_id,
|
||||
self.user_id,
|
||||
self.api_key_id,
|
||||
&candidate.candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
generated_candidate_id,
|
||||
self.required_capabilities,
|
||||
extra_data,
|
||||
self.created_at_unix_ms,
|
||||
self.error_context,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
fn build_attempt(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
candidate_id: String,
|
||||
) -> Self::Attempt {
|
||||
LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiSkippedCandidatePersistencePort for GatewaySkippedCandidatePersistencePort<'_> {
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type ExtraData = Value;
|
||||
type Error = Infallible;
|
||||
|
||||
fn should_persist_skipped_candidate(&self, candidate: &Self::Skipped) -> bool {
|
||||
should_persist_skipped_local_candidate(candidate)
|
||||
}
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Skipped) -> Option<Self::ExtraData> {
|
||||
ai_candidate_extra_data_with_ranking(
|
||||
candidate.extra_data.clone(),
|
||||
candidate.ranking.as_ref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn generate_candidate_id(&self) -> String {
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
|
||||
async fn persist_skipped_candidate(
|
||||
&self,
|
||||
candidate: &Self::Skipped,
|
||||
candidate_index: u32,
|
||||
generated_candidate_id: &str,
|
||||
extra_data: Option<Self::ExtraData>,
|
||||
) -> Result<(), Self::Error> {
|
||||
persist_skipped_local_execution_candidate(
|
||||
self.state,
|
||||
self.trace_id,
|
||||
self.user_id,
|
||||
self.api_key_id,
|
||||
&candidate.candidate,
|
||||
candidate_index,
|
||||
generated_candidate_id,
|
||||
self.required_capabilities,
|
||||
candidate.skip_reason,
|
||||
extra_data,
|
||||
self.error_context,
|
||||
self.record_runtime_miss_diagnostic,
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn materialize_local_execution_candidates_with_serving<F, G>(
|
||||
state: PlannerAppState<'_>,
|
||||
trace_id: &str,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
|
||||
resolution_mode: LocalCandidateResolutionMode,
|
||||
build_available_extra_data: F,
|
||||
decorate_skipped_candidate: G,
|
||||
) -> AiCandidateMaterializationOutcome<LocalExecutionCandidateAttempt>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
|
||||
{
|
||||
let port = GatewayLocalCandidateMaterializationPort {
|
||||
state,
|
||||
trace_id,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
persistence_policy,
|
||||
resolution_mode,
|
||||
build_available_extra_data,
|
||||
decorate_skipped_candidate,
|
||||
};
|
||||
|
||||
match run_ai_candidate_materialization(&port, candidates, preselection_skipped).await {
|
||||
Ok(outcome) => outcome,
|
||||
Err(error) => match error {},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn remember_first_local_candidate_affinity(
|
||||
state: PlannerAppState<'_>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
@@ -72,16 +369,14 @@ pub(crate) fn remember_first_local_candidate_affinity(
|
||||
}
|
||||
|
||||
fn should_persist_available_local_candidate(eligible: &EligibleLocalExecutionCandidate) -> bool {
|
||||
eligible
|
||||
.orchestration
|
||||
.pool_key_index
|
||||
.is_none_or(|index| index == 0)
|
||||
ai_should_persist_available_candidate_for_pool_key(eligible.orchestration.pool_key_index)
|
||||
}
|
||||
|
||||
fn should_persist_skipped_local_candidate(candidate: &SkippedLocalExecutionCandidate) -> bool {
|
||||
candidate.transport.as_ref().is_none_or(|transport| {
|
||||
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_none()
|
||||
})
|
||||
let is_pool_candidate = candidate.transport.as_ref().is_some_and(|transport| {
|
||||
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_some()
|
||||
});
|
||||
ai_should_persist_skipped_candidate_for_pool_membership(is_pool_candidate)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
@@ -96,90 +391,23 @@ pub(crate) async fn persist_available_local_execution_candidates<F>(
|
||||
build_extra_data: F,
|
||||
) -> Vec<LocalExecutionCandidateAttempt>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value>,
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
{
|
||||
let created_at_unix_ms = current_unix_ms();
|
||||
let total_attempts = candidates
|
||||
.iter()
|
||||
.map(|eligible| local_attempt_slot_count(&eligible.transport) as usize)
|
||||
.sum();
|
||||
let mut materialized = Vec::with_capacity(total_attempts);
|
||||
let port = GatewayAvailableCandidatePersistencePort {
|
||||
state,
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
required_capabilities,
|
||||
error_context,
|
||||
created_at_unix_ms: current_unix_ms(),
|
||||
build_extra_data,
|
||||
};
|
||||
|
||||
for (candidate_index, eligible) in candidates.into_iter().enumerate() {
|
||||
let candidate_index = candidate_index as u32;
|
||||
let attempt_slots = local_attempt_slot_count(&eligible.transport);
|
||||
let pool_key_index = eligible.orchestration.pool_key_index;
|
||||
let extra_data = local_candidate_extra_data_with_ranking(
|
||||
build_extra_data(&eligible),
|
||||
eligible.ranking.as_ref(),
|
||||
);
|
||||
let mut owned_eligible = Some(eligible);
|
||||
|
||||
for retry_index in 0..attempt_slots {
|
||||
let eligible = owned_eligible
|
||||
.as_ref()
|
||||
.expect("eligible candidate should remain available until final retry");
|
||||
let attempt_identity = ExecutionAttemptIdentity::new(candidate_index, retry_index)
|
||||
.with_pool_key_index(pool_key_index);
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let candidate_id = if should_persist_available_local_candidate(eligible) {
|
||||
state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
&eligible.candidate,
|
||||
attempt_identity.candidate_index,
|
||||
attempt_identity.retry_index,
|
||||
&generated_candidate_id,
|
||||
required_capabilities,
|
||||
extra_data.clone(),
|
||||
created_at_unix_ms,
|
||||
error_context,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
generated_candidate_id
|
||||
};
|
||||
|
||||
let eligible = if retry_index + 1 == attempt_slots {
|
||||
owned_eligible
|
||||
.take()
|
||||
.expect("final retry should consume owned eligible candidate")
|
||||
} else {
|
||||
eligible.clone()
|
||||
};
|
||||
materialized.push(LocalExecutionCandidateAttempt {
|
||||
eligible,
|
||||
candidate_index: attempt_identity.candidate_index,
|
||||
retry_index: attempt_identity.retry_index,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
match run_ai_available_candidate_persistence(&port, candidates).await {
|
||||
Ok(attempts) => attempts,
|
||||
Err(error) => match error {},
|
||||
}
|
||||
|
||||
materialized
|
||||
}
|
||||
|
||||
fn local_candidate_extra_data_with_ranking(
|
||||
extra_data: Option<Value>,
|
||||
ranking: Option<&SchedulerRankingOutcome>,
|
||||
) -> Option<Value> {
|
||||
let Some(ranking) = ranking else {
|
||||
return extra_data;
|
||||
};
|
||||
|
||||
let mut object = match extra_data {
|
||||
Some(Value::Object(object)) => object,
|
||||
Some(value) => {
|
||||
let mut object = serde_json::Map::new();
|
||||
object.insert("extra".to_string(), value);
|
||||
object
|
||||
}
|
||||
None => serde_json::Map::new(),
|
||||
};
|
||||
append_ranking_metadata_to_object(&mut object, ranking);
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_available_local_execution_candidates_with_context<F>(
|
||||
@@ -190,7 +418,7 @@ pub(crate) async fn persist_available_local_execution_candidates_with_context<F>
|
||||
build_extra_data: F,
|
||||
) -> Vec<LocalExecutionCandidateAttempt>
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value>,
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
{
|
||||
persist_available_local_execution_candidates(
|
||||
state,
|
||||
@@ -330,31 +558,21 @@ pub(crate) async fn persist_skipped_local_execution_candidates(
|
||||
error_context: &'static str,
|
||||
record_runtime_miss_diagnostic: bool,
|
||||
) {
|
||||
let mut next_candidate_index = starting_candidate_index;
|
||||
for skipped_candidate in skipped_candidates {
|
||||
if !should_persist_skipped_local_candidate(&skipped_candidate) {
|
||||
continue;
|
||||
}
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
persist_skipped_local_execution_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
&skipped_candidate.candidate,
|
||||
next_candidate_index,
|
||||
&generated_candidate_id,
|
||||
required_capabilities,
|
||||
skipped_candidate.skip_reason,
|
||||
local_candidate_extra_data_with_ranking(
|
||||
skipped_candidate.extra_data,
|
||||
skipped_candidate.ranking.as_ref(),
|
||||
),
|
||||
error_context,
|
||||
record_runtime_miss_diagnostic,
|
||||
)
|
||||
.await;
|
||||
next_candidate_index = next_candidate_index.saturating_add(1);
|
||||
let port = GatewaySkippedCandidatePersistencePort {
|
||||
state,
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
required_capabilities,
|
||||
error_context,
|
||||
record_runtime_miss_diagnostic,
|
||||
};
|
||||
|
||||
match run_ai_skipped_candidate_persistence(&port, starting_candidate_index, skipped_candidates)
|
||||
.await
|
||||
{
|
||||
Ok(()) => {}
|
||||
Err(error) => match error {},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -423,8 +641,8 @@ mod tests {
|
||||
fn sample_transport(
|
||||
key_id: &str,
|
||||
provider_config: Option<serde_json::Value>,
|
||||
) -> Arc<crate::ai_pipeline::GatewayProviderTransportSnapshot> {
|
||||
Arc::new(crate::ai_pipeline::GatewayProviderTransportSnapshot {
|
||||
) -> Arc<crate::ai_serving::GatewayProviderTransportSnapshot> {
|
||||
Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "provider-1".to_string(),
|
||||
320
apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs
Normal file
320
apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs
Normal file
@@ -0,0 +1,320 @@
|
||||
use aether_ai_serving::{
|
||||
append_ai_execution_contract_fields_to_value, append_ai_ranking_metadata_to_object,
|
||||
build_ai_candidate_metadata_from_candidate,
|
||||
};
|
||||
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::transport::append_transport_diagnostics_to_value;
|
||||
use crate::ai_serving::GatewayProviderTransportSnapshot;
|
||||
use crate::ai_serving::{ConversionMode, ExecutionStrategy};
|
||||
|
||||
pub(crate) struct LocalExecutionCandidateMetadataParts<'a> {
|
||||
pub(crate) eligible: &'a EligibleLocalExecutionCandidate,
|
||||
pub(crate) provider_api_format: &'a str,
|
||||
pub(crate) client_api_format: &'a str,
|
||||
pub(crate) extra_fields: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub(crate) fn append_ranking_metadata_to_object(
|
||||
object: &mut Map<String, Value>,
|
||||
ranking: &SchedulerRankingOutcome,
|
||||
) {
|
||||
append_ai_ranking_metadata_to_object(object, ranking);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_metadata(
|
||||
parts: LocalExecutionCandidateMetadataParts<'_>,
|
||||
) -> Value {
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
&parts.eligible.candidate,
|
||||
Some(parts.eligible.transport.as_ref()),
|
||||
parts.provider_api_format,
|
||||
parts.client_api_format,
|
||||
parts.extra_fields,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_metadata_for_candidate(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
extra_fields: Map<String, Value>,
|
||||
) -> Value {
|
||||
append_transport_diagnostics_to_value(
|
||||
build_ai_candidate_metadata_from_candidate(
|
||||
candidate,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
extra_fields,
|
||||
),
|
||||
transport,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_contract_metadata(
|
||||
parts: LocalExecutionCandidateMetadataParts<'_>,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
append_ai_execution_contract_fields_to_value(
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
&parts.eligible.candidate,
|
||||
Some(parts.eligible.transport.as_ref()),
|
||||
parts.provider_api_format,
|
||||
parts.client_api_format,
|
||||
parts.extra_fields,
|
||||
),
|
||||
execution_strategy.as_str(),
|
||||
conversion_mode.as_str(),
|
||||
parts.client_api_format,
|
||||
provider_contract,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: Option<&GatewayProviderTransportSnapshot>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
extra_fields: Map<String, Value>,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
append_ai_execution_contract_fields_to_value(
|
||||
build_local_execution_candidate_metadata_for_candidate(
|
||||
candidate,
|
||||
transport,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
extra_fields,
|
||||
),
|
||||
execution_strategy.as_str(),
|
||||
conversion_mode.as_str(),
|
||||
client_api_format,
|
||||
provider_contract,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
build_local_execution_candidate_metadata_for_candidate,
|
||||
};
|
||||
use crate::ai_serving::transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider,
|
||||
};
|
||||
use crate::ai_serving::{ConversionMode, ExecutionStrategy, GatewayProviderTransportSnapshot};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn sample_candidate() -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
SchedulerMinimalCandidateSelectionCandidate {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_name: "RightCode".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 22,
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
endpoint_api_format: "openai:responses".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
key_name: "codex".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_internal_priority: 10,
|
||||
key_global_priority_for_format: None,
|
||||
key_capabilities: None,
|
||||
model_id: "model-1".to_string(),
|
||||
global_model_id: "global-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
selected_provider_model_name: "gpt-5.4".to_string(),
|
||||
mapping_matched_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "RightCode".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: Some(json!({"enabled": true, "mode": "node", "node_id": "proxy-node-1"})),
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "openai:responses".to_string(),
|
||||
api_family: None,
|
||||
endpoint_kind: None,
|
||||
is_active: true,
|
||||
base_url: "https://example.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: Some("/v1/responses".to_string()),
|
||||
config: None,
|
||||
format_acceptance_config: Some(json!({
|
||||
"enabled": true,
|
||||
"accept_formats": ["claude:messages"]
|
||||
})),
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "codex".to_string(),
|
||||
auth_type: "oauth".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: Some(json!({
|
||||
"tls_profile": "chrome_136",
|
||||
"user_agent": "Mozilla/5.0"
|
||||
})),
|
||||
decrypted_api_key: "sk-test".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_claude_code_transport_without_auth() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-cc-1".to_string(),
|
||||
name: "NekoCode".to_string(),
|
||||
provider_type: "claude_code".to_string(),
|
||||
website: Some("https://nekocode.ai".to_string()),
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-cc-1".to_string(),
|
||||
provider_id: "provider-cc-1".to_string(),
|
||||
api_format: "claude:messages".to_string(),
|
||||
api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://api.anthropic.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-cc-1".to_string(),
|
||||
provider_id: "provider-cc-1".to_string(),
|
||||
name: "CC-特价-0.4".to_string(),
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:messages".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_contract_metadata_includes_transport_diagnostics() {
|
||||
let metadata = build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
Some(&sample_transport()),
|
||||
"openai:responses",
|
||||
"claude:messages",
|
||||
serde_json::Map::new(),
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
"openai:responses",
|
||||
);
|
||||
|
||||
assert_eq!(metadata["transport_diagnostics"]["provider_type"], "codex");
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["fingerprint"]["tls_profile"],
|
||||
"chrome_136"
|
||||
);
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["resolved_tls_profile"],
|
||||
"chrome_136"
|
||||
);
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["conversion_enabled"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
assert!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["transport_unsupported_reason"]
|
||||
.is_null()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_metadata_marks_missing_transport_snapshot() {
|
||||
let metadata = build_local_execution_candidate_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
None,
|
||||
"openai:responses",
|
||||
"openai:responses",
|
||||
serde_json::Map::new(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["transport_snapshot_available"],
|
||||
Value::Bool(false)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_metadata_uses_same_format_provider_specific_transport_reason() {
|
||||
let metadata = build_local_execution_candidate_metadata_for_candidate(
|
||||
&sample_candidate(),
|
||||
Some(&sample_claude_code_transport_without_auth()),
|
||||
"claude:messages",
|
||||
"claude:messages",
|
||||
serde_json::Map::new(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata["transport_diagnostics"]["request_pair"]["transport_unsupported_reason"],
|
||||
Value::String("transport_auth_unavailable".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,16 +1,15 @@
|
||||
use aether_ai_serving::{
|
||||
prepare_ai_header_authenticated_candidate, resolve_ai_candidate_mapped_model,
|
||||
AiPreparedHeaderAuthenticatedCandidate,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
use crate::ai_serving::{
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct PreparedHeaderAuthenticatedCandidate {
|
||||
pub(crate) auth_header: String,
|
||||
pub(crate) auth_value: String,
|
||||
pub(crate) mapped_model: String,
|
||||
}
|
||||
pub(crate) type PreparedHeaderAuthenticatedCandidate = AiPreparedHeaderAuthenticatedCandidate;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct OauthPreparationContext<'a> {
|
||||
@@ -36,27 +35,29 @@ pub(crate) async fn prepare_header_authenticated_candidate(
|
||||
None
|
||||
};
|
||||
|
||||
let Some((auth_header, auth_value)) = direct_auth.or(oauth_auth) else {
|
||||
return Err("transport_auth_unavailable");
|
||||
};
|
||||
let mapped_model = resolve_candidate_mapped_model(candidate)?;
|
||||
prepare_ai_header_authenticated_candidate(
|
||||
direct_auth,
|
||||
oauth_auth,
|
||||
candidate.selected_provider_model_name.as_str(),
|
||||
)
|
||||
}
|
||||
|
||||
Ok(PreparedHeaderAuthenticatedCandidate {
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
})
|
||||
pub(crate) fn prepare_header_authenticated_candidate_from_auth(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
auth_header: String,
|
||||
auth_value: String,
|
||||
) -> Result<PreparedHeaderAuthenticatedCandidate, &'static str> {
|
||||
prepare_ai_header_authenticated_candidate(
|
||||
Some((auth_header, auth_value)),
|
||||
None,
|
||||
candidate.selected_provider_model_name.as_str(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_candidate_mapped_model(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> Result<String, &'static str> {
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
return Err("mapped_model_missing");
|
||||
}
|
||||
|
||||
Ok(mapped_model)
|
||||
resolve_ai_candidate_mapped_model(candidate.selected_provider_model_name.as_str())
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_candidate_oauth_auth(
|
||||
@@ -92,7 +93,7 @@ mod tests {
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::{prepare_header_authenticated_candidate, OauthPreparationContext};
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::ai_serving::PlannerAppState;
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
@@ -1,19 +1,21 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_serving::{
|
||||
ai_ranking_context, build_ai_rankable_candidate, run_ai_candidate_ranking,
|
||||
AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig,
|
||||
AiRankingSchedulingMode,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
request_candidate_api_format_preference, GatewayAuthApiKeySnapshot, PlannerAppState,
|
||||
};
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
||||
use crate::scheduler::config::{
|
||||
read_scheduler_ordering_config, SchedulerOrderingConfig, SchedulerSchedulingMode,
|
||||
};
|
||||
use aether_scheduler_core::{
|
||||
apply_scheduler_candidate_ranking, matches_affinity_target,
|
||||
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerRankableCandidate,
|
||||
SchedulerRankingContext, SchedulerRankingMode,
|
||||
matches_affinity_target, SchedulerAffinityTarget, SchedulerMinimalCandidateSelectionCandidate,
|
||||
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingOutcome,
|
||||
};
|
||||
|
||||
use super::candidate_affinity_cache::read_cached_scheduler_affinity_target;
|
||||
@@ -22,6 +24,92 @@ use super::candidate_transport_ranking_facts::{
|
||||
resolve_cached_transport_ranking_facts, CandidateTransportRankingFacts,
|
||||
};
|
||||
|
||||
struct GatewayLocalCandidateRankingPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
requested_model: Option<&'a str>,
|
||||
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidateRankingPort for GatewayLocalCandidateRankingPort<'_> {
|
||||
type Candidate = EligibleLocalExecutionCandidate;
|
||||
type AffinityTarget = SchedulerAffinityTarget;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn affinity_requested_model(&self, candidates: &[Self::Candidate]) -> Option<String> {
|
||||
self.requested_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| {
|
||||
candidates
|
||||
.first()
|
||||
.map(|candidate| candidate.candidate.global_model_name.clone())
|
||||
})
|
||||
}
|
||||
|
||||
async fn read_cached_affinity_target(
|
||||
&self,
|
||||
normalized_client_api_format: &str,
|
||||
affinity_requested_model: Option<&str>,
|
||||
) -> Result<Option<Self::AffinityTarget>, Self::Error> {
|
||||
Ok(read_cached_scheduler_affinity_target(
|
||||
self.state,
|
||||
self.auth_snapshot,
|
||||
normalized_client_api_format,
|
||||
affinity_requested_model,
|
||||
))
|
||||
}
|
||||
|
||||
fn cached_affinity_matches(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
target: &Self::AffinityTarget,
|
||||
) -> bool {
|
||||
cached_affinity_matches_local_execution_scope(candidate, target)
|
||||
}
|
||||
|
||||
async fn build_rankable_candidate(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
original_index: usize,
|
||||
normalized_client_api_format: &str,
|
||||
cached_affinity_match: bool,
|
||||
) -> Result<SchedulerRankableCandidate, Self::Error> {
|
||||
let ranking_facts = resolve_transport_ranking_facts_for_candidate(
|
||||
self.state,
|
||||
&candidate.candidate,
|
||||
candidate.transport.as_ref(),
|
||||
self.ordering_config,
|
||||
)
|
||||
.await;
|
||||
Ok(build_ai_rankable_candidate(AiRankableCandidateParts {
|
||||
candidate: &candidate.candidate,
|
||||
original_index,
|
||||
normalized_client_api_format,
|
||||
provider_api_format: candidate.provider_api_format.as_str(),
|
||||
required_capabilities: self.required_capabilities,
|
||||
cached_affinity_match,
|
||||
tunnel_bucket: ranking_facts.tunnel_bucket,
|
||||
keep_priority_on_conversion: ranking_facts.keep_priority_on_conversion,
|
||||
}))
|
||||
}
|
||||
|
||||
fn ranking_context(&self) -> SchedulerRankingContext {
|
||||
ai_ranking_context(ai_ranking_context_config(self.ordering_config))
|
||||
}
|
||||
|
||||
fn apply_ranking_outcome(
|
||||
&self,
|
||||
candidate: &mut Self::Candidate,
|
||||
outcome: SchedulerRankingOutcome,
|
||||
) {
|
||||
candidate.ranking = Some(outcome);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn rank_eligible_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<EligibleLocalExecutionCandidate>,
|
||||
@@ -31,89 +119,35 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Vec<EligibleLocalExecutionCandidate> {
|
||||
let ordering_config = read_scheduler_ordering_config_or_default(state).await;
|
||||
let mut candidates = candidates;
|
||||
let affinity_requested_model = requested_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| {
|
||||
candidates
|
||||
.first()
|
||||
.map(|candidate| candidate.candidate.global_model_name.as_str())
|
||||
});
|
||||
let cached_affinity_target = read_cached_scheduler_affinity_target(
|
||||
let port = GatewayLocalCandidateRankingPort {
|
||||
state,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
normalized_client_api_format,
|
||||
affinity_requested_model,
|
||||
);
|
||||
let mut rankables = Vec::with_capacity(candidates.len());
|
||||
let mut ordering_cache = BTreeMap::new();
|
||||
required_capabilities,
|
||||
ordering_config,
|
||||
};
|
||||
|
||||
for (original_index, eligible) in candidates.iter().enumerate() {
|
||||
let ranking_facts = resolve_cached_transport_ranking_facts(
|
||||
state,
|
||||
&mut ordering_cache,
|
||||
&eligible.candidate,
|
||||
eligible.transport.as_ref(),
|
||||
ordering_config,
|
||||
)
|
||||
.await;
|
||||
rankables.push(rankable_candidate_from_candidate(
|
||||
&eligible.candidate,
|
||||
original_index,
|
||||
ranking_facts,
|
||||
normalized_client_api_format,
|
||||
eligible.provider_api_format.as_str(),
|
||||
required_capabilities,
|
||||
cached_affinity_target.as_ref().is_some_and(|target| {
|
||||
cached_affinity_matches_local_execution_scope(eligible, target)
|
||||
}),
|
||||
));
|
||||
match run_ai_candidate_ranking(&port, candidates, normalized_client_api_format).await {
|
||||
Ok(candidates) => candidates,
|
||||
Err(error) => match error {},
|
||||
}
|
||||
|
||||
drop(ordering_cache);
|
||||
let outcomes = apply_scheduler_candidate_ranking(
|
||||
&mut candidates,
|
||||
&rankables,
|
||||
planner_ranking_context(ordering_config),
|
||||
);
|
||||
for outcome in outcomes {
|
||||
let ranking_index = outcome.ranking_index;
|
||||
if let Some(candidate) = candidates.get_mut(ranking_index) {
|
||||
candidate.ranking = Some(outcome);
|
||||
}
|
||||
}
|
||||
candidates
|
||||
}
|
||||
|
||||
fn rankable_candidate_from_candidate(
|
||||
async fn resolve_transport_ranking_facts_for_candidate(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
original_index: usize,
|
||||
ranking_facts: CandidateTransportRankingFacts,
|
||||
normalized_client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
cached_affinity_match: bool,
|
||||
) -> SchedulerRankableCandidate {
|
||||
let is_same_format = api_format_matches(provider_api_format, normalized_client_api_format);
|
||||
let mut rankable = SchedulerRankableCandidate::from_candidate(candidate, original_index);
|
||||
// The scheduler order is the upstream tie-breaker; pipeline only adds transport facts.
|
||||
rankable.provider_id.clear();
|
||||
rankable.endpoint_id.clear();
|
||||
rankable.key_id.clear();
|
||||
rankable.selected_provider_model_name.clear();
|
||||
|
||||
rankable
|
||||
.with_capability_priority(requested_capability_priority_for_candidate(
|
||||
required_capabilities,
|
||||
candidate,
|
||||
))
|
||||
.with_cached_affinity_match(cached_affinity_match)
|
||||
.with_tunnel_bucket(ranking_facts.tunnel_bucket)
|
||||
.with_format_state(
|
||||
!is_same_format && !ranking_facts.keep_priority_on_conversion,
|
||||
candidate_api_format_preference(normalized_client_api_format, provider_api_format),
|
||||
)
|
||||
transport: &crate::ai_serving::GatewayProviderTransportSnapshot,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> CandidateTransportRankingFacts {
|
||||
let mut ordering_cache = BTreeMap::new();
|
||||
resolve_cached_transport_ranking_facts(
|
||||
state,
|
||||
&mut ordering_cache,
|
||||
candidate,
|
||||
transport,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn cached_affinity_matches_local_execution_scope(
|
||||
@@ -133,36 +167,21 @@ fn local_execution_candidate_uses_pool(eligible: &EligibleLocalExecutionCandidat
|
||||
.is_some()
|
||||
}
|
||||
|
||||
fn planner_ranking_context(ordering_config: SchedulerOrderingConfig) -> SchedulerRankingContext {
|
||||
SchedulerRankingContext {
|
||||
fn ai_ranking_context_config(ordering_config: SchedulerOrderingConfig) -> AiRankingContextConfig {
|
||||
AiRankingContextConfig {
|
||||
priority_mode: ordering_config.priority_mode,
|
||||
ranking_mode: planner_ranking_mode(ordering_config.scheduling_mode),
|
||||
include_health: false,
|
||||
load_balance_seed: 0,
|
||||
scheduling_mode: ai_ranking_scheduling_mode(ordering_config.scheduling_mode),
|
||||
}
|
||||
}
|
||||
|
||||
fn planner_ranking_mode(mode: SchedulerSchedulingMode) -> SchedulerRankingMode {
|
||||
fn ai_ranking_scheduling_mode(mode: SchedulerSchedulingMode) -> AiRankingSchedulingMode {
|
||||
match mode {
|
||||
SchedulerSchedulingMode::FixedOrder => SchedulerRankingMode::FixedOrder,
|
||||
SchedulerSchedulingMode::CacheAffinity => SchedulerRankingMode::CacheAffinity,
|
||||
SchedulerSchedulingMode::LoadBalance => SchedulerRankingMode::LoadBalance,
|
||||
SchedulerSchedulingMode::FixedOrder => AiRankingSchedulingMode::FixedOrder,
|
||||
SchedulerSchedulingMode::CacheAffinity => AiRankingSchedulingMode::CacheAffinity,
|
||||
SchedulerSchedulingMode::LoadBalance => AiRankingSchedulingMode::LoadBalance,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_api_format_alias(value: &str) -> String {
|
||||
crate::ai_pipeline::normalize_api_format_alias(value)
|
||||
}
|
||||
|
||||
fn api_format_matches(left: &str, right: &str) -> bool {
|
||||
normalize_api_format_alias(left) == normalize_api_format_alias(right)
|
||||
}
|
||||
|
||||
fn candidate_api_format_preference(client_api_format: &str, provider_api_format: &str) -> (u8, u8) {
|
||||
request_candidate_api_format_preference(client_api_format, provider_api_format)
|
||||
.unwrap_or((u8::MAX, u8::MAX))
|
||||
}
|
||||
|
||||
async fn read_scheduler_ordering_config_or_default(
|
||||
state: PlannerAppState<'_>,
|
||||
) -> SchedulerOrderingConfig {
|
||||
@@ -185,6 +204,9 @@ mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_ai_serving::{
|
||||
ai_ranking_context, build_ai_rankable_candidate, AiRankableCandidateParts,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
@@ -197,7 +219,7 @@ mod tests {
|
||||
use super::super::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
|
||||
use super::super::candidate_transport_ranking_facts::resolve_cached_candidate_transport_ranking_facts;
|
||||
use super::{PlannerAppState, SchedulerMinimalCandidateSelectionCandidate};
|
||||
use crate::ai_pipeline::planner::candidate_resolution::resolve_and_rank_local_execution_candidates;
|
||||
use crate::ai_serving::planner::candidate_resolution::resolve_and_rank_local_execution_candidates;
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::tunnel::TunnelAttachmentRecord;
|
||||
@@ -224,22 +246,23 @@ mod tests {
|
||||
ordering_config,
|
||||
)
|
||||
.await;
|
||||
rankables.push(super::rankable_candidate_from_candidate(
|
||||
rankables.push(build_ai_rankable_candidate(AiRankableCandidateParts {
|
||||
candidate,
|
||||
original_index,
|
||||
ranking_facts,
|
||||
normalized_client_api_format.as_str(),
|
||||
candidate.endpoint_api_format.as_str(),
|
||||
normalized_client_api_format: normalized_client_api_format.as_str(),
|
||||
provider_api_format: candidate.endpoint_api_format.as_str(),
|
||||
required_capabilities,
|
||||
false,
|
||||
));
|
||||
cached_affinity_match: false,
|
||||
tunnel_bucket: ranking_facts.tunnel_bucket,
|
||||
keep_priority_on_conversion: ranking_facts.keep_priority_on_conversion,
|
||||
}));
|
||||
}
|
||||
|
||||
drop(ordering_cache);
|
||||
apply_scheduler_candidate_ranking(
|
||||
&mut candidates,
|
||||
&rankables,
|
||||
super::planner_ranking_context(ordering_config),
|
||||
ai_ranking_context(super::ai_ranking_context_config(ordering_config)),
|
||||
);
|
||||
candidates
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_ai_serving::{
|
||||
run_ai_candidate_resolution, AiCandidateResolutionMode, AiCandidateResolutionPort,
|
||||
AiCandidateResolutionRequest,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use std::convert::Infallible;
|
||||
use tracing::warn;
|
||||
|
||||
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
|
||||
|
||||
use crate::ai_serving::{
|
||||
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
|
||||
CandidateTransportPolicyFacts, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
|
||||
PlannerAppState,
|
||||
};
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
|
||||
use super::candidate_ranking::rank_eligible_local_execution_candidates;
|
||||
use super::pool_scheduler::apply_local_execution_pool_scheduler;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct EligibleLocalExecutionCandidate {
|
||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(crate) provider_api_format: String,
|
||||
pub(crate) orchestration: LocalExecutionCandidateMetadata,
|
||||
pub(crate) ranking: Option<SchedulerRankingOutcome>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct SkippedLocalExecutionCandidate {
|
||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) skip_reason: &'static str,
|
||||
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
|
||||
pub(crate) ranking: Option<SchedulerRankingOutcome>,
|
||||
pub(crate) extra_data: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl SkippedLocalExecutionCandidate {
|
||||
pub(crate) fn transport_ref(&self) -> Option<&GatewayProviderTransportSnapshot> {
|
||||
self.transport.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
struct GatewayLocalCandidateResolutionPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
requested_model: Option<&'a str>,
|
||||
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
sticky_session_token: Option<&'a str>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
|
||||
type Candidate = SchedulerMinimalCandidateSelectionCandidate;
|
||||
type Transport = GatewayProviderTransportSnapshot;
|
||||
type Eligible = EligibleLocalExecutionCandidate;
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type Error = Infallible;
|
||||
|
||||
async fn read_candidate_transport(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
) -> Result<Option<Self::Transport>, Self::Error> {
|
||||
Ok(read_candidate_transport_snapshot(self.state, candidate).await)
|
||||
}
|
||||
|
||||
fn build_missing_transport_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
) -> Self::Skipped {
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason: "transport_snapshot_missing",
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn candidate_common_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
transport: &Self::Transport,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<&'static str> {
|
||||
candidate_common_transport_skip_reason(
|
||||
transport,
|
||||
candidate_transport_policy_facts(candidate),
|
||||
requested_model,
|
||||
)
|
||||
}
|
||||
|
||||
fn candidate_transport_pair_skip_reason(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
transport: &Self::Transport,
|
||||
normalized_client_api_format: &str,
|
||||
requested_model: &str,
|
||||
) -> Option<&'static str> {
|
||||
let _ = (candidate, requested_model);
|
||||
candidate_transport_pair_skip_reason(transport, normalized_client_api_format)
|
||||
}
|
||||
|
||||
fn build_skipped_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
transport: Self::Transport,
|
||||
skip_reason: &'static str,
|
||||
) -> Self::Skipped {
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason,
|
||||
transport: Some(Arc::new(transport)),
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_eligible_candidate(
|
||||
&self,
|
||||
candidate: Self::Candidate,
|
||||
transport: Self::Transport,
|
||||
) -> Self::Eligible {
|
||||
let provider_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
|
||||
EligibleLocalExecutionCandidate {
|
||||
candidate,
|
||||
transport: Arc::new(transport),
|
||||
provider_api_format,
|
||||
orchestration: LocalExecutionCandidateMetadata::default(),
|
||||
ranking: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn rank_eligible_candidates(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
normalized_client_api_format: &str,
|
||||
) -> Result<Vec<Self::Eligible>, Self::Error> {
|
||||
Ok(rank_eligible_local_execution_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
normalized_client_api_format,
|
||||
self.requested_model,
|
||||
self.auth_snapshot,
|
||||
self.required_capabilities,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
async fn apply_pool_scheduler(
|
||||
&self,
|
||||
candidates: Vec<Self::Eligible>,
|
||||
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
|
||||
Ok(
|
||||
apply_local_execution_pool_scheduler(self.state, candidates, self.sticky_session_token)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_and_rank_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let requested_model = requested_model.trim();
|
||||
resolve_and_rank_local_execution_candidates_with_mode(
|
||||
state,
|
||||
candidates,
|
||||
client_api_format,
|
||||
Some(requested_model),
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
AiCandidateResolutionMode::Standard,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transport_pair_gate(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let requested_model = requested_model.map(str::trim);
|
||||
resolve_and_rank_local_execution_candidates_with_mode(
|
||||
state,
|
||||
candidates,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
AiCandidateResolutionMode::WithoutTransportPairGate,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_and_rank_local_execution_candidates_with_mode(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
requested_model: Option<&str>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
sticky_session_token: Option<&str>,
|
||||
mode: AiCandidateResolutionMode,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let port = GatewayLocalCandidateResolutionPort {
|
||||
state,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
sticky_session_token,
|
||||
};
|
||||
|
||||
let request = AiCandidateResolutionRequest {
|
||||
client_api_format,
|
||||
requested_model,
|
||||
mode,
|
||||
};
|
||||
|
||||
match run_ai_candidate_resolution(&port, candidates, request).await {
|
||||
Ok(outcome) => (outcome.eligible_candidates, outcome.skipped_candidates),
|
||||
Err(error) => match error {},
|
||||
}
|
||||
}
|
||||
|
||||
fn candidate_transport_policy_facts(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> CandidateTransportPolicyFacts<'_> {
|
||||
CandidateTransportPolicyFacts {
|
||||
endpoint_api_format: candidate.endpoint_api_format.as_str(),
|
||||
global_model_name: candidate.global_model_name.as_str(),
|
||||
selected_provider_model_name: candidate.selected_provider_model_name.as_str(),
|
||||
mapping_matched_model: candidate.mapping_matched_model.as_deref(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_candidate_transport_snapshot(
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> Option<GatewayProviderTransportSnapshot> {
|
||||
match state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(transport)) => Some(transport),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "candidate_resolution_transport_load_failed",
|
||||
log_type = "event",
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
error = ?error,
|
||||
"failed to load provider transport while evaluating local candidate eligibility"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
218
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
Normal file
218
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
Normal file
@@ -0,0 +1,218 @@
|
||||
use aether_ai_serving::{
|
||||
run_ai_candidate_preselection, AiCandidatePreselectionOutcome, AiCandidatePreselectionPort,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::scheduler::candidate::SchedulerSkippedCandidate;
|
||||
use crate::GatewayError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum LocalCandidatePreselectionKeyMode {
|
||||
ProviderEndpointKeyModel,
|
||||
ProviderEndpointKeyModelAndApiFormat,
|
||||
}
|
||||
|
||||
struct GatewayLocalCandidatePreselectionPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
client_api_format: &'a str,
|
||||
requested_model: &'a str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
auth_snapshot: &'a GatewayAuthApiKeySnapshot,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
|
||||
type Candidate = SchedulerMinimalCandidateSelectionCandidate;
|
||||
type Skipped = SkippedLocalExecutionCandidate;
|
||||
type Error = GatewayError;
|
||||
|
||||
fn candidate_api_formats(&self) -> Vec<String> {
|
||||
crate::ai_serving::request_candidate_api_formats(
|
||||
self.client_api_format,
|
||||
self.require_streaming,
|
||||
)
|
||||
.into_iter()
|
||||
.map(str::to_string)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn candidate_api_format_matches_client(&self, candidate_api_format: &str) -> bool {
|
||||
if self.use_api_format_alias_match {
|
||||
crate::ai_serving::api_format_alias_matches(
|
||||
candidate_api_format,
|
||||
self.client_api_format,
|
||||
)
|
||||
} else {
|
||||
candidate_api_format == self.client_api_format
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_candidates_for_api_format(
|
||||
&self,
|
||||
candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> Result<(Vec<Self::Candidate>, Vec<Self::Skipped>), Self::Error> {
|
||||
let auth_snapshot = matches_client_format.then_some(self.auth_snapshot);
|
||||
let (candidates, skipped_candidates) = self
|
||||
.state
|
||||
.list_selectable_candidates_with_skip_reasons(
|
||||
candidate_api_format,
|
||||
self.requested_model,
|
||||
self.require_streaming,
|
||||
self.required_capabilities,
|
||||
auth_snapshot,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok((
|
||||
candidates,
|
||||
skipped_candidates
|
||||
.into_iter()
|
||||
.map(skipped_local_execution_candidate_from_scheduler_skip)
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
|
||||
fn candidate_allowed(
|
||||
&self,
|
||||
candidate: &Self::Candidate,
|
||||
_candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> bool {
|
||||
matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
candidate,
|
||||
)
|
||||
}
|
||||
|
||||
fn skipped_candidate_allowed(
|
||||
&self,
|
||||
skipped_candidate: &Self::Skipped,
|
||||
_candidate_api_format: &str,
|
||||
matches_client_format: bool,
|
||||
) -> bool {
|
||||
matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
)
|
||||
}
|
||||
|
||||
fn candidate_key(&self, candidate: &Self::Candidate) -> String {
|
||||
local_candidate_preselection_key(candidate, self.key_mode)
|
||||
}
|
||||
|
||||
fn skipped_candidate_key(&self, skipped_candidate: &Self::Skipped) -> String {
|
||||
local_candidate_preselection_key(&skipped_candidate.candidate, self.key_mode)
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn preselect_local_execution_candidates_with_serving(
|
||||
state: PlannerAppState<'_>,
|
||||
client_api_format: &str,
|
||||
requested_model: &str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
) -> Result<
|
||||
AiCandidatePreselectionOutcome<
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
SkippedLocalExecutionCandidate,
|
||||
>,
|
||||
GatewayError,
|
||||
> {
|
||||
let port = GatewayLocalCandidatePreselectionPort {
|
||||
state,
|
||||
client_api_format,
|
||||
requested_model,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
use_api_format_alias_match,
|
||||
key_mode,
|
||||
};
|
||||
|
||||
run_ai_candidate_preselection(&port).await
|
||||
}
|
||||
|
||||
fn skipped_local_execution_candidate_from_scheduler_skip(
|
||||
skipped_candidate: SchedulerSkippedCandidate,
|
||||
) -> SkippedLocalExecutionCandidate {
|
||||
SkippedLocalExecutionCandidate {
|
||||
candidate: skipped_candidate.candidate,
|
||||
skip_reason: skipped_candidate.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn local_candidate_preselection_key(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
mode: LocalCandidatePreselectionKeyMode,
|
||||
) -> String {
|
||||
match mode {
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel => format!(
|
||||
"{}:{}:{}:{}:{}",
|
||||
candidate.provider_id,
|
||||
candidate.endpoint_id,
|
||||
candidate.key_id,
|
||||
candidate.model_id,
|
||||
candidate.selected_provider_model_name,
|
||||
),
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat => format!(
|
||||
"{}:{}:{}:{}:{}:{}",
|
||||
candidate.provider_id,
|
||||
candidate.endpoint_id,
|
||||
candidate.key_id,
|
||||
candidate.model_id,
|
||||
candidate.selected_provider_model_name,
|
||||
candidate.endpoint_api_format,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn auth_snapshot_allows_cross_format_candidate(
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
requested_model: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> bool {
|
||||
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
|
||||
let provider_allowed = allowed_providers.iter().any(|value| {
|
||||
aether_scheduler_core::provider_matches_allowed_value(
|
||||
value,
|
||||
&candidate.provider_id,
|
||||
&candidate.provider_name,
|
||||
&candidate.provider_type,
|
||||
)
|
||||
});
|
||||
if !provider_allowed {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(allowed_models) = auth_snapshot.effective_allowed_models() {
|
||||
let model_allowed = allowed_models
|
||||
.iter()
|
||||
.any(|value| value == requested_model || value == &candidate.global_model_name);
|
||||
if !model_allowed {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
@@ -5,7 +5,7 @@ use aether_scheduler_core::{
|
||||
};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
|
||||
use crate::ai_serving::{GatewayProviderTransportSnapshot, PlannerAppState};
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
|
||||
use super::candidate_resolution::read_candidate_transport_snapshot;
|
||||
131
apps/aether-gateway/src/ai_serving/planner/common.rs
Normal file
131
apps/aether-gateway/src/ai_serving/planner/common.rs
Normal file
@@ -0,0 +1,131 @@
|
||||
use axum::body::Bytes;
|
||||
|
||||
use crate::ai_serving::{
|
||||
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl,
|
||||
is_json_request, parse_direct_request_body as parse_direct_request_body_impl,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
|
||||
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
|
||||
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
|
||||
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
|
||||
pub(crate) use aether_ai_serving::AiRequestedModelFamily as RequestedModelFamily;
|
||||
|
||||
pub(crate) fn parse_direct_request_body(
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &Bytes,
|
||||
) -> Option<(serde_json::Value, Option<String>)> {
|
||||
parse_direct_request_body_impl(is_json_request(&parts.headers), body_bytes.as_ref())
|
||||
}
|
||||
|
||||
pub(crate) fn force_upstream_streaming_for_provider(
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
) -> bool {
|
||||
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> {
|
||||
aether_ai_serving::extract_ai_standard_requested_model(body_json)
|
||||
}
|
||||
|
||||
pub(crate) fn extract_requested_model_from_request(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
family: RequestedModelFamily,
|
||||
) -> Option<String> {
|
||||
aether_ai_serving::extract_ai_requested_model_from_request_path(
|
||||
parts.uri.path(),
|
||||
body_json,
|
||||
family,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
extract_requested_model_from_request, extract_standard_requested_model,
|
||||
force_upstream_streaming_for_provider, RequestedModelFamily,
|
||||
};
|
||||
use axum::http::Request;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn forces_streaming_for_codex_openai_responses() {
|
||||
assert!(force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_force_streaming_for_compact_or_other_provider_types() {
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
));
|
||||
assert!(!force_upstream_streaming_for_provider(
|
||||
"openai",
|
||||
"openai:responses"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_standard_requested_model_from_request_body() {
|
||||
let requested_model =
|
||||
extract_standard_requested_model(&json!({ "model": " claude-sonnet-4 " }));
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_family_helper_delegates_standard_model_extraction() {
|
||||
let request = Request::builder()
|
||||
.uri("https://example.test/v1/chat/completions")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
&parts,
|
||||
&json!({ "model": " claude-sonnet-4 " }),
|
||||
RequestedModelFamily::Standard,
|
||||
);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_gemini_requested_model_from_request_path() {
|
||||
let request = Request::builder()
|
||||
.uri("https://example.test/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let requested_model =
|
||||
extract_requested_model_from_request(&parts, &json!({}), RequestedModelFamily::Gemini);
|
||||
|
||||
assert_eq!(requested_model.as_deref(), Some("gemini-2.5-pro"));
|
||||
}
|
||||
}
|
||||
@@ -1,33 +1,32 @@
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
use crate::ai_serving::planner::common::{
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
|
||||
GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
|
||||
GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
|
||||
GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
|
||||
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
|
||||
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
|
||||
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
|
||||
OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_openai_chat_stream_plan_from_decision, build_openai_chat_sync_plan_from_decision,
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision, build_passthrough_stream_plan_from_decision,
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
build_standard_sync_plan_from_decision, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
build_standard_sync_plan_from_decision,
|
||||
};
|
||||
use crate::ai_pipeline::planner::route::{
|
||||
use crate::ai_serving::planner::route::{
|
||||
resolve_execution_runtime_stream_plan_kind as resolve_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind as resolve_sync_plan_kind,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{
|
||||
AppState, GatewayControlPlanResponse, GatewayControlSyncDecisionResponse, GatewayError,
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AiExecutionPlanPayload, AppState, GatewayError};
|
||||
use aether_ai_serving::{
|
||||
build_ai_stream_execution_plan_payload, build_ai_sync_execution_plan_payload,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_plan_payload_impl(
|
||||
@@ -38,7 +37,7 @@ pub(crate) async fn maybe_build_sync_plan_payload_impl(
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_sync_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -66,7 +65,7 @@ pub(crate) async fn maybe_build_stream_plan_payload_impl(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -90,8 +89,8 @@ fn build_sync_plan_payload_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
mut payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
mut payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let auth_context = payload.auth_context.take();
|
||||
let plan_and_report = match plan_kind {
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND => {
|
||||
@@ -124,15 +123,16 @@ fn build_sync_plan_payload_from_decision(
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(plan_and_report.map(|value| build_sync_plan_response(plan_kind, value, auth_context)))
|
||||
Ok(plan_and_report
|
||||
.map(|value| build_ai_sync_execution_plan_payload(plan_kind, value, auth_context)))
|
||||
}
|
||||
|
||||
fn build_stream_plan_payload_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
mut payload: GatewayControlSyncDecisionResponse,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
mut payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
let auth_context = payload.auth_context.take();
|
||||
let plan_and_report = match plan_kind {
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND => {
|
||||
@@ -159,35 +159,6 @@ fn build_stream_plan_payload_from_decision(
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(plan_and_report.map(|value| build_stream_plan_response(plan_kind, value, auth_context)))
|
||||
}
|
||||
|
||||
fn build_sync_plan_response(
|
||||
plan_kind: &str,
|
||||
value: LocalSyncPlanAndReport,
|
||||
auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
) -> GatewayControlPlanResponse {
|
||||
GatewayControlPlanResponse {
|
||||
action: EXECUTION_RUNTIME_SYNC_ACTION.to_string(),
|
||||
plan_kind: Some(plan_kind.to_string()),
|
||||
plan: Some(value.plan),
|
||||
report_kind: value.report_kind,
|
||||
report_context: value.report_context,
|
||||
auth_context,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_stream_plan_response(
|
||||
plan_kind: &str,
|
||||
value: LocalStreamPlanAndReport,
|
||||
auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
) -> GatewayControlPlanResponse {
|
||||
GatewayControlPlanResponse {
|
||||
action: EXECUTION_RUNTIME_STREAM_ACTION.to_string(),
|
||||
plan_kind: Some(plan_kind.to_string()),
|
||||
plan: Some(value.plan),
|
||||
report_kind: value.report_kind,
|
||||
report_context: value.report_context,
|
||||
auth_context,
|
||||
}
|
||||
Ok(plan_and_report
|
||||
.map(|value| build_ai_stream_execution_plan_payload(plan_kind, value, auth_context)))
|
||||
}
|
||||
188
apps/aether-gateway/src/ai_serving/planner/decision/stream.rs
Normal file
188
apps/aether-gateway/src/ai_serving/planner/decision/stream.rs
Normal file
@@ -0,0 +1,188 @@
|
||||
use aether_ai_serving::{
|
||||
build_ai_execution_decision_from_plan, run_ai_stream_decision_path,
|
||||
AiExecutionDecisionFromPlanParts, AiStreamDecisionPathPort, AiStreamDecisionStep,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::route::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
};
|
||||
use crate::ai_serving::{resolve_decision_execution_runtime_auth_context, GatewayControlDecision};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if !is_matching_stream_request(plan_kind, parts, body_json, body_base64) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let port = GatewayStreamDecisionPathPort {
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
plan_kind,
|
||||
};
|
||||
|
||||
run_ai_stream_decision_path(&port).await
|
||||
}
|
||||
|
||||
struct GatewayStreamDecisionPathPort<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
decision: &'a GatewayControlDecision,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_base64: Option<&'a str>,
|
||||
plan_kind: &'a str,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiStreamDecisionPathPort for GatewayStreamDecisionPathPort<'_> {
|
||||
type Decision = AiExecutionDecision;
|
||||
type Error = GatewayError;
|
||||
|
||||
async fn build_stream_decision_step(
|
||||
&self,
|
||||
step: AiStreamDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error> {
|
||||
match step {
|
||||
AiStreamDecisionStep::LocalVideoContent => {
|
||||
maybe_build_local_video_task_content_stream_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalImage => {
|
||||
super::maybe_build_stream_local_image_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.body_base64,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalOpenAiChat => {
|
||||
super::maybe_build_stream_local_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalOpenAiResponses => {
|
||||
super::maybe_build_stream_local_openai_responses_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalStandardFamily => {
|
||||
super::maybe_build_stream_local_standard_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalSameFormatProvider => {
|
||||
super::maybe_build_stream_local_same_format_provider_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiStreamDecisionStep::LocalGeminiFiles => {
|
||||
super::maybe_build_stream_local_gemini_files_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_task_content_stream_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if plan_kind != OPENAI_VIDEO_CONTENT_PLAN_KIND
|
||||
|| decision.route_family.as_deref() != Some("openai")
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
|
||||
.await?;
|
||||
|
||||
let Some(action) = state.video_tasks.prepare_openai_content_stream_action(
|
||||
parts.uri.path(),
|
||||
parts.uri.query(),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) = action else {
|
||||
return Ok(None);
|
||||
};
|
||||
let plan = *plan;
|
||||
|
||||
Ok(Some(build_ai_execution_decision_from_plan(
|
||||
AiExecutionDecisionFromPlanParts {
|
||||
action: EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(plan_kind.to_string()),
|
||||
request_id: None,
|
||||
upstream_base_url: None,
|
||||
include_auth_pair: false,
|
||||
plan,
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
auth_context: resolve_decision_execution_runtime_auth_context(decision),
|
||||
},
|
||||
)))
|
||||
}
|
||||
257
apps/aether-gateway/src/ai_serving/planner/decision/sync.rs
Normal file
257
apps/aether-gateway/src/ai_serving/planner/decision/sync.rs
Normal file
@@ -0,0 +1,257 @@
|
||||
use aether_ai_serving::{
|
||||
build_ai_execution_decision_from_plan, infer_ai_upstream_base_url, run_ai_sync_decision_path,
|
||||
AiExecutionDecisionFromPlanParts, AiSyncDecisionPathPort, AiSyncDecisionStep,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DELETE_PLAN_KIND,
|
||||
GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::route::resolve_execution_runtime_sync_plan_kind;
|
||||
use crate::ai_serving::{
|
||||
build_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
|
||||
GatewayControlDecision,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_sync_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let port = GatewaySyncDecisionPathPort {
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
plan_kind,
|
||||
};
|
||||
|
||||
run_ai_sync_decision_path(&port).await
|
||||
}
|
||||
|
||||
struct GatewaySyncDecisionPathPort<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
decision: &'a GatewayControlDecision,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_base64: Option<&'a str>,
|
||||
body_is_empty: bool,
|
||||
plan_kind: &'a str,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiSyncDecisionPathPort for GatewaySyncDecisionPathPort<'_> {
|
||||
type Decision = AiExecutionDecision;
|
||||
type Error = GatewayError;
|
||||
|
||||
fn sync_decision_step_enabled(&self, step: AiSyncDecisionStep) -> bool {
|
||||
if step == AiSyncDecisionStep::LocalGeminiFiles {
|
||||
return matches!(
|
||||
self.plan_kind,
|
||||
GEMINI_FILES_LIST_PLAN_KIND
|
||||
| GEMINI_FILES_GET_PLAN_KIND
|
||||
| GEMINI_FILES_DELETE_PLAN_KIND
|
||||
);
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn build_sync_decision_step(
|
||||
&self,
|
||||
step: AiSyncDecisionStep,
|
||||
) -> Result<Option<Self::Decision>, Self::Error> {
|
||||
match step {
|
||||
AiSyncDecisionStep::VideoTaskFollowUp => {
|
||||
maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalVideo => {
|
||||
super::maybe_build_sync_local_video_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalImage => {
|
||||
super::maybe_build_sync_local_image_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.body_base64,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalOpenAiChat => {
|
||||
super::maybe_build_sync_local_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalOpenAiResponses => {
|
||||
super::maybe_build_sync_local_openai_responses_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalStandardFamily => {
|
||||
super::maybe_build_sync_local_standard_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalSameFormatProvider => {
|
||||
super::maybe_build_sync_local_same_format_provider_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.body_json,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AiSyncDecisionStep::LocalGeminiFiles => {
|
||||
super::maybe_build_sync_local_gemini_files_decision_payload(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
self.body_base64,
|
||||
self.body_is_empty,
|
||||
self.trace_id,
|
||||
self.decision,
|
||||
self.plan_kind,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
if !matches!(
|
||||
plan_kind,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
| OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND
|
||||
| GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
|
||||
.await?;
|
||||
|
||||
let auth_context = resolve_execution_runtime_auth_context(
|
||||
state,
|
||||
decision,
|
||||
&parts.headers,
|
||||
&parts.uri,
|
||||
trace_id,
|
||||
)
|
||||
.await?;
|
||||
let Some(auth_context) = auth_context else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(follow_up) = state.video_tasks.prepare_follow_up_sync_plan(
|
||||
plan_kind,
|
||||
parts.uri.path(),
|
||||
Some(body_json),
|
||||
Some(&auth_context),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let aether_video_tasks_core::LocalVideoTaskFollowUpPlan {
|
||||
plan,
|
||||
report_kind,
|
||||
report_context,
|
||||
} = follow_up;
|
||||
let upstream_base_url = infer_ai_upstream_base_url(&plan.url);
|
||||
|
||||
debug!(
|
||||
event_name = "local_video_follow_up_sync_decision_payload_built",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
request_id = %trace_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
plan_kind,
|
||||
downstream_path = %parts.uri.path(),
|
||||
provider_api_format = %plan.provider_api_format,
|
||||
client_api_format = %plan.client_api_format,
|
||||
upstream_base_url = ?upstream_base_url,
|
||||
upstream_url = %plan.url,
|
||||
"gateway built local video follow-up sync decision payload"
|
||||
);
|
||||
|
||||
Ok(Some(build_ai_execution_decision_from_plan(
|
||||
AiExecutionDecisionFromPlanParts {
|
||||
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(plan_kind.to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
upstream_base_url,
|
||||
include_auth_pair: true,
|
||||
plan,
|
||||
report_kind,
|
||||
report_context,
|
||||
auth_context: Some(build_execution_runtime_auth_context(&auth_context)),
|
||||
},
|
||||
)))
|
||||
}
|
||||
127
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
Normal file
127
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
Normal file
@@ -0,0 +1,127 @@
|
||||
use aether_ai_serving::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::ai_serving::{ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct ResolvedLocalDecisionAuthInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalRequestedModelDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalAuthenticatedDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
struct GatewayAuthenticatedDecisionInputPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
now_unix_secs: u64,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort<'_> {
|
||||
type AuthContext = ExecutionRuntimeAuthContext;
|
||||
type AuthSnapshot = GatewayAuthApiKeySnapshot;
|
||||
type RequiredCapabilities = serde_json::Value;
|
||||
type ResolvedInput = ResolvedLocalDecisionAuthInput;
|
||||
type Error = GatewayError;
|
||||
|
||||
async fn read_auth_snapshot(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
) -> Result<Option<Self::AuthSnapshot>, Self::Error> {
|
||||
self.state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
self.now_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_required_capabilities(
|
||||
&self,
|
||||
auth_context: &Self::AuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&Self::RequiredCapabilities>,
|
||||
) -> Result<Option<Self::RequiredCapabilities>, Self::Error> {
|
||||
Ok(self
|
||||
.state
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
requested_model,
|
||||
explicit_required_capabilities,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
|
||||
fn build_resolved_input(
|
||||
&self,
|
||||
auth_context: Self::AuthContext,
|
||||
auth_snapshot: Self::AuthSnapshot,
|
||||
required_capabilities: Option<Self::RequiredCapabilities>,
|
||||
) -> Self::ResolvedInput {
|
||||
ResolvedLocalDecisionAuthInput {
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
required_capabilities,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_requested_model_decision_input(
|
||||
resolved_input: ResolvedLocalDecisionAuthInput,
|
||||
requested_model: String,
|
||||
) -> LocalRequestedModelDecisionInput {
|
||||
LocalRequestedModelDecisionInput {
|
||||
auth_context: resolved_input.auth_context,
|
||||
requested_model,
|
||||
auth_snapshot: resolved_input.auth_snapshot,
|
||||
required_capabilities: resolved_input.required_capabilities,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_authenticated_decision_input(
|
||||
resolved_input: ResolvedLocalDecisionAuthInput,
|
||||
) -> LocalAuthenticatedDecisionInput {
|
||||
LocalAuthenticatedDecisionInput {
|
||||
auth_context: resolved_input.auth_context,
|
||||
auth_snapshot: resolved_input.auth_snapshot,
|
||||
required_capabilities: resolved_input.required_capabilities,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_local_authenticated_decision_input(
|
||||
state: &AppState,
|
||||
auth_context: ExecutionRuntimeAuthContext,
|
||||
requested_model: Option<&str>,
|
||||
explicit_required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
|
||||
let port = GatewayAuthenticatedDecisionInputPort {
|
||||
state: PlannerAppState::new(state),
|
||||
now_unix_secs: current_unix_secs(),
|
||||
};
|
||||
|
||||
run_ai_authenticated_decision_input(
|
||||
&port,
|
||||
auth_context,
|
||||
requested_model,
|
||||
explicit_required_capabilities,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
use aether_ai_serving::ai_candidate_persistence_policy_spec;
|
||||
pub(crate) use aether_ai_serving::AiCandidatePersistencePolicyKind as LocalCandidatePersistencePolicyKind;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
LocalAvailableCandidatePersistenceContext, LocalSkippedCandidatePersistenceContext,
|
||||
};
|
||||
use crate::ai_serving::ExecutionRuntimeAuthContext;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct LocalCandidatePersistencePolicy<'a> {
|
||||
pub(crate) available: LocalAvailableCandidatePersistenceContext<'a>,
|
||||
pub(crate) skipped: LocalSkippedCandidatePersistenceContext<'a>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_candidate_persistence_policy<'a>(
|
||||
auth_context: &'a ExecutionRuntimeAuthContext,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
kind: LocalCandidatePersistencePolicyKind,
|
||||
) -> LocalCandidatePersistencePolicy<'a> {
|
||||
let spec = ai_candidate_persistence_policy_spec(kind);
|
||||
|
||||
LocalCandidatePersistencePolicy {
|
||||
available: LocalAvailableCandidatePersistenceContext {
|
||||
user_id: &auth_context.user_id,
|
||||
api_key_id: &auth_context.api_key_id,
|
||||
required_capabilities,
|
||||
error_context: spec.available_error_context,
|
||||
},
|
||||
skipped: LocalSkippedCandidatePersistenceContext {
|
||||
user_id: &auth_context.user_id,
|
||||
api_key_id: &auth_context.api_key_id,
|
||||
required_capabilities,
|
||||
error_context: spec.skipped_error_context,
|
||||
record_runtime_miss_diagnostic: spec.record_runtime_miss_diagnostic,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,4 @@
|
||||
use crate::ai_pipeline::contracts::{
|
||||
GatewayControlPlanResponse, GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_serving::{AiExecutionDecision, AiExecutionPlanPayload, GatewayControlDecision};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
mod candidate_affinity_cache;
|
||||
@@ -15,10 +12,8 @@ mod candidate_transport_ranking_facts;
|
||||
mod common;
|
||||
mod decision;
|
||||
mod decision_input;
|
||||
mod failure_diagnostic;
|
||||
mod materialization_policy;
|
||||
mod passthrough;
|
||||
mod payload_metadata;
|
||||
mod plan_builders;
|
||||
mod pool_scheduler;
|
||||
mod report_context;
|
||||
@@ -29,10 +24,6 @@ mod specialized;
|
||||
mod standard;
|
||||
mod state;
|
||||
|
||||
pub(crate) use self::candidate_resolution::extract_pool_sticky_session_token;
|
||||
pub(crate) use self::failure_diagnostic::{
|
||||
CandidateFailureDiagnostic, CandidateFailureDiagnosticKind,
|
||||
};
|
||||
pub(crate) use self::passthrough::{
|
||||
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
|
||||
};
|
||||
@@ -41,7 +32,7 @@ pub(crate) use self::plan_builders::{
|
||||
build_openai_responses_stream_plan_from_decision,
|
||||
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
|
||||
pub(crate) use self::specialized::{
|
||||
@@ -64,6 +55,11 @@ pub(crate) use self::state::{
|
||||
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||
PlannerAppState,
|
||||
};
|
||||
pub(crate) use aether_ai_serving::extract_ai_pool_sticky_session_token as extract_pool_sticky_session_token;
|
||||
pub(crate) use aether_ai_serving::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
CandidateFailureDiagnostic, CandidateFailureDiagnosticKind,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_decision_payload(
|
||||
state: &AppState,
|
||||
@@ -73,7 +69,7 @@ pub(crate) async fn maybe_build_sync_decision_payload(
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
decision::maybe_build_sync_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
@@ -93,7 +89,7 @@ pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
decision::maybe_build_stream_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
@@ -113,7 +109,7 @@ pub(crate) async fn maybe_build_sync_plan_payload(
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
decision::maybe_build_sync_plan_payload_impl(
|
||||
state,
|
||||
parts,
|
||||
@@ -133,7 +129,7 @@ pub(crate) async fn maybe_build_stream_plan_payload(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionPlanPayload>, GatewayError> {
|
||||
decision::maybe_build_stream_plan_payload_impl(
|
||||
state,
|
||||
parts,
|
||||
@@ -8,6 +8,5 @@ pub(crate) use self::provider::{
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
resolve_same_format_provider_transport_unsupported_reason_for_trace,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::transport::provider_types::provider_type_supports_local_same_format_transport;
|
||||
pub(crate) use crate::ai_serving::transport::provider_types::provider_type_supports_local_same_format_transport;
|
||||
@@ -0,0 +1,96 @@
|
||||
use aether_contracts::RequestBody;
|
||||
|
||||
use super::{
|
||||
augment_sync_report_context, build_ai_execution_plan_from_decision,
|
||||
resolve_ai_passthrough_sync_request_body, take_ai_decision_plan_core, take_non_empty_string,
|
||||
AiExecutionPlanFromDecisionParts, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
pub(crate) fn build_passthrough_sync_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let ignored_provider_request_body = serde_json::Value::Null;
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context.take(),
|
||||
&provider_request_headers,
|
||||
&ignored_provider_request_body,
|
||||
)?;
|
||||
let request_body = resolve_ai_passthrough_sync_request_body(
|
||||
payload.provider_request_body.take(),
|
||||
payload.provider_request_body_base64.take(),
|
||||
);
|
||||
let provider_request_method = take_non_empty_string(&mut payload.provider_request_method);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: provider_request_method.unwrap_or_else(|| parts.method.to_string()),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
body: request_body,
|
||||
stream: false,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiSyncAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn build_passthrough_stream_plan_from_decision(
|
||||
parts: &http::request::Parts,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let mut payload = payload;
|
||||
let Some(core) = take_ai_decision_plan_core(&mut payload) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
|
||||
let content_type = payload
|
||||
.content_type
|
||||
.take()
|
||||
.or_else(|| provider_request_headers.get("content-type").cloned());
|
||||
let plan = build_ai_execution_plan_from_decision(
|
||||
&mut payload,
|
||||
AiExecutionPlanFromDecisionParts {
|
||||
core,
|
||||
method: parts.method.to_string(),
|
||||
url: upstream_url,
|
||||
headers: provider_request_headers,
|
||||
content_type,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: true,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(Some(AiStreamAttempt {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context: payload.report_context,
|
||||
}))
|
||||
}
|
||||
@@ -11,53 +11,50 @@ use serde_json::{json, Value};
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
use crate::ai_serving::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::transport::antigravity::{
|
||||
use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt};
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
build_antigravity_v1internal_url, classify_local_antigravity_request_support,
|
||||
AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport,
|
||||
AntigravityRequestSideSupport, AntigravityRequestUrlAction,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
use crate::ai_serving::transport::auth::{
|
||||
build_openai_passthrough_headers, resolve_local_gemini_auth, resolve_local_standard_auth,
|
||||
};
|
||||
use crate::ai_pipeline::transport::claude_code::{
|
||||
use crate::ai_serving::transport::claude_code::{
|
||||
build_claude_code_messages_url, build_claude_code_passthrough_headers,
|
||||
sanitize_claude_code_request_body, supports_local_claude_code_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::transport::kiro::{
|
||||
use crate::ai_serving::transport::kiro::{
|
||||
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
|
||||
build_kiro_provider_request_body, supports_local_kiro_request_transport_with_network,
|
||||
KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_pipeline::transport::policy::{
|
||||
use crate::ai_serving::transport::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::transport::url::{
|
||||
use crate::ai_serving::transport::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::transport::vertex::{
|
||||
use crate::ai_serving::transport::vertex::{
|
||||
build_vertex_api_key_gemini_content_url, resolve_local_vertex_api_key_query_auth,
|
||||
supports_local_vertex_api_key_gemini_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::transport::{
|
||||
use crate::ai_serving::transport::{
|
||||
apply_local_body_rules, apply_local_header_rules, build_passthrough_headers,
|
||||
ensure_upstream_auth_header, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
use crate::ai_serving::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, GatewayControlDecision,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
|
||||
GatewayError,
|
||||
append_execution_contract_fields_to_value, AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
|
||||
mod family;
|
||||
@@ -67,9 +64,8 @@ mod request;
|
||||
pub(crate) use self::family::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
resolve_local_same_format_provider_decision_input,
|
||||
resolve_same_format_provider_transport_unsupported_reason_for_trace,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
resolve_local_same_format_provider_decision_input, LocalSameFormatProviderFamily,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
pub(crate) use self::family::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
@@ -1,11 +1,11 @@
|
||||
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_pipeline::planner::runtime_miss::{
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
|
||||
use super::candidates::{
|
||||
@@ -21,7 +21,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -88,7 +88,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -1,30 +1,25 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||
persist_available_local_execution_candidates_with_context,
|
||||
persist_skipped_local_execution_candidates_with_context,
|
||||
remember_first_local_candidate_affinity,
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_metadata::{
|
||||
use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_contract_metadata,
|
||||
build_local_execution_candidate_contract_metadata_for_candidate,
|
||||
LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_resolution::{
|
||||
extract_pool_sticky_session_token, resolve_and_rank_local_execution_candidates,
|
||||
SkippedLocalExecutionCandidate,
|
||||
};
|
||||
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_pipeline::planner::decision_input::{
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_pipeline::planner::materialization_policy::{
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision, PlannerAppState,
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::{
|
||||
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
@@ -107,63 +102,32 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let (candidates, skipped_candidates) = resolve_and_rank_local_execution_candidates(
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
planner_state,
|
||||
candidates,
|
||||
trace_id,
|
||||
spec_metadata.api_format,
|
||||
&input.requested_model,
|
||||
Some(&input.requested_model),
|
||||
Some(&input.auth_snapshot),
|
||||
input.required_capabilities.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
)
|
||||
.await;
|
||||
let skipped_candidates = preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.chain(skipped_candidates)
|
||||
.map(|mut skipped_candidate| {
|
||||
let provider_api_format = skipped_candidate
|
||||
.transport
|
||||
.as_ref()
|
||||
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| spec_metadata.api_format.to_string());
|
||||
skipped_candidate.extra_data = Some(
|
||||
build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
provider_api_format.as_str(),
|
||||
spec_metadata.api_format,
|
||||
serde_json::Map::new(),
|
||||
ExecutionStrategy::LocalSameFormat,
|
||||
ConversionMode::None,
|
||||
provider_api_format.as_str(),
|
||||
),
|
||||
);
|
||||
skipped_candidate
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let candidate_count = candidates.len() + skipped_candidates.len();
|
||||
|
||||
remember_first_local_candidate_affinity(
|
||||
planner_state,
|
||||
Some(&input.auth_snapshot),
|
||||
spec_metadata.api_format,
|
||||
Some(&input.requested_model),
|
||||
&candidates,
|
||||
);
|
||||
let available_candidate_count = candidates.len() as u32;
|
||||
let attempts = persist_available_local_execution_candidates_with_context(
|
||||
planner_state,
|
||||
trace_id,
|
||||
persistence_policy.available,
|
||||
persistence_policy,
|
||||
candidates,
|
||||
preselection_skipped
|
||||
.into_iter()
|
||||
.map(|item| SkippedLocalExecutionCandidate {
|
||||
candidate: item.candidate,
|
||||
skip_reason: item.skip_reason,
|
||||
transport: None,
|
||||
ranking: None,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect(),
|
||||
LocalCandidateResolutionMode::Standard,
|
||||
|eligible| {
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
);
|
||||
Some(build_local_execution_candidate_contract_metadata(
|
||||
LocalExecutionCandidateMetadataParts {
|
||||
eligible,
|
||||
@@ -171,22 +135,37 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
client_api_format: spec_metadata.api_format,
|
||||
extra_fields: serde_json::Map::new(),
|
||||
},
|
||||
ExecutionStrategy::LocalSameFormat,
|
||||
ConversionMode::None,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
))
|
||||
},
|
||||
|mut skipped_candidate| {
|
||||
let provider_api_format = skipped_candidate
|
||||
.transport
|
||||
.as_ref()
|
||||
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| spec_metadata.api_format.to_string());
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
spec_metadata.api_format,
|
||||
provider_api_format.as_str(),
|
||||
);
|
||||
skipped_candidate.extra_data = Some(
|
||||
build_local_execution_candidate_contract_metadata_for_candidate(
|
||||
&skipped_candidate.candidate,
|
||||
skipped_candidate.transport_ref(),
|
||||
provider_api_format.as_str(),
|
||||
spec_metadata.api_format,
|
||||
serde_json::Map::new(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
provider_api_format.as_str(),
|
||||
),
|
||||
);
|
||||
skipped_candidate
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
persist_skipped_local_execution_candidates_with_context(
|
||||
state,
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
available_candidate_count,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok((attempts, candidate_count))
|
||||
Ok((outcome.attempts, outcome.candidate_count))
|
||||
}
|
||||
@@ -12,7 +12,6 @@ pub(crate) use self::candidates::{
|
||||
resolve_local_same_format_provider_decision_input,
|
||||
};
|
||||
pub(crate) use self::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
|
||||
pub(crate) use self::request::resolve_same_format_provider_transport_unsupported_reason_for_trace;
|
||||
pub(crate) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalSameFormatProviderCandidateAttempt;
|
||||
pub(crate) use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput as LocalSameFormatProviderDecisionInput;
|
||||
pub(crate) use crate::ai_pipeline::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};
|
||||
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalSameFormatProviderCandidateAttempt;
|
||||
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalSameFormatProviderDecisionInput;
|
||||
pub(crate) use crate::ai_serving::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};
|
||||
@@ -1,28 +1,28 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_materialization::{
|
||||
use crate::ai_serving::ai_local_execution_contract_for_formats;
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_metadata::build_request_trace_proxy_value;
|
||||
use crate::ai_pipeline::planner::materialization_policy::{
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
use crate::ai_pipeline::planner::payload_metadata::{
|
||||
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_pipeline::planner::report_context::{
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_pipeline::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_pipeline::transport::{
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::planner::CandidateFailureDiagnostic;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AppState,
|
||||
GatewayControlSyncDecisionResponse,
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
@@ -40,7 +40,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
attempt: LocalSameFormatProviderCandidateAttempt,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let LocalSameFormatProviderCandidateAttempt {
|
||||
eligible,
|
||||
@@ -49,6 +49,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
..
|
||||
} = &attempt;
|
||||
let candidate = &eligible.candidate;
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
|
||||
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
|
||||
state, parts, trace_id, body_json, input, &attempt, spec,
|
||||
)
|
||||
@@ -74,7 +76,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
if resolved.is_kiro {
|
||||
extra_fields.insert(
|
||||
"envelope_name".to_string(),
|
||||
json!(crate::ai_pipeline::transport::kiro::KIRO_ENVELOPE_NAME),
|
||||
json!(crate::ai_serving::transport::kiro::KIRO_ENVELOPE_NAME),
|
||||
);
|
||||
} else if resolved.is_antigravity {
|
||||
extra_fields.insert(
|
||||
@@ -109,7 +111,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_pipeline::request_origin_from_parts(parts)),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
@@ -121,8 +123,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
}),
|
||||
ExecutionStrategy::LocalSameFormat,
|
||||
ConversionMode::None,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.api_format,
|
||||
),
|
||||
@@ -142,12 +144,12 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
provider_request_body,
|
||||
} = resolved;
|
||||
|
||||
Some(build_local_execution_decision_response(
|
||||
LocalExecutionDecisionResponseParts {
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy: ExecutionStrategy::LocalSameFormat,
|
||||
conversion_mode: ConversionMode::None,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
@@ -3,18 +3,15 @@ use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::transport::antigravity::{
|
||||
use crate::ai_serving::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
|
||||
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_complete_passthrough_headers, build_complete_passthrough_headers_with_auth,
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_headers, SameFormatProviderHeadersInput,
|
||||
};
|
||||
use crate::ai_pipeline::transport::claude_code::build_claude_code_passthrough_headers;
|
||||
use crate::ai_pipeline::transport::kiro::{build_kiro_provider_headers, KiroProviderHeadersInput};
|
||||
use crate::ai_pipeline::transport::{apply_local_header_rules, ensure_upstream_auth_header};
|
||||
use crate::ai_pipeline::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::AppState;
|
||||
|
||||
mod policy;
|
||||
@@ -30,51 +27,7 @@ use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
use crate::ai_pipeline::planner::standard::same_format_provider_request_body_failure_extra_data;
|
||||
|
||||
pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trace(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
) -> Option<&'static str> {
|
||||
let provider_api_format =
|
||||
match crate::ai_pipeline::normalize_api_format_alias(provider_api_format).as_str() {
|
||||
"openai:chat" => "openai:chat",
|
||||
"openai:responses" => "openai:responses",
|
||||
"openai:responses:compact" => "openai:responses:compact",
|
||||
"claude:messages" => "claude:messages",
|
||||
"gemini:generate_content" => "gemini:generate_content",
|
||||
_ => return Some("transport_api_format_unsupported"),
|
||||
};
|
||||
let behavior = policy::classify_same_format_provider_request_behavior(
|
||||
transport,
|
||||
crate::ai_pipeline::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata {
|
||||
api_format: provider_api_format,
|
||||
require_streaming: false,
|
||||
requested_model_family: None,
|
||||
decision_kind: "trace_candidate_metadata",
|
||||
report_kind: Some("trace_candidate_metadata"),
|
||||
},
|
||||
);
|
||||
if !behavior.is_antigravity
|
||||
&& !behavior.is_claude_code
|
||||
&& !behavior.is_vertex
|
||||
&& !behavior.is_kiro
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let family = if provider_api_format.starts_with("gemini:") {
|
||||
crate::ai_pipeline::LocalSameFormatProviderFamily::Gemini
|
||||
} else {
|
||||
crate::ai_pipeline::LocalSameFormatProviderFamily::Standard
|
||||
};
|
||||
policy::same_format_provider_transport_unsupported_reason(
|
||||
&behavior,
|
||||
transport,
|
||||
family,
|
||||
provider_api_format,
|
||||
)
|
||||
}
|
||||
use crate::ai_serving::planner::standard::same_format_provider_request_body_failure_extra_data;
|
||||
|
||||
pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
@@ -228,74 +181,28 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(provider_request_headers) = (if let Some(kiro_auth) = prepared.kiro_auth.as_ref() {
|
||||
build_kiro_provider_headers(KiroProviderHeadersInput {
|
||||
let extra_headers = antigravity_auth
|
||||
.as_ref()
|
||||
.map(build_antigravity_static_identity_headers)
|
||||
.unwrap_or_default();
|
||||
let Some(provider_request_headers) =
|
||||
build_same_format_provider_headers(SameFormatProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
|
||||
auth_header: prepared.auth_header.as_deref().unwrap_or_default(),
|
||||
auth_value: prepared.auth_value.as_deref().unwrap_or_default(),
|
||||
auth_config: &kiro_auth.auth_config,
|
||||
machine_id: kiro_auth.machine_id.as_str(),
|
||||
behavior: prepared.behavior,
|
||||
auth_header: prepared.auth_header.as_deref(),
|
||||
auth_value: prepared.auth_value.as_deref(),
|
||||
extra_headers: &extra_headers,
|
||||
key_fingerprint: prepared.transport.key.fingerprint.as_ref(),
|
||||
kiro_auth_config: prepared.kiro_auth.as_ref().map(|auth| &auth.auth_config),
|
||||
kiro_machine_id: prepared
|
||||
.kiro_auth
|
||||
.as_ref()
|
||||
.map(|auth| auth.machine_id.as_str()),
|
||||
})
|
||||
} else {
|
||||
let extra_headers = antigravity_auth
|
||||
.as_ref()
|
||||
.map(build_antigravity_static_identity_headers)
|
||||
.unwrap_or_default();
|
||||
let mut provider_request_headers = if prepared.is_claude_code {
|
||||
build_claude_code_passthrough_headers(
|
||||
&parts.headers,
|
||||
prepared.auth_header.as_deref().unwrap_or_default(),
|
||||
prepared.auth_value.as_deref().unwrap_or_default(),
|
||||
&extra_headers,
|
||||
prepared.upstream_is_stream,
|
||||
prepared.transport.key.fingerprint.as_ref(),
|
||||
)
|
||||
} else if prepared.is_vertex {
|
||||
build_complete_passthrough_headers(
|
||||
&parts.headers,
|
||||
&extra_headers,
|
||||
Some("application/json"),
|
||||
)
|
||||
} else {
|
||||
build_complete_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
prepared.auth_header.as_deref().unwrap_or_default(),
|
||||
prepared.auth_value.as_deref().unwrap_or_default(),
|
||||
&extra_headers,
|
||||
Some("application/json"),
|
||||
)
|
||||
};
|
||||
let protected_headers = prepared
|
||||
.auth_header
|
||||
.as_deref()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.map(|value| vec![value, "content-type"])
|
||||
.unwrap_or_else(|| vec!["content-type"]);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
prepared.transport.endpoint.header_rules.as_ref(),
|
||||
&protected_headers,
|
||||
&provider_request_body,
|
||||
Some(body_json),
|
||||
) {
|
||||
None
|
||||
} else {
|
||||
if let (Some(auth_header), Some(auth_value)) = (
|
||||
prepared.auth_header.as_deref(),
|
||||
prepared.auth_value.as_deref(),
|
||||
) {
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
|
||||
}
|
||||
if prepared.upstream_is_stream {
|
||||
provider_request_headers
|
||||
.insert("accept".to_string(), "text/event-stream".to_string());
|
||||
}
|
||||
Some(provider_request_headers)
|
||||
}
|
||||
}) else {
|
||||
else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
@@ -0,0 +1,86 @@
|
||||
use crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata;
|
||||
use crate::ai_serving::transport::{
|
||||
classify_same_format_provider_request_behavior as classify_same_format_provider_request_behavior_impl,
|
||||
resolve_same_format_provider_direct_auth as resolve_same_format_provider_direct_auth_impl,
|
||||
same_format_provider_transport_supported as same_format_provider_transport_supported_impl,
|
||||
same_format_provider_transport_unsupported_reason as same_format_provider_transport_unsupported_reason_impl,
|
||||
should_try_same_format_provider_oauth_auth as should_try_same_format_provider_oauth_auth_impl,
|
||||
GatewayProviderTransportSnapshot, SameFormatProviderFamily, SameFormatProviderRequestBehavior,
|
||||
SameFormatProviderRequestBehaviorParams,
|
||||
};
|
||||
|
||||
use super::super::LocalSameFormatProviderFamily;
|
||||
|
||||
pub(super) fn classify_same_format_provider_request_behavior(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
spec_metadata: LocalExecutionSurfaceSpecMetadata,
|
||||
) -> SameFormatProviderRequestBehavior {
|
||||
classify_same_format_provider_request_behavior_impl(
|
||||
transport,
|
||||
SameFormatProviderRequestBehaviorParams {
|
||||
require_streaming: spec_metadata.require_streaming,
|
||||
report_kind: spec_metadata
|
||||
.report_kind
|
||||
.expect("same-format provider specs should declare report kind"),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_supported(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
same_format_provider_transport_supported_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn same_format_provider_transport_unsupported_reason(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
api_format: &str,
|
||||
) -> Option<&'static str> {
|
||||
same_format_provider_transport_unsupported_reason_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
api_format,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn should_try_same_format_provider_oauth_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> bool {
|
||||
should_try_same_format_provider_oauth_auth_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn resolve_same_format_provider_direct_auth(
|
||||
behavior: &SameFormatProviderRequestBehavior,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
family: LocalSameFormatProviderFamily,
|
||||
) -> Option<(String, String)> {
|
||||
resolve_same_format_provider_direct_auth_impl(
|
||||
behavior,
|
||||
transport,
|
||||
same_format_provider_family(family),
|
||||
)
|
||||
}
|
||||
|
||||
fn same_format_provider_family(family: LocalSameFormatProviderFamily) -> SameFormatProviderFamily {
|
||||
match family {
|
||||
LocalSameFormatProviderFamily::Standard => SameFormatProviderFamily::Standard,
|
||||
LocalSameFormatProviderFamily::Gemini => SameFormatProviderFamily::Gemini,
|
||||
}
|
||||
}
|
||||
@@ -1,13 +1,14 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_preparation::{
|
||||
use crate::ai_serving::planner::candidate_preparation::{
|
||||
resolve_candidate_mapped_model, resolve_candidate_oauth_auth, OauthPreparationContext,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_pipeline::transport::kiro::KiroRequestAuth;
|
||||
use crate::ai_pipeline::transport::vertex::resolve_local_vertex_api_key_query_auth;
|
||||
use crate::ai_pipeline::{
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||
use crate::ai_serving::transport::kiro::KiroRequestAuth;
|
||||
use crate::ai_serving::transport::vertex::resolve_local_vertex_api_key_query_auth;
|
||||
use crate::ai_serving::transport::SameFormatProviderRequestBehavior;
|
||||
use crate::ai_serving::{
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
use crate::AppState;
|
||||
@@ -22,6 +23,7 @@ use super::policy::{
|
||||
|
||||
pub(super) struct PreparedSameFormatProviderCandidate {
|
||||
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
|
||||
pub(super) behavior: SameFormatProviderRequestBehavior,
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) is_claude_code: bool,
|
||||
pub(super) is_vertex: bool,
|
||||
@@ -158,6 +160,7 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
|
||||
|
||||
Some(PreparedSameFormatProviderCandidate {
|
||||
transport,
|
||||
behavior,
|
||||
is_antigravity: behavior.is_antigravity,
|
||||
is_claude_code: behavior.is_claude_code,
|
||||
is_vertex: behavior.is_vertex,
|
||||
@@ -1,15 +1,15 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_pipeline::planner::runtime_miss::{
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::runtime_miss::{
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::{
|
||||
use crate::ai_serving::planner::spec_metadata::{
|
||||
build_stream_plan_from_requested_model_family, build_sync_plan_from_requested_model_family,
|
||||
local_same_format_provider_spec_metadata,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
pub(crate) use crate::ai_serving::{
|
||||
resolve_local_same_format_stream_spec as resolve_stream_spec,
|
||||
resolve_local_same_format_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
@@ -17,8 +17,8 @@ pub(crate) use crate::ai_pipeline::{
|
||||
use super::{
|
||||
materialize_local_same_format_provider_candidate_attempts,
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate,
|
||||
resolve_local_same_format_provider_decision_input, AppState, GatewayControlDecision,
|
||||
GatewayError, LocalSameFormatProviderSpec, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
resolve_local_same_format_provider_decision_input, AiStreamAttempt, AiSyncAttempt, AppState,
|
||||
GatewayControlDecision, GatewayError, LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
@@ -28,7 +28,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
@@ -113,7 +113,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let requested_model_family = spec_metadata
|
||||
.requested_model_family
|
||||
@@ -0,0 +1,36 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_request_body as build_same_format_provider_request_body_impl,
|
||||
SameFormatProviderFamily, SameFormatProviderRequestBodyInput,
|
||||
};
|
||||
|
||||
pub(crate) fn build_same_format_provider_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
body_rules: Option<&Value>,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
|
||||
is_claude_code: bool,
|
||||
) -> Option<Value> {
|
||||
build_same_format_provider_request_body_impl(SameFormatProviderRequestBodyInput {
|
||||
body_json,
|
||||
mapped_model,
|
||||
family: same_format_provider_family(spec.family),
|
||||
body_rules,
|
||||
upstream_is_stream,
|
||||
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
|
||||
is_claude_code,
|
||||
})
|
||||
}
|
||||
|
||||
fn same_format_provider_family(
|
||||
family: super::super::LocalSameFormatProviderFamily,
|
||||
) -> SameFormatProviderFamily {
|
||||
match family {
|
||||
super::super::LocalSameFormatProviderFamily::Standard => SameFormatProviderFamily::Standard,
|
||||
super::super::LocalSameFormatProviderFamily::Gemini => SameFormatProviderFamily::Gemini,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
use crate::ai_serving::transport::{
|
||||
build_same_format_provider_upstream_url as build_same_format_provider_upstream_url_impl,
|
||||
SameFormatProviderUpstreamUrlParams,
|
||||
};
|
||||
use crate::ai_serving::GatewayProviderTransportSnapshot;
|
||||
|
||||
use super::super::LocalSameFormatProviderSpec;
|
||||
|
||||
pub(crate) fn build_same_format_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
|
||||
) -> Option<String> {
|
||||
build_same_format_provider_upstream_url_impl(
|
||||
transport,
|
||||
SameFormatProviderUpstreamUrlParams {
|
||||
provider_api_format: spec.api_format,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
request_query: parts.uri.query(),
|
||||
kiro_api_region: kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -1,9 +1,16 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::ai_pipeline::augment_sync_report_context as augment_sync_report_context_impl;
|
||||
pub(crate) use crate::ai_pipeline::contracts::generic_decision_missing_exact_provider_request;
|
||||
pub(crate) use crate::ai_pipeline::{LocalStreamPlanAndReport, LocalSyncPlanAndReport};
|
||||
use crate::{GatewayControlSyncDecisionResponse, GatewayError};
|
||||
pub(crate) use aether_ai_serving::{
|
||||
build_ai_execution_plan_from_decision, resolve_ai_passthrough_sync_request_body,
|
||||
take_ai_decision_plan_core, take_ai_non_empty_string as take_non_empty_string,
|
||||
take_ai_upstream_auth_pair, AiExecutionPlanFromDecisionParts,
|
||||
};
|
||||
|
||||
use crate::ai_serving::augment_sync_report_context as augment_sync_report_context_impl;
|
||||
pub(crate) use crate::ai_serving::{
|
||||
generic_decision_missing_exact_provider_request, AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::{AiExecutionDecision, GatewayError};
|
||||
|
||||
#[path = "standard/gemini/plan_builders.rs"]
|
||||
mod gemini_builders;
|
||||
@@ -41,7 +48,3 @@ pub(super) fn augment_sync_report_context(
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(super) fn take_non_empty_string(value: &mut Option<String>) -> Option<String> {
|
||||
value.take().filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
@@ -1,21 +1,24 @@
|
||||
use std::cmp::Ordering;
|
||||
use std::collections::{btree_map::Entry, BTreeMap, BTreeSet};
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
|
||||
|
||||
use aether_ai_serving::{
|
||||
run_ai_pool_scheduler, AiPoolCandidateFacts, AiPoolCandidateInput,
|
||||
AiPoolCandidateOrchestration, AiPoolCatalogKeyContext, AiPoolRuntimeState,
|
||||
AiPoolSchedulingConfig, AiPoolSchedulingPreset,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::{Map, Value};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_resolution::{
|
||||
use crate::ai_serving::planner::candidate_resolution::{
|
||||
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
|
||||
};
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::ai_serving::PlannerAppState;
|
||||
use crate::clock::current_unix_ms;
|
||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
||||
use crate::handlers::shared::provider_pool::read_admin_provider_pool_runtime_state;
|
||||
use crate::handlers::shared::provider_pool::{
|
||||
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
|
||||
AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
|
||||
};
|
||||
use crate::handlers::shared::{
|
||||
parse_catalog_auth_config_json, provider_key_health_summary,
|
||||
@@ -24,39 +27,9 @@ use crate::handlers::shared::{
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
use crate::provider_key_auth::provider_key_auth_semantics;
|
||||
|
||||
const POOL_ACCOUNT_BLOCKED_SKIP_REASON: &str = "pool_account_blocked";
|
||||
const POOL_ACCOUNT_EXHAUSTED_SKIP_REASON: &str = "pool_account_exhausted";
|
||||
const POOL_COOLDOWN_SKIP_REASON: &str = "pool_cooldown";
|
||||
const POOL_COST_LIMIT_REACHED_SKIP_REASON: &str = "pool_cost_limit_reached";
|
||||
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
|
||||
struct PoolGroupKey {
|
||||
provider_id: String,
|
||||
endpoint_id: String,
|
||||
model_id: String,
|
||||
selected_provider_model_name: String,
|
||||
provider_api_format: String,
|
||||
singleton_key_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
struct PoolCatalogKeyContext {
|
||||
oauth_plan_type: Option<String>,
|
||||
quota_usage_ratio: Option<f64>,
|
||||
quota_reset_seconds: Option<f64>,
|
||||
account_blocked: bool,
|
||||
quota_exhausted: bool,
|
||||
health_score: Option<f64>,
|
||||
latency_avg_ms: Option<f64>,
|
||||
catalog_lru_score: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct NormalizedPoolPreset {
|
||||
preset: String,
|
||||
mode: Option<String>,
|
||||
}
|
||||
type PoolCatalogKeyContext = AiPoolCatalogKeyContext;
|
||||
|
||||
pub(crate) async fn apply_local_execution_pool_scheduler(
|
||||
state: PlannerAppState<'_>,
|
||||
@@ -319,80 +292,45 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let mut group_order = Vec::new();
|
||||
let mut groups = BTreeMap::<PoolGroupKey, Vec<EligibleLocalExecutionCandidate>>::new();
|
||||
|
||||
for candidate in candidates {
|
||||
let pool_enabled = pool_config_for_candidate(&candidate).is_some();
|
||||
let group_key = pool_group_key(&candidate, pool_enabled);
|
||||
match groups.entry(group_key) {
|
||||
Entry::Vacant(entry) => {
|
||||
group_order.push(entry.key().clone());
|
||||
entry.insert(vec![candidate]);
|
||||
let runtime_by_provider = runtime_by_provider
|
||||
.iter()
|
||||
.map(|(provider_id, runtime)| (provider_id.clone(), ai_pool_runtime_state(runtime)))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let inputs = candidates
|
||||
.into_iter()
|
||||
.map(|candidate| {
|
||||
let key_context = key_context_by_id
|
||||
.get(&candidate.candidate.key_id)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
AiPoolCandidateInput {
|
||||
facts: ai_pool_candidate_facts(&candidate),
|
||||
pool_config: pool_config_for_candidate(&candidate).map(ai_pool_scheduling_config),
|
||||
key_context,
|
||||
candidate,
|
||||
}
|
||||
Entry::Occupied(mut entry) => {
|
||||
entry.get_mut().push(candidate);
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let outcome = run_ai_pool_scheduler(inputs, &runtime_by_provider, pool_sort_seed().as_str());
|
||||
|
||||
let mut reordered = Vec::new();
|
||||
let mut skipped = Vec::new();
|
||||
let default_runtime = AdminProviderPoolRuntimeState::default();
|
||||
let candidates = outcome
|
||||
.candidates
|
||||
.into_iter()
|
||||
.map(|scheduled| apply_ai_pool_orchestration(scheduled.candidate, scheduled.orchestration))
|
||||
.collect::<Vec<_>>();
|
||||
let skipped_candidates = outcome
|
||||
.skipped_candidates
|
||||
.into_iter()
|
||||
.map(|skipped| SkippedLocalExecutionCandidate {
|
||||
candidate: skipped.candidate.candidate,
|
||||
skip_reason: skipped.skip_reason,
|
||||
transport: Some(skipped.candidate.transport),
|
||||
ranking: skipped.candidate.ranking,
|
||||
extra_data: None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
for group_key in group_order {
|
||||
let Some(group) = groups.remove(&group_key) else {
|
||||
continue;
|
||||
};
|
||||
let candidate_group_id = local_execution_candidate_group_id(&group_key);
|
||||
let Some(pool_config) =
|
||||
pool_config_for_candidate(group.first().expect("group should exist"))
|
||||
else {
|
||||
reordered.extend(annotate_local_execution_group_candidates(
|
||||
group,
|
||||
candidate_group_id.as_str(),
|
||||
false,
|
||||
));
|
||||
continue;
|
||||
};
|
||||
let runtime = runtime_by_provider
|
||||
.get(&group_key.provider_id)
|
||||
.unwrap_or(&default_runtime);
|
||||
let (group_candidates, group_skipped) = schedule_pool_group(
|
||||
group,
|
||||
pool_config,
|
||||
runtime,
|
||||
key_context_by_id,
|
||||
candidate_group_id.as_str(),
|
||||
);
|
||||
reordered.extend(group_candidates);
|
||||
skipped.extend(group_skipped);
|
||||
}
|
||||
|
||||
(reordered, skipped)
|
||||
}
|
||||
|
||||
fn pool_group_key(candidate: &EligibleLocalExecutionCandidate, pool_enabled: bool) -> PoolGroupKey {
|
||||
PoolGroupKey {
|
||||
provider_id: candidate.candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.candidate.endpoint_id.clone(),
|
||||
model_id: candidate.candidate.model_id.clone(),
|
||||
selected_provider_model_name: candidate.candidate.selected_provider_model_name.clone(),
|
||||
provider_api_format: candidate.provider_api_format.clone(),
|
||||
singleton_key_id: (!pool_enabled).then(|| candidate.candidate.key_id.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
fn local_execution_candidate_group_id(group_key: &PoolGroupKey) -> String {
|
||||
format!(
|
||||
"provider={}|endpoint={}|model={}|selected_model={}|api_format={}|singleton_key={}",
|
||||
group_key.provider_id,
|
||||
group_key.endpoint_id,
|
||||
group_key.model_id,
|
||||
group_key.selected_provider_model_name,
|
||||
group_key.provider_api_format,
|
||||
group_key.singleton_key_id.as_deref().unwrap_or("*"),
|
||||
)
|
||||
(candidates, skipped_candidates)
|
||||
}
|
||||
|
||||
fn pool_config_for_candidate(
|
||||
@@ -401,659 +339,76 @@ fn pool_config_for_candidate(
|
||||
admin_provider_pool_config_from_config_value(candidate.transport.provider.config.as_ref())
|
||||
}
|
||||
|
||||
fn schedule_pool_group(
|
||||
group: Vec<EligibleLocalExecutionCandidate>,
|
||||
pool_config: AdminProviderPoolConfig,
|
||||
runtime: &AdminProviderPoolRuntimeState,
|
||||
key_context_by_id: &BTreeMap<String, PoolCatalogKeyContext>,
|
||||
candidate_group_id: &str,
|
||||
) -> (
|
||||
Vec<EligibleLocalExecutionCandidate>,
|
||||
Vec<SkippedLocalExecutionCandidate>,
|
||||
) {
|
||||
let provider_type = group
|
||||
.first()
|
||||
.map(|candidate| {
|
||||
candidate
|
||||
.transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let active_presets =
|
||||
normalize_enabled_pool_presets(&pool_config.scheduling_presets, provider_type.as_str());
|
||||
|
||||
let mut available = Vec::new();
|
||||
let mut skipped = Vec::new();
|
||||
|
||||
for (original_index, eligible) in group.into_iter().enumerate() {
|
||||
let EligibleLocalExecutionCandidate {
|
||||
candidate,
|
||||
transport,
|
||||
provider_api_format,
|
||||
orchestration,
|
||||
ranking,
|
||||
} = eligible;
|
||||
let key_id = candidate.key_id.clone();
|
||||
let mut key_context = key_context_by_id.get(&key_id).cloned().unwrap_or_default();
|
||||
key_context.latency_avg_ms = runtime
|
||||
.latency_avg_ms_by_key
|
||||
.get(&key_id)
|
||||
.copied()
|
||||
.or(key_context.latency_avg_ms);
|
||||
|
||||
if key_context.account_blocked {
|
||||
skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason: POOL_ACCOUNT_BLOCKED_SKIP_REASON,
|
||||
transport: Some(transport),
|
||||
ranking,
|
||||
extra_data: None,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if pool_config.skip_exhausted_accounts && key_context.quota_exhausted {
|
||||
skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason: POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
|
||||
transport: Some(transport),
|
||||
ranking,
|
||||
extra_data: None,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if runtime.cooldown_reason_by_key.contains_key(&key_id) {
|
||||
skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason: POOL_COOLDOWN_SKIP_REASON,
|
||||
transport: Some(transport),
|
||||
ranking,
|
||||
extra_data: None,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if pool_config
|
||||
.cost_limit_per_key_tokens
|
||||
.is_some_and(|limit| runtime_cost_usage(runtime, key_id.as_str()) >= limit)
|
||||
{
|
||||
skipped.push(SkippedLocalExecutionCandidate {
|
||||
candidate,
|
||||
skip_reason: POOL_COST_LIMIT_REACHED_SKIP_REASON,
|
||||
transport: Some(transport),
|
||||
ranking,
|
||||
extra_data: None,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
let lru_score =
|
||||
runtime_lru_score(runtime, key_id.as_str()).or(key_context.catalog_lru_score);
|
||||
|
||||
available.push(PoolGroupCandidateOrdering {
|
||||
eligible: EligibleLocalExecutionCandidate {
|
||||
candidate,
|
||||
transport,
|
||||
provider_api_format,
|
||||
orchestration,
|
||||
ranking,
|
||||
},
|
||||
key_context,
|
||||
original_index,
|
||||
lru_score,
|
||||
cost_usage: runtime_cost_usage(runtime, key_id.as_str()),
|
||||
});
|
||||
}
|
||||
|
||||
if available.is_empty() {
|
||||
return (Vec::new(), skipped);
|
||||
}
|
||||
|
||||
let sticky_candidate = runtime
|
||||
.sticky_bound_key_id
|
||||
.as_ref()
|
||||
.and_then(|sticky_key_id| {
|
||||
available
|
||||
.iter()
|
||||
.position(|item| item.eligible.candidate.key_id == *sticky_key_id)
|
||||
})
|
||||
.map(|index| available.remove(index));
|
||||
|
||||
if !active_presets.is_empty() {
|
||||
let sort_vectors = build_pool_sort_vectors(
|
||||
&available,
|
||||
&active_presets,
|
||||
pool_config.lru_enabled,
|
||||
group_sort_seed(
|
||||
provider_type.as_str(),
|
||||
available.first().map(|item| &item.eligible.candidate),
|
||||
)
|
||||
.as_str(),
|
||||
pool_config.cost_limit_per_key_tokens,
|
||||
);
|
||||
available.sort_by(|left, right| {
|
||||
sort_vectors
|
||||
.get(&left.eligible.candidate.key_id)
|
||||
.cmp(&sort_vectors.get(&right.eligible.candidate.key_id))
|
||||
.then(left.original_index.cmp(&right.original_index))
|
||||
});
|
||||
} else if pool_config.lru_enabled {
|
||||
let lru_ranks = lru_rank_indices(&available, false);
|
||||
available.sort_by(|left, right| {
|
||||
lru_ranks
|
||||
.get(&left.eligible.candidate.key_id)
|
||||
.cmp(&lru_ranks.get(&right.eligible.candidate.key_id))
|
||||
.then(left.original_index.cmp(&right.original_index))
|
||||
});
|
||||
}
|
||||
|
||||
let mut ordered = Vec::new();
|
||||
if let Some(sticky_candidate) = sticky_candidate {
|
||||
ordered.push(sticky_candidate.eligible);
|
||||
}
|
||||
ordered.extend(available.into_iter().map(|item| item.eligible));
|
||||
|
||||
(
|
||||
annotate_local_execution_group_candidates(ordered, candidate_group_id, true),
|
||||
skipped,
|
||||
)
|
||||
}
|
||||
|
||||
fn annotate_local_execution_group_candidates(
|
||||
candidates: Vec<EligibleLocalExecutionCandidate>,
|
||||
candidate_group_id: &str,
|
||||
pool_enabled: bool,
|
||||
) -> Vec<EligibleLocalExecutionCandidate> {
|
||||
candidates
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, mut candidate)| {
|
||||
candidate.orchestration = LocalExecutionCandidateMetadata {
|
||||
candidate_group_id: Some(candidate_group_id.to_string()),
|
||||
pool_key_index: pool_enabled.then_some(index as u32),
|
||||
};
|
||||
candidate
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PoolGroupCandidateOrdering {
|
||||
eligible: EligibleLocalExecutionCandidate,
|
||||
key_context: PoolCatalogKeyContext,
|
||||
original_index: usize,
|
||||
lru_score: Option<f64>,
|
||||
cost_usage: u64,
|
||||
}
|
||||
|
||||
fn build_pool_sort_vectors(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
presets: &[NormalizedPoolPreset],
|
||||
lru_enabled: bool,
|
||||
load_balance_seed: &str,
|
||||
cost_limit_per_key_tokens: Option<u64>,
|
||||
) -> BTreeMap<String, Vec<usize>> {
|
||||
let mut vectors = BTreeMap::<String, Vec<usize>>::new();
|
||||
let lru_ranks = lru_rank_indices(items, false);
|
||||
let cache_affinity_ranks = lru_rank_indices(items, true);
|
||||
|
||||
for preset in presets {
|
||||
let ranks = match preset.preset.as_str() {
|
||||
"cache_affinity" => cache_affinity_ranks.clone(),
|
||||
"priority_first" => priority_first_ranks(items, &lru_ranks),
|
||||
"single_account" => single_account_ranks(items),
|
||||
"plus_first" => plan_ranks(items, &lru_ranks, Some("plus_only")),
|
||||
"pro_first" => plan_ranks(items, &lru_ranks, Some("pro_only")),
|
||||
"free_first" => plan_ranks(items, &lru_ranks, Some("free_only")),
|
||||
"team_first" => plan_ranks(items, &lru_ranks, Some("team_only")),
|
||||
"health_first" => health_first_ranks(items, &lru_ranks),
|
||||
"latency_first" => latency_first_ranks(items, &lru_ranks),
|
||||
"cost_first" => cost_first_ranks(items, &lru_ranks, cost_limit_per_key_tokens),
|
||||
"quota_balanced" => quota_balanced_ranks(items, &lru_ranks, cost_limit_per_key_tokens),
|
||||
"recent_refresh" => recent_refresh_ranks(items, &lru_ranks),
|
||||
"load_balance" => load_balance_ranks(items, load_balance_seed),
|
||||
_ => continue,
|
||||
};
|
||||
for item in items {
|
||||
let key_id = item.eligible.candidate.key_id.clone();
|
||||
vectors
|
||||
.entry(key_id.clone())
|
||||
.or_default()
|
||||
.push(*ranks.get(&key_id).unwrap_or(&0));
|
||||
}
|
||||
}
|
||||
|
||||
if lru_enabled {
|
||||
for item in items {
|
||||
let key_id = item.eligible.candidate.key_id.clone();
|
||||
vectors
|
||||
.entry(key_id.clone())
|
||||
.or_default()
|
||||
.push(*lru_ranks.get(&key_id).unwrap_or(&0));
|
||||
}
|
||||
}
|
||||
|
||||
vectors
|
||||
}
|
||||
|
||||
fn lru_rank_indices(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
descending: bool,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = collect_metric_scores(items, |item| item.lru_score);
|
||||
rank_indices_from_score_map(items, &scores, descending)
|
||||
}
|
||||
|
||||
fn priority_first_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
lru_ranks: &BTreeMap<String, usize>,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = collect_metric_scores(items, |item| {
|
||||
Some(f64::from(item.eligible.candidate.key_internal_priority))
|
||||
});
|
||||
if !score_map_has_variation(&scores) {
|
||||
return lru_ranks.clone();
|
||||
}
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn single_account_ranks(items: &[PoolGroupCandidateOrdering]) -> BTreeMap<String, usize> {
|
||||
let n = items.len().saturating_sub(1).max(1) as f64;
|
||||
let priority_scores = collect_metric_scores(items, |item| {
|
||||
Some(f64::from(item.eligible.candidate.key_internal_priority))
|
||||
});
|
||||
let priority_ranks = rank_indices_from_score_map(items, &priority_scores, false);
|
||||
let lru_desc_ranks = lru_rank_indices(items, true);
|
||||
let combined_scores = items
|
||||
.iter()
|
||||
.map(|item| {
|
||||
let key_id = item.eligible.candidate.key_id.clone();
|
||||
let priority_rank = *priority_ranks.get(&key_id).unwrap_or(&0) as f64 / n;
|
||||
let lru_rank = *lru_desc_ranks.get(&key_id).unwrap_or(&0) as f64 / n;
|
||||
(key_id, Some((priority_rank * 0.75) + (lru_rank * 0.25)))
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
rank_indices_from_score_map(items, &combined_scores, false)
|
||||
}
|
||||
|
||||
fn plan_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
lru_ranks: &BTreeMap<String, usize>,
|
||||
mode: Option<&str>,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = items
|
||||
.iter()
|
||||
.map(|item| {
|
||||
(
|
||||
item.eligible.candidate.key_id.clone(),
|
||||
Some(plan_priority_score(
|
||||
item.key_context.oauth_plan_type.as_deref(),
|
||||
mode,
|
||||
)),
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
if !score_map_has_variation(&scores) {
|
||||
return lru_ranks.clone();
|
||||
}
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn health_first_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
lru_ranks: &BTreeMap<String, usize>,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = collect_metric_scores(items, |item| {
|
||||
item.key_context
|
||||
.health_score
|
||||
.map(|score| 1.0 - score.clamp(0.0, 1.0))
|
||||
});
|
||||
if !score_map_has_signal(&scores) {
|
||||
return lru_ranks.clone();
|
||||
}
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn latency_first_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
lru_ranks: &BTreeMap<String, usize>,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = collect_metric_scores(items, |item| item.key_context.latency_avg_ms);
|
||||
if !score_map_has_signal(&scores) {
|
||||
return lru_ranks.clone();
|
||||
}
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn cost_first_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
lru_ranks: &BTreeMap<String, usize>,
|
||||
cost_limit_per_key_tokens: Option<u64>,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = collect_metric_scores(items, |item| {
|
||||
cost_penalty(item, cost_limit_per_key_tokens).or(item.key_context.quota_usage_ratio)
|
||||
});
|
||||
if !score_map_has_signal(&scores) {
|
||||
return lru_ranks.clone();
|
||||
}
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn quota_balanced_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
lru_ranks: &BTreeMap<String, usize>,
|
||||
cost_limit_per_key_tokens: Option<u64>,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = collect_metric_scores(items, |item| {
|
||||
item.key_context
|
||||
.quota_usage_ratio
|
||||
.or_else(|| cost_penalty(item, cost_limit_per_key_tokens))
|
||||
});
|
||||
if !score_map_has_signal(&scores) {
|
||||
return lru_ranks.clone();
|
||||
}
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn recent_refresh_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
lru_ranks: &BTreeMap<String, usize>,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = collect_metric_scores(items, |item| item.key_context.quota_reset_seconds);
|
||||
if !score_map_has_signal(&scores) {
|
||||
return lru_ranks.clone();
|
||||
}
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn load_balance_ranks(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
load_balance_seed: &str,
|
||||
) -> BTreeMap<String, usize> {
|
||||
let scores = items
|
||||
.iter()
|
||||
.map(|item| {
|
||||
let key_id = item.eligible.candidate.key_id.clone();
|
||||
(
|
||||
key_id.clone(),
|
||||
Some(stable_hash_score(
|
||||
format!("{load_balance_seed}:{key_id}").as_str(),
|
||||
)),
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
rank_indices_from_score_map(items, &scores, false)
|
||||
}
|
||||
|
||||
fn group_sort_seed(
|
||||
provider_type: &str,
|
||||
candidate: Option<&aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
|
||||
) -> String {
|
||||
fn pool_sort_seed() -> String {
|
||||
let now_ms = current_unix_ms();
|
||||
let sequence = LOAD_BALANCE_SEQUENCE.fetch_add(1, AtomicOrdering::Relaxed);
|
||||
match candidate {
|
||||
Some(candidate) => format!(
|
||||
"{provider_type}:{}:{}:{}:{}:{now_ms}:{sequence}",
|
||||
candidate.provider_id,
|
||||
candidate.endpoint_id,
|
||||
candidate.model_id,
|
||||
candidate.selected_provider_model_name,
|
||||
),
|
||||
None => format!("{provider_type}:{now_ms}:{sequence}"),
|
||||
format!("{now_ms}:{sequence}")
|
||||
}
|
||||
|
||||
fn ai_pool_candidate_facts(candidate: &EligibleLocalExecutionCandidate) -> AiPoolCandidateFacts {
|
||||
AiPoolCandidateFacts {
|
||||
provider_id: candidate.candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.candidate.endpoint_id.clone(),
|
||||
model_id: candidate.candidate.model_id.clone(),
|
||||
selected_provider_model_name: candidate.candidate.selected_provider_model_name.clone(),
|
||||
provider_api_format: candidate.provider_api_format.clone(),
|
||||
provider_type: candidate.transport.provider.provider_type.clone(),
|
||||
key_id: candidate.candidate.key_id.clone(),
|
||||
key_internal_priority: candidate.candidate.key_internal_priority,
|
||||
}
|
||||
}
|
||||
|
||||
fn stable_hash_score(seed: &str) -> f64 {
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
seed.hash(&mut hasher);
|
||||
let value = hasher.finish();
|
||||
value as f64 / u64::MAX as f64
|
||||
}
|
||||
|
||||
fn collect_metric_scores<F>(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
mut score_for: F,
|
||||
) -> BTreeMap<String, Option<f64>>
|
||||
where
|
||||
F: FnMut(&PoolGroupCandidateOrdering) -> Option<f64>,
|
||||
{
|
||||
items
|
||||
.iter()
|
||||
.map(|item| (item.eligible.candidate.key_id.clone(), score_for(item)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn score_map_has_signal(scores: &BTreeMap<String, Option<f64>>) -> bool {
|
||||
scores.values().flatten().any(|value| value.is_finite())
|
||||
}
|
||||
|
||||
fn score_map_has_variation(scores: &BTreeMap<String, Option<f64>>) -> bool {
|
||||
let mut values = scores
|
||||
.values()
|
||||
.flatten()
|
||||
.filter(|value| value.is_finite())
|
||||
.map(|value| value.to_bits())
|
||||
.collect::<BTreeSet<_>>();
|
||||
values.len() > 1
|
||||
}
|
||||
|
||||
fn rank_indices_from_score_map(
|
||||
items: &[PoolGroupCandidateOrdering],
|
||||
scores: &BTreeMap<String, Option<f64>>,
|
||||
descending: bool,
|
||||
) -> BTreeMap<String, usize> {
|
||||
if !score_map_has_signal(scores) {
|
||||
return items
|
||||
.iter()
|
||||
.map(|item| (item.eligible.candidate.key_id.clone(), 0))
|
||||
.collect();
|
||||
}
|
||||
|
||||
let mut decorated = items
|
||||
.iter()
|
||||
.map(|item| {
|
||||
let key_id = item.eligible.candidate.key_id.clone();
|
||||
let score = scores
|
||||
.get(&key_id)
|
||||
.copied()
|
||||
.flatten()
|
||||
.filter(|value| value.is_finite());
|
||||
let sortable = score.map(|value| if descending { -value } else { value });
|
||||
(
|
||||
score.is_none(),
|
||||
sortable.unwrap_or(f64::INFINITY),
|
||||
item.original_index,
|
||||
key_id,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
decorated.sort_by(|left, right| {
|
||||
left.0
|
||||
.cmp(&right.0)
|
||||
.then_with(|| left.1.partial_cmp(&right.1).unwrap_or(Ordering::Equal))
|
||||
.then(left.2.cmp(&right.2))
|
||||
});
|
||||
|
||||
decorated
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(rank, (_, _, _, key_id))| (key_id, rank))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn cost_penalty(
|
||||
item: &PoolGroupCandidateOrdering,
|
||||
cost_limit_per_key_tokens: Option<u64>,
|
||||
) -> Option<f64> {
|
||||
if item.cost_usage == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(limit) = cost_limit_per_key_tokens.filter(|limit| *limit > 0) {
|
||||
return Some((item.cost_usage as f64 / limit as f64).clamp(0.0, 1.0));
|
||||
}
|
||||
|
||||
let used = item.cost_usage as f64;
|
||||
Some((used / (used + 10_000.0)).clamp(0.0, 1.0))
|
||||
}
|
||||
|
||||
fn plan_priority_score(plan_type: Option<&str>, mode: Option<&str>) -> f64 {
|
||||
match mode.unwrap_or("both").trim().to_ascii_lowercase().as_str() {
|
||||
"free_only" => match plan_type {
|
||||
Some("free") => 0.0,
|
||||
Some("team") => 0.5,
|
||||
Some("enterprise" | "business") => 0.2,
|
||||
Some("plus" | "pro") => 0.6,
|
||||
Some(_) => 0.7,
|
||||
None => 0.8,
|
||||
},
|
||||
"team_only" => match plan_type {
|
||||
Some("team") => 0.0,
|
||||
Some("free") => 0.5,
|
||||
Some("enterprise" | "business") => 0.2,
|
||||
Some("plus" | "pro") => 0.6,
|
||||
Some(_) => 0.7,
|
||||
None => 0.8,
|
||||
},
|
||||
"plus_only" => match plan_type {
|
||||
Some("plus" | "pro") => 0.0,
|
||||
Some("enterprise" | "business") => 0.3,
|
||||
Some("free" | "team") => 0.7,
|
||||
Some(_) => 0.7,
|
||||
None => 0.8,
|
||||
},
|
||||
"pro_only" => match plan_type {
|
||||
Some("pro") => 0.0,
|
||||
Some("plus") => 0.3,
|
||||
Some("enterprise" | "business") => 0.4,
|
||||
Some("free" | "team") => 0.7,
|
||||
Some(_) => 0.7,
|
||||
None => 0.8,
|
||||
},
|
||||
_ => match plan_type {
|
||||
Some("free" | "team") => 0.0,
|
||||
Some("enterprise" | "business") => 0.2,
|
||||
Some("plus" | "pro") => 0.6,
|
||||
Some(_) => 0.7,
|
||||
None => 0.8,
|
||||
},
|
||||
fn ai_pool_scheduling_config(config: AdminProviderPoolConfig) -> AiPoolSchedulingConfig {
|
||||
AiPoolSchedulingConfig {
|
||||
scheduling_presets: config
|
||||
.scheduling_presets
|
||||
.into_iter()
|
||||
.map(|preset| AiPoolSchedulingPreset {
|
||||
preset: preset.preset,
|
||||
enabled: preset.enabled,
|
||||
mode: preset.mode,
|
||||
})
|
||||
.collect(),
|
||||
lru_enabled: config.lru_enabled,
|
||||
skip_exhausted_accounts: config.skip_exhausted_accounts,
|
||||
cost_limit_per_key_tokens: config.cost_limit_per_key_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_enabled_pool_presets(
|
||||
scheduling_presets: &[AdminProviderPoolSchedulingPreset],
|
||||
provider_type: &str,
|
||||
) -> Vec<NormalizedPoolPreset> {
|
||||
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||
let mut entries = Vec::<(usize, String, bool, Option<String>)>::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
|
||||
for (index, item) in scheduling_presets.iter().enumerate() {
|
||||
let preset = item.preset.trim().to_ascii_lowercase();
|
||||
if preset.is_empty() || !seen.insert(preset.clone()) {
|
||||
continue;
|
||||
}
|
||||
entries.push((index, preset, item.enabled, item.mode.clone()));
|
||||
}
|
||||
|
||||
if provider_type == "codex"
|
||||
&& !entries.is_empty()
|
||||
&& entries
|
||||
.iter()
|
||||
.all(|(_, preset, _, _)| preset != "recent_refresh")
|
||||
{
|
||||
entries.push((entries.len(), "recent_refresh".to_string(), true, None));
|
||||
}
|
||||
|
||||
let mut group_anchor_index = BTreeMap::<String, usize>::new();
|
||||
for (index, preset, _, _) in &entries {
|
||||
let Some(mutex_group) = pool_preset_mutex_group(preset) else {
|
||||
continue;
|
||||
};
|
||||
group_anchor_index
|
||||
.entry(mutex_group.to_string())
|
||||
.or_insert(*index);
|
||||
}
|
||||
|
||||
let mut ordered_enabled = Vec::<(usize, usize, String, Option<String>)>::new();
|
||||
let mut group_enabled = BTreeMap::<String, (usize, usize, String, Option<String>)>::new();
|
||||
|
||||
for (index, preset, enabled, mode) in entries {
|
||||
if !enabled
|
||||
|| preset == "lru"
|
||||
|| !pool_preset_supported_for_provider(&preset, &provider_type)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(mutex_group) = pool_preset_mutex_group(&preset) else {
|
||||
ordered_enabled.push((index, index, preset, mode));
|
||||
continue;
|
||||
};
|
||||
let anchor = group_anchor_index
|
||||
.get(mutex_group)
|
||||
.copied()
|
||||
.unwrap_or(index);
|
||||
let existing = group_enabled.get(mutex_group);
|
||||
if existing.is_none_or(|current| index < current.1) {
|
||||
group_enabled.insert(mutex_group.to_string(), (anchor, index, preset, mode));
|
||||
}
|
||||
}
|
||||
|
||||
ordered_enabled.extend(group_enabled.into_values());
|
||||
ordered_enabled.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
|
||||
ordered_enabled
|
||||
.into_iter()
|
||||
.map(|(_, _, preset, mode)| NormalizedPoolPreset { preset, mode })
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn pool_preset_supported_for_provider(preset: &str, provider_type: &str) -> bool {
|
||||
match preset {
|
||||
"free_first" | "plus_first" | "pro_first" | "recent_refresh" | "team_first" => {
|
||||
matches!(provider_type, "codex" | "kiro")
|
||||
}
|
||||
_ => true,
|
||||
fn ai_pool_runtime_state(runtime: &AdminProviderPoolRuntimeState) -> AiPoolRuntimeState {
|
||||
AiPoolRuntimeState {
|
||||
sticky_bound_key_id: runtime.sticky_bound_key_id.clone(),
|
||||
cooldown_reason_by_key: runtime.cooldown_reason_by_key.clone(),
|
||||
cost_window_usage_by_key: runtime.cost_window_usage_by_key.clone(),
|
||||
latency_avg_ms_by_key: runtime.latency_avg_ms_by_key.clone(),
|
||||
lru_score_by_key: runtime.lru_score_by_key.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn pool_preset_mutex_group(preset: &str) -> Option<&'static str> {
|
||||
match preset {
|
||||
"lru" | "cache_affinity" | "load_balance" | "single_account" => Some("distribution_mode"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn runtime_lru_score(runtime: &AdminProviderPoolRuntimeState, key_id: &str) -> Option<f64> {
|
||||
runtime.lru_score_by_key.get(key_id).copied()
|
||||
}
|
||||
|
||||
fn runtime_cost_usage(runtime: &AdminProviderPoolRuntimeState, key_id: &str) -> u64 {
|
||||
runtime
|
||||
.cost_window_usage_by_key
|
||||
.get(key_id)
|
||||
.copied()
|
||||
.unwrap_or_default()
|
||||
fn apply_ai_pool_orchestration(
|
||||
mut candidate: EligibleLocalExecutionCandidate,
|
||||
orchestration: AiPoolCandidateOrchestration,
|
||||
) -> EligibleLocalExecutionCandidate {
|
||||
candidate.orchestration = LocalExecutionCandidateMetadata {
|
||||
candidate_group_id: orchestration.candidate_group_id,
|
||||
pool_key_index: orchestration.pool_key_index,
|
||||
};
|
||||
candidate
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
|
||||
normalize_enabled_pool_presets, PoolCatalogKeyContext,
|
||||
PoolCatalogKeyContext,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use crate::ai_serving::PlannerAppState;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::handlers::shared::provider_pool::{
|
||||
AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
|
||||
};
|
||||
use crate::handlers::shared::provider_pool::AdminProviderPoolRuntimeState;
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
use crate::AppState;
|
||||
use aether_ai_serving::{normalize_enabled_ai_pool_presets, AiPoolSchedulingPreset};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_provider_transport::snapshot::{
|
||||
@@ -1723,24 +1078,24 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn normalizes_distribution_mutex_group_to_first_enabled_member() {
|
||||
let presets = normalize_enabled_pool_presets(
|
||||
let presets = normalize_enabled_ai_pool_presets(
|
||||
&[
|
||||
AdminProviderPoolSchedulingPreset {
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "lru".to_string(),
|
||||
enabled: false,
|
||||
mode: None,
|
||||
},
|
||||
AdminProviderPoolSchedulingPreset {
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "single_account".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
AdminProviderPoolSchedulingPreset {
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "cache_affinity".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
AdminProviderPoolSchedulingPreset {
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "priority_first".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
@@ -1749,13 +1104,7 @@ mod tests {
|
||||
"openai",
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
presets
|
||||
.iter()
|
||||
.map(|item| item.preset.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["single_account", "priority_first"]
|
||||
);
|
||||
assert_eq!(presets, ["single_account", "priority_first"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1843,7 +1192,7 @@ mod tests {
|
||||
provider_api_format: "openai:chat".to_string(),
|
||||
orchestration: LocalExecutionCandidateMetadata::default(),
|
||||
ranking: None,
|
||||
transport: Arc::new(crate::ai_pipeline::GatewayProviderTransportSnapshot {
|
||||
transport: Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: provider_id.to_string(),
|
||||
name: provider_id.to_string(),
|
||||
218
apps/aether-gateway/src/ai_serving/planner/report_context.rs
Normal file
218
apps/aether-gateway/src/ai_serving/planner/report_context.rs
Normal file
@@ -0,0 +1,218 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_serving::{
|
||||
build_ai_execution_report_context,
|
||||
insert_provider_stream_event_api_format as insert_ai_provider_stream_event_api_format,
|
||||
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
|
||||
AiExecutionReportContextParts, AiRequestOrigin,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerRankingOutcome;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin};
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
|
||||
pub(crate) request_id: &'a str,
|
||||
pub(crate) candidate_id: &'a str,
|
||||
pub(crate) attempt_identity: ExecutionAttemptIdentity,
|
||||
pub(crate) model: &'a str,
|
||||
pub(crate) provider_name: &'a str,
|
||||
pub(crate) provider_id: &'a str,
|
||||
pub(crate) endpoint_id: &'a str,
|
||||
pub(crate) key_id: &'a str,
|
||||
pub(crate) key_name: Option<&'a str>,
|
||||
pub(crate) model_id: Option<&'a str>,
|
||||
pub(crate) global_model_id: Option<&'a str>,
|
||||
pub(crate) global_model_name: Option<&'a str>,
|
||||
pub(crate) provider_api_format: &'a str,
|
||||
pub(crate) client_api_format: &'a str,
|
||||
pub(crate) mapped_model: Option<&'a str>,
|
||||
pub(crate) candidate_group_id: Option<&'a str>,
|
||||
pub(crate) ranking: Option<&'a SchedulerRankingOutcome>,
|
||||
pub(crate) upstream_url: Option<&'a str>,
|
||||
pub(crate) header_rules: Option<&'a Value>,
|
||||
pub(crate) body_rules: Option<&'a Value>,
|
||||
pub(crate) provider_request_method: Option<Value>,
|
||||
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
||||
pub(crate) original_headers: &'a http::HeaderMap,
|
||||
pub(crate) request_origin: Option<RequestOrigin>,
|
||||
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>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_execution_report_context(
|
||||
parts: LocalExecutionReportContextParts<'_>,
|
||||
) -> Value {
|
||||
let RequestOrigin {
|
||||
client_ip,
|
||||
user_agent,
|
||||
} = parts
|
||||
.request_origin
|
||||
.unwrap_or_else(|| request_origin_from_headers(parts.original_headers));
|
||||
let original_headers = crate::ai_serving::collect_control_headers(parts.original_headers);
|
||||
let original_request_body = crate::ai_serving::build_report_context_original_request_echo(
|
||||
parts.original_request_body_json,
|
||||
parts.original_request_body_base64,
|
||||
);
|
||||
|
||||
build_ai_execution_report_context(AiExecutionReportContextParts {
|
||||
auth_context: parts.auth_context,
|
||||
request_id: parts.request_id,
|
||||
candidate_id: parts.candidate_id,
|
||||
candidate_index: parts.attempt_identity.candidate_index,
|
||||
retry_index: parts.attempt_identity.retry_index,
|
||||
pool_key_index: parts.attempt_identity.pool_key_index,
|
||||
model: parts.model,
|
||||
provider_name: parts.provider_name,
|
||||
provider_id: parts.provider_id,
|
||||
endpoint_id: parts.endpoint_id,
|
||||
key_id: parts.key_id,
|
||||
key_name: parts.key_name,
|
||||
model_id: parts.model_id,
|
||||
global_model_id: parts.global_model_id,
|
||||
global_model_name: parts.global_model_name,
|
||||
provider_api_format: parts.provider_api_format,
|
||||
client_api_format: parts.client_api_format,
|
||||
mapped_model: parts.mapped_model,
|
||||
candidate_group_id: parts.candidate_group_id,
|
||||
ranking: parts.ranking,
|
||||
upstream_url: parts.upstream_url,
|
||||
header_rules: parts.header_rules,
|
||||
body_rules: parts.body_rules,
|
||||
provider_request_method: parts.provider_request_method,
|
||||
provider_request_headers: parts.provider_request_headers,
|
||||
original_headers: &original_headers,
|
||||
original_request_body,
|
||||
request_origin: AiRequestOrigin {
|
||||
client_ip,
|
||||
user_agent,
|
||||
},
|
||||
client_requested_stream: parts.client_requested_stream,
|
||||
upstream_is_stream: parts.upstream_is_stream,
|
||||
has_envelope: parts.has_envelope,
|
||||
needs_conversion: parts.needs_conversion,
|
||||
extra_fields: parts.extra_fields,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn provider_stream_event_api_format_for_provider_type(
|
||||
provider_type: &str,
|
||||
) -> Option<&'static str> {
|
||||
ai_provider_stream_event_api_format_for_provider_type(provider_type)
|
||||
}
|
||||
|
||||
pub(crate) fn insert_provider_stream_event_api_format(
|
||||
extra_fields: &mut Map<String, Value>,
|
||||
provider_type: &str,
|
||||
) {
|
||||
insert_ai_provider_stream_event_api_format(extra_fields, provider_type);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use super::{
|
||||
build_local_execution_report_context, provider_stream_event_api_format_for_provider_type,
|
||||
LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_serving::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_serving::RequestOrigin;
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
#[test]
|
||||
fn codex_provider_uses_openai_responses_stream_event_format() {
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("codex"),
|
||||
Some("openai:responses")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("CODEX"),
|
||||
Some("openai:responses")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ordinary_providers_do_not_override_stream_event_format() {
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("openai"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
provider_stream_event_api_format_for_provider_type("anthropic"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_execution_report_context_records_request_origin() {
|
||||
let auth_context = ExecutionRuntimeAuthContext {
|
||||
user_id: "user-1".to_string(),
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
balance_remaining: None,
|
||||
access_allowed: true,
|
||||
api_key_is_standalone: false,
|
||||
};
|
||||
let original_headers = http::HeaderMap::new();
|
||||
let provider_request_headers = BTreeMap::new();
|
||||
|
||||
let report_context =
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &auth_context,
|
||||
request_id: "trace-1",
|
||||
candidate_id: "candidate-1",
|
||||
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
|
||||
model: "gpt-5",
|
||||
provider_name: "OpenAI",
|
||||
provider_id: "provider-1",
|
||||
endpoint_id: "endpoint-1",
|
||||
key_id: "key-1",
|
||||
key_name: None,
|
||||
model_id: None,
|
||||
global_model_id: None,
|
||||
global_model_name: None,
|
||||
provider_api_format: "openai:chat",
|
||||
client_api_format: "openai:chat",
|
||||
mapped_model: None,
|
||||
candidate_group_id: None,
|
||||
ranking: None,
|
||||
upstream_url: None,
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
provider_request_method: None,
|
||||
provider_request_headers: Some(&provider_request_headers),
|
||||
original_headers: &original_headers,
|
||||
request_origin: Some(RequestOrigin {
|
||||
client_ip: Some("203.0.113.8".to_string()),
|
||||
user_agent: Some("Claude-Code/1.0".to_string()),
|
||||
}),
|
||||
original_request_body_json: Some(&json!({"model": "gpt-5"})),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields: Map::new(),
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
report_context["client_ip"],
|
||||
Value::String("203.0.113.8".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
report_context["user_agent"],
|
||||
Value::String("Claude-Code/1.0".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,12 +1,10 @@
|
||||
use super::specialized::is_openai_image_stream_request;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
is_matching_stream_request as is_matching_stream_request_impl,
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
is_matching_stream_http_request as is_matching_stream_http_request_impl,
|
||||
resolve_execution_runtime_stream_plan_kind as resolve_execution_runtime_stream_plan_kind_impl,
|
||||
resolve_execution_runtime_sync_plan_kind as resolve_execution_runtime_sync_plan_kind_impl,
|
||||
supports_stream_scheduler_decision_kind as supports_stream_scheduler_decision_kind_impl,
|
||||
supports_sync_scheduler_decision_kind as supports_sync_scheduler_decision_kind_impl,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
};
|
||||
|
||||
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
|
||||
@@ -41,10 +39,7 @@ pub(crate) fn is_matching_stream_request(
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> bool {
|
||||
if plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND {
|
||||
return is_openai_image_stream_request(parts, body_json, body_base64);
|
||||
}
|
||||
is_matching_stream_request_impl(plan_kind, parts.uri.path(), body_json)
|
||||
is_matching_stream_http_request_impl(plan_kind, parts, body_json, body_base64)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
@@ -65,7 +60,7 @@ mod tests {
|
||||
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
|
||||
supports_sync_scheduler_decision_kind,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
|
||||
fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
|
||||
GatewayControlDecision {
|
||||
@@ -83,7 +78,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_openai_chat_plan_kinds_via_pipeline_crate() {
|
||||
fn resolves_openai_chat_plan_kinds_via_surface_crate() {
|
||||
let request = Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri("/v1/chat/completions")
|
||||
@@ -103,7 +98,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_matching_uses_pipeline_route_logic() {
|
||||
fn stream_matching_uses_surface_route_logic() {
|
||||
let request = Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri("/v1/chat/completions")
|
||||
256
apps/aether-gateway/src/ai_serving/planner/runtime_miss.rs
Normal file
256
apps/aether-gateway/src/ai_serving/planner/runtime_miss.rs
Normal file
@@ -0,0 +1,256 @@
|
||||
use aether_ai_serving::{
|
||||
apply_ai_runtime_candidate_evaluation_progress,
|
||||
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal,
|
||||
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic,
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic,
|
||||
apply_ai_runtime_candidate_terminal_reason, build_ai_runtime_candidate_evaluation_diagnostic,
|
||||
build_ai_runtime_execution_exhausted_diagnostic, record_ai_runtime_candidate_skip_reason,
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic,
|
||||
set_ai_runtime_candidate_evaluation_diagnostic, set_ai_runtime_execution_exhausted_diagnostic,
|
||||
set_ai_runtime_miss_diagnostic_reason, AiRuntimeMissDiagnosticFields,
|
||||
AiRuntimeMissDiagnosticPort,
|
||||
};
|
||||
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
|
||||
|
||||
struct GatewayRuntimeMissDiagnosticPort<'a> {
|
||||
state: Option<&'a AppState>,
|
||||
}
|
||||
|
||||
impl AiRuntimeMissDiagnosticFields for LocalExecutionRuntimeMissDiagnostic {
|
||||
fn set_reason(&mut self, reason: String) {
|
||||
self.reason = reason;
|
||||
}
|
||||
|
||||
fn set_candidate_count(&mut self, candidate_count: usize) {
|
||||
self.candidate_count = Some(candidate_count);
|
||||
}
|
||||
|
||||
fn candidate_count(&self) -> Option<usize> {
|
||||
self.candidate_count
|
||||
}
|
||||
|
||||
fn skipped_candidate_count(&self) -> Option<usize> {
|
||||
self.skipped_candidate_count
|
||||
}
|
||||
|
||||
fn skip_reason_count(&self, skip_reason: &str) -> usize {
|
||||
self.skip_reasons.get(skip_reason).copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
fn skip_reason_len(&self) -> usize {
|
||||
self.skip_reasons.len()
|
||||
}
|
||||
|
||||
fn record_skip_reason(&mut self, skip_reason: &'static str) {
|
||||
*self
|
||||
.skip_reasons
|
||||
.entry(skip_reason.to_string())
|
||||
.or_insert(0) += 1;
|
||||
*self.skipped_candidate_count.get_or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
impl AiRuntimeMissDiagnosticPort for GatewayRuntimeMissDiagnosticPort<'_> {
|
||||
type Decision = GatewayControlDecision;
|
||||
type Diagnostic = LocalExecutionRuntimeMissDiagnostic;
|
||||
|
||||
fn build_runtime_miss_diagnostic(
|
||||
&self,
|
||||
decision: &Self::Decision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> Self::Diagnostic {
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: reason.to_string(),
|
||||
route_family: decision.route_family.clone(),
|
||||
route_kind: decision.route_kind.clone(),
|
||||
public_path: Some(decision.public_path.clone()),
|
||||
plan_kind: Some(plan_kind.to_string()),
|
||||
requested_model: requested_model.map(ToOwned::to_owned),
|
||||
candidate_count: None,
|
||||
skipped_candidate_count: None,
|
||||
skip_reasons: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_candidate_count(&self, diagnostic: &mut Self::Diagnostic, candidate_count: usize) {
|
||||
AiRuntimeMissDiagnosticFields::set_candidate_count(diagnostic, candidate_count);
|
||||
}
|
||||
|
||||
fn apply_candidate_evaluation_progress(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic(diagnostic, candidate_count);
|
||||
}
|
||||
|
||||
fn apply_candidate_terminal_plan_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
no_plan_reason: &'static str,
|
||||
) {
|
||||
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(diagnostic, no_plan_reason);
|
||||
}
|
||||
|
||||
fn record_candidate_skip_reason(
|
||||
&self,
|
||||
diagnostic: &mut Self::Diagnostic,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
record_ai_runtime_candidate_skip_reason_on_diagnostic(diagnostic, skip_reason);
|
||||
}
|
||||
|
||||
fn set_runtime_miss_diagnostic(&self, trace_id: &str, diagnostic: Self::Diagnostic) {
|
||||
self.state
|
||||
.expect("runtime miss diagnostic setter requires gateway state")
|
||||
.set_local_execution_runtime_miss_diagnostic(trace_id, diagnostic);
|
||||
}
|
||||
|
||||
fn mutate_runtime_miss_diagnostic<F>(&self, trace_id: &str, apply: F)
|
||||
where
|
||||
F: FnOnce(&mut Self::Diagnostic) + Send,
|
||||
{
|
||||
self.state
|
||||
.expect("runtime miss diagnostic mutator requires gateway state")
|
||||
.mutate_local_execution_runtime_miss_diagnostic(trace_id, apply);
|
||||
}
|
||||
|
||||
fn runtime_miss_diagnostic_has_candidate_signal(&self, trace_id: &str) -> bool {
|
||||
self.state
|
||||
.expect("runtime miss diagnostic signal check requires gateway state")
|
||||
.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_miss_diagnostic_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
set_ai_runtime_miss_diagnostic_reason(
|
||||
&port,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
reason,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_runtime_execution_exhausted_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: None };
|
||||
build_ai_runtime_execution_exhausted_diagnostic(
|
||||
&port,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_execution_exhausted_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
set_ai_runtime_execution_exhausted_diagnostic(
|
||||
&port,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_runtime_candidate_evaluation_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: None };
|
||||
build_ai_runtime_candidate_evaluation_diagnostic(
|
||||
&port,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_runtime_candidate_evaluation_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
requested_model: Option<&str>,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
set_ai_runtime_candidate_evaluation_diagnostic(
|
||||
&port,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
requested_model,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_evaluation_progress(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
apply_ai_runtime_candidate_evaluation_progress(&port, trace_id, candidate_count);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
&port,
|
||||
trace_id,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_runtime_candidate_terminal_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
no_plan_reason: &'static str,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
apply_ai_runtime_candidate_terminal_reason(&port, trace_id, no_plan_reason);
|
||||
}
|
||||
|
||||
pub(crate) fn record_local_runtime_candidate_skip_reason(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let port = GatewayRuntimeMissDiagnosticPort { state: Some(state) };
|
||||
record_ai_runtime_candidate_skip_reason(&port, trace_id, skip_reason);
|
||||
}
|
||||
53
apps/aether-gateway/src/ai_serving/planner/spec_metadata.rs
Normal file
53
apps/aether-gateway/src/ai_serving/planner/spec_metadata.rs
Normal file
@@ -0,0 +1,53 @@
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::AiExecutionDecision;
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) use aether_ai_serving::{
|
||||
ai_gemini_files_spec_metadata as local_gemini_files_spec_metadata,
|
||||
ai_openai_image_spec_metadata as local_openai_image_spec_metadata,
|
||||
ai_openai_responses_spec_metadata as local_openai_responses_spec_metadata,
|
||||
ai_requested_model_family_for_same_format_provider as requested_model_family_for_same_format_provider,
|
||||
ai_requested_model_family_for_standard_source as requested_model_family_for_standard_source,
|
||||
ai_requested_model_family_for_video_create as requested_model_family_for_video_create,
|
||||
ai_same_format_provider_spec_metadata as local_same_format_provider_spec_metadata,
|
||||
ai_standard_spec_metadata as local_standard_spec_metadata,
|
||||
ai_video_create_spec_metadata as local_video_create_spec_metadata,
|
||||
AiExecutionSurfaceSpecMetadata as LocalExecutionSurfaceSpecMetadata,
|
||||
AiRequestedModelFamily as RequestedModelFamily,
|
||||
};
|
||||
|
||||
pub(crate) fn build_sync_plan_from_requested_model_family(
|
||||
family: RequestedModelFamily,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
match family {
|
||||
RequestedModelFamily::Standard => {
|
||||
build_standard_sync_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
RequestedModelFamily::Gemini => {
|
||||
build_gemini_sync_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_stream_plan_from_requested_model_family(
|
||||
family: RequestedModelFamily,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
payload: AiExecutionDecision,
|
||||
) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
match family {
|
||||
RequestedModelFamily::Standard => {
|
||||
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
|
||||
}
|
||||
RequestedModelFamily::Gemini => {
|
||||
build_gemini_stream_plan_from_decision(parts, body_json, payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4,17 +4,17 @@ mod support;
|
||||
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
use crate::ai_serving::planner::plan_builders::{
|
||||
build_passthrough_stream_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
resolve_gemini_files_stream_spec as resolve_stream_spec,
|
||||
resolve_gemini_files_sync_spec as resolve_sync_spec, LocalGeminiFilesSpec,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use self::decision::maybe_build_local_gemini_files_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
@@ -30,7 +30,7 @@ pub(crate) async fn build_local_gemini_files_sync_plan_and_reports_for_kind(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
@@ -54,7 +54,7 @@ pub(crate) async fn build_local_gemini_files_stream_plan_and_reports_for_kind(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
@@ -71,7 +71,7 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -111,7 +111,7 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -155,7 +155,7 @@ async fn build_local_sync_plan_and_reports(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
|
||||
) -> Result<Vec<AiSyncAttempt>, GatewayError> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
|
||||
else {
|
||||
@@ -206,7 +206,7 @@ async fn build_local_stream_plan_and_reports(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
|
||||
else {
|
||||
@@ -1,18 +1,18 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_pipeline::planner::candidate_metadata::build_request_trace_proxy_value;
|
||||
use crate::ai_pipeline::planner::payload_metadata::{
|
||||
build_local_execution_decision_response, LocalExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_pipeline::planner::report_context::{
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_pipeline::transport::{
|
||||
use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_serving::planner::{
|
||||
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
|
||||
};
|
||||
use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{AiExecutionDecision, AppState};
|
||||
|
||||
use super::request::resolve_local_gemini_files_candidate_payload_parts;
|
||||
use super::support::{
|
||||
@@ -31,7 +31,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
attempt: LocalGeminiFilesCandidateAttempt,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
) -> Option<AiExecutionDecision> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
@@ -54,6 +54,10 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
} = attempt;
|
||||
let candidate = eligible.candidate;
|
||||
let transport = resolved.transport;
|
||||
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
);
|
||||
let proxy = planner_state
|
||||
.app()
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||
@@ -90,7 +94,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
provider_request_method: None,
|
||||
provider_request_headers: None,
|
||||
original_headers: &parts.headers,
|
||||
request_origin: Some(crate::ai_pipeline::request_origin_from_parts(parts)),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
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,
|
||||
@@ -110,12 +114,12 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
file_name: _,
|
||||
} = resolved;
|
||||
|
||||
Some(build_local_execution_decision_response(
|
||||
LocalExecutionDecisionResponseParts {
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy: ExecutionStrategy::LocalSameFormat,
|
||||
conversion_mode: ConversionMode::None,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
@@ -1,17 +1,14 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_pipeline::contracts::GEMINI_FILES_UPLOAD_PLAN_KIND;
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_passthrough_headers_with_auth, resolve_local_gemini_auth,
|
||||
use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata;
|
||||
use crate::ai_serving::transport::{
|
||||
build_gemini_files_headers, build_gemini_files_request_body, build_gemini_files_upstream_url,
|
||||
gemini_files_transport_unsupported_reason, resolve_gemini_files_auth, GeminiFilesHeadersInput,
|
||||
GeminiFilesRequestBodyError,
|
||||
};
|
||||
use crate::ai_pipeline::transport::local_gemini_transport_unsupported_reason_with_network;
|
||||
use crate::ai_pipeline::transport::url::build_gemini_files_passthrough_url;
|
||||
use crate::ai_pipeline::transport::{apply_local_body_rules, apply_local_header_rules};
|
||||
use crate::ai_pipeline::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::ai_serving::GEMINI_FILES_UPLOAD_PLAN_KIND;
|
||||
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
|
||||
use crate::AppState;
|
||||
|
||||
use super::support::{
|
||||
@@ -49,10 +46,9 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
|
||||
if let Some(skip_reason) = local_gemini_transport_unsupported_reason_with_network(
|
||||
transport,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
) {
|
||||
if let Some(skip_reason) =
|
||||
gemini_files_transport_unsupported_reason(transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
|
||||
{
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
@@ -66,7 +62,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_local_gemini_auth(transport) else {
|
||||
let Some((auth_header, auth_value)) = resolve_gemini_files_auth(transport) else {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
@@ -80,18 +76,9 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
return None;
|
||||
};
|
||||
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
let passthrough_path = custom_path.unwrap_or(parts.uri.path());
|
||||
let Some(upstream_url) = build_gemini_files_passthrough_url(
|
||||
&transport.endpoint.base_url,
|
||||
passthrough_path,
|
||||
parts.uri.query(),
|
||||
) else {
|
||||
let Some(upstream_url) =
|
||||
build_gemini_files_upstream_url(transport, parts.uri.path(), parts.uri.query())
|
||||
else {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
@@ -110,46 +97,33 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_body = if spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND
|
||||
&& !body_is_empty
|
||||
&& body_base64.is_none()
|
||||
{
|
||||
Some(body_json.clone())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let provider_request_body_base64 =
|
||||
if spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND {
|
||||
body_base64
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if provider_request_body_base64.is_some() && transport.endpoint.body_rules.is_some() {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_body_rules_unsupported_for_binary_upload",
|
||||
CandidateFailureDiagnostic::body_rules_unsupported_for_binary_upload(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"gemini_files_binary_upload",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
if let Some(body) = provider_request_body.as_mut() {
|
||||
if !apply_local_body_rules(
|
||||
body,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(body_json),
|
||||
) {
|
||||
let body_parts = match build_gemini_files_request_body(
|
||||
body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) {
|
||||
Ok(parts) => parts,
|
||||
Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload) => {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
attempt.candidate_index,
|
||||
&attempt.candidate_id,
|
||||
"transport_body_rules_unsupported_for_binary_upload",
|
||||
CandidateFailureDiagnostic::body_rules_unsupported_for_binary_upload(
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"gemini_files_binary_upload",
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(GeminiFilesRequestBodyError::BodyRulesApplyFailed) => {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
@@ -167,31 +141,18 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
);
|
||||
let null_original_request_body = serde_json::Value::Null;
|
||||
let base64_original_request_body = provider_request_body_base64
|
||||
.as_ref()
|
||||
.map(|body_bytes_b64| json!({ "body_bytes_b64": body_bytes_b64 }));
|
||||
let original_request_body = base64_original_request_body
|
||||
.as_ref()
|
||||
.or_else(|| (!body_is_empty).then_some(body_json))
|
||||
.unwrap_or(&null_original_request_body);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
&[&auth_header, "content-type"],
|
||||
provider_request_body
|
||||
.as_ref()
|
||||
.unwrap_or(original_request_body),
|
||||
Some(original_request_body),
|
||||
) {
|
||||
let Some(provider_request_headers) = build_gemini_files_headers(GeminiFilesHeadersInput {
|
||||
headers: &parts.headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
provider_request_body: body_parts.provider_request_body.as_ref(),
|
||||
provider_request_body_base64: body_parts.provider_request_body_base64.as_deref(),
|
||||
original_request_body_json: body_json,
|
||||
original_body_is_empty: body_is_empty,
|
||||
}) else {
|
||||
mark_skipped_local_gemini_files_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
@@ -208,7 +169,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let file_name = parts
|
||||
.uri
|
||||
@@ -222,8 +183,8 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
provider_request_body_base64,
|
||||
provider_request_body: body_parts.provider_request_body,
|
||||
provider_request_body_base64: body_parts.provider_request_body_base64,
|
||||
upstream_url,
|
||||
file_name,
|
||||
})
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user