mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: 大规模模块拆分与重组，新增 aether-admin crate
- 新建独立 aether-admin crate 承载 admin 相关共享契约与纯辅助函数 - 拆分 ai_pipeline 下 kiro/private_envelope/conversion/planner 等大文件为子模块目录 - 重组 admin handlers 各业务域（billing/oauth/provider/system/users 等）为目录结构，移除 shared.rs/builders.rs 等反模式 - 移除 ai_pipeline runtime adapters 旧实现（claude/openai/gemini/kiro/vertex/antigravity 等），改由 provider transport 统一承载 - 移除 control_facade/execution_facade/auth_snapshot_facade 等冗余 facade 层 - 拆分 query/billing 与 query/monitoring 模块、state/runtime/payments 与 security 模块 - 扩展架构测试覆盖 admin_billing/admin_model/admin_users 等新模块 - 删除 docs/architecture/refactor-execution-plan.md 已完成的执行计划文档
This commit is contained in:
fawney19
@@ -7,6 +7,7 @@ repository.workspace = true
|
||||
description = "Rust ingress gateway for Aether phase 3a transparent proxy"
|
||||
|
||||
[dependencies]
|
||||
aether-admin.workspace = true
|
||||
aether-ai-pipeline.workspace = true
|
||||
aether-billing.workspace = true
|
||||
aether-cache.workspace = true
|
||||
|
||||
109
apps/aether-gateway/src/admin_api.rs
Normal file
109
apps/aether-gateway/src/admin_api.rs
Normal file
@@ -0,0 +1,109 @@
|
||||
pub(crate) use crate::handlers::admin::{
|
||||
admin_provider_ops_local_action_response, build_internal_control_error_response,
|
||||
maybe_build_local_admin_pool_response, maybe_build_local_admin_response, AdminAppState,
|
||||
AdminRequestContext, AdminRouteRequest, AdminRouteResponse, AdminRouteResult,
|
||||
AdminStatsTimeRange, AdminStatsUsageFilter,
|
||||
};
|
||||
|
||||
use crate::handlers::admin::{
|
||||
admin_stats_bad_request_response as admin_stats_bad_request_response_impl,
|
||||
list_usage_for_optional_range as list_usage_for_optional_range_impl,
|
||||
parse_bounded_u32 as parse_bounded_u32_impl, round_to as round_to_impl,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
response::Response,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_security_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::handlers::admin::maybe_build_local_admin_security_response(
|
||||
state,
|
||||
request_context,
|
||||
request_body,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_endpoint_health_status_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
lookback_hours: u64,
|
||||
) -> Option<serde_json::Value> {
|
||||
crate::handlers::admin::build_admin_endpoint_health_status_payload(state, lookback_hours).await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_video_tasks_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::handlers::admin::maybe_build_local_admin_video_tasks_response(state, request_context)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_usage_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::handlers::admin::maybe_build_local_admin_usage_response(
|
||||
state,
|
||||
request_context,
|
||||
request_body,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn admin_stats_bad_request_response(detail: String) -> Response<Body> {
|
||||
admin_stats_bad_request_response_impl(detail)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_usage_for_optional_range(
|
||||
state: &AdminAppState<'_>,
|
||||
time_range: Option<&AdminStatsTimeRange>,
|
||||
filters: &AdminStatsUsageFilter,
|
||||
) -> Result<Vec<aether_data_contracts::repository::usage::StoredRequestUsageAudit>, GatewayError> {
|
||||
list_usage_for_optional_range_impl(state, time_range, filters).await
|
||||
}
|
||||
|
||||
pub(crate) fn parse_bounded_u32(
|
||||
field: &str,
|
||||
value: &str,
|
||||
min: u32,
|
||||
max: u32,
|
||||
) -> Result<u32, String> {
|
||||
parse_bounded_u32_impl(field, value, min, max)
|
||||
}
|
||||
|
||||
pub(crate) fn round_to(value: f64, decimals: u32) -> f64 {
|
||||
round_to_impl(value, decimals)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_provider_oauth_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::handlers::admin::maybe_build_local_admin_provider_oauth_response(
|
||||
state,
|
||||
request_context,
|
||||
request_body,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_providers_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::handlers::admin::maybe_build_local_admin_providers_response(
|
||||
state,
|
||||
request_context,
|
||||
request_body,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -1,11 +1,8 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
const CONTEXT_WINDOW_TOKENS: f64 = 200_000.0;
|
||||
const MAX_THINKING_BUFFER: usize = 1024 * 1024;
|
||||
const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
|
||||
const MAX_BUFFER_SIZE: usize = MAX_MESSAGE_SIZE;
|
||||
const MAX_ERRORS: usize = 5;
|
||||
const QUOTE_CHARS: &str = "`\"'\\#!@$%^&*()-_=+[]{};:<>,.?/";
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct KiroToClaudeCliStreamState {
|
||||
@@ -68,8 +65,6 @@ enum FrameParseError {
|
||||
mod decoder;
|
||||
#[path = "stream/state.rs"]
|
||||
mod state;
|
||||
#[path = "stream/util.rs"]
|
||||
mod util;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "stream/tests.rs"]
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::util::crc32;
|
||||
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> {
|
||||
|
||||
@@ -1,605 +1,8 @@
|
||||
use serde_json::{json, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
use super::util::{
|
||||
encode_events, estimate_tokens, find_real_thinking_end_tag,
|
||||
find_real_thinking_end_tag_at_buffer_end, find_real_thinking_start_tag,
|
||||
};
|
||||
use super::{
|
||||
AwsEventFrame, EventStreamDecoder, KiroClaudeStreamState, KiroToClaudeCliStreamState,
|
||||
CONTEXT_WINDOW_TOKENS, MAX_THINKING_BUFFER,
|
||||
};
|
||||
|
||||
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 {
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_initial_bytes(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
let mut events = vec![json!({
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": self.message_id,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"model": self.model,
|
||||
"stop_reason": Value::Null,
|
||||
"stop_sequence": Value::Null,
|
||||
"usage": {
|
||||
"input_tokens": self.estimated_input_tokens as u64,
|
||||
"output_tokens": 1,
|
||||
},
|
||||
}
|
||||
})];
|
||||
if !self.thinking_enabled {
|
||||
events.extend(self.ensure_text_block_open());
|
||||
}
|
||||
encode_events(events)
|
||||
}
|
||||
|
||||
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_events(vec![json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": error_type,
|
||||
"message": message,
|
||||
}
|
||||
})])
|
||||
}
|
||||
|
||||
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()),
|
||||
}
|
||||
}
|
||||
|
||||
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(((percentage * CONTEXT_WINDOW_TOKENS) / 100.0) as usize);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
encode_events(events)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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_tokens(content);
|
||||
|
||||
if !self.thinking_enabled {
|
||||
return self.emit_text_delta(content);
|
||||
}
|
||||
|
||||
self.thinking_buffer.push_str(content);
|
||||
if self.thinking_buffer.len() > 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_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 self.thinking_buffer.len() > keep {
|
||||
let split = self.thinking_buffer.len() - keep;
|
||||
let safe = self.thinking_buffer[..split].to_string();
|
||||
if !safe.trim().is_empty() {
|
||||
events.extend(self.emit_text_delta(&safe));
|
||||
self.thinking_buffer = self.thinking_buffer[split..].to_string();
|
||||
}
|
||||
}
|
||||
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_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 self.thinking_buffer.len() > keep {
|
||||
let split = self.thinking_buffer.len() - keep;
|
||||
let safe = self.thinking_buffer[..split].to_string();
|
||||
if !safe.is_empty() {
|
||||
events.extend(self.emit_thinking_delta(&safe));
|
||||
self.thinking_buffer = self.thinking_buffer[split..].to_string();
|
||||
}
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
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_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_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
|
||||
}
|
||||
|
||||
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_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_events(flush_events)?;
|
||||
for idx in self
|
||||
.open_blocks
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
{
|
||||
output.extend(encode_events(self.close_block(idx))?);
|
||||
}
|
||||
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_events(self.close_block(idx))?);
|
||||
}
|
||||
output.extend(self.final_message_bytes()?);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
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_events(vec![
|
||||
json!({
|
||||
"type": "message_delta",
|
||||
"delta": {
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": Value::Null,
|
||||
},
|
||||
"usage": {
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": self.output_tokens as u64,
|
||||
}
|
||||
}),
|
||||
json!({"type": "message_stop"}),
|
||||
])
|
||||
}
|
||||
|
||||
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": ""}
|
||||
})]
|
||||
}
|
||||
|
||||
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": ""}
|
||||
})]
|
||||
}
|
||||
|
||||
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})]
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
#[path = "state/blocks.rs"]
|
||||
mod blocks;
|
||||
#[path = "state/events.rs"]
|
||||
mod events;
|
||||
#[path = "state/finalize.rs"]
|
||||
mod finalize;
|
||||
#[path = "state/lifecycle.rs"]
|
||||
mod lifecycle;
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,303 @@
|
||||
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;
|
||||
|
||||
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 self.thinking_buffer.len() > keep {
|
||||
let split = self.thinking_buffer.len() - keep;
|
||||
let safe = self.thinking_buffer[..split].to_string();
|
||||
if !safe.trim().is_empty() {
|
||||
events.extend(self.emit_text_delta(&safe));
|
||||
self.thinking_buffer = self.thinking_buffer[split..].to_string();
|
||||
}
|
||||
}
|
||||
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 self.thinking_buffer.len() > keep {
|
||||
let split = self.thinking_buffer.len() - keep;
|
||||
let safe = self.thinking_buffer[..split].to_string();
|
||||
if !safe.is_empty() {
|
||||
events.extend(self.emit_thinking_delta(&safe));
|
||||
self.thinking_buffer = self.thinking_buffer[split..].to_string();
|
||||
}
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
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()))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
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,4 +1,4 @@
|
||||
use super::util::crc32;
|
||||
use aether_ai_pipeline::adaptation::kiro_stream::kiro_crc32 as crc32;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::KiroToClaudeCliStreamState;
|
||||
|
||||
@@ -1,118 +0,0 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
use super::QUOTE_CHARS;
|
||||
|
||||
pub(super) fn encode_events(events: Vec<Value>) -> Result<Vec<u8>, GatewayError> {
|
||||
let mut output = Vec::new();
|
||||
for event in events {
|
||||
output.extend(encode_sse_event(&event)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub(super) fn encode_sse_event(event: &Value) -> Result<Vec<u8>, GatewayError> {
|
||||
let encoded =
|
||||
serde_json::to_string(event).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if let Some(event_type) = event.get("type").and_then(Value::as_str) {
|
||||
Ok(format!("event: {event_type}\ndata: {encoded}\n\n").into_bytes())
|
||||
} else {
|
||||
Ok(format!("data: {encoded}\n\n").into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn estimate_tokens(text: &str) -> usize {
|
||||
if text.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
let mut chinese = 0usize;
|
||||
let mut other = 0usize;
|
||||
for ch in text.chars() {
|
||||
if ('\u{4e00}'..='\u{9fff}').contains(&ch) {
|
||||
chinese += 1;
|
||||
} else {
|
||||
other += 1;
|
||||
}
|
||||
}
|
||||
let chinese_tokens = (chinese * 2).div_ceil(3);
|
||||
let other_tokens = other.div_ceil(4);
|
||||
(chinese_tokens + other_tokens).max(1)
|
||||
}
|
||||
|
||||
pub(super) fn is_quote_char(buffer: &str, pos: usize) -> bool {
|
||||
buffer
|
||||
.as_bytes()
|
||||
.get(pos)
|
||||
.map(|byte| QUOTE_CHARS.as_bytes().contains(byte))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(super) fn find_real_thinking_start_tag(buffer: &str) -> Option<usize> {
|
||||
let tag = "<thinking>";
|
||||
let mut search = 0usize;
|
||||
loop {
|
||||
let pos = buffer[search..].find(tag).map(|value| value + search)?;
|
||||
let has_before = pos > 0 && is_quote_char(buffer, pos - 1);
|
||||
let after_pos = pos + tag.len();
|
||||
let has_after = is_quote_char(buffer, after_pos);
|
||||
if !has_before && !has_after {
|
||||
return Some(pos);
|
||||
}
|
||||
search = pos + 1;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn find_real_thinking_end_tag(buffer: &str) -> Option<usize> {
|
||||
let tag = "</thinking>";
|
||||
let mut search = 0usize;
|
||||
loop {
|
||||
let pos = buffer[search..].find(tag).map(|value| value + search)?;
|
||||
let has_before = pos > 0 && is_quote_char(buffer, pos - 1);
|
||||
let after_pos = pos + tag.len();
|
||||
let has_after = is_quote_char(buffer, after_pos);
|
||||
if has_before || has_after {
|
||||
search = pos + 1;
|
||||
continue;
|
||||
}
|
||||
let after = &buffer[after_pos..];
|
||||
if after.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
if after.starts_with("\n\n") {
|
||||
return Some(pos);
|
||||
}
|
||||
search = pos + 1;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn find_real_thinking_end_tag_at_buffer_end(buffer: &str) -> Option<usize> {
|
||||
let tag = "</thinking>";
|
||||
let mut search = 0usize;
|
||||
loop {
|
||||
let pos = buffer[search..].find(tag).map(|value| value + search)?;
|
||||
let has_before = pos > 0 && is_quote_char(buffer, pos - 1);
|
||||
let after_pos = pos + tag.len();
|
||||
let has_after = is_quote_char(buffer, after_pos);
|
||||
if has_before || has_after {
|
||||
search = pos + 1;
|
||||
continue;
|
||||
}
|
||||
if buffer[after_pos..].trim().is_empty() {
|
||||
return Some(pos);
|
||||
}
|
||||
search = pos + 1;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn crc32(data: &[u8]) -> u32 {
|
||||
let mut crc = 0xffff_ffffu32;
|
||||
for &byte in data {
|
||||
crc ^= byte as u32;
|
||||
for _ in 0..8 {
|
||||
let mask = if crc & 1 == 1 { 0xedb8_8320 } else { 0 };
|
||||
crc = (crc >> 1) ^ mask;
|
||||
}
|
||||
}
|
||||
!crc
|
||||
}
|
||||
@@ -1,7 +1,13 @@
|
||||
pub(crate) mod kiro;
|
||||
pub(crate) mod private_envelope;
|
||||
pub(crate) mod surfaces;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
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,
|
||||
provider_adaptation_should_unwrap_stream_envelope, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
|
||||
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
|
||||
};
|
||||
pub(crate) use kiro::KiroToClaudeCliStreamState;
|
||||
pub(crate) use private_envelope::{
|
||||
maybe_build_provider_private_stream_normalizer,
|
||||
@@ -9,10 +15,3 @@ pub(crate) use private_envelope::{
|
||||
normalize_provider_private_report_context, normalize_provider_private_response_value,
|
||||
provider_private_response_allows_sync_finalize, transform_provider_private_stream_line,
|
||||
};
|
||||
pub(crate) use surfaces::{
|
||||
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,
|
||||
provider_adaptation_should_unwrap_stream_envelope, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
|
||||
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
|
||||
};
|
||||
|
||||
@@ -1,530 +1,17 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use base64::Engine as _;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::adaptation::surfaces::{
|
||||
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_descriptor_for_envelope,
|
||||
provider_adaptation_should_unwrap_stream_envelope, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
|
||||
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_pipeline::runtime::adapters::kiro::{KiroToClaudeCliStreamState, KIRO_ENVELOPE_NAME};
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
enum ProviderPrivateStreamNormalizeMode {
|
||||
EnvelopeUnwrap,
|
||||
KiroToClaudeCli(KiroToClaudeCliStreamState),
|
||||
}
|
||||
|
||||
pub(crate) struct ProviderPrivateStreamNormalizer {
|
||||
report_context: Value,
|
||||
buffered: Vec<u8>,
|
||||
mode: ProviderPrivateStreamNormalizeMode,
|
||||
}
|
||||
|
||||
pub(crate) fn provider_private_response_allows_sync_finalize(report_context: &Value) -> bool {
|
||||
let has_envelope = report_context
|
||||
.get("has_envelope")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
if !has_envelope {
|
||||
return true;
|
||||
}
|
||||
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_allows_sync_finalize_envelope(envelope_name, provider_api_format)
|
||||
|| matches!(envelope_name, "claude:cli")
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_provider_private_report_context(
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let report_context = report_context?;
|
||||
if !report_context
|
||||
.get("has_envelope")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Some(report_context.clone());
|
||||
}
|
||||
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();
|
||||
if provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format).is_none() {
|
||||
return Some(report_context.clone());
|
||||
}
|
||||
Some(clear_private_envelope_context(report_context))
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_provider_private_stream_normalizer(
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<ProviderPrivateStreamNormalizer> {
|
||||
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: report_context.clone(),
|
||||
buffered: Vec::new(),
|
||||
mode,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_provider_private_response_value(
|
||||
data: Value,
|
||||
report_context: &Value,
|
||||
) -> Result<Option<Value>, GatewayError> {
|
||||
if !report_context
|
||||
.get("has_envelope")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Ok(Some(data));
|
||||
}
|
||||
let mut unwrapped = match report_context.get("envelope_name").and_then(Value::as_str) {
|
||||
Some("claude:cli") | Some(KIRO_ENVELOPE_NAME) => data,
|
||||
Some(GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME) => {
|
||||
if let Some(response) = data
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.filter(|response| !response.contains_key("response"))
|
||||
{
|
||||
Value::Object(response.clone())
|
||||
} else {
|
||||
data
|
||||
}
|
||||
}
|
||||
Some(ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME) => {
|
||||
if let Some(response) = data
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.filter(|response| !response.contains_key("response"))
|
||||
{
|
||||
let mut unwrapped = response.clone();
|
||||
if let Some(response_id) = data.get("responseId").cloned() {
|
||||
unwrapped.insert("_v1internal_response_id".to_string(), response_id);
|
||||
}
|
||||
Value::Object(unwrapped)
|
||||
} else {
|
||||
data
|
||||
}
|
||||
}
|
||||
_ => return Ok(None),
|
||||
};
|
||||
postprocess_private_response_value(&mut unwrapped, report_context);
|
||||
Ok(Some(unwrapped))
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_normalize_provider_private_sync_report_payload(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<GatewaySyncReportRequest>, GatewayError> {
|
||||
let Some(report_context) = payload.report_context.as_ref() else {
|
||||
return Ok(Some(payload.clone()));
|
||||
};
|
||||
if !report_context
|
||||
.get("has_envelope")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Ok(Some(payload.clone()));
|
||||
}
|
||||
if !provider_private_response_allows_sync_finalize(report_context) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut normalized = payload.clone();
|
||||
normalized.report_context = normalize_provider_private_report_context(Some(report_context));
|
||||
|
||||
if let Some(body_json) = payload.body_json.clone() {
|
||||
normalized.body_json =
|
||||
normalize_provider_private_response_value(body_json, report_context)?;
|
||||
if normalized.body_json.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(body_base64) = payload.body_base64.as_deref() {
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let Some(normalized_bytes) =
|
||||
normalize_provider_private_stream_bytes(report_context, &body_bytes)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if stream_body_contains_error_event(&normalized_bytes) {
|
||||
return Ok(None);
|
||||
}
|
||||
normalized.body_base64 = (!normalized_bytes.is_empty())
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(normalized_bytes));
|
||||
}
|
||||
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
|
||||
pub(crate) fn transform_provider_private_stream_line(
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let Ok(text) = std::str::from_utf8(&line) else {
|
||||
return Ok(line);
|
||||
};
|
||||
let trimmed = text.trim_matches('\r').trim();
|
||||
if trimmed.is_empty() || trimmed.starts_with(':') || trimmed.starts_with("event:") {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(data_line) = trimmed.strip_prefix("data:") else {
|
||||
return Ok(line);
|
||||
};
|
||||
let data_line = data_line.trim();
|
||||
if data_line.is_empty() || data_line == "[DONE]" {
|
||||
return Ok(line);
|
||||
}
|
||||
|
||||
let body: Value = match serde_json::from_str(data_line) {
|
||||
Ok(value) => value,
|
||||
Err(_) => return Ok(line),
|
||||
};
|
||||
|
||||
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();
|
||||
if !provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format) {
|
||||
return Ok(line);
|
||||
}
|
||||
let unwrapped = match envelope_name {
|
||||
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME => body.get("response").cloned().unwrap_or(body),
|
||||
ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME => {
|
||||
let mut response = body.get("response").cloned().unwrap_or(body.clone());
|
||||
if let Some(response_id) = body.get("responseId").cloned() {
|
||||
if let Some(object) = response.as_object_mut() {
|
||||
object
|
||||
.entry("_v1internal_response_id".to_string())
|
||||
.or_insert(response_id);
|
||||
}
|
||||
}
|
||||
inject_antigravity_stream_tool_ids(&mut response);
|
||||
response
|
||||
}
|
||||
_ => body,
|
||||
};
|
||||
|
||||
let mut out = b"data: ".to_vec();
|
||||
out.extend(
|
||||
serde_json::to_vec(&unwrapped).map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
out.extend_from_slice(b"\n\n");
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
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,
|
||||
)?);
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn clear_private_envelope_context(report_context: &Value) -> Value {
|
||||
let mut normalized = report_context.clone();
|
||||
if let Some(object) = normalized.as_object_mut() {
|
||||
object.insert("has_envelope".to_string(), Value::Bool(false));
|
||||
object.remove("envelope_name");
|
||||
}
|
||||
normalized
|
||||
}
|
||||
|
||||
fn normalize_provider_private_stream_bytes(
|
||||
report_context: &Value,
|
||||
body: &[u8],
|
||||
) -> Result<Option<Vec<u8>>, GatewayError> {
|
||||
let Some(mut normalizer) = maybe_build_provider_private_stream_normalizer(Some(report_context))
|
||||
else {
|
||||
return Ok(Some(body.to_vec()));
|
||||
};
|
||||
let mut normalized = normalizer.push_chunk(body)?;
|
||||
normalized.extend(normalizer.finish()?);
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
|
||||
fn local_finalize_response_model(report_context: &Value) -> &str {
|
||||
report_context
|
||||
.get("mapped_model")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| report_context.get("model").and_then(Value::as_str))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn inject_antigravity_stream_tool_ids(value: &mut Value) {
|
||||
let Some(candidates) = value.get_mut("candidates").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
|
||||
for candidate in candidates {
|
||||
let Some(parts) = candidate
|
||||
.get_mut("content")
|
||||
.and_then(Value::as_object_mut)
|
||||
.and_then(|content| content.get_mut("parts"))
|
||||
.and_then(Value::as_array_mut)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let mut counters: BTreeMap<String, usize> = BTreeMap::new();
|
||||
for part in parts {
|
||||
let Some(function_call) = part.get_mut("functionCall").and_then(Value::as_object_mut)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let has_id = function_call
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if has_id {
|
||||
continue;
|
||||
}
|
||||
let name = function_call
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
let index = counters.entry(name.clone()).or_insert(0);
|
||||
function_call.insert(
|
||||
"id".to_string(),
|
||||
Value::String(format!("call_{name}_{index}")),
|
||||
);
|
||||
*index += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn inject_antigravity_sync_tool_ids(response: &mut Value, model: &str) {
|
||||
if !model.to_ascii_lowercase().contains("claude") {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(candidates) = response.get_mut("candidates").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
|
||||
for candidate in candidates {
|
||||
let Some(parts) = candidate
|
||||
.get_mut("content")
|
||||
.and_then(Value::as_object_mut)
|
||||
.and_then(|content| content.get_mut("parts"))
|
||||
.and_then(Value::as_array_mut)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let mut name_counters: BTreeMap<String, usize> = BTreeMap::new();
|
||||
for part in parts {
|
||||
let function_call = if let Some(function_call) =
|
||||
part.get_mut("functionCall").and_then(Value::as_object_mut)
|
||||
{
|
||||
function_call
|
||||
} else if let Some(function_call) =
|
||||
part.get_mut("function_call").and_then(Value::as_object_mut)
|
||||
{
|
||||
function_call
|
||||
} else {
|
||||
continue;
|
||||
};
|
||||
let has_id = function_call
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
if has_id {
|
||||
continue;
|
||||
}
|
||||
let function_name = function_call
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
let count = name_counters.entry(function_name.clone()).or_insert(0);
|
||||
function_call.insert(
|
||||
"id".to_string(),
|
||||
Value::String(format!("call_{function_name}_{count}")),
|
||||
);
|
||||
*count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn postprocess_private_response_value(data: &mut Value, report_context: &Value) {
|
||||
if !matches!(
|
||||
report_context.get("envelope_name").and_then(Value::as_str),
|
||||
Some(ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
if let Some(object) = data.as_object_mut() {
|
||||
if !object.contains_key("_v1internal_response_id") {
|
||||
if let Some(response_id) = object.remove("responseId") {
|
||||
object.insert("_v1internal_response_id".to_string(), response_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
inject_antigravity_sync_tool_ids(data, local_finalize_response_model(report_context));
|
||||
}
|
||||
|
||||
fn stream_body_contains_error_event(body: &[u8]) -> bool {
|
||||
let Ok(text) = std::str::from_utf8(body) else {
|
||||
return false;
|
||||
};
|
||||
let mut current_event_type: Option<String> = None;
|
||||
for raw_line in text.lines() {
|
||||
let line = raw_line.trim_matches('\r').trim();
|
||||
if line.is_empty() || line.starts_with(':') {
|
||||
continue;
|
||||
}
|
||||
if let Some(event_name) = line.strip_prefix("event:") {
|
||||
current_event_type = Some(event_name.trim().to_string());
|
||||
continue;
|
||||
}
|
||||
let data_line = if let Some(rest) = line.strip_prefix("data:") {
|
||||
rest.trim()
|
||||
} else {
|
||||
line
|
||||
};
|
||||
if data_line.is_empty() || data_line == "[DONE]" {
|
||||
continue;
|
||||
}
|
||||
let Ok(mut event) = serde_json::from_str::<Value>(data_line) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(event_object) = event.as_object_mut() {
|
||||
if !event_object.contains_key("type") {
|
||||
if let Some(event_name) = current_event_type.take() {
|
||||
event_object.insert("type".to_string(), Value::String(event_name));
|
||||
}
|
||||
}
|
||||
}
|
||||
if event
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("error"))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
current_event_type = None;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
#[path = "private_envelope/stream.rs"]
|
||||
mod stream;
|
||||
#[path = "private_envelope/sync.rs"]
|
||||
mod sync;
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
#[path = "private_envelope/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
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:cli",
|
||||
});
|
||||
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:cli",
|
||||
"client_api_format": "gemini:cli",
|
||||
"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\""));
|
||||
}
|
||||
}
|
||||
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,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
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 {
|
||||
report_context: Value,
|
||||
buffered: Vec<u8>,
|
||||
mode: ProviderPrivateStreamNormalizeMode,
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_provider_private_stream_normalizer(
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<ProviderPrivateStreamNormalizer> {
|
||||
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: report_context.clone(),
|
||||
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()))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
use base64::Engine as _;
|
||||
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,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_normalize_provider_private_sync_report_payload(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<GatewaySyncReportRequest>, GatewayError> {
|
||||
let Some(report_context) = payload.report_context.as_ref() else {
|
||||
return Ok(Some(payload.clone()));
|
||||
};
|
||||
if !report_context
|
||||
.get("has_envelope")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Ok(Some(payload.clone()));
|
||||
}
|
||||
if !provider_private_response_allows_sync_finalize(report_context) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut normalized = payload.clone();
|
||||
normalized.report_context = normalize_provider_private_report_context(Some(report_context));
|
||||
|
||||
if let Some(body_json) = payload.body_json.clone() {
|
||||
normalized.body_json = normalize_provider_private_response_value(body_json, report_context);
|
||||
if normalized.body_json.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(body_base64) = payload.body_base64.as_deref() {
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let Some(normalized_bytes) =
|
||||
normalize_provider_private_stream_bytes(report_context, &body_bytes)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if stream_body_contains_error_event(&normalized_bytes) {
|
||||
return Ok(None);
|
||||
}
|
||||
normalized.body_base64 = (!normalized_bytes.is_empty())
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(normalized_bytes));
|
||||
}
|
||||
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
|
||||
fn normalize_provider_private_stream_bytes(
|
||||
report_context: &Value,
|
||||
body: &[u8],
|
||||
) -> Result<Option<Vec<u8>>, GatewayError> {
|
||||
let Some(mut normalizer): Option<ProviderPrivateStreamNormalizer> =
|
||||
maybe_build_provider_private_stream_normalizer(Some(report_context))
|
||||
else {
|
||||
return Ok(Some(body.to_vec()));
|
||||
};
|
||||
let mut normalized = normalizer.push_chunk(body)?;
|
||||
normalized.extend(normalizer.finish()?);
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
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:cli",
|
||||
});
|
||||
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:cli",
|
||||
"client_api_format": "gemini:cli",
|
||||
"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,8 +0,0 @@
|
||||
pub(crate) use aether_ai_pipeline::adaptation::surfaces::{
|
||||
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,
|
||||
provider_adaptation_should_unwrap_stream_envelope, ProviderAdaptationDescriptor,
|
||||
ProviderAdaptationSurface, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
|
||||
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
|
||||
};
|
||||
@@ -1,128 +1,14 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, ProxySnapshot};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::control_facade::{
|
||||
collect_control_headers, resolve_execution_runtime_auth_context, GatewayControlAuthContext,
|
||||
GatewayControlDecision,
|
||||
use crate::ai_pipeline::{
|
||||
generic_decision_missing_exact_provider_request as generic_decision_missing_exact_provider_request_impl,
|
||||
GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct GatewayControlPlanRequest {
|
||||
pub(crate) trace_id: String,
|
||||
pub(crate) method: String,
|
||||
pub(crate) path: String,
|
||||
pub(crate) query_string: Option<String>,
|
||||
pub(crate) headers: BTreeMap<String, String>,
|
||||
pub(crate) body_json: serde_json::Value,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) body_base64: Option<String>,
|
||||
pub(crate) auth_context: Option<GatewayControlAuthContext>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub(crate) struct GatewayControlPlanResponse {
|
||||
pub(crate) action: String,
|
||||
#[serde(default)]
|
||||
pub(crate) plan_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) plan: Option<ExecutionPlan>,
|
||||
#[serde(default)]
|
||||
pub(crate) report_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) report_context: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub(crate) auth_context: Option<GatewayControlAuthContext>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub(crate) struct GatewayControlSyncDecisionResponse {
|
||||
pub(crate) action: String,
|
||||
#[serde(default)]
|
||||
pub(crate) decision_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) execution_strategy: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) conversion_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) request_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) candidate_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) endpoint_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) key_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) upstream_base_url: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) upstream_url: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_request_method: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) auth_header: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) auth_value: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_api_format: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) client_api_format: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_contract: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) client_contract: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) model_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) mapped_model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) prompt_cache_key: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) extra_headers: BTreeMap<String, String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_request_headers: BTreeMap<String, String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_request_body: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_request_body_base64: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) content_type: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) proxy: Option<ProxySnapshot>,
|
||||
#[serde(default)]
|
||||
pub(crate) tls_profile: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) timeouts: Option<ExecutionTimeouts>,
|
||||
#[serde(default)]
|
||||
pub(crate) upstream_is_stream: bool,
|
||||
#[serde(default)]
|
||||
pub(crate) report_kind: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) report_context: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub(crate) auth_context: Option<GatewayControlAuthContext>,
|
||||
}
|
||||
|
||||
fn decision_has_exact_provider_request(payload: &GatewayControlSyncDecisionResponse) -> bool {
|
||||
!payload.provider_request_headers.is_empty()
|
||||
&& (payload.provider_request_body.is_some()
|
||||
|| payload
|
||||
.provider_request_body_base64
|
||||
.as_ref()
|
||||
.map(|value| !value.trim().is_empty())
|
||||
.unwrap_or(false))
|
||||
}
|
||||
|
||||
pub(crate) fn generic_decision_missing_exact_provider_request(
|
||||
payload: &GatewayControlSyncDecisionResponse,
|
||||
) -> bool {
|
||||
if decision_has_exact_provider_request(payload) {
|
||||
if !generic_decision_missing_exact_provider_request_impl(payload) {
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -134,32 +20,3 @@ pub(crate) fn generic_decision_missing_exact_provider_request(
|
||||
);
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) async fn build_gateway_plan_request(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: serde_json::Value,
|
||||
body_base64: Option<String>,
|
||||
) -> Result<GatewayControlPlanRequest, GatewayError> {
|
||||
let auth_context = resolve_execution_runtime_auth_context(
|
||||
state,
|
||||
decision,
|
||||
&parts.headers,
|
||||
&parts.uri,
|
||||
trace_id,
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(GatewayControlPlanRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
method: parts.method.to_string(),
|
||||
path: parts.uri.path().to_string(),
|
||||
query_string: parts.uri.query().map(ToOwned::to_owned),
|
||||
headers: collect_control_headers(&parts.headers),
|
||||
body_json,
|
||||
body_base64,
|
||||
auth_context,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
pub(crate) mod control_payloads;
|
||||
|
||||
pub(crate) use aether_ai_pipeline::contracts::{
|
||||
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,
|
||||
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,
|
||||
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,
|
||||
@@ -32,7 +34,4 @@ pub(crate) use aether_ai_pipeline::contracts::{
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
pub(crate) use control_payloads::{
|
||||
build_gateway_plan_request, generic_decision_missing_exact_provider_request,
|
||||
GatewayControlPlanRequest, GatewayControlPlanResponse, GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
pub(crate) use control_payloads::generic_decision_missing_exact_provider_request;
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
use axum::http::Uri;
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(crate) use crate::control::{GatewayControlAuthContext, GatewayControlDecision};
|
||||
|
||||
pub(crate) async fn resolve_execution_runtime_auth_context(
|
||||
state: &AppState,
|
||||
decision: &GatewayControlDecision,
|
||||
headers: &http::HeaderMap,
|
||||
uri: &Uri,
|
||||
trace_id: &str,
|
||||
) -> Result<Option<GatewayControlAuthContext>, GatewayError> {
|
||||
crate::control::resolve_execution_runtime_auth_context(state, decision, headers, uri, trace_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn collect_control_headers(
|
||||
headers: &http::HeaderMap,
|
||||
) -> std::collections::BTreeMap<String, String> {
|
||||
crate::headers::collect_control_headers(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
|
||||
crate::headers::is_json_request(headers)
|
||||
}
|
||||
@@ -1,16 +1,70 @@
|
||||
pub(crate) mod registry;
|
||||
pub(crate) mod request;
|
||||
pub(crate) mod response;
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) use aether_ai_pipeline::conversion::core_success_background_report_kind;
|
||||
pub(crate) use aether_ai_pipeline::conversion::{
|
||||
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 registry::{
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
request_conversion_direct_auth, request_conversion_kind,
|
||||
request_conversion_transport_supported, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
|
||||
SyncCliResponseConversionKind,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
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("claude:chat", "openai:chat"),
|
||||
Some(RequestConversionKind::ToOpenAIChat)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:chat", "claude:chat"),
|
||||
Some(RequestConversionKind::ToClaudeStandard)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:cli", "openai:compact"),
|
||||
Some(RequestConversionKind::ToOpenAICompact)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:compact", "gemini:cli"),
|
||||
Some(RequestConversionKind::ToGeminiStandard)
|
||||
);
|
||||
assert_eq!(request_conversion_kind("claude:chat", "claude:chat"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_response_conversion_registry_supports_bidirectional_standard_matrix() {
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("openai:chat", "claude:chat"),
|
||||
Some(SyncChatResponseConversionKind::ToClaudeChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("claude:chat", "gemini:chat"),
|
||||
Some(SyncChatResponseConversionKind::ToGeminiChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("gemini:chat", "openai:chat"),
|
||||
Some(SyncChatResponseConversionKind::ToOpenAIChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("openai:cli", "gemini:cli"),
|
||||
Some(SyncCliResponseConversionKind::ToGeminiCli)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("claude:cli", "openai:compact"),
|
||||
Some(SyncCliResponseConversionKind::ToOpenAIFamilyCli)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("gemini:cli", "claude:cli"),
|
||||
Some(SyncCliResponseConversionKind::ToClaudeCli)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,116 +0,0 @@
|
||||
use crate::ai_pipeline::provider_transport_facade::auth::{
|
||||
resolve_local_gemini_auth, resolve_local_openai_chat_auth, resolve_local_standard_auth,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::policy::{
|
||||
supports_local_openai_chat_transport, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::{
|
||||
supports_local_gemini_transport_with_network, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
pub(crate) use aether_ai_pipeline::conversion::{
|
||||
request_conversion_kind, sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
|
||||
RequestConversionKind, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
|
||||
};
|
||||
|
||||
pub(crate) fn request_conversion_transport_supported(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
_kind: RequestConversionKind,
|
||||
) -> bool {
|
||||
match transport
|
||||
.endpoint
|
||||
.api_format
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
.as_str()
|
||||
{
|
||||
"openai:chat" => supports_local_openai_chat_transport(transport),
|
||||
"openai:cli" => supports_local_standard_transport_with_network(transport, "openai:cli"),
|
||||
"openai:compact" => {
|
||||
supports_local_standard_transport_with_network(transport, "openai:compact")
|
||||
}
|
||||
"claude:chat" => supports_local_standard_transport_with_network(transport, "claude:chat"),
|
||||
"claude:cli" => supports_local_standard_transport_with_network(transport, "claude:cli"),
|
||||
"gemini:chat" => supports_local_gemini_transport_with_network(transport, "gemini:chat"),
|
||||
"gemini:cli" => supports_local_gemini_transport_with_network(transport, "gemini:cli"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn request_conversion_direct_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
_kind: RequestConversionKind,
|
||||
) -> Option<(String, String)> {
|
||||
match transport
|
||||
.endpoint
|
||||
.api_format
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
.as_str()
|
||||
{
|
||||
"openai:chat" => resolve_local_openai_chat_auth(transport),
|
||||
"gemini:chat" | "gemini:cli" => resolve_local_gemini_auth(transport),
|
||||
"openai:cli" | "openai:compact" | "claude:chat" | "claude:cli" => {
|
||||
resolve_local_standard_auth(transport)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
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("claude:chat", "openai:chat"),
|
||||
Some(RequestConversionKind::ToOpenAIChat)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:chat", "claude:chat"),
|
||||
Some(RequestConversionKind::ToClaudeStandard)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:cli", "openai:compact"),
|
||||
Some(RequestConversionKind::ToOpenAICompact)
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:compact", "gemini:cli"),
|
||||
Some(RequestConversionKind::ToGeminiStandard)
|
||||
);
|
||||
assert_eq!(request_conversion_kind("claude:chat", "claude:chat"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_response_conversion_registry_supports_bidirectional_standard_matrix() {
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("openai:chat", "claude:chat"),
|
||||
Some(SyncChatResponseConversionKind::ToClaudeChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("claude:chat", "gemini:chat"),
|
||||
Some(SyncChatResponseConversionKind::ToGeminiChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_chat_response_conversion_kind("gemini:chat", "openai:chat"),
|
||||
Some(SyncChatResponseConversionKind::ToOpenAIChat)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("openai:cli", "gemini:cli"),
|
||||
Some(SyncCliResponseConversionKind::ToGeminiCli)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("claude:cli", "openai:compact"),
|
||||
Some(SyncCliResponseConversionKind::ToOpenAIFamilyCli)
|
||||
);
|
||||
assert_eq!(
|
||||
sync_cli_response_conversion_kind("gemini:cli", "claude:cli"),
|
||||
Some(SyncCliResponseConversionKind::ToClaudeCli)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
pub(crate) use aether_ai_pipeline::conversion::request::{
|
||||
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
|
||||
convert_openai_chat_request_to_openai_cli_request, extract_openai_text_content,
|
||||
normalize_claude_request_to_openai_chat_request,
|
||||
normalize_gemini_request_to_openai_chat_request,
|
||||
normalize_openai_cli_request_to_openai_chat_request, parse_openai_tool_result_content,
|
||||
};
|
||||
@@ -1,7 +0,0 @@
|
||||
pub(crate) use aether_ai_pipeline::conversion::response::{
|
||||
build_openai_cli_response, convert_claude_chat_response_to_openai_chat,
|
||||
convert_claude_cli_response_to_openai_cli, convert_gemini_chat_response_to_openai_chat,
|
||||
convert_gemini_cli_response_to_openai_cli, convert_openai_chat_response_to_claude_chat,
|
||||
convert_openai_chat_response_to_gemini_chat, convert_openai_chat_response_to_openai_cli,
|
||||
convert_openai_cli_response_to_openai_chat,
|
||||
};
|
||||
@@ -1,14 +0,0 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
|
||||
|
||||
pub(crate) fn maybe_build_local_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &crate::ai_pipeline::control_facade::GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::execution_runtime::maybe_build_local_sync_finalize_response(trace_id, decision, payload)
|
||||
}
|
||||
@@ -4,12 +4,18 @@ use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use serde_json::Value;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::adaptation::private_envelope::{
|
||||
use crate::ai_pipeline::{
|
||||
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,
|
||||
};
|
||||
use crate::ai_pipeline::contracts::core_success_background_report_kind;
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::api::response::build_client_response_from_parts;
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
@@ -26,7 +32,7 @@ pub(crate) fn build_local_success_outcome(
|
||||
) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
|
||||
let headers = payload.headers.clone();
|
||||
let background_report =
|
||||
map_local_finalize_to_success_report(payload, body_json.clone(), headers.clone());
|
||||
build_local_success_background_report_impl(payload, body_json.clone(), headers.clone());
|
||||
build_local_success_outcome_with_report(
|
||||
trace_id,
|
||||
decision,
|
||||
@@ -42,15 +48,11 @@ pub(crate) fn build_local_success_outcome_with_report(
|
||||
decision: &GatewayControlDecision,
|
||||
status_code: u16,
|
||||
body_json: Value,
|
||||
mut headers: BTreeMap<String, String>,
|
||||
headers: BTreeMap<String, String>,
|
||||
background_report: Option<GatewaySyncReportRequest>,
|
||||
) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
|
||||
headers.remove("content-encoding");
|
||||
headers.remove("content-length");
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
let body_bytes =
|
||||
serde_json::to_vec(&body_json).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
let (body_bytes, headers) = prepare_local_success_response_parts_impl(&headers, &body_json)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let response = build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
@@ -71,30 +73,11 @@ pub(crate) fn build_local_success_outcome_with_conversion_report(
|
||||
client_body_json: Value,
|
||||
provider_body_json: Value,
|
||||
) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
|
||||
let Some(report_kind) =
|
||||
map_local_finalize_kind_to_success_report_kind(payload.report_kind.as_str())
|
||||
else {
|
||||
return build_local_success_outcome_with_report(
|
||||
trace_id,
|
||||
decision,
|
||||
payload.status_code,
|
||||
client_body_json,
|
||||
payload.headers.clone(),
|
||||
None,
|
||||
);
|
||||
};
|
||||
|
||||
let report_payload = GatewaySyncReportRequest {
|
||||
trace_id: payload.trace_id.clone(),
|
||||
report_kind: report_kind.to_string(),
|
||||
report_context: payload.report_context.clone(),
|
||||
status_code: payload.status_code,
|
||||
headers: payload.headers.clone(),
|
||||
body_json: Some(provider_body_json),
|
||||
client_body_json: Some(client_body_json.clone()),
|
||||
body_base64: None,
|
||||
telemetry: payload.telemetry.clone(),
|
||||
};
|
||||
let report_payload = build_local_success_conversion_background_report_impl(
|
||||
payload,
|
||||
client_body_json.clone(),
|
||||
provider_body_json,
|
||||
);
|
||||
|
||||
build_local_success_outcome_with_report(
|
||||
trace_id,
|
||||
@@ -102,42 +85,6 @@ pub(crate) fn build_local_success_outcome_with_conversion_report(
|
||||
payload.status_code,
|
||||
client_body_json,
|
||||
payload.headers.clone(),
|
||||
Some(report_payload),
|
||||
report_payload,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_local_finalize_to_success_report(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
body_json: Value,
|
||||
headers: BTreeMap<String, String>,
|
||||
) -> Option<GatewaySyncReportRequest> {
|
||||
let report_kind = map_local_finalize_kind_to_success_report_kind(payload.report_kind.as_str())?;
|
||||
|
||||
Some(GatewaySyncReportRequest {
|
||||
trace_id: payload.trace_id.clone(),
|
||||
report_kind: report_kind.to_string(),
|
||||
report_context: payload.report_context.clone(),
|
||||
status_code: payload.status_code,
|
||||
headers,
|
||||
body_json: Some(body_json),
|
||||
client_body_json: None,
|
||||
body_base64: None,
|
||||
telemetry: payload.telemetry.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn map_local_finalize_kind_to_success_report_kind(report_kind: &str) -> Option<&'static str> {
|
||||
core_success_background_report_kind(report_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn canonicalize_tool_arguments(value: Option<Value>) -> String {
|
||||
match value {
|
||||
Some(Value::String(text)) => text,
|
||||
Some(other) => serde_json::to_string(&other).unwrap_or_else(|_| "null".to_string()),
|
||||
None => "{}".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_generated_tool_call_id(index: usize) -> String {
|
||||
format!("call_auto_{index}")
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
#[path = "stream_rewrite.rs"]
|
||||
|
||||
@@ -1,12 +1,10 @@
|
||||
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::standard::StreamingStandardConversionState;
|
||||
use crate::ai_pipeline::runtime::adapters::kiro::KiroToClaudeCliStreamState;
|
||||
use crate::ai_pipeline::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
|
||||
use crate::GatewayError;
|
||||
use aether_ai_pipeline::finalize::{
|
||||
resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode,
|
||||
};
|
||||
|
||||
enum RewriteMode {
|
||||
EnvelopeUnwrap,
|
||||
@@ -81,7 +79,8 @@ impl LocalStreamRewriter {
|
||||
|
||||
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),
|
||||
RewriteMode::EnvelopeUnwrap => transform_envelope_line(&self.report_context, line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string())),
|
||||
RewriteMode::Standard(state) => state.transform_line(&self.report_context, line),
|
||||
RewriteMode::KiroToClaudeCli(_) => Ok(Vec::new()),
|
||||
}
|
||||
|
||||
@@ -1,13 +1,9 @@
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{build_generated_tool_call_id, canonicalize_tool_arguments};
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
|
||||
pub(crate) use crate::ai_pipeline::conversion::response::{
|
||||
convert_claude_chat_response_to_openai_chat, convert_claude_cli_response_to_openai_cli,
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_cli_response_to_openai_cli,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::finalize::common::{
|
||||
build_generated_tool_call_id, build_local_success_outcome,
|
||||
build_local_success_outcome_with_conversion_report, canonicalize_tool_arguments,
|
||||
build_local_success_outcome, build_local_success_outcome_with_conversion_report,
|
||||
local_finalize_allows_envelope, unwrap_local_finalize_response_value,
|
||||
LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
@@ -15,10 +11,14 @@ pub(crate) use crate::ai_pipeline::finalize::standard::{
|
||||
maybe_build_standard_sync_finalize_product_from_normalized_payload,
|
||||
StandardSyncFinalizeNormalizedProduct,
|
||||
};
|
||||
pub(crate) use aether_ai_pipeline::finalize::sync_products::{
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
|
||||
aggregate_openai_chat_stream_sync_response, aggregate_openai_cli_stream_sync_response,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
convert_claude_chat_response_to_openai_chat, convert_claude_cli_response_to_openai_cli,
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_cli_response_to_openai_cli,
|
||||
};
|
||||
|
||||
pub(crate) fn maybe_build_local_core_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
@@ -51,7 +51,7 @@ pub(crate) fn maybe_build_local_core_sync_finalize_response(
|
||||
|
||||
match product {
|
||||
StandardSyncFinalizeNormalizedProduct::SuccessBody(body_json) => {
|
||||
let Some(body_json) = unwrap_local_finalize_response_value(body_json, report_context)?
|
||||
let Some(body_json) = unwrap_local_finalize_response_value(body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -61,7 +61,7 @@ pub(crate) fn maybe_build_local_core_sync_finalize_response(
|
||||
}
|
||||
StandardSyncFinalizeNormalizedProduct::CrossFormat(product) => {
|
||||
let Some(provider_body_json) =
|
||||
unwrap_local_finalize_response_value(product.provider_body_json, report_context)?
|
||||
unwrap_local_finalize_response_value(product.provider_body_json, report_context)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
@@ -1,12 +1,4 @@
|
||||
pub(crate) mod common;
|
||||
pub(crate) mod internal;
|
||||
pub(crate) mod sse;
|
||||
pub(crate) mod standard;
|
||||
pub(crate) use crate::ai_pipeline::execution_facade::maybe_build_local_sync_finalize_response;
|
||||
pub(crate) use crate::api::response::{build_client_response, build_client_response_from_parts};
|
||||
pub(crate) use common::build_local_success_outcome;
|
||||
pub(crate) use internal::{
|
||||
maybe_build_stream_response_rewriter, maybe_build_sync_finalize_outcome,
|
||||
maybe_compile_sync_finalize_response, LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
|
||||
pub(crate) mod internal;
|
||||
|
||||
@@ -1,14 +1,15 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::{
|
||||
encode_done_sse, encode_json_sse as encode_json_sse_impl, map_claude_stop_reason,
|
||||
PipelineFinalizeError,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_ai_pipeline::finalize::{self, PipelineFinalizeError};
|
||||
|
||||
pub(crate) use finalize::sse::{encode_done_sse, map_claude_stop_reason};
|
||||
|
||||
fn map_error(err: PipelineFinalizeError) -> GatewayError {
|
||||
err.into()
|
||||
}
|
||||
|
||||
pub(crate) fn encode_json_sse(event: Option<&str>, value: &Value) -> Result<Vec<u8>, GatewayError> {
|
||||
finalize::sse::encode_json_sse(event, value).map_err(map_error)
|
||||
encode_json_sse_impl(event, value).map_err(map_error)
|
||||
}
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
pub(super) mod stream;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::conversion::response::{
|
||||
convert_claude_chat_response_to_openai_chat, convert_claude_cli_response_to_openai_cli,
|
||||
convert_openai_chat_response_to_claude_chat,
|
||||
};
|
||||
@@ -1,3 +0,0 @@
|
||||
pub(crate) use aether_ai_pipeline::finalize::standard::claude::stream::{
|
||||
ClaudeClientEmitter, ClaudeProviderState,
|
||||
};
|
||||
@@ -1,6 +0,0 @@
|
||||
pub(super) mod stream;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::conversion::response::{
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_cli_response_to_openai_cli,
|
||||
convert_openai_chat_response_to_gemini_chat,
|
||||
};
|
||||
@@ -1,3 +0,0 @@
|
||||
pub(crate) use aether_ai_pipeline::finalize::standard::gemini::stream::{
|
||||
GeminiClientEmitter, GeminiProviderState,
|
||||
};
|
||||
@@ -1,18 +1,16 @@
|
||||
//! Standard finalize surface for standard contract sync/stream compilation.
|
||||
|
||||
mod claude;
|
||||
mod gemini;
|
||||
mod openai;
|
||||
#[path = "stream_core/mod.rs"]
|
||||
mod stream;
|
||||
|
||||
pub(crate) use crate::ai_pipeline::conversion::response::{
|
||||
build_openai_cli_response, convert_openai_chat_response_to_openai_cli,
|
||||
convert_openai_cli_response_to_openai_chat,
|
||||
};
|
||||
pub(crate) use aether_ai_pipeline::finalize::sync_products::{
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response,
|
||||
convert_standard_chat_response, convert_standard_cli_response,
|
||||
build_openai_cli_response, convert_claude_chat_response_to_openai_chat,
|
||||
convert_claude_cli_response_to_openai_cli, convert_gemini_chat_response_to_openai_chat,
|
||||
convert_gemini_cli_response_to_openai_cli, convert_openai_chat_response_to_claude_chat,
|
||||
convert_openai_chat_response_to_gemini_chat, convert_openai_chat_response_to_openai_cli,
|
||||
convert_openai_cli_response_to_openai_chat, convert_standard_chat_response,
|
||||
convert_standard_cli_response,
|
||||
maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload,
|
||||
maybe_build_openai_cli_cross_format_sync_product_from_normalized_payload,
|
||||
maybe_build_openai_cli_same_family_sync_body_from_normalized_payload,
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
pub(super) mod stream;
|
||||
@@ -1,4 +0,0 @@
|
||||
pub(crate) use aether_ai_pipeline::finalize::standard::openai::stream::{
|
||||
OpenAIChatClientEmitter, OpenAIChatProviderState, OpenAICliClientEmitter,
|
||||
OpenAICliProviderState,
|
||||
};
|
||||
@@ -1 +0,0 @@
|
||||
pub(crate) use aether_ai_pipeline::finalize::standard::stream_core::common::*;
|
||||
@@ -1,7 +1,6 @@
|
||||
//! Standard finalize streaming conversion helpers.
|
||||
pub use aether_ai_pipeline::finalize::standard::stream_core::CanonicalStreamFrame;
|
||||
pub(crate) use crate::ai_pipeline::CanonicalStreamFrame;
|
||||
|
||||
pub(crate) mod common;
|
||||
mod orchestrator;
|
||||
|
||||
pub(crate) use orchestrator::StreamingStandardConversionState;
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
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::surfaces::provider_adaptation_should_unwrap_stream_envelope;
|
||||
use crate::ai_pipeline::{
|
||||
provider_adaptation_should_unwrap_stream_envelope, StreamingStandardFormatMatrix,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_ai_pipeline::finalize::standard::stream_core::StreamingStandardFormatMatrix;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct StreamingStandardConversionState {
|
||||
@@ -17,7 +18,8 @@ impl StreamingStandardConversionState {
|
||||
line: Vec<u8>,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
let line = if should_unwrap_envelope(report_context) {
|
||||
transform_envelope_line(report_context, line)?
|
||||
transform_envelope_line(report_context, line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
} else {
|
||||
line
|
||||
};
|
||||
|
||||
@@ -10,8 +10,8 @@ use super::{
|
||||
convert_gemini_chat_response_to_openai_chat, convert_gemini_cli_response_to_openai_cli,
|
||||
maybe_build_local_core_sync_finalize_response,
|
||||
};
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::conversion::response::{
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
convert_openai_chat_response_to_openai_cli, convert_openai_cli_response_to_openai_chat,
|
||||
};
|
||||
use crate::usage::GatewaySyncReportRequest;
|
||||
|
||||
@@ -1,9 +1,99 @@
|
||||
pub(crate) mod adaptation;
|
||||
pub(crate) mod contracts;
|
||||
pub(crate) mod control_facade;
|
||||
pub(crate) mod conversion;
|
||||
pub(crate) mod execution_facade;
|
||||
pub(crate) mod finalize;
|
||||
pub(crate) mod planner;
|
||||
pub(crate) mod provider_transport_facade;
|
||||
pub(crate) mod runtime;
|
||||
mod adaptation;
|
||||
mod contracts;
|
||||
mod conversion;
|
||||
mod finalize;
|
||||
mod planner;
|
||||
mod pure;
|
||||
pub(crate) mod transport;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Response, Uri};
|
||||
|
||||
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
||||
|
||||
use self::contracts::ExecutionRuntimeAuthContext;
|
||||
|
||||
pub(crate) use self::adaptation::maybe_build_provider_private_stream_normalizer;
|
||||
pub(crate) use self::finalize::common::LocalCoreSyncFinalizeOutcome;
|
||||
pub(crate) use self::finalize::internal::{
|
||||
maybe_build_stream_response_rewriter, maybe_build_sync_finalize_outcome,
|
||||
maybe_compile_sync_finalize_response,
|
||||
};
|
||||
pub(crate) use self::planner::{
|
||||
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,
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_cli_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_cli_sync_plan_and_reports_for_kind,
|
||||
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
|
||||
build_local_video_sync_plan_and_reports_for_kind, build_openai_cli_stream_plan_from_decision,
|
||||
build_openai_cli_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
build_standard_family_stream_plan_and_reports, build_standard_family_sync_plan_and_reports,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
|
||||
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic, GatewayAuthApiKeySnapshot,
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
pub(crate) use self::pure::*;
|
||||
pub(crate) use crate::control::GatewayControlDecision;
|
||||
pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
|
||||
|
||||
pub(crate) async fn resolve_execution_runtime_auth_context(
|
||||
state: &AppState,
|
||||
decision: &GatewayControlDecision,
|
||||
headers: &http::HeaderMap,
|
||||
uri: &Uri,
|
||||
trace_id: &str,
|
||||
) -> Result<Option<crate::control::GatewayControlAuthContext>, GatewayError> {
|
||||
crate::control::resolve_execution_runtime_auth_context(state, decision, headers, uri, trace_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn collect_control_headers(
|
||||
headers: &http::HeaderMap,
|
||||
) -> std::collections::BTreeMap<String, String> {
|
||||
crate::headers::collect_control_headers(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
|
||||
crate::headers::is_json_request(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn build_execution_runtime_auth_context(
|
||||
auth_context: &crate::control::GatewayControlAuthContext,
|
||||
) -> ExecutionRuntimeAuthContext {
|
||||
ExecutionRuntimeAuthContext {
|
||||
user_id: auth_context.user_id.clone(),
|
||||
api_key_id: auth_context.api_key_id.clone(),
|
||||
balance_remaining: auth_context.balance_remaining,
|
||||
access_allowed: auth_context.access_allowed,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_decision_execution_runtime_auth_context(
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<ExecutionRuntimeAuthContext> {
|
||||
decision
|
||||
.auth_context
|
||||
.as_ref()
|
||||
.map(build_execution_runtime_auth_context)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_local_decision_execution_runtime_auth_context(
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<ExecutionRuntimeAuthContext> {
|
||||
resolve_decision_execution_runtime_auth_context(decision).filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::execution_runtime::maybe_build_local_sync_finalize_response(trace_id, decision, payload)
|
||||
}
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
pub(crate) use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(crate) async fn read_auth_api_key_snapshot(
|
||||
state: &AppState,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<GatewayAuthApiKeySnapshot>, GatewayError> {
|
||||
state
|
||||
.data
|
||||
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::transport_facade::read_provider_transport_snapshot;
|
||||
use crate::ai_pipeline::provider_transport_facade::resolve_transport_proxy_snapshot;
|
||||
use crate::ai_pipeline::transport::resolve_transport_proxy_snapshot;
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::AppState;
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
@@ -13,7 +13,7 @@ enum TunnelOwnerAffinityBucket {
|
||||
}
|
||||
|
||||
pub(crate) async fn prefer_local_tunnel_owner_candidates(
|
||||
state: &AppState,
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
|
||||
let mut ranked = Vec::with_capacity(candidates.len());
|
||||
@@ -29,16 +29,16 @@ pub(crate) async fn prefer_local_tunnel_owner_candidates(
|
||||
}
|
||||
|
||||
async fn resolve_candidate_tunnel_owner_affinity(
|
||||
state: &AppState,
|
||||
state: PlannerAppState<'_>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> TunnelOwnerAffinityBucket {
|
||||
let transport = match read_provider_transport_snapshot(
|
||||
state,
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(transport)) => transport,
|
||||
Ok(None) => return TunnelOwnerAffinityBucket::Neutral,
|
||||
@@ -71,16 +71,17 @@ async fn resolve_candidate_tunnel_owner_affinity(
|
||||
return TunnelOwnerAffinityBucket::Neutral;
|
||||
};
|
||||
|
||||
if state.tunnel.has_local_proxy(node_id) {
|
||||
if state.app().tunnel.has_local_proxy(node_id) {
|
||||
return TunnelOwnerAffinityBucket::LocalTunnel;
|
||||
}
|
||||
|
||||
match state
|
||||
.app()
|
||||
.tunnel
|
||||
.lookup_attachment_owner(state.data.as_ref(), node_id)
|
||||
.lookup_attachment_owner(state.app().data.as_ref(), node_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some(owner)) if owner.gateway_instance_id == state.tunnel.local_instance_id() => {
|
||||
Ok(Some(owner)) if owner.gateway_instance_id == state.app().tunnel.local_instance_id() => {
|
||||
TunnelOwnerAffinityBucket::LocalTunnel
|
||||
}
|
||||
Ok(Some(_)) => TunnelOwnerAffinityBucket::RemoteTunnel,
|
||||
@@ -109,7 +110,8 @@ mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
prefer_local_tunnel_owner_candidates, AppState, SchedulerMinimalCandidateSelectionCandidate,
|
||||
prefer_local_tunnel_owner_candidates, AppState, PlannerAppState,
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::tunnel::TunnelAttachmentRecord;
|
||||
@@ -258,7 +260,7 @@ mod tests {
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
|
||||
|
||||
let reordered = prefer_local_tunnel_owner_candidates(
|
||||
&state,
|
||||
PlannerAppState::new(&state),
|
||||
vec![
|
||||
sample_candidate("endpoint-remote", "key-remote"),
|
||||
sample_candidate("endpoint-local", "key-local"),
|
||||
@@ -287,7 +289,7 @@ mod tests {
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
|
||||
|
||||
let reordered = prefer_local_tunnel_owner_candidates(
|
||||
&state,
|
||||
PlannerAppState::new(&state),
|
||||
vec![
|
||||
sample_candidate("endpoint-a", "key-a"),
|
||||
sample_candidate("endpoint-b", "key-b"),
|
||||
|
||||
@@ -1,60 +0,0 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn persist_available_local_candidate(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
extra_data: Option<Value>,
|
||||
created_at_unix_secs: u64,
|
||||
error_context: &'static str,
|
||||
) -> String {
|
||||
crate::request_candidate_runtime::persist_available_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
extra_data,
|
||||
created_at_unix_secs,
|
||||
error_context,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn persist_skipped_local_candidate(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &str,
|
||||
finished_at_unix_secs: u64,
|
||||
error_context: &'static str,
|
||||
) {
|
||||
crate::request_candidate_runtime::persist_skipped_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
finished_at_unix_secs,
|
||||
error_context,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -1,7 +1,6 @@
|
||||
use axum::body::Bytes;
|
||||
|
||||
use crate::ai_pipeline::control_facade::is_json_request;
|
||||
pub(crate) use aether_ai_pipeline::contracts::{
|
||||
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,
|
||||
@@ -15,25 +14,23 @@ pub(crate) use aether_ai_pipeline::contracts::{
|
||||
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::{
|
||||
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) fn parse_direct_request_body(
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &Bytes,
|
||||
) -> Option<(serde_json::Value, Option<String>)> {
|
||||
aether_ai_pipeline::planner::common::parse_direct_request_body(
|
||||
is_json_request(&parts.headers),
|
||||
body_bytes.as_ref(),
|
||||
)
|
||||
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 {
|
||||
provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& provider_api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("openai:cli")
|
||||
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use crate::ai_pipeline::control_facade::{GatewayControlAuthContext, GatewayControlDecision};
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::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,
|
||||
@@ -19,6 +19,11 @@ use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::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,
|
||||
};
|
||||
@@ -32,7 +37,7 @@ pub(crate) async fn maybe_build_sync_plan_payload_impl(
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
let Some(plan_kind) = super::resolve_sync_plan_kind(parts, decision) else {
|
||||
let Some(plan_kind) = resolve_sync_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(payload) = super::maybe_build_sync_decision_payload(
|
||||
@@ -59,7 +64,7 @@ pub(crate) async fn maybe_build_stream_plan_payload_impl(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
let Some(plan_kind) = super::resolve_stream_plan_kind(parts, decision) else {
|
||||
let Some(plan_kind) = resolve_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(payload) =
|
||||
@@ -147,7 +152,7 @@ fn build_stream_plan_payload_from_decision(
|
||||
fn build_sync_plan_response(
|
||||
plan_kind: &str,
|
||||
value: LocalSyncPlanAndReport,
|
||||
auth_context: Option<GatewayControlAuthContext>,
|
||||
auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
) -> GatewayControlPlanResponse {
|
||||
GatewayControlPlanResponse {
|
||||
action: EXECUTION_RUNTIME_SYNC_ACTION.to_string(),
|
||||
@@ -162,7 +167,7 @@ fn build_sync_plan_response(
|
||||
fn build_stream_plan_response(
|
||||
plan_kind: &str,
|
||||
value: LocalStreamPlanAndReport,
|
||||
auth_context: Option<GatewayControlAuthContext>,
|
||||
auth_context: Option<ExecutionRuntimeAuthContext>,
|
||||
) -> GatewayControlPlanResponse {
|
||||
GatewayControlPlanResponse {
|
||||
action: EXECUTION_RUNTIME_STREAM_ACTION.to_string(),
|
||||
|
||||
@@ -23,7 +23,3 @@ pub(crate) use super::standard::{
|
||||
maybe_build_sync_local_openai_cli_decision_payload,
|
||||
maybe_build_sync_local_standard_decision_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::planner::{
|
||||
resolve_execution_runtime_stream_plan_kind as resolve_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind as resolve_sync_plan_kind,
|
||||
};
|
||||
|
||||
@@ -1,13 +1,15 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::planner::{
|
||||
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(
|
||||
@@ -151,6 +153,6 @@ async fn maybe_build_local_video_task_content_stream_decision_payload(
|
||||
upstream_is_stream: true,
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
auth_context: decision.auth_context.clone(),
|
||||
auth_context: resolve_decision_execution_runtime_auth_context(decision),
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -2,17 +2,17 @@ use std::collections::BTreeMap;
|
||||
|
||||
use url::Url;
|
||||
|
||||
use crate::ai_pipeline::control_facade::{
|
||||
resolve_execution_runtime_auth_context, GatewayControlDecision,
|
||||
};
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
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::resolve_execution_runtime_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(
|
||||
@@ -190,7 +190,7 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
upstream_is_stream: false,
|
||||
report_kind: follow_up.report_kind,
|
||||
report_context: follow_up.report_context,
|
||||
auth_context: Some(auth_context),
|
||||
auth_context: Some(build_execution_runtime_auth_context(&auth_context)),
|
||||
}))
|
||||
}
|
||||
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) async fn mark_unused_local_candidate_items<T, FPlan, FContext>(
|
||||
state: &AppState,
|
||||
remaining: Vec<T>,
|
||||
plan: FPlan,
|
||||
report_context: FContext,
|
||||
) where
|
||||
FPlan: Fn(&T) -> &ExecutionPlan,
|
||||
FContext: Fn(&T) -> Option<&Value>,
|
||||
{
|
||||
crate::executor::mark_unused_local_candidate_items(state, remaining, plan, report_context).await
|
||||
}
|
||||
@@ -1,27 +1,45 @@
|
||||
use crate::ai_pipeline::contracts::{
|
||||
GatewayControlPlanResponse, GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(crate) mod auth_snapshot_facade;
|
||||
pub(crate) mod candidate_affinity;
|
||||
pub(crate) mod candidate_runtime_facade;
|
||||
pub(crate) mod common;
|
||||
mod candidate_affinity;
|
||||
mod common;
|
||||
mod decision;
|
||||
pub(crate) mod executor_facade;
|
||||
pub(crate) mod passthrough;
|
||||
pub(crate) mod plan_builders;
|
||||
mod passthrough;
|
||||
mod plan_builders;
|
||||
mod route;
|
||||
pub(crate) mod scheduler_facade;
|
||||
pub(crate) mod specialized;
|
||||
pub(crate) mod standard;
|
||||
pub(crate) mod transport_facade;
|
||||
mod specialized;
|
||||
mod standard;
|
||||
mod state;
|
||||
|
||||
pub(crate) use self::route::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
|
||||
supports_sync_scheduler_decision_kind,
|
||||
pub(crate) use self::passthrough::{
|
||||
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
|
||||
};
|
||||
pub(crate) use self::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_openai_cli_stream_plan_from_decision, build_openai_cli_sync_plan_from_decision,
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
build_standard_sync_plan_from_decision, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
pub(crate) use self::specialized::{
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_video_sync_plan_and_reports_for_kind,
|
||||
};
|
||||
pub(crate) use self::standard::{
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_cli_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_cli_sync_plan_and_reports_for_kind,
|
||||
build_local_stream_plan_and_reports as build_standard_family_stream_plan_and_reports,
|
||||
build_local_sync_plan_and_reports as build_standard_family_sync_plan_and_reports,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use self::state::{
|
||||
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||
PlannerAppState,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_decision_payload(
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
//! Requests that can stay in the same public/provider contract family.
|
||||
|
||||
pub(crate) mod provider;
|
||||
mod provider;
|
||||
|
||||
pub(crate) use self::provider::{
|
||||
build_local_stream_plan_and_reports as build_local_same_format_stream_plan_and_reports,
|
||||
build_local_sync_plan_and_reports as build_local_same_format_sync_plan_and_reports,
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::provider_transport_facade::provider_types::provider_type_supports_local_same_format_transport;
|
||||
pub(crate) use crate::ai_pipeline::transport::provider_types::provider_type_supports_local_same_format_transport;
|
||||
|
||||
@@ -11,8 +11,6 @@ use serde_json::{json, Value};
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::control_facade::{collect_control_headers, GatewayControlDecision};
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
@@ -20,48 +18,51 @@ use crate::ai_pipeline::planner::common::{
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::antigravity::{
|
||||
use crate::ai_pipeline::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::provider_transport_facade::auth::{
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_openai_passthrough_headers, resolve_local_gemini_auth, resolve_local_standard_auth,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::claude_code::{
|
||||
use crate::ai_pipeline::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::provider_transport_facade::kiro::{
|
||||
use crate::ai_pipeline::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::provider_transport_facade::policy::{
|
||||
use crate::ai_pipeline::transport::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::url::{
|
||||
use crate::ai_pipeline::transport::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::vertex::{
|
||||
use crate::ai_pipeline::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::provider_transport_facade::{
|
||||
use crate::ai_pipeline::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::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, GatewayControlDecision,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
|
||||
GatewayError,
|
||||
};
|
||||
|
||||
pub(crate) mod family;
|
||||
pub(crate) mod plans;
|
||||
mod family;
|
||||
mod plans;
|
||||
mod request;
|
||||
|
||||
pub(super) use self::family::{
|
||||
@@ -74,6 +75,9 @@ pub(crate) use self::family::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
maybe_build_sync_local_same_format_provider_decision_payload,
|
||||
};
|
||||
pub(crate) use self::plans::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
|
||||
};
|
||||
use self::request::{
|
||||
build_same_format_provider_request_body, build_same_format_upstream_url,
|
||||
extract_gemini_model_from_path,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
|
||||
|
||||
@@ -2,16 +2,15 @@ use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::read_auth_api_key_snapshot;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_available_local_candidate;
|
||||
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision, PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
|
||||
|
||||
use super::types::{
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
};
|
||||
@@ -24,9 +23,8 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<LocalSameFormatProviderDecisionInput> {
|
||||
let Some(auth_context) = decision.auth_context.clone().filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
}) else {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
@@ -42,13 +40,13 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
}
|
||||
};
|
||||
|
||||
let auth_snapshot = match read_auth_api_key_snapshot(
|
||||
state,
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
let auth_snapshot = match planner_state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => return None,
|
||||
@@ -76,16 +74,17 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<LocalSameFormatProviderCandidateAttempt>, GatewayError> {
|
||||
let candidates = list_selectable_candidates(
|
||||
state,
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
spec.require_streaming,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidates = planner_state
|
||||
.list_selectable_candidates(
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
spec.require_streaming,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
@@ -109,19 +108,19 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
spec.api_format,
|
||||
);
|
||||
|
||||
let candidate_id = persist_available_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local same-format decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
let candidate_id = planner_state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local same-format decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalSameFormatProviderCandidateAttempt {
|
||||
candidate,
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
|
||||
|
||||
mod build;
|
||||
mod candidates;
|
||||
mod payload;
|
||||
mod types;
|
||||
|
||||
pub(crate) use self::build::{
|
||||
maybe_build_stream_local_same_format_provider_decision_payload,
|
||||
@@ -12,4 +14,18 @@ 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::types::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};
|
||||
pub(crate) use crate::ai_pipeline::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalSameFormatProviderDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalSameFormatProviderCandidateAttempt {
|
||||
pub(crate) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) candidate_index: u32,
|
||||
pub(crate) candidate_id: String,
|
||||
}
|
||||
|
||||
@@ -1,41 +1,21 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::control_facade::collect_control_headers;
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_skipped_local_candidate;
|
||||
use crate::ai_pipeline::planner::transport_facade::{
|
||||
read_provider_transport_snapshot, resolve_local_oauth_request_auth,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::antigravity::{
|
||||
use crate::ai_pipeline::transport::antigravity::{
|
||||
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
|
||||
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
|
||||
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::auth::{
|
||||
build_openai_passthrough_headers, resolve_local_gemini_auth, resolve_local_standard_auth,
|
||||
use crate::ai_pipeline::transport::auth::build_openai_passthrough_headers;
|
||||
use crate::ai_pipeline::transport::claude_code::build_claude_code_passthrough_headers;
|
||||
use crate::ai_pipeline::transport::kiro::{build_kiro_provider_headers, KIRO_ENVELOPE_NAME};
|
||||
use crate::ai_pipeline::transport::{
|
||||
apply_local_header_rules, ensure_upstream_auth_header, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::claude_code::{
|
||||
build_claude_code_passthrough_headers, supports_local_claude_code_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::kiro::{
|
||||
build_kiro_provider_headers, supports_local_kiro_request_transport_with_network,
|
||||
KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::vertex::{
|
||||
resolve_local_vertex_api_key_query_auth,
|
||||
supports_local_vertex_api_key_gemini_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::{
|
||||
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,
|
||||
use crate::ai_pipeline::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, PlannerAppState,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{
|
||||
@@ -44,9 +24,15 @@ use crate::{
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::types::{
|
||||
#[path = "payload/prepare.rs"]
|
||||
mod prepare;
|
||||
|
||||
use self::prepare::{
|
||||
prepare_local_same_format_provider_candidate, PreparedSameFormatProviderCandidate,
|
||||
};
|
||||
use super::{
|
||||
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
@@ -58,205 +44,35 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
attempt: LocalSameFormatProviderCandidateAttempt,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let LocalSameFormatProviderCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
} = attempt;
|
||||
|
||||
let transport = match read_provider_transport_snapshot(
|
||||
state,
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
let PreparedSameFormatProviderCandidate {
|
||||
transport,
|
||||
is_antigravity,
|
||||
is_claude_code,
|
||||
is_vertex,
|
||||
is_kiro,
|
||||
kiro_auth,
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
report_kind,
|
||||
upstream_is_stream,
|
||||
} = prepare_local_same_format_provider_candidate(
|
||||
planner_state.app(),
|
||||
trace_id,
|
||||
input,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format decision provider transport read failed"
|
||||
);
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_read_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
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 = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("vertex_ai");
|
||||
let transport_supported = match spec.family {
|
||||
_ if transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("kiro") =>
|
||||
{
|
||||
supports_local_kiro_request_transport_with_network(&transport)
|
||||
}
|
||||
_ if is_antigravity => true,
|
||||
_ if is_claude_code => {
|
||||
supports_local_claude_code_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
_ if is_vertex => supports_local_vertex_api_key_gemini_transport_with_network(&transport),
|
||||
LocalSameFormatProviderFamily::Standard => {
|
||||
supports_local_standard_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
LocalSameFormatProviderFamily::Gemini => {
|
||||
supports_local_gemini_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
};
|
||||
if !transport_supported {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let is_kiro = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("kiro");
|
||||
let vertex_query_auth = if is_vertex {
|
||||
resolve_local_vertex_api_key_query_auth(&transport)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let should_try_oauth_auth = is_kiro
|
||||
|| matches!(spec.family, LocalSameFormatProviderFamily::Standard)
|
||||
&& resolve_local_standard_auth(&transport).is_none()
|
||||
|| matches!(spec.family, LocalSameFormatProviderFamily::Gemini)
|
||||
&& !is_vertex
|
||||
&& resolve_local_gemini_auth(&transport).is_none();
|
||||
let oauth_auth = if should_try_oauth_auth {
|
||||
match resolve_local_oauth_request_auth(state, &transport).await {
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(auth))) => {
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth))
|
||||
}
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => {
|
||||
Some(LocalResolvedOAuthRequestAuth::Header { name, value })
|
||||
}
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec.api_format,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
error = ?err,
|
||||
"gateway local same-format oauth auth resolution failed"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let kiro_auth = match oauth_auth.as_ref() {
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth),
|
||||
_ => None,
|
||||
};
|
||||
let auth = if let Some(auth) = kiro_auth.as_ref() {
|
||||
Some((auth.name.to_string(), auth.value.clone()))
|
||||
} else if let Some(LocalResolvedOAuthRequestAuth::Header { name, value }) = oauth_auth.as_ref()
|
||||
{
|
||||
Some((name.clone(), value.clone()))
|
||||
} else if is_vertex {
|
||||
None
|
||||
} else {
|
||||
match spec.family {
|
||||
LocalSameFormatProviderFamily::Standard => resolve_local_standard_auth(&transport),
|
||||
LocalSameFormatProviderFamily::Gemini => resolve_local_gemini_auth(&transport),
|
||||
}
|
||||
};
|
||||
let (auth_header, auth_value) = match auth {
|
||||
Some((name, value)) => (Some(name), Some(value)),
|
||||
None if is_vertex && vertex_query_auth.is_some() => (None, None),
|
||||
None => {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if is_vertex && vertex_query_auth.is_none() {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
.await?;
|
||||
|
||||
let Some(base_provider_request_body) =
|
||||
super::super::request::build_same_format_provider_request_body(
|
||||
@@ -264,8 +80,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
&mapped_model,
|
||||
spec,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
is_kiro || is_antigravity || spec.require_streaming,
|
||||
kiro_auth,
|
||||
upstream_is_stream,
|
||||
kiro_auth.as_ref(),
|
||||
is_claude_code,
|
||||
)
|
||||
else {
|
||||
@@ -332,18 +148,6 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
} else {
|
||||
base_provider_request_body
|
||||
};
|
||||
let upstream_is_stream = is_kiro || is_antigravity || spec.require_streaming;
|
||||
let report_kind = if is_kiro && !spec.require_streaming {
|
||||
"claude_cli_sync_finalize"
|
||||
} else if is_antigravity && !spec.require_streaming {
|
||||
match spec.api_format {
|
||||
"gemini:chat" => "gemini_chat_sync_finalize",
|
||||
"gemini:cli" => "gemini_cli_sync_finalize",
|
||||
_ => spec.report_kind,
|
||||
}
|
||||
} else {
|
||||
spec.report_kind
|
||||
};
|
||||
|
||||
let Some(upstream_url) = super::super::request::build_same_format_upstream_url(
|
||||
parts,
|
||||
@@ -351,7 +155,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
&mapped_model,
|
||||
spec,
|
||||
upstream_is_stream,
|
||||
kiro_auth,
|
||||
kiro_auth.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
@@ -392,7 +196,11 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
transport.key.fingerprint.as_ref(),
|
||||
)
|
||||
} else if is_vertex {
|
||||
build_passthrough_headers(&parts.headers, &extra_headers, Some("application/json"))
|
||||
crate::ai_pipeline::transport::build_passthrough_headers(
|
||||
&parts.headers,
|
||||
&extra_headers,
|
||||
Some("application/json"),
|
||||
)
|
||||
} else {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
@@ -440,13 +248,16 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let prompt_cache_key = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let proxy = resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &transport).await;
|
||||
let proxy =
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), &transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
let report_context = append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
@@ -538,17 +349,17 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate(
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
persist_skipped_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local same-format decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
PlannerAppState::new(state)
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local same-format decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,270 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::transport::auth::{resolve_local_gemini_auth, resolve_local_standard_auth};
|
||||
use crate::ai_pipeline::transport::claude_code::supports_local_claude_code_transport_with_network;
|
||||
use crate::ai_pipeline::transport::kiro::{
|
||||
supports_local_kiro_request_transport_with_network, KiroRequestAuth,
|
||||
};
|
||||
use crate::ai_pipeline::transport::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::transport::vertex::{
|
||||
resolve_local_vertex_api_key_query_auth,
|
||||
supports_local_vertex_api_key_gemini_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
use super::super::{
|
||||
LocalSameFormatProviderDecisionInput, LocalSameFormatProviderFamily,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
pub(super) struct PreparedSameFormatProviderCandidate {
|
||||
pub(super) transport: GatewayProviderTransportSnapshot,
|
||||
pub(super) is_antigravity: bool,
|
||||
pub(super) is_claude_code: bool,
|
||||
pub(super) is_vertex: bool,
|
||||
pub(super) is_kiro: bool,
|
||||
pub(super) kiro_auth: Option<KiroRequestAuth>,
|
||||
pub(super) auth_header: Option<String>,
|
||||
pub(super) auth_value: Option<String>,
|
||||
pub(super) mapped_model: String,
|
||||
pub(super) report_kind: &'static str,
|
||||
pub(super) upstream_is_stream: bool,
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_local_same_format_provider_candidate(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<PreparedSameFormatProviderCandidate> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let transport = match planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
super::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_snapshot_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format decision provider transport read failed"
|
||||
);
|
||||
super::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_snapshot_read_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
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 = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("vertex_ai");
|
||||
let is_kiro = transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("kiro");
|
||||
let transport_supported = if is_kiro {
|
||||
supports_local_kiro_request_transport_with_network(&transport)
|
||||
} else if is_antigravity {
|
||||
true
|
||||
} else if is_claude_code {
|
||||
supports_local_claude_code_transport_with_network(&transport, spec.api_format)
|
||||
} else if is_vertex {
|
||||
supports_local_vertex_api_key_gemini_transport_with_network(&transport)
|
||||
} else {
|
||||
match spec.family {
|
||||
LocalSameFormatProviderFamily::Standard => {
|
||||
supports_local_standard_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
LocalSameFormatProviderFamily::Gemini => {
|
||||
supports_local_gemini_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
}
|
||||
};
|
||||
if !transport_supported {
|
||||
super::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let vertex_query_auth = if is_vertex {
|
||||
resolve_local_vertex_api_key_query_auth(&transport)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let should_try_oauth_auth = is_kiro
|
||||
|| matches!(spec.family, LocalSameFormatProviderFamily::Standard)
|
||||
&& resolve_local_standard_auth(&transport).is_none()
|
||||
|| matches!(spec.family, LocalSameFormatProviderFamily::Gemini)
|
||||
&& !is_vertex
|
||||
&& resolve_local_gemini_auth(&transport).is_none();
|
||||
let oauth_auth = if should_try_oauth_auth {
|
||||
match planner_state
|
||||
.resolve_local_oauth_request_auth(&transport)
|
||||
.await
|
||||
{
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(auth))) => {
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth))
|
||||
}
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => {
|
||||
Some(LocalResolvedOAuthRequestAuth::Header { name, value })
|
||||
}
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec.api_format,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
error = ?err,
|
||||
"gateway local same-format oauth auth resolution failed"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let kiro_auth = match oauth_auth.as_ref() {
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth.clone()),
|
||||
_ => None,
|
||||
};
|
||||
let auth = if let Some(kiro_auth) = kiro_auth.as_ref() {
|
||||
Some((kiro_auth.name.to_string(), kiro_auth.value.clone()))
|
||||
} else if let Some(LocalResolvedOAuthRequestAuth::Header { name, value }) = oauth_auth.as_ref()
|
||||
{
|
||||
Some((name.clone(), value.clone()))
|
||||
} else if is_vertex {
|
||||
None
|
||||
} else {
|
||||
match spec.family {
|
||||
LocalSameFormatProviderFamily::Standard => resolve_local_standard_auth(&transport),
|
||||
LocalSameFormatProviderFamily::Gemini => resolve_local_gemini_auth(&transport),
|
||||
}
|
||||
};
|
||||
let (auth_header, auth_value) = match auth {
|
||||
Some((name, value)) => (Some(name), Some(value)),
|
||||
None if is_vertex && vertex_query_auth.is_some() => (None, None),
|
||||
None => {
|
||||
super::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if is_vertex && vertex_query_auth.is_none() {
|
||||
super::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
super::mark_skipped_local_same_format_provider_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let upstream_is_stream = is_kiro || is_antigravity || spec.require_streaming;
|
||||
let report_kind = if is_kiro && !spec.require_streaming {
|
||||
"claude_cli_sync_finalize"
|
||||
} else if is_antigravity && !spec.require_streaming {
|
||||
match spec.api_format {
|
||||
"gemini:chat" => "gemini_chat_sync_finalize",
|
||||
"gemini:cli" => "gemini_cli_sync_finalize",
|
||||
_ => spec.report_kind,
|
||||
}
|
||||
} else {
|
||||
spec.report_kind
|
||||
};
|
||||
|
||||
Some(PreparedSameFormatProviderCandidate {
|
||||
transport,
|
||||
is_antigravity,
|
||||
is_claude_code,
|
||||
is_vertex,
|
||||
is_kiro,
|
||||
kiro_auth,
|
||||
auth_header,
|
||||
auth_value,
|
||||
mapped_model,
|
||||
report_kind,
|
||||
upstream_is_stream,
|
||||
})
|
||||
}
|
||||
@@ -1,20 +0,0 @@
|
||||
use crate::ai_pipeline::control_facade::GatewayControlAuthContext;
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
|
||||
|
||||
pub(crate) use aether_ai_pipeline::planner::passthrough::provider::{
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalSameFormatProviderDecisionInput {
|
||||
pub(crate) auth_context: GatewayControlAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalSameFormatProviderCandidateAttempt {
|
||||
pub(crate) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) candidate_index: u32,
|
||||
pub(crate) candidate_id: String,
|
||||
}
|
||||
@@ -1,7 +1,8 @@
|
||||
use tracing::warn;
|
||||
|
||||
pub(crate) use aether_ai_pipeline::planner::passthrough::provider::{
|
||||
resolve_stream_spec, resolve_sync_spec,
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
resolve_local_same_format_stream_spec as resolve_stream_spec,
|
||||
resolve_local_same_format_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
|
||||
use super::{
|
||||
|
||||
@@ -3,7 +3,7 @@ use std::collections::BTreeMap;
|
||||
use serde_json::Value;
|
||||
use url::form_urlencoded;
|
||||
|
||||
use crate::ai_pipeline::planner::transport_facade::GatewayProviderTransportSnapshot;
|
||||
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
|
||||
|
||||
use super::{
|
||||
apply_local_body_rules, build_antigravity_v1internal_url, build_claude_code_messages_url,
|
||||
@@ -20,7 +20,7 @@ pub(super) fn build_same_format_provider_request_body(
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
body_rules: Option<&Value>,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_pipeline::provider_transport_facade::kiro::KiroRequestAuth>,
|
||||
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
|
||||
is_claude_code: bool,
|
||||
) -> Option<Value> {
|
||||
if let Some(kiro_auth) = kiro_auth {
|
||||
@@ -66,7 +66,7 @@ pub(super) fn build_same_format_upstream_url(
|
||||
mapped_model: &str,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
upstream_is_stream: bool,
|
||||
kiro_auth: Option<&crate::ai_pipeline::provider_transport_facade::kiro::KiroRequestAuth>,
|
||||
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
|
||||
) -> Option<String> {
|
||||
if let Some(kiro_auth) = kiro_auth {
|
||||
return build_kiro_generate_assistant_response_url(
|
||||
|
||||
@@ -1,22 +1,10 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::ExecutionPlan;
|
||||
|
||||
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) struct LocalSyncPlanAndReport {
|
||||
pub(crate) plan: ExecutionPlan,
|
||||
pub(crate) report_kind: Option<String>,
|
||||
pub(crate) report_context: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
pub(crate) struct LocalStreamPlanAndReport {
|
||||
pub(crate) plan: ExecutionPlan,
|
||||
pub(crate) report_kind: Option<String>,
|
||||
pub(crate) report_context: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[path = "standard/gemini/plan_builders.rs"]
|
||||
mod gemini_builders;
|
||||
#[path = "standard/openai/plan_builders.rs"]
|
||||
@@ -45,21 +33,10 @@ pub(super) fn augment_sync_report_context(
|
||||
provider_request_headers: &BTreeMap<String, String>,
|
||||
provider_request_body: &serde_json::Value,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
let mut report_context = match report_context {
|
||||
Some(serde_json::Value::Object(map)) => map,
|
||||
Some(_) => serde_json::Map::new(),
|
||||
None => serde_json::Map::new(),
|
||||
};
|
||||
|
||||
report_context.insert(
|
||||
"provider_request_headers".to_string(),
|
||||
serde_json::to_value(provider_request_headers)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?,
|
||||
);
|
||||
report_context.insert(
|
||||
"provider_request_body".to_string(),
|
||||
provider_request_body.clone(),
|
||||
);
|
||||
|
||||
Ok(Some(serde_json::Value::Object(report_context)))
|
||||
augment_sync_report_context_impl(
|
||||
report_context,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
@@ -1,10 +1,17 @@
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
is_matching_stream_request as is_matching_stream_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,
|
||||
};
|
||||
|
||||
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
aether_ai_pipeline::planner::route::resolve_execution_runtime_stream_plan_kind(
|
||||
resolve_execution_runtime_stream_plan_kind_impl(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
@@ -17,7 +24,7 @@ pub(crate) fn resolve_execution_runtime_sync_plan_kind(
|
||||
parts: &http::request::Parts,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<&'static str> {
|
||||
aether_ai_pipeline::planner::route::resolve_execution_runtime_sync_plan_kind(
|
||||
resolve_execution_runtime_sync_plan_kind_impl(
|
||||
decision.route_class.as_deref(),
|
||||
decision.route_family.as_deref(),
|
||||
decision.route_kind.as_deref(),
|
||||
@@ -31,19 +38,15 @@ pub(crate) fn is_matching_stream_request(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
) -> bool {
|
||||
aether_ai_pipeline::planner::route::is_matching_stream_request(
|
||||
plan_kind,
|
||||
parts.uri.path(),
|
||||
body_json,
|
||||
)
|
||||
is_matching_stream_request_impl(plan_kind, parts.uri.path(), body_json)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
aether_ai_pipeline::planner::route::supports_sync_scheduler_decision_kind(plan_kind)
|
||||
supports_sync_scheduler_decision_kind_impl(plan_kind)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
aether_ai_pipeline::planner::route::supports_stream_scheduler_decision_kind(plan_kind)
|
||||
supports_stream_scheduler_decision_kind_impl(plan_kind)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -55,7 +58,7 @@ mod tests {
|
||||
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
|
||||
supports_sync_scheduler_decision_kind,
|
||||
};
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
|
||||
fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
|
||||
GatewayControlDecision {
|
||||
|
||||
@@ -1,44 +0,0 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(crate) async fn list_selectable_candidates(
|
||||
state: &AppState,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
require_streaming: bool,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
crate::scheduler::candidate::list_selectable_candidates(
|
||||
state.data.as_ref(),
|
||||
state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
state: &AppState,
|
||||
candidate_api_format: &str,
|
||||
required_capability: &str,
|
||||
require_streaming: bool,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
crate::scheduler::candidate::list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
state.data.as_ref(),
|
||||
state,
|
||||
candidate_api_format,
|
||||
required_capability,
|
||||
require_streaming,
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -1,64 +1,24 @@
|
||||
use std::collections::BTreeMap;
|
||||
mod decision;
|
||||
mod support;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::control_facade::{
|
||||
collect_control_headers, GatewayControlAuthContext, GatewayControlDecision,
|
||||
};
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::{
|
||||
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_runtime_facade::{
|
||||
persist_available_local_candidate, persist_skipped_local_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
use crate::ai_pipeline::planner::executor_facade::mark_unused_local_candidate_items;
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_passthrough_stream_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates_for_required_capability_without_requested_model;
|
||||
use crate::ai_pipeline::planner::transport_facade::read_provider_transport_snapshot;
|
||||
use crate::ai_pipeline::provider_transport_facade::auth::{
|
||||
build_passthrough_headers_with_auth, resolve_local_gemini_auth,
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_gemini_files_stream_spec as resolve_stream_spec,
|
||||
resolve_gemini_files_sync_spec as resolve_sync_spec, LocalGeminiFilesSpec,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::policy::supports_local_gemini_transport_with_network;
|
||||
use crate::ai_pipeline::provider_transport_facade::url::build_gemini_files_passthrough_url;
|
||||
use crate::ai_pipeline::provider_transport_facade::{
|
||||
apply_local_body_rules, apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
use aether_ai_pipeline::contracts::GEMINI_FILES_UPLOAD_PLAN_KIND;
|
||||
use aether_ai_pipeline::planner::specialized::files::{
|
||||
resolve_stream_spec, resolve_sync_spec, LocalGeminiFilesSpec,
|
||||
|
||||
use self::decision::maybe_build_local_gemini_files_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
materialize_local_gemini_files_candidate_attempts, resolve_local_gemini_files_decision_input,
|
||||
};
|
||||
|
||||
const GEMINI_FILES_CANDIDATE_API_FORMAT: &str = "gemini:chat";
|
||||
const GEMINI_FILES_CLIENT_API_FORMAT: &str = "gemini:files";
|
||||
const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct LocalGeminiFilesDecisionInput {
|
||||
auth_context: GatewayControlAuthContext,
|
||||
auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct LocalGeminiFilesCandidateAttempt {
|
||||
candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: String,
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_gemini_files_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -287,442 +247,3 @@ async fn build_local_stream_plan_and_reports(
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
async fn resolve_local_gemini_files_decision_input(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<LocalGeminiFilesDecisionInput> {
|
||||
let Some(auth_context) = decision.auth_context.clone().filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
}) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let auth_snapshot = match read_auth_api_key_snapshot(
|
||||
state,
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local gemini files decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalGeminiFilesDecisionInput {
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
async fn materialize_local_gemini_files_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
) -> Result<Vec<LocalGeminiFilesCandidateAttempt>, GatewayError> {
|
||||
let candidates = list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
state,
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
GEMINI_FILES_REQUIRED_CAPABILITY,
|
||||
false,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let extra_data = json!({
|
||||
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"candidate_api_format": GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"global_model_id": candidate.global_model_id.clone(),
|
||||
"global_model_name": candidate.global_model_name.clone(),
|
||||
"model_id": candidate.model_id.clone(),
|
||||
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
|
||||
"mapping_matched_model": candidate.mapping_matched_model.clone(),
|
||||
"provider_name": candidate.provider_name.clone(),
|
||||
"key_name": candidate.key_name.clone(),
|
||||
});
|
||||
|
||||
let candidate_id = persist_available_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local gemini files request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalGeminiFilesCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index: candidate_index as u32,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(attempts)
|
||||
}
|
||||
|
||||
async fn maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &str,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
attempt: LocalGeminiFilesCandidateAttempt,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let LocalGeminiFilesCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
} = attempt;
|
||||
|
||||
let transport = match read_provider_transport_snapshot(
|
||||
state,
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local gemini files provider transport read failed"
|
||||
);
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_read_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
if !supports_local_gemini_transport_with_network(&transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
|
||||
{
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_local_gemini_auth(&transport) else {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
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 upstream_url =
|
||||
if spec.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND || custom_path.is_some() {
|
||||
build_gemini_files_passthrough_url(
|
||||
&transport.endpoint.base_url,
|
||||
passthrough_path,
|
||||
parts.uri.query(),
|
||||
)
|
||||
} else {
|
||||
build_gemini_files_passthrough_url(
|
||||
&transport.endpoint.base_url,
|
||||
passthrough_path,
|
||||
parts.uri.query(),
|
||||
)
|
||||
};
|
||||
let Some(upstream_url) = upstream_url else {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_body = if spec.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.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND {
|
||||
body_base64
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let original_request_body = if let Some(body_bytes_b64) = provider_request_body_base64.clone() {
|
||||
json!({"body_bytes_b64": body_bytes_b64})
|
||||
} else if !body_is_empty {
|
||||
body_json.clone()
|
||||
} else {
|
||||
serde_json::Value::Null
|
||||
};
|
||||
if provider_request_body_base64.is_some() && transport.endpoint.body_rules.is_some() {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_body_rules_unsupported_for_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),
|
||||
) {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_body_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
let mut provider_request_headers = build_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
);
|
||||
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),
|
||||
) {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
let file_name = parts
|
||||
.uri
|
||||
.path()
|
||||
.trim_start_matches("/v1beta/")
|
||||
.trim()
|
||||
.to_string();
|
||||
let proxy = resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &transport).await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: if spec.require_streaming {
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string()
|
||||
} else {
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
|
||||
},
|
||||
decision_kind: Some(spec.decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.clone()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url),
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
client_api_format: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
provider_contract: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
client_contract: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
model_name: Some("gemini-files".to_string()),
|
||||
mapped_model: Some(candidate.selected_provider_model_name.clone()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
provider_request_body_base64,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: spec.require_streaming,
|
||||
report_kind: spec.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": "gemini-files",
|
||||
"provider_name": transport.provider.name,
|
||||
"provider_id": candidate.provider_id,
|
||||
"endpoint_id": candidate.endpoint_id,
|
||||
"key_id": candidate.key_id,
|
||||
"file_key_id": candidate.key_id,
|
||||
"file_name": file_name,
|
||||
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": original_request_body,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
})),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
|
||||
async fn mark_skipped_local_gemini_files_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
persist_skipped_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local gemini files failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn mark_unused_local_files_candidates<T>(state: &AppState, remaining: Vec<T>)
|
||||
where
|
||||
T: LocalGeminiFilesPlanAndReport,
|
||||
{
|
||||
mark_unused_local_candidate_items(
|
||||
state,
|
||||
remaining,
|
||||
|item| item.plan(),
|
||||
|item| item.report_context(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
trait LocalGeminiFilesPlanAndReport {
|
||||
fn plan(&self) -> &aether_contracts::ExecutionPlan;
|
||||
|
||||
fn report_context(&self) -> Option<&serde_json::Value>;
|
||||
}
|
||||
|
||||
impl LocalGeminiFilesPlanAndReport for LocalSyncPlanAndReport {
|
||||
fn plan(&self) -> &aether_contracts::ExecutionPlan {
|
||||
&self.plan
|
||||
}
|
||||
|
||||
fn report_context(&self) -> Option<&serde_json::Value> {
|
||||
self.report_context.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl LocalGeminiFilesPlanAndReport for LocalStreamPlanAndReport {
|
||||
fn plan(&self) -> &aether_contracts::ExecutionPlan {
|
||||
&self.plan
|
||||
}
|
||||
|
||||
fn report_context(&self) -> Option<&serde_json::Value> {
|
||||
self.report_context.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,304 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::contracts::GEMINI_FILES_UPLOAD_PLAN_KIND;
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_passthrough_headers_with_auth, resolve_local_gemini_auth,
|
||||
};
|
||||
use crate::ai_pipeline::transport::policy::supports_local_gemini_transport_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, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, PlannerAppState,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse};
|
||||
|
||||
use super::support::{
|
||||
mark_skipped_local_gemini_files_candidate, LocalGeminiFilesCandidateAttempt,
|
||||
LocalGeminiFilesDecisionInput, GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
};
|
||||
use super::LocalGeminiFilesSpec;
|
||||
|
||||
pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &str,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
attempt: LocalGeminiFilesCandidateAttempt,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let LocalGeminiFilesCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
} = attempt;
|
||||
|
||||
let transport = match planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local gemini files provider transport read failed"
|
||||
);
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_read_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
if !supports_local_gemini_transport_with_network(&transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
|
||||
{
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_local_gemini_auth(&transport) else {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
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 {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_body = if spec.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.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND {
|
||||
body_base64
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let original_request_body = if let Some(body_bytes_b64) = provider_request_body_base64.clone() {
|
||||
json!({"body_bytes_b64": body_bytes_b64})
|
||||
} else if !body_is_empty {
|
||||
body_json.clone()
|
||||
} else {
|
||||
serde_json::Value::Null
|
||||
};
|
||||
if provider_request_body_base64.is_some() && transport.endpoint.body_rules.is_some() {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_body_rules_unsupported_for_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),
|
||||
) {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_body_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
let mut provider_request_headers = build_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
);
|
||||
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),
|
||||
) {
|
||||
mark_skipped_local_gemini_files_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
let file_name = parts
|
||||
.uri
|
||||
.path()
|
||||
.trim_start_matches("/v1beta/")
|
||||
.trim()
|
||||
.to_string();
|
||||
let proxy =
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), &transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: if spec.require_streaming {
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string()
|
||||
} else {
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
|
||||
},
|
||||
decision_kind: Some(spec.decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.clone()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url),
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
client_api_format: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
provider_contract: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
client_contract: Some(GEMINI_FILES_CLIENT_API_FORMAT.to_string()),
|
||||
model_name: Some("gemini-files".to_string()),
|
||||
mapped_model: Some(candidate.selected_provider_model_name.clone()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
provider_request_body_base64,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: spec.require_streaming,
|
||||
report_kind: spec.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": "gemini-files",
|
||||
"provider_name": transport.provider.name,
|
||||
"provider_id": candidate.provider_id,
|
||||
"endpoint_id": candidate.endpoint_id,
|
||||
"key_id": candidate.key_id,
|
||||
"file_key_id": candidate.key_id,
|
||||
"file_name": file_name,
|
||||
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": original_request_body,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
})),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
|
||||
};
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(super) const GEMINI_FILES_CANDIDATE_API_FORMAT: &str = "gemini:chat";
|
||||
pub(super) const GEMINI_FILES_CLIENT_API_FORMAT: &str = "gemini:files";
|
||||
pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalGeminiFilesDecisionInput {
|
||||
pub(super) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalGeminiFilesCandidateAttempt {
|
||||
pub(super) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(super) candidate_index: u32,
|
||||
pub(super) candidate_id: String,
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_local_gemini_files_decision_input(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<LocalGeminiFilesDecisionInput> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let auth_snapshot = match planner_state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local gemini files decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalGeminiFilesDecisionInput {
|
||||
auth_context,
|
||||
auth_snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
) -> Result<Vec<LocalGeminiFilesCandidateAttempt>, GatewayError> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidates = planner_state
|
||||
.list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
GEMINI_FILES_REQUIRED_CAPABILITY,
|
||||
false,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let extra_data = json!({
|
||||
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"candidate_api_format": GEMINI_FILES_CANDIDATE_API_FORMAT,
|
||||
"global_model_id": candidate.global_model_id.clone(),
|
||||
"global_model_name": candidate.global_model_name.clone(),
|
||||
"model_id": candidate.model_id.clone(),
|
||||
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
|
||||
"mapping_matched_model": candidate.mapping_matched_model.clone(),
|
||||
"provider_name": candidate.provider_name.clone(),
|
||||
"key_name": candidate.key_name.clone(),
|
||||
});
|
||||
|
||||
let candidate_id = planner_state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local gemini files request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalGeminiFilesCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index: candidate_index as u32,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(attempts)
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_gemini_files_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
PlannerAppState::new(state)
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local gemini files failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -1,10 +1,14 @@
|
||||
//! Non-matrix AI surfaces such as files and video.
|
||||
|
||||
pub(crate) mod files;
|
||||
pub(crate) mod video;
|
||||
mod files;
|
||||
mod video;
|
||||
|
||||
pub(crate) use self::files::{
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
maybe_build_stream_local_gemini_files_decision_payload,
|
||||
maybe_build_sync_local_gemini_files_decision_payload,
|
||||
};
|
||||
pub(crate) use self::video::maybe_build_sync_local_video_decision_payload;
|
||||
pub(crate) use self::video::{
|
||||
build_local_video_sync_plan_and_reports_for_kind, maybe_build_sync_local_video_decision_payload,
|
||||
};
|
||||
|
||||
@@ -1,63 +1,23 @@
|
||||
use std::collections::BTreeMap;
|
||||
mod decision;
|
||||
mod support;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::control_facade::{
|
||||
collect_control_headers, GatewayControlAuthContext, GatewayControlDecision,
|
||||
};
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::{
|
||||
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_runtime_facade::{
|
||||
persist_available_local_candidate, persist_skipped_local_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION;
|
||||
use crate::ai_pipeline::planner::executor_facade::mark_unused_local_candidate_items;
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_passthrough_sync_plan_from_decision, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
|
||||
use crate::ai_pipeline::planner::transport_facade::{
|
||||
read_provider_transport_snapshot, GatewayProviderTransportSnapshot,
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_video_sync_spec as resolve_sync_spec, LocalVideoCreateFamily,
|
||||
LocalVideoCreateSpec,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::auth::{
|
||||
build_passthrough_headers_with_auth, resolve_local_gemini_auth, resolve_local_openai_chat_auth,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::url::{
|
||||
build_gemini_video_predict_long_running_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::{
|
||||
apply_local_body_rules, apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
use aether_ai_pipeline::planner::specialized::video::{
|
||||
resolve_sync_spec, LocalVideoCreateFamily, LocalVideoCreateSpec,
|
||||
|
||||
use self::decision::maybe_build_local_video_create_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
list_local_video_create_candidate_attempts, resolve_local_video_create_decision_input,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct LocalVideoCreateDecisionInput {
|
||||
auth_context: GatewayControlAuthContext,
|
||||
requested_model: String,
|
||||
auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct LocalVideoCreateCandidateAttempt {
|
||||
candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: String,
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_video_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -93,36 +53,17 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let candidates = match list_selectable_candidates(
|
||||
state,
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision scheduler selection failed"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let attempts = materialize_local_video_create_candidate_attempts(
|
||||
let Some(attempts) = list_local_video_create_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
candidates,
|
||||
spec.api_format,
|
||||
spec.decision_kind,
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
@@ -153,36 +94,17 @@ async fn build_local_sync_plan_and_reports(
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let candidates = match list_selectable_candidates(
|
||||
state,
|
||||
spec.api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision scheduler selection failed"
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
|
||||
let attempts = materialize_local_video_create_candidate_attempts(
|
||||
let Some(attempts) = list_local_video_create_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
candidates,
|
||||
spec.api_format,
|
||||
spec.decision_kind,
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
@@ -210,448 +132,3 @@ async fn build_local_sync_plan_and_reports(
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
async fn resolve_local_video_create_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<LocalVideoCreateDecisionInput> {
|
||||
let Some(auth_context) = decision.auth_context.clone().filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
}) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = match spec.family {
|
||||
LocalVideoCreateFamily::OpenAi => body_json
|
||||
.get("model")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)?,
|
||||
LocalVideoCreateFamily::Gemini => extract_gemini_video_model_from_path(parts.uri.path())?,
|
||||
};
|
||||
|
||||
let auth_snapshot = match read_auth_api_key_snapshot(
|
||||
state,
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalVideoCreateDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
attempt: LocalVideoCreateCandidateAttempt,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let LocalVideoCreateCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
} = attempt;
|
||||
let transport = match read_provider_transport_snapshot(
|
||||
state,
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision provider transport read failed"
|
||||
);
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_read_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let transport_supported = match spec.family {
|
||||
LocalVideoCreateFamily::OpenAi => {
|
||||
supports_local_standard_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
LocalVideoCreateFamily::Gemini => {
|
||||
supports_local_gemini_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
};
|
||||
if !transport_supported {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let auth = match spec.family {
|
||||
LocalVideoCreateFamily::OpenAi => resolve_local_openai_chat_auth(&transport),
|
||||
LocalVideoCreateFamily::Gemini => resolve_local_gemini_auth(&transport),
|
||||
};
|
||||
let Some((auth_header, auth_value)) = auth else {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let upstream_url = build_video_upstream_url(parts, &transport, &mapped_model, spec.family);
|
||||
let Some(upstream_url) = upstream_url else {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(provider_request_body) = build_provider_request_body(
|
||||
body_json,
|
||||
spec.family,
|
||||
&mapped_model,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"provider_request_body_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mut provider_request_headers = build_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
&[&auth_header, "content-type"],
|
||||
&provider_request_body,
|
||||
Some(body_json),
|
||||
) {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
let proxy = resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &transport).await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(spec.decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.clone()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url),
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some(spec.api_format.to_string()),
|
||||
client_api_format: Some(spec.api_format.to_string()),
|
||||
provider_contract: Some(spec.api_format.to_string()),
|
||||
client_contract: Some(spec.api_format.to_string()),
|
||||
model_name: Some(input.requested_model.clone()),
|
||||
mapped_model: Some(mapped_model.clone()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some(spec.report_kind.to_string()),
|
||||
report_context: Some(json!({
|
||||
"user_id": input.auth_context.user_id.clone(),
|
||||
"api_key_id": input.auth_context.api_key_id.clone(),
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": input.requested_model.clone(),
|
||||
"provider_name": transport.provider.name.clone(),
|
||||
"provider_id": candidate.provider_id.clone(),
|
||||
"endpoint_id": candidate.endpoint_id.clone(),
|
||||
"key_id": candidate.key_id.clone(),
|
||||
"provider_api_format": spec.api_format,
|
||||
"client_api_format": spec.api_format,
|
||||
"mapped_model": mapped_model,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
})),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
|
||||
fn build_provider_request_body(
|
||||
body_json: &serde_json::Value,
|
||||
family: LocalVideoCreateFamily,
|
||||
mapped_model: &str,
|
||||
body_rules: Option<&serde_json::Value>,
|
||||
) -> Option<serde_json::Value> {
|
||||
let mut provider_request_body = match family {
|
||||
LocalVideoCreateFamily::OpenAi => {
|
||||
let mut provider_request_body = body_json.as_object().cloned().unwrap_or_default();
|
||||
provider_request_body
|
||||
.insert("model".to_string(), Value::String(mapped_model.to_string()));
|
||||
serde_json::Value::Object(provider_request_body)
|
||||
}
|
||||
LocalVideoCreateFamily::Gemini => body_json.clone(),
|
||||
};
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
fn build_video_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
family: LocalVideoCreateFamily,
|
||||
) -> Option<String> {
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
if let Some(path) = custom_path {
|
||||
let blocked_keys = match family {
|
||||
LocalVideoCreateFamily::OpenAi => &[][..],
|
||||
LocalVideoCreateFamily::Gemini => &["key"][..],
|
||||
};
|
||||
return build_passthrough_path_url(
|
||||
&transport.endpoint.base_url,
|
||||
path,
|
||||
parts.uri.query(),
|
||||
blocked_keys,
|
||||
);
|
||||
}
|
||||
|
||||
match family {
|
||||
LocalVideoCreateFamily::OpenAi => build_passthrough_path_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.path(),
|
||||
parts.uri.query(),
|
||||
&[],
|
||||
),
|
||||
LocalVideoCreateFamily::Gemini => build_gemini_video_predict_long_running_url(
|
||||
&transport.endpoint.base_url,
|
||||
mapped_model,
|
||||
parts.uri.query(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
async fn materialize_local_video_create_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
api_format: &str,
|
||||
) -> Vec<LocalVideoCreateCandidateAttempt> {
|
||||
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let extra_data = json!({
|
||||
"provider_api_format": api_format,
|
||||
"client_api_format": api_format,
|
||||
"global_model_id": candidate.global_model_id.clone(),
|
||||
"global_model_name": candidate.global_model_name.clone(),
|
||||
"model_id": candidate.model_id.clone(),
|
||||
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
|
||||
"mapping_matched_model": candidate.mapping_matched_model.clone(),
|
||||
"provider_name": candidate.provider_name.clone(),
|
||||
"key_name": candidate.key_name.clone(),
|
||||
});
|
||||
|
||||
let candidate_id = persist_available_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local video decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalVideoCreateCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index: candidate_index as u32,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
|
||||
attempts
|
||||
}
|
||||
|
||||
async fn mark_skipped_local_video_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
persist_skipped_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local video decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn mark_unused_local_video_candidates(
|
||||
state: &AppState,
|
||||
remaining: Vec<LocalSyncPlanAndReport>,
|
||||
) {
|
||||
mark_unused_local_candidate_items(
|
||||
state,
|
||||
remaining,
|
||||
|item| &item.plan,
|
||||
|item| item.report_context.as_ref(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
fn extract_gemini_video_model_from_path(path: &str) -> Option<String> {
|
||||
let suffix = path.strip_prefix("/v1beta/models/")?;
|
||||
let model = suffix.split(':').next()?.trim();
|
||||
if model.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(model.to_string())
|
||||
}
|
||||
|
||||
@@ -0,0 +1,332 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION;
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_passthrough_headers_with_auth, resolve_local_gemini_auth, resolve_local_openai_chat_auth,
|
||||
};
|
||||
use crate::ai_pipeline::transport::policy::{
|
||||
supports_local_gemini_transport_with_network, supports_local_standard_transport_with_network,
|
||||
};
|
||||
use crate::ai_pipeline::transport::url::{
|
||||
build_gemini_video_predict_long_running_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::transport::{
|
||||
apply_local_body_rules, apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, GatewayProviderTransportSnapshot,
|
||||
PlannerAppState,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse};
|
||||
|
||||
use super::support::{
|
||||
mark_skipped_local_video_candidate, LocalVideoCreateCandidateAttempt,
|
||||
LocalVideoCreateDecisionInput,
|
||||
};
|
||||
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
|
||||
|
||||
pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
attempt: LocalVideoCreateCandidateAttempt,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let LocalVideoCreateCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
} = attempt;
|
||||
let transport = match planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision provider transport read failed"
|
||||
);
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_snapshot_read_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let transport_supported = match spec.family {
|
||||
LocalVideoCreateFamily::OpenAi => {
|
||||
supports_local_standard_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
LocalVideoCreateFamily::Gemini => {
|
||||
supports_local_gemini_transport_with_network(&transport, spec.api_format)
|
||||
}
|
||||
};
|
||||
if !transport_supported {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let auth = match spec.family {
|
||||
LocalVideoCreateFamily::OpenAi => resolve_local_openai_chat_auth(&transport),
|
||||
LocalVideoCreateFamily::Gemini => resolve_local_gemini_auth(&transport),
|
||||
};
|
||||
let Some((auth_header, auth_value)) = auth else {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some(upstream_url) =
|
||||
build_video_upstream_url(parts, &transport, &mapped_model, spec.family)
|
||||
else {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(provider_request_body) = build_provider_request_body(
|
||||
body_json,
|
||||
spec.family,
|
||||
&mapped_model,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"provider_request_body_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mut provider_request_headers = build_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
&[&auth_header, "content-type"],
|
||||
&provider_request_body,
|
||||
Some(body_json),
|
||||
) {
|
||||
mark_skipped_local_video_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
&candidate,
|
||||
candidate_index,
|
||||
&candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
let proxy =
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), &transport)
|
||||
.await;
|
||||
let tls_profile = resolve_transport_tls_profile(&transport);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
|
||||
decision_kind: Some(spec.decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.clone()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url),
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some(spec.api_format.to_string()),
|
||||
client_api_format: Some(spec.api_format.to_string()),
|
||||
provider_contract: Some(spec.api_format.to_string()),
|
||||
client_contract: Some(spec.api_format.to_string()),
|
||||
model_name: Some(input.requested_model.clone()),
|
||||
mapped_model: Some(mapped_model.clone()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some(spec.report_kind.to_string()),
|
||||
report_context: Some(json!({
|
||||
"user_id": input.auth_context.user_id.clone(),
|
||||
"api_key_id": input.auth_context.api_key_id.clone(),
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": input.requested_model.clone(),
|
||||
"provider_name": transport.provider.name.clone(),
|
||||
"provider_id": candidate.provider_id.clone(),
|
||||
"endpoint_id": candidate.endpoint_id.clone(),
|
||||
"key_id": candidate.key_id.clone(),
|
||||
"provider_api_format": spec.api_format,
|
||||
"client_api_format": spec.api_format,
|
||||
"mapped_model": mapped_model,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
})),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
|
||||
fn build_provider_request_body(
|
||||
body_json: &serde_json::Value,
|
||||
family: LocalVideoCreateFamily,
|
||||
mapped_model: &str,
|
||||
body_rules: Option<&serde_json::Value>,
|
||||
) -> Option<serde_json::Value> {
|
||||
let mut provider_request_body = match family {
|
||||
LocalVideoCreateFamily::OpenAi => {
|
||||
let mut provider_request_body = body_json.as_object().cloned().unwrap_or_default();
|
||||
provider_request_body
|
||||
.insert("model".to_string(), Value::String(mapped_model.to_string()));
|
||||
serde_json::Value::Object(provider_request_body)
|
||||
}
|
||||
LocalVideoCreateFamily::Gemini => body_json.clone(),
|
||||
};
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
fn build_video_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
family: LocalVideoCreateFamily,
|
||||
) -> Option<String> {
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
if let Some(path) = custom_path {
|
||||
let blocked_keys = match family {
|
||||
LocalVideoCreateFamily::OpenAi => &[][..],
|
||||
LocalVideoCreateFamily::Gemini => &["key"][..],
|
||||
};
|
||||
return build_passthrough_path_url(
|
||||
&transport.endpoint.base_url,
|
||||
path,
|
||||
parts.uri.query(),
|
||||
blocked_keys,
|
||||
);
|
||||
}
|
||||
|
||||
match family {
|
||||
LocalVideoCreateFamily::OpenAi => build_passthrough_path_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.path(),
|
||||
parts.uri.query(),
|
||||
&[],
|
||||
),
|
||||
LocalVideoCreateFamily::Gemini => build_gemini_video_predict_long_running_url(
|
||||
&transport.endpoint.base_url,
|
||||
mapped_model,
|
||||
parts.uri.query(),
|
||||
),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision,
|
||||
};
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::AppState;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalVideoCreateDecisionInput {
|
||||
pub(super) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalVideoCreateCandidateAttempt {
|
||||
pub(super) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(super) candidate_index: u32,
|
||||
pub(super) candidate_id: String,
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_local_video_create_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<LocalVideoCreateDecisionInput> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = match spec.family {
|
||||
LocalVideoCreateFamily::OpenAi => body_json
|
||||
.get("model")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)?,
|
||||
LocalVideoCreateFamily::Gemini => extract_gemini_video_model_from_path(parts.uri.path())?,
|
||||
};
|
||||
|
||||
let auth_snapshot = match planner_state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalVideoCreateDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn list_local_video_create_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
api_format: &str,
|
||||
decision_kind: &str,
|
||||
) -> Option<Vec<LocalVideoCreateCandidateAttempt>> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidates = match planner_state
|
||||
.list_selectable_candidates(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
false,
|
||||
Some(&input.auth_snapshot),
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision scheduler selection failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(
|
||||
materialize_local_video_create_candidate_attempts(
|
||||
planner_state,
|
||||
trace_id,
|
||||
input,
|
||||
candidates,
|
||||
api_format,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn materialize_local_video_create_candidate_attempts(
|
||||
state: PlannerAppState<'_>,
|
||||
trace_id: &str,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
api_format: &str,
|
||||
) -> Vec<LocalVideoCreateCandidateAttempt> {
|
||||
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let extra_data = json!({
|
||||
"provider_api_format": api_format,
|
||||
"client_api_format": api_format,
|
||||
"global_model_id": candidate.global_model_id.clone(),
|
||||
"global_model_name": candidate.global_model_name.clone(),
|
||||
"model_id": candidate.model_id.clone(),
|
||||
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
|
||||
"mapping_matched_model": candidate.mapping_matched_model.clone(),
|
||||
"provider_name": candidate.provider_name.clone(),
|
||||
"key_name": candidate.key_name.clone(),
|
||||
});
|
||||
|
||||
let candidate_id = state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local video decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalVideoCreateCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index: candidate_index as u32,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
|
||||
attempts
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_video_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
PlannerAppState::new(state)
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local video decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
fn extract_gemini_video_model_from_path(path: &str) -> Option<String> {
|
||||
let suffix = path.strip_prefix("/v1beta/models/")?;
|
||||
let model = suffix.split(':').next()?.trim();
|
||||
if model.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(model.to_string())
|
||||
}
|
||||
@@ -1,15 +1,14 @@
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
use aether_ai_pipeline::planner::standard::claude::{
|
||||
resolve_stream_spec as resolve_pipeline_stream_spec,
|
||||
resolve_sync_spec as resolve_pipeline_sync_spec,
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_claude_stream_spec as resolve_pipeline_stream_spec,
|
||||
resolve_claude_sync_spec as resolve_pipeline_sync_spec,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
use super::family::{
|
||||
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::conversion::request::normalize_claude_request_to_openai_chat_request;
|
||||
pub(crate) use crate::ai_pipeline::normalize_claude_request_to_openai_chat_request;
|
||||
|
||||
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
|
||||
resolve_pipeline_sync_spec(plan_kind)
|
||||
|
||||
@@ -1,433 +1,7 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt::Write;
|
||||
|
||||
use crate::ai_pipeline::provider_transport_facade::body_rules_handle_path;
|
||||
use serde_json::{json, Value};
|
||||
use sha1::{Digest as Sha1Digest, Sha1};
|
||||
use sha2::{Digest as Sha2Digest, Sha256};
|
||||
use uuid::Uuid;
|
||||
|
||||
const CODEX_PROMPT_CACHE_NAMESPACE_VERSION: &str = "v3";
|
||||
const UUID_NAMESPACE_OID_BYTES: [u8; 16] = [
|
||||
0x6b, 0xa7, 0xb8, 0x12, 0x9d, 0xad, 0x11, 0xd1, 0x80, 0xb4, 0x00, 0xc0, 0x4f, 0xd4, 0x30, 0xc8,
|
||||
];
|
||||
|
||||
fn is_codex_openai_cli_request(provider_type: &str, provider_api_format: &str) -> bool {
|
||||
provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& matches!(
|
||||
provider_api_format.trim().to_ascii_lowercase().as_str(),
|
||||
"openai:cli" | "openai:compact"
|
||||
)
|
||||
}
|
||||
|
||||
fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String> {
|
||||
let normalized = user_api_key_id.trim();
|
||||
if normalized.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let namespace = format!(
|
||||
"aether:codex:prompt-cache:{CODEX_PROMPT_CACHE_NAMESPACE_VERSION}:user:{normalized}"
|
||||
);
|
||||
let mut hasher = Sha1::new();
|
||||
hasher.update(UUID_NAMESPACE_OID_BYTES);
|
||||
hasher.update(namespace.as_bytes());
|
||||
|
||||
let digest = hasher.finalize();
|
||||
let mut bytes = [0u8; 16];
|
||||
bytes.copy_from_slice(&digest[..16]);
|
||||
bytes[6] = (bytes[6] & 0x0f) | 0x50;
|
||||
bytes[8] = (bytes[8] & 0x3f) | 0x80;
|
||||
Some(Uuid::from_bytes(bytes).to_string())
|
||||
}
|
||||
|
||||
fn maybe_inject_codex_prompt_cache_key(
|
||||
provider_request_body: &mut Value,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
user_api_key_id: Option<&str>,
|
||||
) {
|
||||
if !is_codex_openai_cli_request(provider_type, provider_api_format) {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(body_object) = provider_request_body.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
|
||||
let existing = body_object
|
||||
.get("prompt_cache_key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !existing.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(prompt_cache_key) = user_api_key_id.and_then(build_stable_codex_prompt_cache_key)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
body_object.insert(
|
||||
"prompt_cache_key".to_string(),
|
||||
Value::String(prompt_cache_key),
|
||||
);
|
||||
}
|
||||
|
||||
fn build_short_codex_header_id(seed: &str) -> Option<String> {
|
||||
let normalized = seed.trim();
|
||||
if normalized.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let digest = Sha256::digest(normalized.as_bytes());
|
||||
let mut short_id = String::with_capacity(16);
|
||||
for byte in digest.iter().take(8) {
|
||||
let _ = write!(&mut short_id, "{byte:02x}");
|
||||
}
|
||||
Some(short_id)
|
||||
}
|
||||
|
||||
fn header_map_has_non_empty_value(headers: &http::HeaderMap, header_name: &str) -> bool {
|
||||
let target = header_name.trim().to_ascii_lowercase();
|
||||
if target.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
headers.iter().any(|(name, value)| {
|
||||
if name.as_str().trim().to_ascii_lowercase() != target {
|
||||
return false;
|
||||
}
|
||||
value
|
||||
.to_str()
|
||||
.ok()
|
||||
.map(str::trim)
|
||||
.map(|value| !value.is_empty())
|
||||
.unwrap_or(false)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_codex_account_id(decrypted_auth_config_raw: Option<&str>) -> Option<String> {
|
||||
let raw = decrypted_auth_config_raw?.trim();
|
||||
if raw.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
|
||||
value
|
||||
.get("account_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn apply_codex_openai_cli_special_headers(
|
||||
provider_request_headers: &mut BTreeMap<String, String>,
|
||||
provider_request_body: &Value,
|
||||
original_headers: &http::HeaderMap,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
request_id: Option<&str>,
|
||||
decrypted_auth_config_raw: Option<&str>,
|
||||
) {
|
||||
if !is_codex_openai_cli_request(provider_type, provider_api_format) {
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(account_id) = extract_codex_account_id(decrypted_auth_config_raw) {
|
||||
provider_request_headers.insert("chatgpt-account-id".to_string(), account_id);
|
||||
}
|
||||
if !provider_request_headers
|
||||
.get("x-client-request-id")
|
||||
.map(|value| !value.trim().is_empty())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
if let Some(request_id) = request_id.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
provider_request_headers
|
||||
.insert("x-client-request-id".to_string(), request_id.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let prompt_cache_key = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
let Some(short_id) = prompt_cache_key.and_then(build_short_codex_header_id) else {
|
||||
return;
|
||||
};
|
||||
|
||||
if !header_map_has_non_empty_value(original_headers, "session_id") {
|
||||
provider_request_headers.insert("session_id".to_string(), short_id.clone());
|
||||
}
|
||||
|
||||
if provider_api_format.trim().to_ascii_lowercase() != "openai:compact"
|
||||
&& !header_map_has_non_empty_value(original_headers, "conversation_id")
|
||||
{
|
||||
provider_request_headers.insert("conversation_id".to_string(), short_id);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn apply_codex_openai_cli_special_body_edits(
|
||||
provider_request_body: &mut Value,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) {
|
||||
if !is_codex_openai_cli_request(provider_type, provider_api_format) {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(body_object) = provider_request_body.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
|
||||
if !body_rules_handle_path(body_rules, "max_output_tokens") {
|
||||
body_object.remove("max_output_tokens");
|
||||
}
|
||||
if !body_rules_handle_path(body_rules, "temperature") {
|
||||
body_object.remove("temperature");
|
||||
}
|
||||
if !body_rules_handle_path(body_rules, "top_p") {
|
||||
body_object.remove("top_p");
|
||||
}
|
||||
if !body_rules_handle_path(body_rules, "metadata") {
|
||||
body_object.remove("metadata");
|
||||
}
|
||||
if !body_rules_handle_path(body_rules, "store") {
|
||||
body_object.insert("store".to_string(), json!(false));
|
||||
}
|
||||
if !body_rules_handle_path(body_rules, "instructions")
|
||||
&& !body_object.contains_key("instructions")
|
||||
{
|
||||
body_object.insert("instructions".to_string(), json!("You are GPT-5."));
|
||||
}
|
||||
|
||||
maybe_inject_codex_prompt_cache_key(
|
||||
provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
user_api_key_id,
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
#[path = "codex/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
use super::{
|
||||
apply_codex_openai_cli_special_body_edits, apply_codex_openai_cli_special_headers,
|
||||
};
|
||||
use http::{HeaderMap, HeaderValue};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"max_output_tokens": 128,
|
||||
"temperature": 0.3,
|
||||
"top_p": 0.9,
|
||||
"metadata": {"client": "desktop"},
|
||||
"store": true
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(&mut body, "codex", "openai:cli", None, None);
|
||||
|
||||
assert!(body.get("max_output_tokens").is_none());
|
||||
assert!(body.get("temperature").is_none());
|
||||
assert!(body.get("top_p").is_none());
|
||||
assert!(body.get("metadata").is_none());
|
||||
assert_eq!(body["store"], false);
|
||||
assert_eq!(body["instructions"], "You are GPT-5.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defers_to_user_body_rules_for_handled_fields() {
|
||||
let body_rules = json!([
|
||||
{"action":"set","path":"store","value":true},
|
||||
{"action":"set","path":"instructions","value":"Keep custom"},
|
||||
{"action":"set","path":"metadata","value":{"client":"desktop","mode":"custom"}},
|
||||
{"action":"set","path":"top_p","value":0.5}
|
||||
]);
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"max_output_tokens": 128,
|
||||
"metadata": {"client": "desktop", "mode": "custom"},
|
||||
"store": true,
|
||||
"instructions": "Keep custom",
|
||||
"top_p": 0.5
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:compact",
|
||||
Some(&body_rules),
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(body.get("max_output_tokens").is_none());
|
||||
assert_eq!(body["store"], true);
|
||||
assert_eq!(body["instructions"], "Keep custom");
|
||||
assert_eq!(body["metadata"]["mode"], "custom");
|
||||
assert_eq!(body["top_p"], 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_stable_prompt_cache_key_for_codex_requests() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:cli",
|
||||
None,
|
||||
Some("key-123"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_existing_prompt_cache_key_for_codex_requests() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
"prompt_cache_key": "existing-key",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:cli",
|
||||
None,
|
||||
Some("key-123"),
|
||||
);
|
||||
|
||||
assert_eq!(body["prompt_cache_key"], "existing-key");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_chatgpt_account_id_and_session_headers_for_codex_requests() {
|
||||
let mut headers = BTreeMap::new();
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&HeaderMap::new(),
|
||||
"codex",
|
||||
"openai:cli",
|
||||
Some("trace-codex-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id"),
|
||||
Some(&"acc-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"trace-codex-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("conversation_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn respects_existing_codex_request_and_session_headers() {
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert(
|
||||
"x-client-request-id".to_string(),
|
||||
"kept-by-rule-request".to_string(),
|
||||
);
|
||||
headers.insert("session_id".to_string(), "kept-by-rule".to_string());
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
let mut original_headers = HeaderMap::new();
|
||||
original_headers.insert(
|
||||
"x-client-request-id",
|
||||
HeaderValue::from_static("user-specified-request"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"session_id",
|
||||
HeaderValue::from_static("user-specified-session"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"conversation_id",
|
||||
HeaderValue::from_static("user-specified-conversation"),
|
||||
);
|
||||
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&original_headers,
|
||||
"codex",
|
||||
"openai:cli",
|
||||
Some("trace-codex-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"kept-by-rule-request".to_string())
|
||||
);
|
||||
assert_eq!(headers.get("session_id"), Some(&"kept-by-rule".to_string()));
|
||||
assert!(headers.get("conversation_id").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_conversation_id_for_compact_codex_requests() {
|
||||
let mut headers = BTreeMap::new();
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&HeaderMap::new(),
|
||||
"codex",
|
||||
"openai:compact",
|
||||
Some("trace-codex-compact-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id"),
|
||||
Some(&"acc-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"trace-codex-compact-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
assert!(headers.get("conversation_id").is_none());
|
||||
}
|
||||
}
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
apply_codex_openai_cli_special_body_edits, apply_codex_openai_cli_special_headers,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{apply_codex_openai_cli_special_body_edits, apply_codex_openai_cli_special_headers};
|
||||
use http::{HeaderMap, HeaderValue};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"max_output_tokens": 128,
|
||||
"temperature": 0.3,
|
||||
"top_p": 0.9,
|
||||
"metadata": {"client": "desktop"},
|
||||
"store": true
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(&mut body, "codex", "openai:cli", None, None);
|
||||
|
||||
assert!(body.get("max_output_tokens").is_none());
|
||||
assert!(body.get("temperature").is_none());
|
||||
assert!(body.get("top_p").is_none());
|
||||
assert!(body.get("metadata").is_none());
|
||||
assert_eq!(body["store"], false);
|
||||
assert_eq!(body["instructions"], "You are GPT-5.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defers_to_user_body_rules_for_handled_fields() {
|
||||
let body_rules = json!([
|
||||
{"action":"set","path":"store","value":true},
|
||||
{"action":"set","path":"instructions","value":"Keep custom"},
|
||||
{"action":"set","path":"metadata","value":{"client":"desktop","mode":"custom"}},
|
||||
{"action":"set","path":"top_p","value":0.5}
|
||||
]);
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"max_output_tokens": 128,
|
||||
"metadata": {"client": "desktop", "mode": "custom"},
|
||||
"store": true,
|
||||
"instructions": "Keep custom",
|
||||
"top_p": 0.5
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:compact",
|
||||
Some(&body_rules),
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(body.get("max_output_tokens").is_none());
|
||||
assert_eq!(body["store"], true);
|
||||
assert_eq!(body["instructions"], "Keep custom");
|
||||
assert_eq!(body["metadata"]["mode"], "custom");
|
||||
assert_eq!(body["top_p"], 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_stable_prompt_cache_key_for_codex_requests() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:cli",
|
||||
None,
|
||||
Some("key-123"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_existing_prompt_cache_key_for_codex_requests() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
"prompt_cache_key": "existing-key",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut body,
|
||||
"codex",
|
||||
"openai:cli",
|
||||
None,
|
||||
Some("key-123"),
|
||||
);
|
||||
|
||||
assert_eq!(body["prompt_cache_key"], "existing-key");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_chatgpt_account_id_and_session_headers_for_codex_requests() {
|
||||
let mut headers = BTreeMap::new();
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&HeaderMap::new(),
|
||||
"codex",
|
||||
"openai:cli",
|
||||
Some("trace-codex-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id"),
|
||||
Some(&"acc-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"trace-codex-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("conversation_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn respects_existing_codex_request_and_session_headers() {
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert(
|
||||
"x-client-request-id".to_string(),
|
||||
"kept-by-rule-request".to_string(),
|
||||
);
|
||||
headers.insert("session_id".to_string(), "kept-by-rule".to_string());
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
let mut original_headers = HeaderMap::new();
|
||||
original_headers.insert(
|
||||
"x-client-request-id",
|
||||
HeaderValue::from_static("user-specified-request"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"session_id",
|
||||
HeaderValue::from_static("user-specified-session"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"conversation_id",
|
||||
HeaderValue::from_static("user-specified-conversation"),
|
||||
);
|
||||
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&original_headers,
|
||||
"codex",
|
||||
"openai:cli",
|
||||
Some("trace-codex-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"kept-by-rule-request".to_string())
|
||||
);
|
||||
assert_eq!(headers.get("session_id"), Some(&"kept-by-rule".to_string()));
|
||||
assert!(headers.get("conversation_id").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_conversation_id_for_compact_codex_requests() {
|
||||
let mut headers = BTreeMap::new();
|
||||
let body = json!({
|
||||
"model": "gpt-5",
|
||||
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut headers,
|
||||
&body,
|
||||
&HeaderMap::new(),
|
||||
"codex",
|
||||
"openai:compact",
|
||||
Some("trace-codex-compact-123"),
|
||||
Some(r#"{"account_id":"acc-123"}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id"),
|
||||
Some(&"acc-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-client-request-id"),
|
||||
Some(&"trace-codex-compact-123".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
);
|
||||
assert!(headers.get("conversation_id").is_none());
|
||||
}
|
||||
@@ -1,18 +1,18 @@
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
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::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
use super::candidates::{
|
||||
materialize_local_standard_candidate_attempts, resolve_local_standard_decision_input,
|
||||
};
|
||||
use super::payload::maybe_build_local_standard_decision_payload_for_candidate;
|
||||
use super::types::{LocalStandardSourceFamily, LocalStandardSpec};
|
||||
use super::{LocalStandardSourceFamily, LocalStandardSpec};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
state: &AppState,
|
||||
|
||||
@@ -5,18 +5,16 @@ use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::{
|
||||
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_available_local_candidate;
|
||||
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_decision_execution_runtime_auth_context, ConversionMode, ExecutionStrategy,
|
||||
GatewayControlDecision,
|
||||
};
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{append_execution_contract_fields_to_value, AppState, GatewayError};
|
||||
|
||||
use super::types::{
|
||||
use super::{
|
||||
LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSourceFamily,
|
||||
LocalStandardSourceMode, LocalStandardSpec,
|
||||
};
|
||||
@@ -29,9 +27,8 @@ pub(super) async fn resolve_local_standard_decision_input(
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<LocalStandardDecisionInput> {
|
||||
let Some(auth_context) = decision.auth_context.clone().filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
}) else {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
@@ -45,13 +42,13 @@ pub(super) async fn resolve_local_standard_decision_input(
|
||||
LocalStandardSourceFamily::Gemini => extract_gemini_model_from_path(parts.uri.path())?,
|
||||
};
|
||||
|
||||
let auth_snapshot = match read_auth_api_key_snapshot(
|
||||
state,
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
let auth_snapshot = match planner_state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => return None,
|
||||
@@ -79,6 +76,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
input: &LocalStandardDecisionInput,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Result<Vec<LocalStandardCandidateAttempt>, GatewayError> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let mut seen_candidates = BTreeSet::new();
|
||||
let mut candidates = Vec::new();
|
||||
for candidate_api_format in candidate_api_formats_for_spec(spec) {
|
||||
@@ -87,15 +85,15 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let mut selected_candidates = list_selectable_candidates(
|
||||
state,
|
||||
candidate_api_format,
|
||||
&input.requested_model,
|
||||
spec.require_streaming,
|
||||
auth_snapshot,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
let mut selected_candidates = planner_state
|
||||
.list_selectable_candidates(
|
||||
candidate_api_format,
|
||||
&input.requested_model,
|
||||
spec.require_streaming,
|
||||
auth_snapshot,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await?;
|
||||
if auth_snapshot.is_none() {
|
||||
selected_candidates.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_candidate(
|
||||
@@ -120,7 +118,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
}
|
||||
}
|
||||
}
|
||||
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
@@ -160,19 +158,19 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
candidate.endpoint_api_format.as_str(),
|
||||
);
|
||||
|
||||
let stored_candidate_id = persist_available_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local standard decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
let stored_candidate_id = planner_state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local standard decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalStandardCandidateAttempt {
|
||||
candidate,
|
||||
|
||||
@@ -1,12 +1,28 @@
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
|
||||
|
||||
mod build;
|
||||
mod candidates;
|
||||
mod payload;
|
||||
mod types;
|
||||
|
||||
pub(crate) use self::build::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
|
||||
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
|
||||
};
|
||||
pub(crate) use self::types::{
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalStandardDecisionInput {
|
||||
pub(super) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalStandardCandidateAttempt {
|
||||
pub(super) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(super) candidate_index: u32,
|
||||
pub(super) candidate_id: String,
|
||||
}
|
||||
|
||||
@@ -3,29 +3,24 @@ use std::collections::BTreeMap;
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::control_facade::collect_control_headers;
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::candidate_runtime_facade::persist_skipped_local_candidate;
|
||||
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
|
||||
use crate::ai_pipeline::planner::standard::apply_codex_openai_cli_special_headers;
|
||||
use crate::ai_pipeline::planner::transport_facade::{
|
||||
read_provider_transport_snapshot, resolve_local_oauth_request_auth,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::auth::{
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::{
|
||||
use crate::ai_pipeline::transport::{
|
||||
apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::{collect_control_headers, ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::{LocalResolvedOAuthRequestAuth, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
|
||||
use super::types::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
|
||||
use super::{LocalStandardCandidateAttempt, LocalStandardDecisionInput, LocalStandardSpec};
|
||||
|
||||
pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
@@ -36,6 +31,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
attempt: LocalStandardCandidateAttempt,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let LocalStandardCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index,
|
||||
@@ -52,13 +48,13 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
return None;
|
||||
};
|
||||
|
||||
let transport = match read_provider_transport_snapshot(
|
||||
state,
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
let transport = match planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
@@ -115,7 +111,10 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
let resolved_auth =
|
||||
crate::ai_pipeline::conversion::request_conversion_direct_auth(&transport, conversion_kind);
|
||||
let oauth_auth = if resolved_auth.is_none() {
|
||||
match resolve_local_oauth_request_auth(state, &transport).await {
|
||||
match planner_state
|
||||
.resolve_local_oauth_request_auth(&transport)
|
||||
.await
|
||||
{
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
|
||||
Ok(None) => None,
|
||||
@@ -342,17 +341,17 @@ pub(super) async fn mark_skipped_local_standard_candidate(
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
persist_skipped_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local standard decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
PlannerAppState::new(state)
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local standard decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
use crate::ai_pipeline::control_facade::GatewayControlAuthContext;
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
|
||||
|
||||
pub(crate) use aether_ai_pipeline::planner::standard::family::{
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalStandardDecisionInput {
|
||||
pub(super) auth_context: GatewayControlAuthContext,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalStandardCandidateAttempt {
|
||||
pub(super) candidate: aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(super) candidate_index: u32,
|
||||
pub(super) candidate_id: String,
|
||||
}
|
||||
@@ -1,15 +1,14 @@
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
use aether_ai_pipeline::planner::standard::gemini::{
|
||||
resolve_stream_spec as resolve_pipeline_stream_spec,
|
||||
resolve_sync_spec as resolve_pipeline_sync_spec,
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_gemini_stream_spec as resolve_pipeline_stream_spec,
|
||||
resolve_gemini_sync_spec as resolve_pipeline_sync_spec,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
use super::family::{
|
||||
maybe_build_stream_via_standard_family_payload, maybe_build_sync_via_standard_family_payload,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::conversion::request::normalize_gemini_request_to_openai_chat_request;
|
||||
pub(crate) use crate::ai_pipeline::normalize_gemini_request_to_openai_chat_request;
|
||||
|
||||
pub(crate) fn resolve_sync_spec(plan_kind: &str) -> Option<super::family::LocalStandardSpec> {
|
||||
resolve_pipeline_sync_spec(plan_kind)
|
||||
|
||||
@@ -4,7 +4,7 @@ use super::{
|
||||
augment_sync_report_context, generic_decision_missing_exact_provider_request,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::ensure_upstream_auth_header;
|
||||
use crate::ai_pipeline::transport::ensure_upstream_auth_header;
|
||||
use crate::{GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) fn build_gemini_sync_plan_from_decision(
|
||||
|
||||
@@ -1,339 +0,0 @@
|
||||
use super::codex::apply_codex_openai_cli_special_body_edits;
|
||||
use crate::ai_pipeline::planner::transport_facade::GatewayProviderTransportSnapshot;
|
||||
use crate::ai_pipeline::provider_transport_facade::apply_local_body_rules;
|
||||
use crate::ai_pipeline::provider_transport_facade::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
|
||||
build_openai_cli_url, build_passthrough_path_url,
|
||||
};
|
||||
use aether_ai_pipeline::planner::matrix::build_standard_request_body_from_canonical;
|
||||
pub(crate) use aether_ai_pipeline::planner::standard::normalize_standard_request_to_openai_chat_request;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
pub(crate) fn build_standard_request_body(
|
||||
body_json: &Value,
|
||||
client_api_format: &str,
|
||||
mapped_model: &str,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
request_path: &str,
|
||||
upstream_is_stream: bool,
|
||||
body_rules: Option<&Value>,
|
||||
user_api_key_id: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let canonical_request = normalize_standard_request_to_openai_chat_request(
|
||||
body_json,
|
||||
client_api_format,
|
||||
request_path,
|
||||
)?;
|
||||
if cfg!(test) {
|
||||
println!("canonical_request: {canonical_request:#?}");
|
||||
}
|
||||
let mut provider_request_body = build_standard_request_body_from_canonical(
|
||||
&canonical_request,
|
||||
mapped_model,
|
||||
provider_api_format,
|
||||
upstream_is_stream,
|
||||
)?;
|
||||
if cfg!(test) {
|
||||
println!("provider_request_body before rules: {provider_request_body:#?}");
|
||||
}
|
||||
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_standard_upstream_url(
|
||||
parts: &http::request::Parts,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
mapped_model: &str,
|
||||
provider_api_format: &str,
|
||||
upstream_is_stream: 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 => match provider_api_format.trim().to_ascii_lowercase().as_str() {
|
||||
"openai:chat" => Some(build_openai_chat_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
"openai:cli" => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
"openai:compact" => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
true,
|
||||
)),
|
||||
"claude:chat" | "claude:cli" => Some(build_claude_messages_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
"gemini:chat" | "gemini:cli" => build_gemini_content_url(
|
||||
&transport.endpoint.base_url,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::build_standard_request_body;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn builds_openai_chat_request_from_claude_chat_source() {
|
||||
let request = json!({
|
||||
"model": "claude-3-7-sonnet",
|
||||
"system": "You are concise.",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Hello from Claude"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 128
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:chat",
|
||||
"gpt-5",
|
||||
"openai",
|
||||
"openai:chat",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude chat should convert to openai chat");
|
||||
|
||||
assert_eq!(converted["model"], "gpt-5");
|
||||
assert_eq!(converted["messages"][0]["role"], "system");
|
||||
assert_eq!(converted["messages"][0]["content"], "You are concise.");
|
||||
assert_eq!(converted["messages"][1]["role"], "user");
|
||||
assert_eq!(converted["messages"][1]["content"], "Hello from Claude");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_claude_chat_request_from_gemini_chat_source() {
|
||||
let request = json!({
|
||||
"systemInstruction": {
|
||||
"parts": [{"text": "Be brief."}]
|
||||
},
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"text": "Hello from Gemini"}]
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"gemini:chat",
|
||||
"claude-sonnet-4-5",
|
||||
"anthropic",
|
||||
"claude:chat",
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("gemini chat should convert to claude chat");
|
||||
|
||||
assert_eq!(converted["model"], "claude-sonnet-4-5");
|
||||
assert_eq!(converted["messages"][0]["role"], "user");
|
||||
assert!(
|
||||
converted["messages"]
|
||||
.to_string()
|
||||
.contains("Hello from Gemini"),
|
||||
"converted claude payload should retain the gemini user text: {converted}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_gemini_cli_request_from_claude_cli_source() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need CLI output"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gemini-2.5-pro",
|
||||
"google",
|
||||
"gemini:cli",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to gemini cli");
|
||||
|
||||
assert_eq!(converted["contents"][0]["role"], "user");
|
||||
assert_eq!(
|
||||
converted["contents"][0]["parts"][0]["text"],
|
||||
"Need CLI output"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_cli_request_from_claude_cli_source_with_forced_stream() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gpt-5",
|
||||
"openai",
|
||||
"openai:cli",
|
||||
"/v1/messages",
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to openai cli");
|
||||
|
||||
assert_eq!(converted["model"], "gpt-5");
|
||||
assert_eq!(converted["input"][0]["role"], "user");
|
||||
assert_eq!(converted["input"][0]["content"][0]["type"], "input_text");
|
||||
assert_eq!(
|
||||
converted["input"][0]["content"][0]["text"],
|
||||
"Need OpenAI CLI output"
|
||||
);
|
||||
assert_eq!(converted["stream"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_cli_requests() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
|
||||
}],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gpt-5",
|
||||
"codex",
|
||||
"openai:cli",
|
||||
"/v1/messages",
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to codex cli");
|
||||
|
||||
assert!(converted.get("metadata").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_codex_defaults_unless_body_rules_handle_the_field() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
|
||||
}],
|
||||
"max_tokens": 64
|
||||
});
|
||||
let body_rules = json!([
|
||||
{"action":"set","path":"store","value":true},
|
||||
{"action":"set","path":"instructions","value":"Custom instructions"},
|
||||
{"action":"set","path":"metadata","value":{"trace_id":"keep-me"}}
|
||||
]);
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gpt-5",
|
||||
"codex",
|
||||
"openai:cli",
|
||||
"/v1/messages",
|
||||
true,
|
||||
Some(&body_rules),
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to codex cli");
|
||||
|
||||
assert_eq!(converted["store"], true);
|
||||
assert_eq!(converted["instructions"], "Custom instructions");
|
||||
assert_eq!(converted["metadata"]["trace_id"], "keep-me");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_codex_prompt_cache_key_for_standard_requests() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
|
||||
}],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gpt-5",
|
||||
"codex",
|
||||
"openai:cli",
|
||||
"/v1/messages",
|
||||
true,
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("claude cli should convert to codex cli");
|
||||
|
||||
assert_eq!(
|
||||
converted["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -3,21 +3,19 @@
|
||||
//! This groups the standard planning surface in one place:
|
||||
//! request-side conversion, matrix registry, and decision payload builders.
|
||||
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) mod claude;
|
||||
mod claude;
|
||||
mod codex;
|
||||
pub(crate) mod family;
|
||||
pub(crate) mod gemini;
|
||||
mod matrix;
|
||||
mod family;
|
||||
mod gemini;
|
||||
mod normalize;
|
||||
pub(crate) mod openai;
|
||||
mod openai;
|
||||
|
||||
pub(crate) use self::codex::apply_codex_openai_cli_special_headers;
|
||||
pub(crate) use self::matrix::{
|
||||
build_standard_request_body, build_standard_upstream_url,
|
||||
normalize_standard_request_to_openai_chat_request,
|
||||
pub(crate) use self::family::{
|
||||
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
|
||||
};
|
||||
pub(crate) use self::normalize::{
|
||||
build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url,
|
||||
@@ -26,17 +24,16 @@ pub(crate) use self::normalize::{
|
||||
build_local_openai_cli_request_body, build_local_openai_cli_upstream_url,
|
||||
};
|
||||
pub(crate) use self::openai::{
|
||||
copy_request_number_field, copy_request_number_field_as,
|
||||
map_openai_reasoning_effort_to_claude_output, map_openai_reasoning_effort_to_gemini_budget,
|
||||
maybe_build_stream_local_decision_payload,
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_cli_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_cli_sync_plan_and_reports_for_kind, copy_request_number_field,
|
||||
copy_request_number_field_as, map_openai_reasoning_effort_to_claude_output,
|
||||
map_openai_reasoning_effort_to_gemini_budget, maybe_build_stream_local_decision_payload,
|
||||
maybe_build_stream_local_openai_cli_decision_payload, maybe_build_sync_local_decision_payload,
|
||||
maybe_build_sync_local_openai_cli_decision_payload, parse_openai_stop_sequences,
|
||||
resolve_openai_chat_max_tokens, value_as_u64,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::conversion::request::{
|
||||
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
|
||||
convert_openai_chat_request_to_openai_cli_request, extract_openai_text_content,
|
||||
normalize_openai_cli_request_to_openai_chat_request, parse_openai_tool_result_content,
|
||||
resolve_openai_chat_max_tokens, set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
value_as_u64,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::conversion::{
|
||||
build_core_error_body_for_client_format, request_conversion_kind,
|
||||
@@ -44,6 +41,13 @@ pub(crate) use crate::ai_pipeline::conversion::{
|
||||
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
|
||||
SyncCliResponseConversionKind,
|
||||
};
|
||||
pub(crate) use crate::ai_pipeline::normalize_standard_request_to_openai_chat_request;
|
||||
pub(crate) use crate::ai_pipeline::{
|
||||
build_standard_request_body, build_standard_upstream_url,
|
||||
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
|
||||
convert_openai_chat_request_to_openai_cli_request, extract_openai_text_content,
|
||||
normalize_openai_cli_request_to_openai_chat_request, parse_openai_tool_result_content,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_standard_decision_payload(
|
||||
state: &AppState,
|
||||
@@ -88,3 +92,179 @@ pub(crate) async fn maybe_build_stream_local_standard_decision_payload(
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::build_standard_request_body;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn builds_openai_chat_request_from_claude_chat_source() {
|
||||
let request = json!({
|
||||
"model": "claude-3-7-sonnet",
|
||||
"system": "You are concise.",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Hello from Claude"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 128
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:chat",
|
||||
"gpt-5",
|
||||
"openai",
|
||||
"openai:chat",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude chat should convert to openai chat");
|
||||
|
||||
assert_eq!(converted["model"], "gpt-5");
|
||||
assert_eq!(converted["messages"][0]["role"], "system");
|
||||
assert_eq!(converted["messages"][0]["content"], "You are concise.");
|
||||
assert_eq!(converted["messages"][1]["role"], "user");
|
||||
assert_eq!(converted["messages"][1]["content"], "Hello from Claude");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_claude_chat_request_from_gemini_chat_source() {
|
||||
let request = json!({
|
||||
"systemInstruction": {
|
||||
"parts": [{"text": "Be brief."}]
|
||||
},
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"text": "Hello from Gemini"}]
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"gemini:chat",
|
||||
"claude-sonnet-4-5",
|
||||
"anthropic",
|
||||
"claude:chat",
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("gemini chat should convert to claude chat");
|
||||
|
||||
assert_eq!(converted["model"], "claude-sonnet-4-5");
|
||||
assert_eq!(converted["messages"][0]["role"], "user");
|
||||
assert!(
|
||||
converted["messages"]
|
||||
.to_string()
|
||||
.contains("Hello from Gemini"),
|
||||
"converted claude payload should retain the gemini user text: {converted}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_gemini_cli_request_from_claude_cli_source() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need CLI output"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gemini-2.5-pro",
|
||||
"google",
|
||||
"gemini:cli",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to gemini cli");
|
||||
|
||||
assert_eq!(converted["contents"][0]["role"], "user");
|
||||
assert_eq!(
|
||||
converted["contents"][0]["parts"][0]["text"],
|
||||
"Need CLI output"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_cli_request_from_claude_cli_source_with_forced_stream() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
|
||||
}
|
||||
],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gpt-5",
|
||||
"openai",
|
||||
"openai:cli",
|
||||
"/v1/messages",
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to openai cli");
|
||||
|
||||
assert_eq!(converted["model"], "gpt-5");
|
||||
assert_eq!(converted["input"][0]["role"], "user");
|
||||
assert_eq!(converted["input"][0]["content"][0]["type"], "input_text");
|
||||
assert_eq!(
|
||||
converted["input"][0]["content"][0]["text"],
|
||||
"Need OpenAI CLI output"
|
||||
);
|
||||
assert_eq!(converted["stream"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_cli_requests() {
|
||||
let request = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Need OpenAI CLI output"}]
|
||||
}],
|
||||
"max_tokens": 64
|
||||
});
|
||||
|
||||
let converted = build_standard_request_body(
|
||||
&request,
|
||||
"claude:cli",
|
||||
"gpt-5",
|
||||
"codex",
|
||||
"openai:cli",
|
||||
"/v1/messages",
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli should convert to codex request");
|
||||
|
||||
assert!(converted.get("metadata").is_none());
|
||||
assert_eq!(converted["store"], false);
|
||||
assert_eq!(converted["instructions"], "You are GPT-5.");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,422 +1,16 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
use url::form_urlencoded;
|
||||
|
||||
use super::codex::apply_codex_openai_cli_special_body_edits;
|
||||
use crate::ai_pipeline::conversion::{request_conversion_kind, RequestConversionKind};
|
||||
use crate::ai_pipeline::planner::transport_facade::GatewayProviderTransportSnapshot;
|
||||
use crate::ai_pipeline::provider_transport_facade::antigravity::{
|
||||
build_antigravity_v1internal_url, AntigravityRequestUrlAction,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::apply_local_body_rules;
|
||||
use crate::ai_pipeline::provider_transport_facade::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
|
||||
build_openai_cli_url, build_passthrough_path_url,
|
||||
};
|
||||
use aether_ai_pipeline::planner::standard::normalize::{
|
||||
build_cross_format_openai_chat_request_body as pipeline_build_cross_format_openai_chat_request_body,
|
||||
build_cross_format_openai_cli_request_body as pipeline_build_cross_format_openai_cli_request_body,
|
||||
build_local_openai_chat_request_body as pipeline_build_local_openai_chat_request_body,
|
||||
build_local_openai_cli_request_body as pipeline_build_local_openai_cli_request_body,
|
||||
};
|
||||
|
||||
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_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
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 => build_gemini_content_url(
|
||||
&transport.endpoint.base_url,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
),
|
||||
RequestConversionKind::ToOpenAIFamilyCli => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
true,
|
||||
)),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_cli_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_cli_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_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_cli_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_cli_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_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_cli_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_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
compact,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_cli_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::ToOpenAIFamilyCli => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
true,
|
||||
)),
|
||||
RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
RequestConversionKind::ToGeminiStandard => build_gemini_content_url(
|
||||
&transport.endpoint.base_url,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[path = "normalize/chat.rs"]
|
||||
mod chat;
|
||||
#[path = "normalize/cli.rs"]
|
||||
mod cli;
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::build_cross_format_openai_cli_request_body;
|
||||
use serde_json::json;
|
||||
#[path = "normalize/tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[test]
|
||||
fn builds_openai_family_cross_format_request_body_from_compact_source() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"openai:compact",
|
||||
"openai:cli",
|
||||
false,
|
||||
"openai",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("compact to openai cli body should build");
|
||||
|
||||
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
|
||||
assert_eq!(provider_request_body["input"][0]["type"], "message");
|
||||
assert_eq!(provider_request_body["input"][0]["role"], "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_cli_requests() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:cli",
|
||||
"openai:cli",
|
||||
true,
|
||||
"codex",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert!(provider_request_body.get("metadata").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_codex_defaults_unless_body_rules_handle_the_field() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"store": true
|
||||
});
|
||||
let body_rules = json!([
|
||||
{"action":"set","path":"store","value":true},
|
||||
{"action":"set","path":"instructions","value":"Custom instructions"},
|
||||
{"action":"set","path":"metadata","value":{"trace_id":"keep-me"}}
|
||||
]);
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:cli",
|
||||
"openai:cli",
|
||||
true,
|
||||
"codex",
|
||||
Some(&body_rules),
|
||||
None,
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert_eq!(provider_request_body["store"], true);
|
||||
assert_eq!(provider_request_body["instructions"], "Custom instructions");
|
||||
assert_eq!(provider_request_body["metadata"]["trace_id"], "keep-me");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_codex_prompt_cache_key_for_openai_cli_cross_format_requests() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:cli",
|
||||
"openai:cli",
|
||||
true,
|
||||
"codex",
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": "hello"
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = super::build_cross_format_openai_chat_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"codex",
|
||||
"openai:cli",
|
||||
false,
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("openai chat to codex request should build");
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
}
|
||||
pub(crate) use self::chat::{
|
||||
build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url,
|
||||
build_local_openai_chat_request_body, build_local_openai_chat_upstream_url,
|
||||
};
|
||||
pub(crate) use self::cli::{
|
||||
build_cross_format_openai_cli_request_body, build_cross_format_openai_cli_upstream_url,
|
||||
build_local_openai_cli_request_body, build_local_openai_cli_upstream_url,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::codex::apply_codex_openai_cli_special_body_edits;
|
||||
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_gemini_content_url, build_openai_chat_url,
|
||||
build_openai_cli_url, build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
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_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
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 => build_gemini_content_url(
|
||||
&transport.endpoint.base_url,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
),
|
||||
RequestConversionKind::ToOpenAIFamilyCli => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
true,
|
||||
)),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
use url::form_urlencoded;
|
||||
|
||||
use super::super::codex::apply_codex_openai_cli_special_body_edits;
|
||||
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_gemini_content_url, build_openai_cli_url,
|
||||
build_passthrough_path_url,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
build_cross_format_openai_cli_request_body as pipeline_build_cross_format_openai_cli_request_body,
|
||||
build_local_openai_cli_request_body as pipeline_build_local_openai_cli_request_body,
|
||||
GatewayProviderTransportSnapshot,
|
||||
};
|
||||
|
||||
pub(crate) fn build_local_openai_cli_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_cli_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_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_cli_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_cli_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_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
provider_type,
|
||||
provider_api_format,
|
||||
body_rules,
|
||||
user_api_key_id,
|
||||
);
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_openai_cli_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_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
compact,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_cross_format_openai_cli_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::ToOpenAIFamilyCli => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
false,
|
||||
)),
|
||||
RequestConversionKind::ToOpenAICompact => Some(build_openai_cli_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
true,
|
||||
)),
|
||||
RequestConversionKind::ToClaudeStandard => Some(build_claude_messages_url(
|
||||
&transport.endpoint.base_url,
|
||||
parts.uri.query(),
|
||||
)),
|
||||
RequestConversionKind::ToGeminiStandard => build_gemini_content_url(
|
||||
&transport.endpoint.base_url,
|
||||
mapped_model,
|
||||
upstream_is_stream,
|
||||
parts.uri.query(),
|
||||
),
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
use serde_json::json;
|
||||
|
||||
use super::build_cross_format_openai_cli_request_body;
|
||||
|
||||
#[test]
|
||||
fn builds_openai_family_cross_format_request_body_from_compact_source() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"openai:compact",
|
||||
"openai:cli",
|
||||
false,
|
||||
"openai",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("compact to openai cli body should build");
|
||||
|
||||
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
|
||||
assert_eq!(provider_request_body["input"][0]["type"], "message");
|
||||
assert_eq!(provider_request_body["input"][0]["role"], "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_cli_requests() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:cli",
|
||||
"openai:cli",
|
||||
true,
|
||||
"codex",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert!(provider_request_body.get("metadata").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_codex_defaults_unless_body_rules_handle_the_field() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
"metadata": {"trace_id": "abc"},
|
||||
"store": true
|
||||
});
|
||||
let body_rules = json!([
|
||||
{"action":"set","path":"store","value":true},
|
||||
{"action":"set","path":"instructions","value":"Custom instructions"},
|
||||
{"action":"set","path":"metadata","value":{"trace_id":"keep-me"}}
|
||||
]);
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:cli",
|
||||
"openai:cli",
|
||||
true,
|
||||
"codex",
|
||||
Some(&body_rules),
|
||||
None,
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert_eq!(provider_request_body["store"], true);
|
||||
assert_eq!(provider_request_body["instructions"], "Custom instructions");
|
||||
assert_eq!(provider_request_body["metadata"]["trace_id"], "keep-me");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_codex_prompt_cache_key_for_openai_cli_cross_format_requests() {
|
||||
let body_json = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hello"}]
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = build_cross_format_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"claude:cli",
|
||||
"openai:cli",
|
||||
true,
|
||||
"codex",
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("claude cli to codex request should build");
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": "hello"
|
||||
}],
|
||||
});
|
||||
|
||||
let provider_request_body = super::build_cross_format_openai_chat_request_body(
|
||||
&body_json,
|
||||
"gpt-5-upstream",
|
||||
"codex",
|
||||
"openai:cli",
|
||||
false,
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("openai chat to codex request should build");
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["prompt_cache_key"],
|
||||
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
|
||||
);
|
||||
}
|
||||
@@ -1,65 +1,23 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::control_facade::{collect_control_headers, GatewayControlAuthContext};
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
|
||||
};
|
||||
use crate::ai_pipeline::execution_facade::{ConversionMode, ExecutionStrategy};
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::GatewayAuthApiKeySnapshot;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::candidate_runtime_facade::{
|
||||
persist_available_local_candidate, persist_skipped_local_candidate,
|
||||
};
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::planner::executor_facade::mark_unused_local_candidate_items;
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::transport_facade::{
|
||||
read_provider_transport_snapshot, resolve_local_oauth_request_auth,
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::auth::{
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header, resolve_local_openai_chat_auth,
|
||||
};
|
||||
use crate::ai_pipeline::provider_transport_facade::policy::supports_local_openai_chat_transport;
|
||||
use crate::ai_pipeline::provider_transport_facade::{
|
||||
apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use crate::ai_pipeline::PlannerAppState;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse};
|
||||
|
||||
use super::plans::current_unix_secs;
|
||||
use crate::ai_pipeline::planner::standard::{
|
||||
apply_codex_openai_cli_special_headers, build_cross_format_openai_chat_request_body,
|
||||
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
|
||||
build_local_openai_chat_upstream_url,
|
||||
#[path = "decision/cross_format.rs"]
|
||||
mod cross_format;
|
||||
#[path = "decision/same_format.rs"]
|
||||
mod same_format;
|
||||
#[path = "decision/support.rs"]
|
||||
mod support;
|
||||
|
||||
use self::cross_format::build_cross_format_local_openai_chat_decision_payload_for_candidate;
|
||||
use self::same_format::build_same_format_local_openai_chat_decision_payload_for_candidate;
|
||||
use self::support::mark_skipped_local_openai_chat_candidate;
|
||||
pub(super) use self::support::{
|
||||
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
|
||||
LocalOpenAiChatDecisionInput,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalOpenAiChatDecisionInput {
|
||||
pub(super) auth_context: GatewayControlAuthContext,
|
||||
pub(super) requested_model: String,
|
||||
pub(super) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LocalOpenAiChatCandidateAttempt {
|
||||
pub(super) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(super) candidate_index: u32,
|
||||
pub(super) candidate_id: String,
|
||||
}
|
||||
|
||||
pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -71,18 +29,19 @@ pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
report_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let LocalOpenAiChatCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
} = attempt;
|
||||
let transport = match read_provider_transport_snapshot(
|
||||
state,
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
let transport = match planner_state
|
||||
.read_provider_transport_snapshot(
|
||||
&candidate.provider_id,
|
||||
&candidate.endpoint_id,
|
||||
&candidate.key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
@@ -169,625 +128,3 @@ pub(super) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn build_same_format_local_openai_chat_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
decision_kind: &str,
|
||||
report_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
if !supports_local_openai_chat_transport(transport) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let oauth_auth = if resolve_local_openai_chat_auth(transport).is_none() {
|
||||
match resolve_local_oauth_request_auth(state, transport).await {
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
error = ?err,
|
||||
"gateway local openai chat oauth auth resolution failed"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_local_openai_chat_auth(transport).or(oauth_auth)
|
||||
else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some(provider_request_body) = build_local_openai_chat_request_body(
|
||||
body_json,
|
||||
&mapped_model,
|
||||
upstream_is_stream,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"provider_request_body_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
&[&auth_header, "content-type"],
|
||||
&provider_request_body,
|
||||
Some(body_json),
|
||||
) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
transport.endpoint.api_format.as_str(),
|
||||
Some(trace_id),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
);
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
|
||||
if upstream_is_stream {
|
||||
provider_request_headers
|
||||
.entry("accept".to_string())
|
||||
.or_insert_with(|| "text/event-stream".to_string());
|
||||
}
|
||||
let proxy = resolve_transport_proxy_snapshot_with_tunnel_affinity(state, transport).await;
|
||||
let tls_profile = resolve_transport_tls_profile(transport);
|
||||
let prompt_cache_key = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: if upstream_is_stream {
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string()
|
||||
} else {
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
|
||||
},
|
||||
decision_kind: Some(decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.to_string()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url.clone()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some(input.requested_model.clone()),
|
||||
mapped_model: Some(mapped_model.clone()),
|
||||
prompt_cache_key,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: provider_request_headers.clone(),
|
||||
provider_request_body: Some(provider_request_body.clone()),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": input.requested_model,
|
||||
"provider_name": transport.provider.name,
|
||||
"provider_id": candidate.provider_id,
|
||||
"endpoint_id": candidate.endpoint_id,
|
||||
"key_id": candidate.key_id,
|
||||
"key_name": candidate.key_name,
|
||||
"provider_api_format": "openai:chat",
|
||||
"client_api_format": "openai:chat",
|
||||
"mapped_model": mapped_model,
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
}),
|
||||
ExecutionStrategy::LocalSameFormat,
|
||||
ConversionMode::None,
|
||||
"openai:chat",
|
||||
"openai:chat",
|
||||
)),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn build_cross_format_local_openai_chat_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
decision_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
|
||||
let Some(conversion_kind) =
|
||||
request_conversion_kind("openai:chat", provider_api_format.as_str())
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
let transport_supported = request_conversion_transport_supported(transport, conversion_kind);
|
||||
if !transport_supported {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let resolve_auth = request_conversion_direct_auth(transport, conversion_kind);
|
||||
let oauth_auth = if resolve_auth.is_none() {
|
||||
match resolve_local_oauth_request_auth(state, transport).await {
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
provider_api_format = %provider_api_format,
|
||||
error = ?err,
|
||||
"gateway local openai chat cross-format oauth auth resolution failed"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_auth.or(oauth_auth) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some(provider_request_body) = build_cross_format_openai_chat_request_body(
|
||||
body_json,
|
||||
&mapped_model,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"provider_request_body_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(upstream_url) = build_cross_format_openai_chat_upstream_url(
|
||||
parts,
|
||||
transport,
|
||||
&mapped_model,
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
&[&auth_header, "content-type"],
|
||||
&provider_request_body,
|
||||
Some(body_json),
|
||||
) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
Some(trace_id),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
);
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
|
||||
if upstream_is_stream {
|
||||
provider_request_headers
|
||||
.entry("accept".to_string())
|
||||
.or_insert_with(|| "text/event-stream".to_string());
|
||||
}
|
||||
|
||||
let report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
"openai_chat_stream_success"
|
||||
} else {
|
||||
"openai_chat_sync_finalize"
|
||||
};
|
||||
let proxy = resolve_transport_proxy_snapshot_with_tunnel_affinity(state, transport).await;
|
||||
let tls_profile = resolve_transport_tls_profile(transport);
|
||||
let prompt_cache_key = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: if upstream_is_stream {
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION.to_string()
|
||||
} else {
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
|
||||
},
|
||||
decision_kind: Some(decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalCrossFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::Bidirectional.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.to_string()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url.clone()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some(provider_api_format.clone()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some(provider_api_format.clone()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some(input.requested_model.clone()),
|
||||
mapped_model: Some(mapped_model.clone()),
|
||||
prompt_cache_key,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: provider_request_headers.clone(),
|
||||
provider_request_body: Some(provider_request_body.clone()),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": input.requested_model,
|
||||
"provider_name": transport.provider.name,
|
||||
"provider_id": candidate.provider_id,
|
||||
"endpoint_id": candidate.endpoint_id,
|
||||
"key_id": candidate.key_id,
|
||||
"key_name": candidate.key_name,
|
||||
"provider_api_format": provider_api_format,
|
||||
"client_api_format": "openai:chat",
|
||||
"mapped_model": mapped_model,
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": true,
|
||||
}),
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
"openai:chat",
|
||||
provider_api_format.as_str(),
|
||||
)),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_openai_chat_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_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;
|
||||
});
|
||||
let terminal_unix_secs = current_unix_secs();
|
||||
persist_skipped_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
terminal_unix_secs,
|
||||
"gateway local openai chat decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn materialize_local_openai_chat_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
) -> Vec<LocalOpenAiChatCandidateAttempt> {
|
||||
let candidates = prefer_local_tunnel_owner_candidates(state, candidates).await;
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
|
||||
let (execution_strategy, conversion_mode) = if provider_api_format == "openai:chat" {
|
||||
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
|
||||
} else {
|
||||
(
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
)
|
||||
};
|
||||
let extra_data = append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
"provider_api_format": provider_api_format,
|
||||
"client_api_format": "openai:chat",
|
||||
"global_model_id": candidate.global_model_id.clone(),
|
||||
"global_model_name": candidate.global_model_name.clone(),
|
||||
"model_id": candidate.model_id.clone(),
|
||||
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
|
||||
"mapping_matched_model": candidate.mapping_matched_model.clone(),
|
||||
"provider_name": candidate.provider_name.clone(),
|
||||
"key_name": candidate.key_name.clone(),
|
||||
}),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
"openai:chat",
|
||||
candidate.endpoint_api_format.trim(),
|
||||
);
|
||||
|
||||
let candidate_id = persist_available_local_candidate(
|
||||
state,
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local openai chat decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalOpenAiChatCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index: candidate_index as u32,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
|
||||
attempts
|
||||
}
|
||||
|
||||
pub(super) async fn mark_unused_local_openai_chat_candidates<T>(state: &AppState, remaining: Vec<T>)
|
||||
where
|
||||
T: LocalOpenAiChatPlanAndReport,
|
||||
{
|
||||
mark_unused_local_candidate_items(
|
||||
state,
|
||||
remaining,
|
||||
|item| item.plan(),
|
||||
|item| item.report_context(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) trait LocalOpenAiChatPlanAndReport {
|
||||
fn plan(&self) -> &aether_contracts::ExecutionPlan;
|
||||
|
||||
fn report_context(&self) -> Option<&serde_json::Value>;
|
||||
}
|
||||
|
||||
impl LocalOpenAiChatPlanAndReport for LocalSyncPlanAndReport {
|
||||
fn plan(&self) -> &aether_contracts::ExecutionPlan {
|
||||
&self.plan
|
||||
}
|
||||
|
||||
fn report_context(&self) -> Option<&serde_json::Value> {
|
||||
self.report_context.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl LocalOpenAiChatPlanAndReport for LocalStreamPlanAndReport {
|
||||
fn plan(&self) -> &aether_contracts::ExecutionPlan {
|
||||
&self.plan
|
||||
}
|
||||
|
||||
fn report_context(&self) -> Option<&serde_json::Value> {
|
||||
self.report_context.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::collect_control_headers;
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_conversion_direct_auth, request_conversion_kind, request_conversion_transport_supported,
|
||||
};
|
||||
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
use crate::ai_pipeline::planner::standard::{
|
||||
apply_codex_openai_cli_special_headers, build_cross_format_openai_chat_request_body,
|
||||
build_cross_format_openai_chat_upstream_url,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
};
|
||||
use crate::ai_pipeline::transport::{
|
||||
apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
|
||||
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
|
||||
use super::support::{mark_skipped_local_openai_chat_candidate, LocalOpenAiChatDecisionInput};
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
decision_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
|
||||
let Some(conversion_kind) =
|
||||
request_conversion_kind("openai:chat", provider_api_format.as_str())
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
if !request_conversion_transport_supported(transport, conversion_kind) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let resolve_auth = request_conversion_direct_auth(transport, conversion_kind);
|
||||
let oauth_auth = if resolve_auth.is_none() {
|
||||
match planner_state
|
||||
.resolve_local_oauth_request_auth(transport)
|
||||
.await
|
||||
{
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
provider_api_format = %provider_api_format,
|
||||
error = ?err,
|
||||
"gateway local openai chat cross-format oauth auth resolution failed"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_auth.or(oauth_auth) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some(provider_request_body) = build_cross_format_openai_chat_request_body(
|
||||
body_json,
|
||||
&mapped_model,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"provider_request_body_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(upstream_url) = build_cross_format_openai_chat_upstream_url(
|
||||
parts,
|
||||
transport,
|
||||
&mapped_model,
|
||||
provider_api_format.as_str(),
|
||||
upstream_is_stream,
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
&[&auth_header, "content-type"],
|
||||
&provider_request_body,
|
||||
Some(body_json),
|
||||
) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
Some(trace_id),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
);
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
|
||||
if upstream_is_stream {
|
||||
provider_request_headers
|
||||
.entry("accept".to_string())
|
||||
.or_insert_with(|| "text/event-stream".to_string());
|
||||
}
|
||||
|
||||
let report_kind = if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
"openai_chat_stream_success"
|
||||
} else {
|
||||
"openai_chat_sync_finalize"
|
||||
};
|
||||
let proxy =
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), transport).await;
|
||||
let tls_profile = resolve_transport_tls_profile(transport);
|
||||
let prompt_cache_key = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: if upstream_is_stream {
|
||||
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_STREAM_DECISION_ACTION
|
||||
.to_string()
|
||||
} else {
|
||||
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
|
||||
},
|
||||
decision_kind: Some(decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalCrossFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::Bidirectional.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.to_string()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url.clone()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some(provider_api_format.clone()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some(provider_api_format.clone()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some(input.requested_model.clone()),
|
||||
mapped_model: Some(mapped_model.clone()),
|
||||
prompt_cache_key,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: provider_request_headers.clone(),
|
||||
provider_request_body: Some(provider_request_body.clone()),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": input.requested_model,
|
||||
"provider_name": transport.provider.name,
|
||||
"provider_id": candidate.provider_id,
|
||||
"endpoint_id": candidate.endpoint_id,
|
||||
"key_id": candidate.key_id,
|
||||
"key_name": candidate.key_name,
|
||||
"provider_api_format": provider_api_format,
|
||||
"client_api_format": "openai:chat",
|
||||
"mapped_model": mapped_model,
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": true,
|
||||
}),
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
"openai:chat",
|
||||
provider_api_format.as_str(),
|
||||
)),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,265 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::standard::{
|
||||
apply_codex_openai_cli_special_headers, build_local_openai_chat_request_body,
|
||||
build_local_openai_chat_upstream_url,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header, resolve_local_openai_chat_auth,
|
||||
};
|
||||
use crate::ai_pipeline::transport::policy::supports_local_openai_chat_transport;
|
||||
use crate::ai_pipeline::transport::{
|
||||
apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
|
||||
};
|
||||
use crate::ai_pipeline::{
|
||||
collect_control_headers, ConversionMode, ExecutionStrategy, PlannerAppState,
|
||||
};
|
||||
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AppState, GatewayControlSyncDecisionResponse,
|
||||
};
|
||||
|
||||
use super::support::{mark_skipped_local_openai_chat_candidate, LocalOpenAiChatDecisionInput};
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn build_same_format_local_openai_chat_decision_payload_for_candidate(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
body_json: &serde_json::Value,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
decision_kind: &str,
|
||||
report_kind: &str,
|
||||
upstream_is_stream: bool,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
if !supports_local_openai_chat_transport(transport) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_unsupported",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let oauth_auth = if resolve_local_openai_chat_auth(transport).is_none() {
|
||||
match planner_state
|
||||
.resolve_local_oauth_request_auth(transport)
|
||||
.await
|
||||
{
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => Some((name, value)),
|
||||
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => None,
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
error = ?err,
|
||||
"gateway local openai chat oauth auth resolution failed"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let Some((auth_header, auth_value)) = resolve_local_openai_chat_auth(transport).or(oauth_auth)
|
||||
else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_auth_unavailable",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
let mapped_model = candidate.selected_provider_model_name.trim().to_string();
|
||||
if mapped_model.is_empty() {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"mapped_model_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
|
||||
let Some(provider_request_body) = build_local_openai_chat_request_body(
|
||||
body_json,
|
||||
&mapped_model,
|
||||
upstream_is_stream,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"provider_request_body_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"upstream_url_missing",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
);
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
&[&auth_header, "content-type"],
|
||||
&provider_request_body,
|
||||
Some(body_json),
|
||||
) {
|
||||
mark_skipped_local_openai_chat_candidate(
|
||||
state,
|
||||
input,
|
||||
trace_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
"transport_header_rules_apply_failed",
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
apply_codex_openai_cli_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
transport.endpoint.api_format.as_str(),
|
||||
Some(trace_id),
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
);
|
||||
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
|
||||
if upstream_is_stream {
|
||||
provider_request_headers
|
||||
.entry("accept".to_string())
|
||||
.or_insert_with(|| "text/event-stream".to_string());
|
||||
}
|
||||
let proxy =
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity(planner_state.app(), transport).await;
|
||||
let tls_profile = resolve_transport_tls_profile(transport);
|
||||
let prompt_cache_key = provider_request_body
|
||||
.get("prompt_cache_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: if upstream_is_stream {
|
||||
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_STREAM_DECISION_ACTION
|
||||
.to_string()
|
||||
} else {
|
||||
crate::ai_pipeline::planner::common::EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string()
|
||||
},
|
||||
decision_kind: Some(decision_kind.to_string()),
|
||||
execution_strategy: Some(ExecutionStrategy::LocalSameFormat.as_str().to_string()),
|
||||
conversion_mode: Some(ConversionMode::None.as_str().to_string()),
|
||||
request_id: Some(trace_id.to_string()),
|
||||
candidate_id: Some(candidate_id.to_string()),
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: Some(candidate.provider_id.clone()),
|
||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||
key_id: Some(candidate.key_id.clone()),
|
||||
upstream_base_url: Some(transport.endpoint.base_url.clone()),
|
||||
upstream_url: Some(upstream_url.clone()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some(input.requested_model.clone()),
|
||||
mapped_model: Some(mapped_model.clone()),
|
||||
prompt_cache_key,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: provider_request_headers.clone(),
|
||||
provider_request_body: Some(provider_request_body.clone()),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
"user_id": input.auth_context.user_id,
|
||||
"api_key_id": input.auth_context.api_key_id,
|
||||
"request_id": trace_id,
|
||||
"candidate_id": candidate_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"model": input.requested_model,
|
||||
"provider_name": transport.provider.name,
|
||||
"provider_id": candidate.provider_id,
|
||||
"endpoint_id": candidate.endpoint_id,
|
||||
"key_id": candidate.key_id,
|
||||
"key_name": candidate.key_name,
|
||||
"provider_api_format": "openai:chat",
|
||||
"client_api_format": "openai:chat",
|
||||
"mapped_model": mapped_model,
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
}),
|
||||
ExecutionStrategy::LocalSameFormat,
|
||||
ConversionMode::None,
|
||||
"openai:chat",
|
||||
"openai:chat",
|
||||
)),
|
||||
auth_context: Some(input.auth_context.clone()),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::GatewayAuthApiKeySnapshot;
|
||||
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{append_execution_contract_fields_to_value, AppState};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalOpenAiChatDecisionInput {
|
||||
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalOpenAiChatCandidateAttempt {
|
||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||
pub(crate) candidate_index: u32,
|
||||
pub(crate) candidate_id: String,
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_skipped_local_openai_chat_candidate(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
trace_id: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
candidate_index: u32,
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
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;
|
||||
});
|
||||
planner_state
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
candidate,
|
||||
candidate_index,
|
||||
candidate_id,
|
||||
skip_reason,
|
||||
current_unix_secs(),
|
||||
"gateway local openai chat decision failed to persist skipped candidate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
) -> Vec<LocalOpenAiChatCandidateAttempt> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidates = prefer_local_tunnel_owner_candidates(planner_state, candidates).await;
|
||||
let created_at_unix_secs = current_unix_secs();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
|
||||
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let provider_api_format = candidate.endpoint_api_format.trim().to_ascii_lowercase();
|
||||
let (execution_strategy, conversion_mode) = if provider_api_format == "openai:chat" {
|
||||
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
|
||||
} else {
|
||||
(
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
)
|
||||
};
|
||||
let extra_data = append_execution_contract_fields_to_value(
|
||||
json!({
|
||||
"provider_api_format": provider_api_format,
|
||||
"client_api_format": "openai:chat",
|
||||
"global_model_id": candidate.global_model_id.clone(),
|
||||
"global_model_name": candidate.global_model_name.clone(),
|
||||
"model_id": candidate.model_id.clone(),
|
||||
"selected_provider_model_name": candidate.selected_provider_model_name.clone(),
|
||||
"mapping_matched_model": candidate.mapping_matched_model.clone(),
|
||||
"provider_name": candidate.provider_name.clone(),
|
||||
"key_name": candidate.key_name.clone(),
|
||||
}),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
"openai:chat",
|
||||
candidate.endpoint_api_format.trim(),
|
||||
);
|
||||
|
||||
let candidate_id = planner_state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
&candidate,
|
||||
candidate_index as u32,
|
||||
&generated_candidate_id,
|
||||
Some(extra_data),
|
||||
created_at_unix_secs,
|
||||
"gateway local openai chat decision request candidate upsert failed",
|
||||
)
|
||||
.await;
|
||||
|
||||
attempts.push(LocalOpenAiChatCandidateAttempt {
|
||||
candidate,
|
||||
candidate_index: candidate_index as u32,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
|
||||
attempts
|
||||
}
|
||||
@@ -1,10 +1,10 @@
|
||||
use serde_json::Value;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::control_facade::GatewayControlDecision;
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{
|
||||
AppState, GatewayControlSyncDecisionResponse, GatewayError, LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
@@ -18,9 +18,8 @@ use self::decision::{
|
||||
};
|
||||
use self::plans::{
|
||||
build_local_openai_chat_miss_diagnostic, build_local_openai_chat_stream_plan_and_reports,
|
||||
build_local_openai_chat_sync_plan_and_reports, current_unix_secs,
|
||||
list_local_openai_chat_candidates, resolve_local_openai_chat_decision_input,
|
||||
set_local_openai_chat_miss_diagnostic,
|
||||
build_local_openai_chat_sync_plan_and_reports, list_local_openai_chat_candidates,
|
||||
resolve_local_openai_chat_decision_input, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports_for_kind(
|
||||
|
||||
@@ -1,496 +1,18 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
#[path = "plans/candidates.rs"]
|
||||
mod candidates;
|
||||
#[path = "plans/diagnostic.rs"]
|
||||
mod diagnostic;
|
||||
#[path = "plans/resolve.rs"]
|
||||
mod resolve;
|
||||
#[path = "plans/stream.rs"]
|
||||
mod stream;
|
||||
#[path = "plans/sync.rs"]
|
||||
mod sync;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use tracing::warn;
|
||||
|
||||
use super::{
|
||||
materialize_local_openai_chat_candidate_attempts,
|
||||
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
|
||||
GatewayError, LocalExecutionRuntimeMissDiagnostic, LocalOpenAiChatDecisionInput,
|
||||
pub(super) use self::candidates::list_local_openai_chat_candidates;
|
||||
pub(super) use self::diagnostic::{
|
||||
build_local_openai_chat_miss_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
use crate::ai_pipeline::planner::auth_snapshot_facade::{
|
||||
read_auth_api_key_snapshot, GatewayAuthApiKeySnapshot,
|
||||
};
|
||||
use crate::ai_pipeline::planner::candidate_affinity::prefer_local_tunnel_owner_candidates;
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_openai_chat_stream_plan_from_decision, build_openai_chat_sync_plan_from_decision,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::scheduler_facade::list_selectable_candidates;
|
||||
|
||||
pub(super) fn build_local_openai_chat_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: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn set_local_openai_chat_miss_diagnostic(
|
||||
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_openai_chat_miss_diagnostic(decision, plan_kind, requested_model, reason),
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_SYNC_PLAN_KIND {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let candidates = match list_local_openai_chat_candidates(state, &input, false).await {
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision scheduler selection failed"
|
||||
);
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"scheduler_selection_failed",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(0),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_list_empty",
|
||||
)
|
||||
},
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(candidates.len()),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
)
|
||||
},
|
||||
);
|
||||
|
||||
let attempts =
|
||||
materialize_local_openai_chat_candidate_attempts(state, trace_id, &input, candidates).await;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
"openai_chat_sync_success",
|
||||
false,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_sync_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
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 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_sync_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
pub(super) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let candidates = match list_local_openai_chat_candidates(state, &input, true).await {
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision scheduler selection failed"
|
||||
);
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"scheduler_selection_failed",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(0),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_list_empty",
|
||||
)
|
||||
},
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(candidates.len()),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
)
|
||||
},
|
||||
);
|
||||
|
||||
let attempts =
|
||||
materialize_local_openai_chat_candidate_attempts(state, trace_id, &input, candidates).await;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_stream_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
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 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_stream_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
pub(super) fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
pub(super) async fn list_local_openai_chat_candidates(
|
||||
state: &AppState,
|
||||
input: &LocalOpenAiChatDecisionInput,
|
||||
require_streaming: bool,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let mut combined = Vec::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
|
||||
let api_formats = if require_streaming {
|
||||
vec!["openai:chat", "claude:chat", "gemini:chat", "openai:cli"]
|
||||
} else {
|
||||
vec![
|
||||
"openai:chat",
|
||||
"claude:chat",
|
||||
"gemini:chat",
|
||||
"openai:cli",
|
||||
"openai:compact",
|
||||
]
|
||||
};
|
||||
|
||||
for api_format in api_formats {
|
||||
let auth_snapshot = if api_format == "openai:chat" {
|
||||
Some(&input.auth_snapshot)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let mut candidates = list_selectable_candidates(
|
||||
state,
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
require_streaming,
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await?;
|
||||
if api_format != "openai:chat" {
|
||||
candidates.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_openai_chat_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
candidate,
|
||||
)
|
||||
});
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
fn auth_snapshot_allows_cross_format_openai_chat_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| {
|
||||
value
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(candidate.provider_id.trim())
|
||||
|| value
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(candidate.provider_name.trim())
|
||||
});
|
||||
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
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_local_openai_chat_decision_input(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
record_miss_diagnostic: bool,
|
||||
) -> Option<LocalOpenAiChatDecisionInput> {
|
||||
let Some(auth_context) = decision.auth_context.clone().filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
}) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
route_class = ?decision.route_class,
|
||||
route_family = ?decision.route_family,
|
||||
route_kind = ?decision.route_kind,
|
||||
"gateway local openai chat decision skipped: missing_auth_context"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
body_json.get("model").and_then(|value| value.as_str()),
|
||||
"missing_auth_context",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(requested_model) = body_json
|
||||
.get("model")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
"gateway local openai chat decision skipped: missing_requested_model"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
None,
|
||||
"missing_requested_model",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
};
|
||||
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let auth_snapshot = match read_auth_api_key_snapshot(
|
||||
state,
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
user_id = %auth_context.user_id,
|
||||
api_key_id = %auth_context.api_key_id,
|
||||
"gateway local openai chat decision skipped: auth_snapshot_missing"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_missing",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat decision auth snapshot read failed"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_read_failed",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalOpenAiChatDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
})
|
||||
}
|
||||
pub(super) use self::resolve::resolve_local_openai_chat_decision_input;
|
||||
pub(super) use self::stream::build_local_openai_chat_stream_plan_and_reports;
|
||||
pub(super) use self::sync::build_local_openai_chat_sync_plan_and_reports;
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
|
||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, 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>, 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 api_formats = if require_streaming {
|
||||
vec!["openai:chat", "claude:chat", "gemini:chat", "openai:cli"]
|
||||
} else {
|
||||
vec![
|
||||
"openai:chat",
|
||||
"claude:chat",
|
||||
"gemini:chat",
|
||||
"openai:cli",
|
||||
"openai:compact",
|
||||
]
|
||||
};
|
||||
|
||||
for api_format in api_formats {
|
||||
let auth_snapshot = if api_format == "openai:chat" {
|
||||
Some(&input.auth_snapshot)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let mut candidates = planner_state
|
||||
.list_selectable_candidates(
|
||||
api_format,
|
||||
&input.requested_model,
|
||||
require_streaming,
|
||||
auth_snapshot,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await?;
|
||||
if api_format != "openai:chat" {
|
||||
candidates.retain(|candidate| {
|
||||
auth_snapshot_allows_cross_format_openai_chat_candidate(
|
||||
&input.auth_snapshot,
|
||||
&input.requested_model,
|
||||
candidate,
|
||||
)
|
||||
});
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
fn auth_snapshot_allows_cross_format_openai_chat_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| {
|
||||
value
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(candidate.provider_id.trim())
|
||||
|| value
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(candidate.provider_name.trim())
|
||||
});
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::super::{GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic};
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) fn build_local_openai_chat_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: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_openai_chat_miss_diagnostic(
|
||||
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_openai_chat_miss_diagnostic(decision, plan_kind, requested_model, reason),
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::{GatewayControlDecision, LocalOpenAiChatDecisionInput};
|
||||
use super::diagnostic::set_local_openai_chat_miss_diagnostic;
|
||||
use crate::ai_pipeline::{resolve_local_decision_execution_runtime_auth_context, PlannerAppState};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
record_miss_diagnostic: bool,
|
||||
) -> Option<LocalOpenAiChatDecisionInput> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
route_class = ?decision.route_class,
|
||||
route_family = ?decision.route_family,
|
||||
route_kind = ?decision.route_kind,
|
||||
"gateway local openai chat decision skipped: missing_auth_context"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
body_json.get("model").and_then(|value| value.as_str()),
|
||||
"missing_auth_context",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(requested_model) = body_json
|
||||
.get("model")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
"gateway local openai chat decision skipped: missing_requested_model"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
None,
|
||||
"missing_requested_model",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
};
|
||||
|
||||
let auth_snapshot = match planner_state
|
||||
.read_auth_api_key_snapshot(
|
||||
&auth_context.user_id,
|
||||
&auth_context.api_key_id,
|
||||
current_unix_secs(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(snapshot)) => snapshot,
|
||||
Ok(None) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
user_id = %auth_context.user_id,
|
||||
api_key_id = %auth_context.api_key_id,
|
||||
"gateway local openai chat decision skipped: auth_snapshot_missing"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_missing",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat decision auth snapshot read failed"
|
||||
);
|
||||
if record_miss_diagnostic {
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_read_failed",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(LocalOpenAiChatDecisionInput {
|
||||
auth_context,
|
||||
requested_model,
|
||||
auth_snapshot,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::{
|
||||
materialize_local_openai_chat_candidate_attempts,
|
||||
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
|
||||
GatewayError, LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
build_local_openai_chat_miss_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
use super::resolve::resolve_local_openai_chat_decision_input;
|
||||
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_openai_chat_stream_plan_from_decision, LocalStreamPlanAndReport,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_STREAM_PLAN_KIND {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let candidates = match list_local_openai_chat_candidates(state, &input, true).await {
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision scheduler selection failed"
|
||||
);
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"scheduler_selection_failed",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(0),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_list_empty",
|
||||
)
|
||||
},
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(candidates.len()),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
)
|
||||
},
|
||||
);
|
||||
|
||||
let attempts =
|
||||
materialize_local_openai_chat_candidate_attempts(state, trace_id, &input, candidates).await;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_stream_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
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 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_stream_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::{
|
||||
materialize_local_openai_chat_candidate_attempts,
|
||||
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
|
||||
GatewayError, LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
build_local_openai_chat_miss_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
use super::resolve::resolve_local_openai_chat_decision_input;
|
||||
use crate::ai_pipeline::planner::common::OPENAI_CHAT_SYNC_PLAN_KIND;
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_openai_chat_sync_plan_from_decision, LocalSyncPlanAndReport,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
|
||||
if plan_kind != OPENAI_CHAT_SYNC_PLAN_KIND {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let candidates = match list_local_openai_chat_candidates(state, &input, false).await {
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision scheduler selection failed"
|
||||
);
|
||||
set_local_openai_chat_miss_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"scheduler_selection_failed",
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(0),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_list_empty",
|
||||
)
|
||||
},
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
candidate_count: Some(candidates.len()),
|
||||
..build_local_openai_chat_miss_diagnostic(
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
)
|
||||
},
|
||||
);
|
||||
|
||||
let attempts =
|
||||
materialize_local_openai_chat_candidate_attempts(state, trace_id, &input, candidates).await;
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
"openai_chat_sync_success",
|
||||
false,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_sync_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
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 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_sync_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user