refactor: purge Node.js, Rust is the one true backend
- Moved boosapi-rs/* to repo root - Deleted all Node.js source (src/server/, src/web/, scripts/) - Deleted node_modules, package.json, tsconfig, vite.config - Updated .gitignore and AGENTS.md for Rust-only stack - Fixed static file path for frontend dist/web/ - Dockerfile stays the same (already Rust-only) - Kept dist/web/ for Rust static serving
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
use crate::proxy::providers::types::ProviderCategory;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProviderCapabilities {
|
||||
pub category: ProviderCategory,
|
||||
pub supports_chat: bool,
|
||||
pub supports_responses: bool,
|
||||
pub supports_images: bool,
|
||||
pub supports_audio: bool,
|
||||
pub supports_streaming: bool,
|
||||
pub supports_token_refresh: bool,
|
||||
}
|
||||
|
||||
pub fn detect_capabilities(platform: &str) -> ProviderCapabilities {
|
||||
let normalized = platform.to_ascii_lowercase();
|
||||
let category = match normalized.as_str() {
|
||||
"openai" => ProviderCategory::OpenAi,
|
||||
"anthropic" | "claude" => ProviderCategory::Anthropic,
|
||||
"gemini" | "google" => ProviderCategory::Gemini,
|
||||
"codex" => ProviderCategory::Codex,
|
||||
"gemini_cli" | "gemini-cli" => ProviderCategory::GeminiCli,
|
||||
"antigravity" => ProviderCategory::Antigravity,
|
||||
_ => ProviderCategory::GenericOpenAi,
|
||||
};
|
||||
|
||||
match category {
|
||||
ProviderCategory::Anthropic => ProviderCapabilities {
|
||||
category,
|
||||
supports_chat: true,
|
||||
supports_responses: false,
|
||||
supports_images: false,
|
||||
supports_audio: false,
|
||||
supports_streaming: true,
|
||||
supports_token_refresh: true,
|
||||
},
|
||||
ProviderCategory::Gemini | ProviderCategory::GeminiCli => ProviderCapabilities {
|
||||
category,
|
||||
supports_chat: true,
|
||||
supports_responses: false,
|
||||
supports_images: true,
|
||||
supports_audio: false,
|
||||
supports_streaming: true,
|
||||
supports_token_refresh: true,
|
||||
},
|
||||
_ => ProviderCapabilities {
|
||||
category,
|
||||
supports_chat: true,
|
||||
supports_responses: true,
|
||||
supports_images: true,
|
||||
supports_audio: true,
|
||||
supports_streaming: true,
|
||||
supports_token_refresh: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub enum CliProfileKind {
|
||||
Codex,
|
||||
ClaudeCode,
|
||||
GeminiCli,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct CliProfile {
|
||||
pub id: &'static str,
|
||||
pub kind: CliProfileKind,
|
||||
pub display_name: &'static str,
|
||||
pub default_model: &'static str,
|
||||
pub env_key_name: &'static str,
|
||||
}
|
||||
|
||||
pub fn builtin_profiles() -> Vec<CliProfile> {
|
||||
vec![
|
||||
CliProfile {
|
||||
id: "codex",
|
||||
kind: CliProfileKind::Codex,
|
||||
display_name: "Codex CLI",
|
||||
default_model: "gpt-5.5",
|
||||
env_key_name: "OPENAI_API_KEY",
|
||||
},
|
||||
CliProfile {
|
||||
id: "claude-code",
|
||||
kind: CliProfileKind::ClaudeCode,
|
||||
display_name: "Claude Code",
|
||||
default_model: "claude-sonnet-4",
|
||||
env_key_name: "ANTHROPIC_API_KEY",
|
||||
},
|
||||
CliProfile {
|
||||
id: "gemini-cli",
|
||||
kind: CliProfileKind::GeminiCli,
|
||||
display_name: "Gemini CLI",
|
||||
default_model: "gemini-2.0-flash",
|
||||
env_key_name: "GEMINI_API_KEY",
|
||||
},
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,411 @@
|
||||
use super::retry_policy::classify_failure;
|
||||
use super::types::{
|
||||
AttemptFailure, AttemptFailureAction, AttemptResult, ExecuteAttemptContext, ExecuteInput,
|
||||
ExecuteResult, FailureInfo, FailureReason, SelectedChannel, SuccessMetrics,
|
||||
};
|
||||
use futures::future::BoxFuture;
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
|
||||
const MAX_ATTEMPTS: usize = 5;
|
||||
|
||||
pub struct ProxyConductor {
|
||||
pub select_channel: Arc<
|
||||
dyn Fn(String, Option<Value>) -> BoxFuture<'static, Option<SelectedChannel>> + Send + Sync,
|
||||
>,
|
||||
pub select_next_channel: Arc<
|
||||
dyn Fn(String, Vec<i64>, Option<Value>) -> BoxFuture<'static, Option<SelectedChannel>>
|
||||
+ Send
|
||||
+ Sync,
|
||||
>,
|
||||
pub record_success: Arc<dyn Fn(i64, SuccessMetrics) -> BoxFuture<'static, ()> + Send + Sync>,
|
||||
pub record_failure: Arc<dyn Fn(i64, FailureInfo) -> BoxFuture<'static, ()> + Send + Sync>,
|
||||
pub refresh_auth: Option<
|
||||
Arc<
|
||||
dyn Fn(SelectedChannel, FailureInfo) -> BoxFuture<'static, Option<SelectedChannel>>
|
||||
+ Send
|
||||
+ Sync,
|
||||
>,
|
||||
>,
|
||||
}
|
||||
|
||||
impl ProxyConductor {
|
||||
pub async fn execute(&self, input: ExecuteInput) -> ExecuteResult {
|
||||
let mut attempts = 0;
|
||||
let mut exclude_channel_ids = Vec::new();
|
||||
let mut selected = match (self.select_channel)(
|
||||
input.requested_model.clone(),
|
||||
input.downstream_policy.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Some(selected) => selected,
|
||||
None => {
|
||||
return ExecuteResult::Failure {
|
||||
reason: FailureReason::NoChannel,
|
||||
selected: None,
|
||||
status: None,
|
||||
raw_error_text: None,
|
||||
attempts,
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
loop {
|
||||
if attempts >= MAX_ATTEMPTS {
|
||||
return exhausted_failure(
|
||||
selected,
|
||||
None,
|
||||
None,
|
||||
attempts,
|
||||
input.on_terminal_failure.as_deref(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let ctx = ExecuteAttemptContext {
|
||||
selected: selected.clone(),
|
||||
attempt_index: attempts,
|
||||
exclude_channel_ids: exclude_channel_ids.clone(),
|
||||
};
|
||||
attempts += 1;
|
||||
|
||||
match (input.attempt_fn)(ctx).await {
|
||||
AttemptResult::Success {
|
||||
response,
|
||||
latency_ms,
|
||||
cost,
|
||||
} => {
|
||||
(self.record_success)(
|
||||
selected.channel_id,
|
||||
SuccessMetrics {
|
||||
latency_ms,
|
||||
cost,
|
||||
metadata: Default::default(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
return ExecuteResult::Success {
|
||||
selected,
|
||||
response,
|
||||
attempts,
|
||||
};
|
||||
}
|
||||
AttemptResult::Failure {
|
||||
action,
|
||||
status,
|
||||
raw_error_text,
|
||||
} => {
|
||||
let classified = if matches!(action, AttemptFailureAction::Stop) {
|
||||
AttemptFailureAction::Stop
|
||||
} else {
|
||||
classify_failure(status, raw_error_text.as_deref())
|
||||
};
|
||||
let info = FailureInfo {
|
||||
action: classified.as_str().to_string(),
|
||||
status,
|
||||
raw_error_text: raw_error_text.clone(),
|
||||
};
|
||||
(self.record_failure)(selected.channel_id, info.clone()).await;
|
||||
|
||||
match classified {
|
||||
AttemptFailureAction::RetrySameChannel => {
|
||||
continue;
|
||||
}
|
||||
AttemptFailureAction::RefreshAuth => {
|
||||
if let Some(refresh_auth) = &self.refresh_auth {
|
||||
if let Some(refreshed) =
|
||||
refresh_auth(selected.clone(), info.clone()).await
|
||||
{
|
||||
selected = refreshed;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
exclude_channel_ids.push(selected.channel_id);
|
||||
match (self.select_next_channel)(
|
||||
input.requested_model.clone(),
|
||||
exclude_channel_ids.clone(),
|
||||
input.downstream_policy.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Some(next) => {
|
||||
selected = next;
|
||||
}
|
||||
None => {
|
||||
return terminal_failure(
|
||||
selected,
|
||||
FailureReason::Failed,
|
||||
status,
|
||||
raw_error_text,
|
||||
attempts,
|
||||
input.on_terminal_failure.as_deref(),
|
||||
&info,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
AttemptFailureAction::Failover => {
|
||||
exclude_channel_ids.push(selected.channel_id);
|
||||
match (self.select_next_channel)(
|
||||
input.requested_model.clone(),
|
||||
exclude_channel_ids.clone(),
|
||||
input.downstream_policy.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Some(next) => {
|
||||
selected = next;
|
||||
}
|
||||
None => {
|
||||
return terminal_failure(
|
||||
selected,
|
||||
FailureReason::Failed,
|
||||
status,
|
||||
raw_error_text,
|
||||
attempts,
|
||||
input.on_terminal_failure.as_deref(),
|
||||
&info,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
AttemptFailureAction::Terminal => {
|
||||
return terminal_failure(
|
||||
selected,
|
||||
FailureReason::Terminal,
|
||||
status,
|
||||
raw_error_text,
|
||||
attempts,
|
||||
input.on_terminal_failure.as_deref(),
|
||||
&info,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
AttemptFailureAction::Stop => {
|
||||
return terminal_failure(
|
||||
selected,
|
||||
FailureReason::Failed,
|
||||
status,
|
||||
raw_error_text,
|
||||
attempts,
|
||||
input.on_terminal_failure.as_deref(),
|
||||
&info,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn terminal_failure(
|
||||
selected: SelectedChannel,
|
||||
reason: FailureReason,
|
||||
status: Option<u16>,
|
||||
raw_error_text: Option<String>,
|
||||
attempts: usize,
|
||||
on_terminal_failure: Option<
|
||||
&(dyn Fn(SelectedChannel, AttemptFailure) -> BoxFuture<'static, ()> + Send + Sync),
|
||||
>,
|
||||
info: &FailureInfo,
|
||||
) -> ExecuteResult {
|
||||
if let Some(callback) = on_terminal_failure {
|
||||
callback(selected.clone(), AttemptFailure::from(info)).await;
|
||||
}
|
||||
ExecuteResult::Failure {
|
||||
reason,
|
||||
selected: Some(selected),
|
||||
status,
|
||||
raw_error_text,
|
||||
attempts,
|
||||
}
|
||||
}
|
||||
|
||||
async fn exhausted_failure(
|
||||
selected: SelectedChannel,
|
||||
status: Option<u16>,
|
||||
raw_error_text: Option<String>,
|
||||
attempts: usize,
|
||||
on_terminal_failure: Option<
|
||||
&(dyn Fn(SelectedChannel, AttemptFailure) -> BoxFuture<'static, ()> + Send + Sync),
|
||||
>,
|
||||
) -> ExecuteResult {
|
||||
let info = FailureInfo {
|
||||
action: AttemptFailureAction::Stop.as_str().to_string(),
|
||||
status,
|
||||
raw_error_text: raw_error_text.clone(),
|
||||
};
|
||||
terminal_failure(
|
||||
selected,
|
||||
FailureReason::Failed,
|
||||
status,
|
||||
raw_error_text,
|
||||
attempts,
|
||||
on_terminal_failure,
|
||||
&info,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use futures::FutureExt;
|
||||
use serde_json::json;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
fn channel(id: i64) -> SelectedChannel {
|
||||
SelectedChannel {
|
||||
channel_id: id,
|
||||
route_id: Some(10),
|
||||
site_id: 20,
|
||||
account_id: Some(30),
|
||||
base_url: "https://example.test/v1".into(),
|
||||
api_key: "key".into(),
|
||||
model: "gpt-5.5".into(),
|
||||
actual_model: "gpt-5.5".into(),
|
||||
token_name: None,
|
||||
token_value: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retries_same_channel_then_succeeds() {
|
||||
let attempts_seen = Arc::new(AtomicUsize::new(0));
|
||||
let recorded_failures = Arc::new(AtomicUsize::new(0));
|
||||
let recorded_successes = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let conductor = ProxyConductor {
|
||||
select_channel: Arc::new(|_, _| async { Some(channel(1)) }.boxed()),
|
||||
select_next_channel: Arc::new(|_, _, _| async { Some(channel(2)) }.boxed()),
|
||||
record_success: {
|
||||
let recorded_successes = Arc::clone(&recorded_successes);
|
||||
Arc::new(move |_, _| {
|
||||
let recorded_successes = Arc::clone(&recorded_successes);
|
||||
async move {
|
||||
recorded_successes.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
.boxed()
|
||||
})
|
||||
},
|
||||
record_failure: {
|
||||
let recorded_failures = Arc::clone(&recorded_failures);
|
||||
Arc::new(move |_, _| {
|
||||
let recorded_failures = Arc::clone(&recorded_failures);
|
||||
async move {
|
||||
recorded_failures.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
.boxed()
|
||||
})
|
||||
},
|
||||
refresh_auth: None,
|
||||
};
|
||||
|
||||
let result = conductor
|
||||
.execute(ExecuteInput {
|
||||
requested_model: "gpt-5.5".into(),
|
||||
downstream_policy: None,
|
||||
attempt_fn: Box::new({
|
||||
let attempts_seen = Arc::clone(&attempts_seen);
|
||||
move |_| {
|
||||
let attempt = attempts_seen.fetch_add(1, Ordering::SeqCst);
|
||||
async move {
|
||||
if attempt == 0 {
|
||||
AttemptResult::Failure {
|
||||
action: AttemptFailureAction::RetrySameChannel,
|
||||
status: Some(429),
|
||||
raw_error_text: Some("rate limited".into()),
|
||||
}
|
||||
} else {
|
||||
AttemptResult::Success {
|
||||
response: json!({"ok": true}),
|
||||
latency_ms: Some(12),
|
||||
cost: Some(0.1),
|
||||
}
|
||||
}
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
}),
|
||||
on_terminal_failure: None,
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
ExecuteResult::Success {
|
||||
selected, attempts, ..
|
||||
} => {
|
||||
assert_eq!(selected.channel_id, 1);
|
||||
assert_eq!(attempts, 2);
|
||||
}
|
||||
other => panic!("expected success, got {other:?}"),
|
||||
}
|
||||
assert_eq!(recorded_failures.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(recorded_successes.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failover_selects_next_channel() {
|
||||
let conductor = ProxyConductor {
|
||||
select_channel: Arc::new(|_, _| async { Some(channel(1)) }.boxed()),
|
||||
select_next_channel: Arc::new(|_, excluded, _| {
|
||||
async move {
|
||||
assert_eq!(excluded, vec![1]);
|
||||
Some(channel(2))
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
record_success: Arc::new(|_, _| async {}.boxed()),
|
||||
record_failure: Arc::new(|_, _| async {}.boxed()),
|
||||
refresh_auth: None,
|
||||
};
|
||||
|
||||
let attempts_seen = Arc::new(AtomicUsize::new(0));
|
||||
let result = conductor
|
||||
.execute(ExecuteInput {
|
||||
requested_model: "gpt-5.5".into(),
|
||||
downstream_policy: None,
|
||||
attempt_fn: Box::new({
|
||||
let attempts_seen = Arc::clone(&attempts_seen);
|
||||
move |ctx| {
|
||||
let attempt = attempts_seen.fetch_add(1, Ordering::SeqCst);
|
||||
async move {
|
||||
if attempt == 0 {
|
||||
assert_eq!(ctx.selected.channel_id, 1);
|
||||
AttemptResult::Failure {
|
||||
action: AttemptFailureAction::Failover,
|
||||
status: Some(500),
|
||||
raw_error_text: None,
|
||||
}
|
||||
} else {
|
||||
assert_eq!(ctx.selected.channel_id, 2);
|
||||
AttemptResult::Success {
|
||||
response: json!({"ok": true}),
|
||||
latency_ms: None,
|
||||
cost: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
}),
|
||||
on_terminal_failure: None,
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
ExecuteResult::Success {
|
||||
selected, attempts, ..
|
||||
} => {
|
||||
assert_eq!(selected.channel_id, 2);
|
||||
assert_eq!(attempts, 2);
|
||||
}
|
||||
other => panic!("expected success, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
pub mod conductor;
|
||||
pub mod retry_policy;
|
||||
pub mod types;
|
||||
|
||||
pub use conductor::ProxyConductor;
|
||||
pub use retry_policy::{
|
||||
classify_failure, is_terminal, should_failover, should_refresh_auth, should_retry_same_channel,
|
||||
};
|
||||
pub use types::{
|
||||
AttemptFailure, AttemptFailureAction, AttemptResult, ExecuteAttemptContext, ExecuteInput,
|
||||
ExecuteResult, FailureInfo, FailureReason, SelectedChannel, SuccessMetrics,
|
||||
};
|
||||
@@ -0,0 +1,82 @@
|
||||
use super::types::AttemptFailureAction;
|
||||
|
||||
pub fn classify_failure(status: Option<u16>, _error_text: Option<&str>) -> AttemptFailureAction {
|
||||
match status {
|
||||
Some(429 | 503) => AttemptFailureAction::RetrySameChannel,
|
||||
Some(401) => AttemptFailureAction::RefreshAuth,
|
||||
Some(code) if (500..600).contains(&code) => AttemptFailureAction::Failover,
|
||||
Some(code) if (400..500).contains(&code) => AttemptFailureAction::Terminal,
|
||||
_ => AttemptFailureAction::Failover,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_terminal(action: AttemptFailureAction) -> bool {
|
||||
matches!(action, AttemptFailureAction::Terminal)
|
||||
}
|
||||
|
||||
pub fn should_retry_same_channel(action: AttemptFailureAction) -> bool {
|
||||
matches!(action, AttemptFailureAction::RetrySameChannel)
|
||||
}
|
||||
|
||||
pub fn should_refresh_auth(action: AttemptFailureAction) -> bool {
|
||||
matches!(action, AttemptFailureAction::RefreshAuth)
|
||||
}
|
||||
|
||||
pub fn should_failover(action: AttemptFailureAction) -> bool {
|
||||
matches!(action, AttemptFailureAction::Failover)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn classifies_retryable_rate_limit_and_unavailable() {
|
||||
assert_eq!(
|
||||
classify_failure(Some(429), None),
|
||||
AttemptFailureAction::RetrySameChannel
|
||||
);
|
||||
assert_eq!(
|
||||
classify_failure(Some(503), Some("busy")),
|
||||
AttemptFailureAction::RetrySameChannel
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_auth_server_and_terminal_failures() {
|
||||
assert_eq!(
|
||||
classify_failure(Some(401), None),
|
||||
AttemptFailureAction::RefreshAuth
|
||||
);
|
||||
assert_eq!(
|
||||
classify_failure(Some(500), None),
|
||||
AttemptFailureAction::Failover
|
||||
);
|
||||
assert_eq!(
|
||||
classify_failure(Some(502), None),
|
||||
AttemptFailureAction::Failover
|
||||
);
|
||||
assert_eq!(
|
||||
classify_failure(Some(400), None),
|
||||
AttemptFailureAction::Terminal
|
||||
);
|
||||
assert_eq!(
|
||||
classify_failure(Some(404), None),
|
||||
AttemptFailureAction::Terminal
|
||||
);
|
||||
assert_eq!(
|
||||
classify_failure(None, Some("timeout")),
|
||||
AttemptFailureAction::Failover
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn helper_predicates_match_actions() {
|
||||
assert!(is_terminal(AttemptFailureAction::Terminal));
|
||||
assert!(should_retry_same_channel(
|
||||
AttemptFailureAction::RetrySameChannel
|
||||
));
|
||||
assert!(should_refresh_auth(AttemptFailureAction::RefreshAuth));
|
||||
assert!(should_failover(AttemptFailureAction::Failover));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
use futures::future::BoxFuture;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SelectedChannel {
|
||||
pub channel_id: i64,
|
||||
pub route_id: Option<i64>,
|
||||
pub site_id: i64,
|
||||
pub account_id: Option<i64>,
|
||||
pub base_url: String,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
pub actual_model: String,
|
||||
pub token_name: Option<String>,
|
||||
pub token_value: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AttemptFailureAction {
|
||||
RetrySameChannel,
|
||||
RefreshAuth,
|
||||
Failover,
|
||||
Terminal,
|
||||
Stop,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum AttemptResult {
|
||||
Success {
|
||||
response: Value,
|
||||
latency_ms: Option<u64>,
|
||||
cost: Option<f64>,
|
||||
},
|
||||
Failure {
|
||||
action: AttemptFailureAction,
|
||||
status: Option<u16>,
|
||||
raw_error_text: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct ExecuteInput {
|
||||
pub requested_model: String,
|
||||
pub downstream_policy: Option<Value>,
|
||||
pub attempt_fn: Box<dyn Fn(ExecuteAttemptContext) -> BoxFuture<'static, AttemptResult> + Send>,
|
||||
pub on_terminal_failure: Option<
|
||||
Box<dyn Fn(SelectedChannel, AttemptFailure) -> BoxFuture<'static, ()> + Send + Sync>,
|
||||
>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ExecuteAttemptContext {
|
||||
pub selected: SelectedChannel,
|
||||
pub attempt_index: usize,
|
||||
pub exclude_channel_ids: Vec<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum ExecuteResult {
|
||||
Success {
|
||||
selected: SelectedChannel,
|
||||
response: Value,
|
||||
attempts: usize,
|
||||
},
|
||||
Failure {
|
||||
reason: FailureReason,
|
||||
selected: Option<SelectedChannel>,
|
||||
status: Option<u16>,
|
||||
raw_error_text: Option<String>,
|
||||
attempts: usize,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FailureReason {
|
||||
NoChannel,
|
||||
Failed,
|
||||
Terminal,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct SuccessMetrics {
|
||||
pub latency_ms: Option<u64>,
|
||||
pub cost: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub metadata: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct FailureInfo {
|
||||
pub action: String,
|
||||
pub status: Option<u16>,
|
||||
pub raw_error_text: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AttemptFailure {
|
||||
pub status: Option<u16>,
|
||||
pub raw_error_text: Option<String>,
|
||||
pub action: String,
|
||||
}
|
||||
|
||||
impl From<&FailureInfo> for AttemptFailure {
|
||||
fn from(value: &FailureInfo) -> Self {
|
||||
Self {
|
||||
status: value.status,
|
||||
raw_error_text: value.raw_error_text.clone(),
|
||||
action: value.action.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AttemptFailureAction {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::RetrySameChannel => "retry_same_channel",
|
||||
Self::RefreshAuth => "refresh_auth",
|
||||
Self::Failover => "failover",
|
||||
Self::Terminal => "terminal",
|
||||
Self::Stop => "stop",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
// Proxy executors — forward requests to upstream APIs
|
||||
//
|
||||
// Each executor:
|
||||
// 1. Transforms the OpenAI-compatible request to the upstream format
|
||||
// 2. Sends the HTTP request
|
||||
// 3. Transforms the response back to OpenAI-compatible format
|
||||
|
||||
use crate::error::{AppError, AppResult};
|
||||
use crate::AppState;
|
||||
use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Resolved upstream info for a given model request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UpstreamRoute {
|
||||
pub base_url: String,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
/// Create a shared reqwest client with sensible defaults.
|
||||
fn http_client() -> Client {
|
||||
Client::builder()
|
||||
.timeout(Duration::from_secs(300))
|
||||
.pool_max_idle_per_host(32)
|
||||
.build()
|
||||
.expect("Failed to build HTTP client")
|
||||
}
|
||||
|
||||
/// Execute a non-streaming chat completions request.
|
||||
pub async fn execute_chat_completions(
|
||||
_state: AppState,
|
||||
route: UpstreamRoute,
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!("{}/chat/completions", route.base_url.trim_end_matches('/'));
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: Value = serde_json::from_str(&text)
|
||||
.map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
|
||||
|
||||
Ok(json)
|
||||
}
|
||||
|
||||
/// Execute a streaming chat completions request.
|
||||
/// Returns a vector of SSE data chunks as strings.
|
||||
pub async fn execute_chat_completions_stream(
|
||||
_state: AppState,
|
||||
route: UpstreamRoute,
|
||||
body: Value,
|
||||
) -> AppResult<Vec<String>> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!("{}/chat/completions", route.base_url.trim_end_matches('/'));
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|_| "unknown error".to_string());
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
// Read the streaming response as chunks
|
||||
let mut chunks = Vec::new();
|
||||
let mut stream = resp.bytes_stream();
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
match chunk_result {
|
||||
Ok(bytes) => {
|
||||
let text = String::from_utf8_lossy(&bytes).to_string();
|
||||
// Parse SSE lines from the chunk
|
||||
for line in text.lines() {
|
||||
if line.starts_with("data: ") {
|
||||
chunks.push(line.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(AppError::UpstreamError(format!("Stream error: {e}")));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
/// Generic non-streaming proxy request to any upstream endpoint path.
|
||||
pub async fn execute_proxy_request(
|
||||
_state: AppState,
|
||||
route: UpstreamRoute,
|
||||
path_suffix: &str,
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!(
|
||||
"{}/{}",
|
||||
route.base_url.trim_end_matches('/'),
|
||||
path_suffix.trim_start_matches('/')
|
||||
);
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: Value = serde_json::from_str(&text)
|
||||
.map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
|
||||
|
||||
Ok(json)
|
||||
}
|
||||
|
||||
/// Execute an image generation request.
|
||||
pub async fn execute_image_generation(
|
||||
_state: AppState,
|
||||
route: UpstreamRoute,
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!(
|
||||
"{}/images/generations",
|
||||
route.base_url.trim_end_matches('/')
|
||||
);
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: Value = serde_json::from_str(&text)
|
||||
.map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
|
||||
|
||||
Ok(json)
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
// Proxy core — OpenAI-compatible request routing
|
||||
//
|
||||
// Architecture:
|
||||
// /v1/chat/completions → auth (downstream key) → route resolution → upstream call → response
|
||||
// /v1/images/generations → ...same flow...
|
||||
// /v1/audio/speech → ...same flow...
|
||||
|
||||
pub mod capabilities;
|
||||
pub mod cli_profiles;
|
||||
pub mod conductor;
|
||||
pub mod executors;
|
||||
pub mod providers;
|
||||
pub mod runtime;
|
||||
pub mod surfaces;
|
||||
pub mod token_router;
|
||||
@@ -0,0 +1,6 @@
|
||||
// Provider profiles & registry
|
||||
pub mod registry;
|
||||
pub mod types;
|
||||
|
||||
pub use registry::ProviderRegistry;
|
||||
pub use types::{ProviderCategory, ProviderProfile, RouteResolution};
|
||||
@@ -0,0 +1,134 @@
|
||||
// Provider registry — holds all configured upstream provider profiles
|
||||
use crate::proxy::providers::types::{ProviderCategory, ProviderProfile, RouteResolution};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Registry of all available provider profiles
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProviderRegistry {
|
||||
providers: HashMap<String, ProviderProfile>,
|
||||
}
|
||||
|
||||
impl ProviderRegistry {
|
||||
/// Create a new registry with default providers
|
||||
pub fn new() -> Self {
|
||||
let mut registry = Self {
|
||||
providers: HashMap::new(),
|
||||
};
|
||||
registry.register_defaults();
|
||||
registry
|
||||
}
|
||||
|
||||
/// Register the built-in default providers
|
||||
fn register_defaults(&mut self) {
|
||||
// OpenAI
|
||||
self.register(ProviderProfile::openai_compatible(
|
||||
"openai",
|
||||
"OpenAI",
|
||||
"https://api.openai.com/v1",
|
||||
));
|
||||
// Anthropic
|
||||
self.register(ProviderProfile::anthropic(
|
||||
"anthropic",
|
||||
"Anthropic",
|
||||
"https://api.anthropic.com/v1",
|
||||
));
|
||||
// Google Gemini
|
||||
self.register(ProviderProfile::gemini(
|
||||
"gemini",
|
||||
"Google Gemini",
|
||||
"https://generativelanguage.googleapis.com/v1beta",
|
||||
));
|
||||
// Codex
|
||||
self.register(ProviderProfile {
|
||||
id: "codex".into(),
|
||||
name: "OpenAI Codex".into(),
|
||||
category: ProviderCategory::Codex,
|
||||
base_url: "https://api.openai.com/v1".into(),
|
||||
default_model: Some("codex".into()),
|
||||
requires_api_key: true,
|
||||
supports_streaming: true,
|
||||
headers: HashMap::new(),
|
||||
});
|
||||
// Claude Code
|
||||
self.register(ProviderProfile {
|
||||
id: "claude".into(),
|
||||
name: "Claude Code".into(),
|
||||
category: ProviderCategory::Claude,
|
||||
base_url: "https://api.anthropic.com/v1".into(),
|
||||
default_model: Some("claude-sonnet-4".into()),
|
||||
requires_api_key: true,
|
||||
supports_streaming: true,
|
||||
headers: HashMap::new(),
|
||||
});
|
||||
// Gemini CLI
|
||||
self.register(ProviderProfile {
|
||||
id: "gemini_cli".into(),
|
||||
name: "Gemini CLI".into(),
|
||||
category: ProviderCategory::GeminiCli,
|
||||
base_url: "https://generativelanguage.googleapis.com/v1beta".into(),
|
||||
default_model: Some("gemini-2.0-flash".into()),
|
||||
requires_api_key: true,
|
||||
supports_streaming: true,
|
||||
headers: HashMap::new(),
|
||||
});
|
||||
}
|
||||
|
||||
/// Register a new provider profile
|
||||
pub fn register(&mut self, profile: ProviderProfile) {
|
||||
self.providers.insert(profile.id.clone(), profile);
|
||||
}
|
||||
|
||||
/// Get a provider by ID
|
||||
pub fn get(&self, id: &str) -> Option<&ProviderProfile> {
|
||||
self.providers.get(id)
|
||||
}
|
||||
|
||||
/// List all registered providers
|
||||
pub fn list(&self) -> Vec<&ProviderProfile> {
|
||||
self.providers.values().collect()
|
||||
}
|
||||
|
||||
/// Remove a provider by ID
|
||||
pub fn remove(&mut self, id: &str) -> Option<ProviderProfile> {
|
||||
self.providers.remove(id)
|
||||
}
|
||||
|
||||
/// Resolve a model name to an upstream route.
|
||||
/// Tries to match the model to a provider category, falls back to first OpenAI-compatible provider.
|
||||
pub fn resolve(&self, model: &str, api_key: &str) -> Option<RouteResolution> {
|
||||
// First try to find a provider that supports this model
|
||||
for provider in self.providers.values() {
|
||||
if provider.supports_model(model) {
|
||||
let base_url = provider.base_url.clone();
|
||||
return Some(RouteResolution {
|
||||
provider: provider.clone(),
|
||||
api_key: api_key.to_string(),
|
||||
model: model.to_string(),
|
||||
base_url,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to the first OpenAI-compatible provider
|
||||
self.providers
|
||||
.values()
|
||||
.find(|p| {
|
||||
matches!(
|
||||
p.category,
|
||||
ProviderCategory::OpenAi | ProviderCategory::GenericOpenAi
|
||||
)
|
||||
})
|
||||
.map(|provider| RouteResolution {
|
||||
provider: provider.clone(),
|
||||
api_key: api_key.to_string(),
|
||||
model: model.to_string(),
|
||||
base_url: provider.base_url.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProviderRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
// Provider profile types — upstream API provider definitions
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Provider category — determines request transformation and routing behavior
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub enum ProviderCategory {
|
||||
/// OpenAI-compatible chat/completions API
|
||||
OpenAi,
|
||||
/// Anthropic Messages API
|
||||
Anthropic,
|
||||
/// Google Gemini API
|
||||
Gemini,
|
||||
/// Codex CLI / Responses API
|
||||
Codex,
|
||||
/// Claude Code CLI
|
||||
Claude,
|
||||
/// Gemini CLI
|
||||
GeminiCli,
|
||||
/// Antigravity / assistant-style API
|
||||
Antigravity,
|
||||
/// Generic OpenAI-compatible proxy (new-api, one-api, etc.)
|
||||
GenericOpenAi,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ProviderCategory {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::OpenAi => write!(f, "openai"),
|
||||
Self::Anthropic => write!(f, "anthropic"),
|
||||
Self::Gemini => write!(f, "gemini"),
|
||||
Self::Codex => write!(f, "codex"),
|
||||
Self::Claude => write!(f, "claude"),
|
||||
Self::GeminiCli => write!(f, "gemini_cli"),
|
||||
Self::Antigravity => write!(f, "antigravity"),
|
||||
Self::GenericOpenAi => write!(f, "generic_openai"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider configuration profile
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ProviderProfile {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub category: ProviderCategory,
|
||||
pub base_url: String,
|
||||
pub default_model: Option<String>,
|
||||
pub requires_api_key: bool,
|
||||
pub supports_streaming: bool,
|
||||
pub headers: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl ProviderProfile {
|
||||
/// Create a standard OpenAI-compatible provider
|
||||
pub fn openai_compatible(id: &str, name: &str, base_url: &str) -> Self {
|
||||
let mut headers = HashMap::new();
|
||||
headers.insert("Content-Type".into(), "application/json".into());
|
||||
Self {
|
||||
id: id.to_string(),
|
||||
name: name.to_string(),
|
||||
category: ProviderCategory::GenericOpenAi,
|
||||
base_url: base_url.trim_end_matches('/').to_string(),
|
||||
default_model: None,
|
||||
requires_api_key: true,
|
||||
supports_streaming: true,
|
||||
headers,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an Anthropic Messages API provider
|
||||
pub fn anthropic(id: &str, name: &str, base_url: &str) -> Self {
|
||||
let mut headers = HashMap::new();
|
||||
headers.insert("Content-Type".into(), "application/json".into());
|
||||
headers.insert("anthropic-version".into(), "2023-06-01".into());
|
||||
Self {
|
||||
id: id.to_string(),
|
||||
name: name.to_string(),
|
||||
category: ProviderCategory::Anthropic,
|
||||
base_url: base_url.trim_end_matches('/').to_string(),
|
||||
default_model: None,
|
||||
requires_api_key: true,
|
||||
supports_streaming: true,
|
||||
headers,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a Gemini API provider
|
||||
pub fn gemini(id: &str, name: &str, base_url: &str) -> Self {
|
||||
let mut headers = HashMap::new();
|
||||
headers.insert("Content-Type".into(), "application/json".into());
|
||||
Self {
|
||||
id: id.to_string(),
|
||||
name: name.to_string(),
|
||||
category: ProviderCategory::Gemini,
|
||||
base_url: base_url.trim_end_matches('/').to_string(),
|
||||
default_model: None,
|
||||
requires_api_key: true,
|
||||
supports_streaming: true,
|
||||
headers,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this provider handles the given model
|
||||
pub fn supports_model(&self, model: &str) -> bool {
|
||||
let model_lower = model.to_lowercase();
|
||||
match self.category {
|
||||
ProviderCategory::OpenAi | ProviderCategory::GenericOpenAi => {
|
||||
model_lower.contains("gpt")
|
||||
|| model_lower.contains("o1")
|
||||
|| model_lower.contains("o3")
|
||||
|| model_lower.contains("dall-e")
|
||||
}
|
||||
ProviderCategory::Anthropic => model_lower.contains("claude"),
|
||||
ProviderCategory::Gemini => model_lower.contains("gemini"),
|
||||
ProviderCategory::Codex => model_lower.contains("codex"),
|
||||
ProviderCategory::Claude => {
|
||||
model_lower.contains("claude") && model_lower.contains("code")
|
||||
}
|
||||
ProviderCategory::GeminiCli => {
|
||||
model_lower.contains("gemini") && model_lower.contains("cli")
|
||||
}
|
||||
ProviderCategory::Antigravity => {
|
||||
model_lower.contains("antigravity") || model_lower.contains("assistant")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of resolving a model to an upstream route
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RouteResolution {
|
||||
pub provider: ProviderProfile,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
pub base_url: String,
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
use crate::proxy::conductor::{FailureInfo, ProxyConductor, SelectedChannel, SuccessMetrics};
|
||||
use crate::proxy::token_router::TokenRouter;
|
||||
use futures::FutureExt;
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ProxyRuntime {
|
||||
pub router: TokenRouter,
|
||||
}
|
||||
|
||||
impl ProxyRuntime {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self {
|
||||
router: TokenRouter::new(db),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn conductor(&self) -> ProxyConductor {
|
||||
let router = self.router.clone();
|
||||
let next_router = self.router.clone();
|
||||
let success_router = self.router.clone();
|
||||
let failure_router = self.router.clone();
|
||||
|
||||
ProxyConductor {
|
||||
select_channel: Arc::new(move |model, _policy| {
|
||||
let router = router.clone();
|
||||
async move { router.resolve_selected(&model).await.ok() }.boxed()
|
||||
}),
|
||||
select_next_channel: Arc::new(move |model, excluded, _policy| {
|
||||
let router = next_router.clone();
|
||||
async move { router.resolve_next_selected(&model, &excluded).await.ok() }.boxed()
|
||||
}),
|
||||
record_success: Arc::new(move |channel_id: i64, metrics: SuccessMetrics| {
|
||||
let router = success_router.clone();
|
||||
async move {
|
||||
let _ = router.record_success(channel_id, metrics).await;
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
record_failure: Arc::new(move |channel_id: i64, info: FailureInfo| {
|
||||
let router = failure_router.clone();
|
||||
async move {
|
||||
let _ = router.record_failure(channel_id, info).await;
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
refresh_auth: Some(Arc::new(|selected: SelectedChannel, _info: FailureInfo| {
|
||||
async move { Some(selected) }.boxed()
|
||||
})),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ProxySurfaceState {
|
||||
sessions: Arc<Mutex<HashMap<String, SurfaceSession>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SurfaceSession {
|
||||
pub id: String,
|
||||
pub model: String,
|
||||
pub channel_id: Option<i64>,
|
||||
pub metadata: Value,
|
||||
}
|
||||
|
||||
impl ProxySurfaceState {
|
||||
pub fn upsert_session(&self, session: SurfaceSession) {
|
||||
self.sessions.lock().insert(session.id.clone(), session);
|
||||
}
|
||||
|
||||
pub fn get_session(&self, id: &str) -> Option<SurfaceSession> {
|
||||
self.sessions.lock().get(id).cloned()
|
||||
}
|
||||
|
||||
pub fn remove_session(&self, id: &str) -> Option<SurfaceSession> {
|
||||
self.sessions.lock().remove(id)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod router;
|
||||
pub mod types;
|
||||
|
||||
pub use router::TokenRouter;
|
||||
pub use types::{RouteDecision, RouteMatch};
|
||||
@@ -0,0 +1,430 @@
|
||||
use super::types::{RouteDecision, RouteMatch};
|
||||
use crate::error::{AppError, AppResult};
|
||||
use crate::proxy::conductor::{FailureInfo, SelectedChannel, SuccessMetrics};
|
||||
use crate::proxy::executors::UpstreamRoute;
|
||||
use regex::Regex;
|
||||
use sqlx::{FromRow, SqlitePool};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct TokenRouter {
|
||||
db: Arc<SqlitePool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct RouteRow {
|
||||
id: i64,
|
||||
model_pattern: String,
|
||||
model_mapping: Option<String>,
|
||||
route_mode: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct ChannelRow {
|
||||
channel_id: i64,
|
||||
route_id: i64,
|
||||
account_id: i64,
|
||||
site_id: i64,
|
||||
base_url: String,
|
||||
site_api_key: Option<String>,
|
||||
account_access_token: String,
|
||||
account_api_token: Option<String>,
|
||||
token_name: Option<String>,
|
||||
token_value: Option<String>,
|
||||
source_model: Option<String>,
|
||||
}
|
||||
|
||||
impl TokenRouter {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
|
||||
pub async fn resolve(&self, model: &str) -> AppResult<UpstreamRoute> {
|
||||
let decision = self.resolve_decision(model, &[]).await?;
|
||||
Ok(UpstreamRoute {
|
||||
base_url: decision.selected.base_url,
|
||||
api_key: decision.selected.api_key,
|
||||
model: decision.selected.actual_model,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn resolve_selected(&self, model: &str) -> AppResult<SelectedChannel> {
|
||||
Ok(self.resolve_decision(model, &[]).await?.selected)
|
||||
}
|
||||
|
||||
pub async fn resolve_next_selected(
|
||||
&self,
|
||||
model: &str,
|
||||
exclude_channel_ids: &[i64],
|
||||
) -> AppResult<SelectedChannel> {
|
||||
Ok(self
|
||||
.resolve_decision(model, exclude_channel_ids)
|
||||
.await?
|
||||
.selected)
|
||||
}
|
||||
|
||||
pub async fn resolve_decision(
|
||||
&self,
|
||||
model: &str,
|
||||
exclude_channel_ids: &[i64],
|
||||
) -> AppResult<RouteDecision> {
|
||||
let route = self.match_route(model).await?.ok_or_else(|| {
|
||||
AppError::BadRequest(format!("No route configured for model '{model}'"))
|
||||
})?;
|
||||
let channel = self
|
||||
.select_channel_for_route(&route.0, model, exclude_channel_ids)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
AppError::BadRequest(format!(
|
||||
"No enabled channel configured for model '{}'",
|
||||
model
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(RouteDecision {
|
||||
requested_model: model.to_string(),
|
||||
matched_pattern: route.0.model_pattern,
|
||||
match_kind: route.1,
|
||||
selected: channel,
|
||||
})
|
||||
}
|
||||
|
||||
async fn match_route(&self, model: &str) -> AppResult<Option<(RouteRow, RouteMatch)>> {
|
||||
let rows: Vec<RouteRow> = sqlx::query_as(
|
||||
"SELECT id, model_pattern, model_mapping, route_mode
|
||||
FROM token_routes
|
||||
WHERE enabled = 1
|
||||
ORDER BY id",
|
||||
)
|
||||
.fetch_all(&*self.db)
|
||||
.await?;
|
||||
|
||||
let mut best: Option<(RouteRow, RouteMatch, usize, usize)> = None;
|
||||
for row in rows {
|
||||
if let Some(kind) = match_pattern(model, &row.model_pattern, row.route_mode.as_deref())
|
||||
{
|
||||
let rank = match kind {
|
||||
RouteMatch::Exact => 3,
|
||||
RouteMatch::Wildcard => 2,
|
||||
RouteMatch::Regex => 1,
|
||||
};
|
||||
let specificity = row.model_pattern.chars().filter(|c| *c != '*').count();
|
||||
let replace = best
|
||||
.as_ref()
|
||||
.map(|(_, _, best_rank, best_specificity)| {
|
||||
rank > *best_rank || (rank == *best_rank && specificity > *best_specificity)
|
||||
})
|
||||
.unwrap_or(true);
|
||||
if replace {
|
||||
best = Some((row, kind, rank, specificity));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(best.map(|(row, kind, _, _)| (row, kind)))
|
||||
}
|
||||
|
||||
async fn select_channel_for_route(
|
||||
&self,
|
||||
route: &RouteRow,
|
||||
requested_model: &str,
|
||||
exclude_channel_ids: &[i64],
|
||||
) -> AppResult<Option<SelectedChannel>> {
|
||||
let rows: Vec<ChannelRow> = sqlx::query_as(
|
||||
"SELECT
|
||||
rc.id AS channel_id,
|
||||
rc.route_id AS route_id,
|
||||
rc.account_id AS account_id,
|
||||
a.site_id AS site_id,
|
||||
s.url AS base_url,
|
||||
s.api_key AS site_api_key,
|
||||
a.access_token AS account_access_token,
|
||||
a.api_token AS account_api_token,
|
||||
at.name AS token_name,
|
||||
at.token AS token_value,
|
||||
rc.source_model AS source_model
|
||||
FROM route_channels rc
|
||||
JOIN accounts a ON a.id = rc.account_id
|
||||
JOIN sites s ON s.id = a.site_id
|
||||
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
|
||||
WHERE rc.route_id = ?
|
||||
AND rc.enabled = 1
|
||||
AND a.status = 'active'
|
||||
AND s.status = 'active'
|
||||
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
|
||||
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
|
||||
ORDER BY
|
||||
COALESCE(rc.priority, 0) DESC,
|
||||
rc.consecutive_fail_count ASC,
|
||||
COALESCE(rc.weight, 10) DESC,
|
||||
rc.last_selected_at IS NOT NULL,
|
||||
rc.id ASC",
|
||||
)
|
||||
.bind(route.id)
|
||||
.fetch_all(&*self.db)
|
||||
.await?;
|
||||
|
||||
let row = rows
|
||||
.into_iter()
|
||||
.find(|row| !exclude_channel_ids.contains(&row.channel_id));
|
||||
|
||||
Ok(row.map(|row| {
|
||||
let actual_model = route
|
||||
.model_mapping
|
||||
.clone()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.or_else(|| {
|
||||
row.source_model
|
||||
.clone()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
})
|
||||
.unwrap_or_else(|| requested_model.to_string());
|
||||
let api_key = row
|
||||
.token_value
|
||||
.clone()
|
||||
.or(row.account_api_token.clone())
|
||||
.unwrap_or(row.account_access_token.clone());
|
||||
|
||||
SelectedChannel {
|
||||
channel_id: row.channel_id,
|
||||
route_id: Some(row.route_id),
|
||||
site_id: row.site_id,
|
||||
account_id: Some(row.account_id),
|
||||
base_url: row.base_url,
|
||||
api_key: api_key
|
||||
.trim()
|
||||
.to_string()
|
||||
.if_empty_then(row.site_api_key.unwrap_or_default()),
|
||||
model: requested_model.to_string(),
|
||||
actual_model,
|
||||
token_name: row.token_name,
|
||||
token_value: row.token_value,
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn record_success(&self, channel_id: i64, metrics: SuccessMetrics) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE route_channels
|
||||
SET success_count = COALESCE(success_count, 0) + 1,
|
||||
total_latency_ms = COALESCE(total_latency_ms, 0) + ?,
|
||||
total_cost = COALESCE(total_cost, 0) + ?,
|
||||
consecutive_fail_count = 0,
|
||||
cooldown_level = 0,
|
||||
cooldown_until = NULL,
|
||||
last_used_at = datetime('now'),
|
||||
last_selected_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(metrics.latency_ms.unwrap_or(0) as i64)
|
||||
.bind(metrics.cost.unwrap_or(0.0))
|
||||
.bind(channel_id)
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn record_failure(&self, channel_id: i64, _info: FailureInfo) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE route_channels
|
||||
SET fail_count = COALESCE(fail_count, 0) + 1,
|
||||
consecutive_fail_count = consecutive_fail_count + 1,
|
||||
cooldown_level = cooldown_level + 1,
|
||||
last_fail_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(channel_id)
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
trait IfEmptyThen {
|
||||
fn if_empty_then(self, fallback: String) -> String;
|
||||
}
|
||||
|
||||
impl IfEmptyThen for String {
|
||||
fn if_empty_then(self, fallback: String) -> String {
|
||||
if self.is_empty() {
|
||||
fallback
|
||||
} else {
|
||||
self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn match_pattern(model: &str, pattern: &str, route_mode: Option<&str>) -> Option<RouteMatch> {
|
||||
match route_mode.unwrap_or("pattern") {
|
||||
"exact" => (model == pattern).then_some(RouteMatch::Exact),
|
||||
"regex" => Regex::new(pattern)
|
||||
.ok()
|
||||
.and_then(|re| re.is_match(model).then_some(RouteMatch::Regex)),
|
||||
_ if pattern == model => Some(RouteMatch::Exact),
|
||||
_ if pattern.starts_with("regex:") => Regex::new(pattern.trim_start_matches("regex:"))
|
||||
.ok()
|
||||
.and_then(|re| re.is_match(model).then_some(RouteMatch::Regex)),
|
||||
_ if pattern.contains('*') => {
|
||||
wildcard_matches(model, pattern).then_some(RouteMatch::Wildcard)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn wildcard_matches(model: &str, pattern: &str) -> bool {
|
||||
let escaped = pattern
|
||||
.split('*')
|
||||
.map(regex::escape)
|
||||
.collect::<Vec<_>>()
|
||||
.join(".*");
|
||||
Regex::new(&format!("^{escaped}$"))
|
||||
.map(|re| re.is_match(model))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn test_pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
sqlx::query(
|
||||
"CREATE TABLE sites (
|
||||
id INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
platform TEXT NOT NULL DEFAULT 'openai',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
api_key TEXT
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE token_routes (
|
||||
id INTEGER PRIMARY KEY,
|
||||
model_pattern TEXT NOT NULL,
|
||||
model_mapping TEXT,
|
||||
route_mode TEXT,
|
||||
enabled INTEGER DEFAULT 1
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE accounts (
|
||||
id INTEGER PRIMARY KEY,
|
||||
site_id INTEGER NOT NULL,
|
||||
access_token TEXT NOT NULL,
|
||||
api_token TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'active'
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE account_tokens (
|
||||
id INTEGER PRIMARY KEY,
|
||||
account_id INTEGER NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
token TEXT NOT NULL,
|
||||
enabled INTEGER DEFAULT 1
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE route_channels (
|
||||
id INTEGER PRIMARY KEY,
|
||||
route_id INTEGER NOT NULL,
|
||||
account_id INTEGER NOT NULL,
|
||||
token_id INTEGER,
|
||||
source_model TEXT,
|
||||
priority INTEGER DEFAULT 0,
|
||||
weight INTEGER DEFAULT 10,
|
||||
enabled INTEGER DEFAULT 1,
|
||||
consecutive_fail_count INTEGER NOT NULL DEFAULT 0,
|
||||
cooldown_level INTEGER NOT NULL DEFAULT 0,
|
||||
cooldown_until TEXT,
|
||||
last_selected_at TEXT,
|
||||
success_count INTEGER DEFAULT 0,
|
||||
fail_count INTEGER DEFAULT 0,
|
||||
total_latency_ms INTEGER DEFAULT 0,
|
||||
total_cost REAL DEFAULT 0,
|
||||
last_used_at TEXT,
|
||||
last_fail_at TEXT
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
sqlx::query("INSERT INTO sites (id, name, url, api_key) VALUES (1, 'OpenAI', 'https://api.example/v1', 'site-key')")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (1, 1, 'access-key', NULL)")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (1, 1, 'primary', 'token-key')")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolves_exact_model_route() {
|
||||
let pool = test_pool().await;
|
||||
sqlx::query("INSERT INTO token_routes (id, model_pattern, model_mapping, route_mode) VALUES (1, 'gpt-5.5', 'upstream-gpt-5.5', 'exact')")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let router = TokenRouter::new(Arc::new(pool));
|
||||
let decision = router.resolve_decision("gpt-5.5", &[]).await.unwrap();
|
||||
assert_eq!(decision.match_kind, RouteMatch::Exact);
|
||||
assert_eq!(decision.selected.actual_model, "upstream-gpt-5.5");
|
||||
assert_eq!(decision.selected.api_key, "token-key");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolves_wildcard_route_and_skips_excluded_channel() {
|
||||
let pool = test_pool().await;
|
||||
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (1, 1, 1, 1, 10)")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (2, 1, 1, 1, 1)")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let router = TokenRouter::new(Arc::new(pool));
|
||||
let decision = router.resolve_decision("gpt-4o", &[1]).await.unwrap();
|
||||
assert_eq!(decision.match_kind, RouteMatch::Wildcard);
|
||||
assert_eq!(decision.selected.channel_id, 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
use crate::proxy::conductor::SelectedChannel;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum RouteMatch {
|
||||
Exact,
|
||||
Wildcard,
|
||||
Regex,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RouteDecision {
|
||||
pub requested_model: String,
|
||||
pub matched_pattern: String,
|
||||
pub match_kind: RouteMatch,
|
||||
pub selected: SelectedChannel,
|
||||
}
|
||||
Reference in New Issue
Block a user