boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: optimize proxy streaming and database access

Browse Source
This commit is contained in:
boos
2026-06-07 00:00:24 +08:00
parent f3d2d50f9a
commit aaba6f152e
13 changed files with 960 additions and 492 deletions
Download Patch File Download Diff File Expand all files Collapse all files
migrations/20260606000002_performance_projections.sql
+51
View File
@@ -0,0 +1,51 @@
-- Performance-oriented projections and indexes for logs/workspace views.
CREATE INDEX IF NOT EXISTS proxy_logs_key_created_status_idx
ON proxy_logs (downstream_api_key_id, created_at DESC, status);
CREATE INDEX IF NOT EXISTS proxy_logs_key_model_created_idx
ON proxy_logs (downstream_api_key_id, model_actual, created_at DESC);
CREATE INDEX IF NOT EXISTS proxy_logs_content_summary_idx
ON proxy_logs (content_summary);
CREATE INDEX IF NOT EXISTS workspace_assets_user_output_created_idx
ON workspace_assets (
user_id,
status,
deleted_at,
source,
kind,
created_at DESC
);
CREATE INDEX IF NOT EXISTS workspace_assets_output_role_created_idx
ON workspace_assets (
json_extract(metadata, '$.workspaceAsset.role'),
created_at DESC
);
CREATE INDEX IF NOT EXISTS workspace_assets_workflow_output_created_idx
ON workspace_assets (
json_extract(metadata, '$.workflowUi.type'),
created_at DESC
);
CREATE TABLE IF NOT EXISTS workspace_asset_daily (
local_day TEXT NOT NULL,
user_id INTEGER NOT NULL DEFAULT 0,
source TEXT NOT NULL,
kind TEXT NOT NULL,
asset_count INTEGER NOT NULL DEFAULT 0,
local_file_count INTEGER NOT NULL DEFAULT 0,
total_size_bytes INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
UNIQUE(local_day, user_id, source, kind)
);
CREATE INDEX IF NOT EXISTS workspace_asset_daily_user_day_idx
ON workspace_asset_daily (user_id, local_day DESC);
CREATE INDEX IF NOT EXISTS workspace_asset_daily_source_day_idx
ON workspace_asset_daily (source, local_day DESC);
migrations_pg/20260606000002_performance_projections.sql
+45
View File
@@ -0,0 +1,45 @@
-- Performance-oriented projections and indexes for logs/workspace views.
CREATE INDEX IF NOT EXISTS proxy_logs_key_created_status_idx
ON proxy_logs (downstream_api_key_id, created_at DESC, status);
CREATE INDEX IF NOT EXISTS proxy_logs_key_model_created_idx
ON proxy_logs (downstream_api_key_id, model_actual, created_at DESC);
CREATE INDEX IF NOT EXISTS proxy_logs_content_summary_idx
ON proxy_logs (content_summary);
CREATE INDEX IF NOT EXISTS workspace_assets_user_output_created_idx
ON workspace_assets (
user_id,
status,
deleted_at,
source,
kind,
created_at DESC
);
CREATE INDEX IF NOT EXISTS workspace_assets_output_role_created_idx
ON workspace_assets (((metadata::jsonb #>> '{workspaceAsset,role}')), created_at DESC);
CREATE INDEX IF NOT EXISTS workspace_assets_workflow_output_created_idx
ON workspace_assets (((metadata::jsonb #>> '{workflowUi,type}')), created_at DESC);
CREATE TABLE IF NOT EXISTS workspace_asset_daily (
local_day TEXT NOT NULL,
user_id BIGINT NOT NULL DEFAULT 0,
source TEXT NOT NULL,
kind TEXT NOT NULL,
asset_count BIGINT NOT NULL DEFAULT 0,
local_file_count BIGINT NOT NULL DEFAULT 0,
total_size_bytes BIGINT NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
UNIQUE(local_day, user_id, source, kind)
);
CREATE INDEX IF NOT EXISTS workspace_asset_daily_user_day_idx
ON workspace_asset_daily (user_id, local_day DESC);
CREATE INDEX IF NOT EXISTS workspace_asset_daily_source_day_idx
ON workspace_asset_daily (source, local_day DESC);
src/db/mod.rs
+17
View File
@@ -46,6 +46,23 @@ pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Re
}
let pool = SqlitePoolOptions::new()
.max_connections(max_connections)
.after_connect(|conn, _meta| {
Box::pin(async move {
sqlx::query("PRAGMA journal_mode = WAL")
.execute(&mut *conn)
.await?;
sqlx::query("PRAGMA synchronous = NORMAL")
.execute(&mut *conn)
.await?;
sqlx::query("PRAGMA busy_timeout = 5000")
.execute(&mut *conn)
.await?;
sqlx::query("PRAGMA foreign_keys = ON")
.execute(&mut *conn)
.await?;
Ok(())
})
})
.connect(database_url)
.await?;
Ok(pool)
src/db/schema_compat.rs
+28
View File
@@ -100,6 +100,34 @@ const SQLITE_COMPAT_DDL: &[&str] = &[
ON proxy_logs (client_app_id, created_at)",
"CREATE INDEX IF NOT EXISTS proxy_logs_client_family_created_at_idx
ON proxy_logs (client_family, created_at)",
"CREATE INDEX IF NOT EXISTS proxy_logs_key_created_status_idx
ON proxy_logs (downstream_api_key_id, created_at DESC, status)",
"CREATE INDEX IF NOT EXISTS proxy_logs_key_model_created_idx
ON proxy_logs (downstream_api_key_id, model_actual, created_at DESC)",
"CREATE INDEX IF NOT EXISTS proxy_logs_content_summary_idx
ON proxy_logs (content_summary)",
"CREATE INDEX IF NOT EXISTS workspace_assets_user_output_created_idx
ON workspace_assets (user_id, status, deleted_at, source, kind, created_at DESC)",
"CREATE INDEX IF NOT EXISTS workspace_assets_output_role_created_idx
ON workspace_assets (json_extract(metadata, '$.workspaceAsset.role'), created_at DESC)",
"CREATE INDEX IF NOT EXISTS workspace_assets_workflow_output_created_idx
ON workspace_assets (json_extract(metadata, '$.workflowUi.type'), created_at DESC)",
"CREATE TABLE IF NOT EXISTS workspace_asset_daily (
local_day TEXT NOT NULL,
user_id INTEGER NOT NULL DEFAULT 0,
source TEXT NOT NULL,
kind TEXT NOT NULL,
asset_count INTEGER NOT NULL DEFAULT 0,
local_file_count INTEGER NOT NULL DEFAULT 0,
total_size_bytes INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
UNIQUE(local_day, user_id, source, kind)
)",
"CREATE INDEX IF NOT EXISTS workspace_asset_daily_user_day_idx
ON workspace_asset_daily (user_id, local_day DESC)",
"CREATE INDEX IF NOT EXISTS workspace_asset_daily_source_day_idx
ON workspace_asset_daily (source, local_day DESC)",
"CREATE INDEX IF NOT EXISTS route_channels_oauth_route_unit_id_idx
ON route_channels (oauth_route_unit_id)",
"CREATE TABLE IF NOT EXISTS site_api_endpoints (
src/lib.rs
+14 -1
View File
@@ -11,12 +11,13 @@ pub mod transformers;
pub mod workflow;
use sqlx::SqlitePool;
use std::sync::Arc;
use std::{sync::Arc, time::Duration};
#[derive(Clone)]
pub struct AppState {
pub config: config::AppConfig,
pub db: Arc<SqlitePool>,
pub http: reqwest::Client,
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
pub runtime_auth: Arc<services::runtime_auth::RuntimeAuthSettings>,
@@ -44,10 +45,22 @@ impl AppState {
runtime_auth.proxy_token(),
comfyui_agent_secret.clone(),
);
let http = reqwest::Client::builder()
.timeout(Duration::from_secs(config.proxy.upstream_timeout_secs.max(30)))
.connect_timeout(Duration::from_secs(15))
.pool_idle_timeout(Duration::from_secs(90))
.pool_max_idle_per_host(64)
.tcp_nodelay(true)
.build()
.unwrap_or_else(|err| {
tracing::warn!(%err, "failed to build shared HTTP client; falling back to reqwest default");
reqwest::Client::new()
});
Self {
config,
db: Arc::new(db),
http,
comfyui,
builtin_keys: Arc::new(builtin_keys),
runtime_auth: Arc::new(runtime_auth),
src/proxy/executors/mod.rs
+83 -49
View File
@@ -14,9 +14,7 @@ use crate::transformers::{
canonical::CanonicalRequest,
};
use crate::AppState;
use reqwest::Client;
use serde_json::{json, Value};
use std::time::Duration;
/// Resolved upstream info for a given model request.
#[derive(Debug, Clone)]
@@ -26,15 +24,6 @@ pub struct UpstreamRoute {
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")
}
fn parse_upstream_json(
text: &str,
upstream_url: &str,
@@ -174,11 +163,11 @@ fn aggregate_chat_sse(text: &str, fallback_model: &str) -> Option<Value> {
/// Execute a non-streaming chat completions request.
pub async fn execute_chat_completions(
_state: AppState,
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, "/chat/completions");
let resp = client
@@ -210,19 +199,29 @@ pub async fn execute_chat_completions(
/// Execute a streaming chat completions request.
/// Returns a vector of SSE data chunks as strings.
pub async fn execute_chat_completions_stream(
_state: AppState,
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<Vec<String>> {
execute_streaming_request(route, "/chat/completions", body).await
execute_streaming_request(state, route, "/chat/completions", body).await
}
/// Execute a streaming chat completions request and return the live upstream response.
pub async fn execute_chat_completions_stream_response(
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<reqwest::Response> {
execute_streaming_response(state, route, "/chat/completions", body).await
}
/// Execute an Anthropic Messages API request against an Anthropic-shaped upstream.
pub async fn execute_anthropic_messages(
state: AppState,
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = format!("{}/messages", route.base_url.trim_end_matches('/'));
let compatibility = AnthropicMessagesCompatibility::default();
let body = format_messages_request_with_compat(&request, &compatibility);
@@ -260,10 +259,21 @@ pub async fn execute_anthropic_messages(
/// Execute a streaming Anthropic Messages request and return Anthropic SSE frames.
pub async fn execute_anthropic_messages_stream(
state: AppState,
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<Vec<String>> {
let client = http_client();
let response = execute_anthropic_messages_stream_response(state, route, request).await?;
collect_anthropic_stream_frames(response).await
}
/// Execute a streaming Anthropic Messages request and return the live upstream response.
pub async fn execute_anthropic_messages_stream_response(
state: AppState,
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let upstream_url = format!("{}/messages", route.base_url.trim_end_matches('/'));
let compatibility = AnthropicMessagesCompatibility::default();
let body = format_messages_request_with_compat(&request, &compatibility);
@@ -292,6 +302,10 @@ pub async fn execute_anthropic_messages_stream(
)));
}
Ok(resp)
}
async fn collect_anthropic_stream_frames(resp: reqwest::Response) -> AppResult<Vec<String>> {
let mut frames = Vec::new();
let mut stream = resp.bytes_stream();
use futures::StreamExt;
@@ -406,32 +420,12 @@ fn denormalize_anthropic_stop_reason(reason: &str) -> &str {
/// Execute a streaming request against any upstream path (e.g. /responses).
/// Returns SSE `data: ...` lines verbatim.
pub async fn execute_streaming_request(
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
body: Value,
) -> AppResult<Vec<String>> {
let client = http_client();
let upstream_url = upstream_url(&route.base_url, path_suffix);
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}"
)));
}
let resp = execute_streaming_response(state, route, path_suffix, body).await?;
let mut chunks = Vec::new();
let mut stream = resp.bytes_stream();
@@ -455,14 +449,50 @@ pub async fn execute_streaming_request(
Ok(chunks)
}
/// Execute a streaming request against any upstream path and return the live response.
///
/// The caller owns draining `bytes_stream()` so bytes can be flushed to the
/// downstream client as soon as the upstream sends them.
pub async fn execute_streaming_response(
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
body: Value,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, path_suffix);
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}"
)));
}
Ok(resp)
}
/// Generic non-streaming proxy request to any upstream endpoint path.
pub async fn execute_proxy_request(
_state: AppState,
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
body: Value,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
@@ -492,8 +522,12 @@ pub async fn execute_proxy_request(
}
/// Generic GET proxy request to an upstream endpoint path.
pub async fn execute_proxy_get(route: UpstreamRoute, path_suffix: &str) -> AppResult<Value> {
let client = http_client();
pub async fn execute_proxy_get(
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
) -> AppResult<Value> {
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
@@ -532,7 +566,7 @@ pub async fn execute_video_task_status(
} else {
format!("/videos/{id}")
};
execute_proxy_get(route, &suffix).await
execute_proxy_get(state, route, &suffix).await
}
/// Execute an OpenAI-shaped video task request, with Volcengine Ark native
@@ -929,7 +963,7 @@ async fn comfy_json(
path: &str,
body: Option<Value>,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let mut request = client.request(method, comfy_native_url(state, path));
if let Some(api_key) = state
.config
@@ -1016,11 +1050,11 @@ fn comfy_output_objects(value: &Value) -> bool {
/// Execute an image generation request.
pub async fn execute_image_generation(
_state: AppState,
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, "/images/generations");
let resp = client
@@ -1073,13 +1107,13 @@ pub struct MultipartField {
/// the route-resolved upstream model and skip any client-side fields
/// that don't belong on the wire.
pub async fn execute_image_edit(
_state: AppState,
state: AppState,
route: UpstreamRoute,
fields: Vec<MultipartField>,
) -> AppResult<Value> {
use reqwest::multipart::{Form, Part};
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, "/images/edits");
let mut form = Form::new();
src/proxy/token_router/router.rs
+126 -53
View File
@@ -2,16 +2,52 @@ use super::types::{RouteDecision, RouteMatch};
use crate::error::{AppError, AppResult};
use crate::proxy::conductor::{FailureInfo, SelectedChannel, SuccessMetrics};
use crate::proxy::executors::UpstreamRoute;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use regex::Regex;
use sqlx::{FromRow, SqlitePool};
use std::sync::Arc;
use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
const ROUTE_CACHE_TTL: Duration = Duration::from_secs(2);
static ROUTE_ROWS_CACHE: Lazy<RwLock<Option<TimedCache<Vec<RouteRow>>>>> =
Lazy::new(|| RwLock::new(None));
static CHANNEL_ROWS_CACHE: Lazy<RwLock<HashMap<i64, TimedCache<Vec<ChannelRow>>>>> =
Lazy::new(|| RwLock::new(HashMap::new()));
#[derive(Debug, Clone)]
struct TimedCache<T> {
loaded_at: Instant,
value: T,
}
impl<T> TimedCache<T> {
fn fresh(&self) -> bool {
self.loaded_at.elapsed() <= ROUTE_CACHE_TTL
}
}
pub fn invalidate_route_cache() {
*ROUTE_ROWS_CACHE.write() = None;
CHANNEL_ROWS_CACHE.write().clear();
}
fn invalidate_channel_cache_for_channel(channel_id: i64) {
CHANNEL_ROWS_CACHE
.write()
.retain(|_, cache| !cache.value.iter().any(|row| row.channel_id == channel_id));
}
#[derive(Clone)]
pub struct TokenRouter {
db: Arc<SqlitePool>,
}
#[derive(Debug, FromRow)]
#[derive(Debug, Clone, FromRow)]
struct RouteRow {
id: i64,
model_pattern: String,
@@ -19,7 +55,7 @@ struct RouteRow {
route_mode: Option<String>,
}
#[derive(Debug, FromRow)]
#[derive(Debug, Clone, FromRow)]
struct ChannelRow {
channel_id: i64,
route_id: i64,
@@ -91,14 +127,7 @@ impl TokenRouter {
}
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 rows = self.route_rows().await?;
let mut best: Option<(RouteRow, RouteMatch, usize, usize)> = None;
for row in rows {
@@ -131,48 +160,7 @@ impl TokenRouter {
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,
COALESCE((
SELECT sae.url
FROM site_api_endpoints sae
WHERE sae.site_id = s.id
AND COALESCE(sae.enabled, 1) = 1
AND (sae.cooldown_until IS NULL OR sae.cooldown_until <= datetime('now'))
ORDER BY COALESCE(sae.sort_order, 0), sae.id
LIMIT 1
), 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,
a.unit_cost AS account_unit_cost,
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 rows = self.channel_rows(route.id).await?;
let row = rows
.into_iter()
@@ -214,6 +202,89 @@ impl TokenRouter {
}))
}
async fn route_rows(&self) -> AppResult<Vec<RouteRow>> {
if let Some(cache) = ROUTE_ROWS_CACHE.read().as_ref() {
if cache.fresh() {
return Ok(cache.value.clone());
}
}
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?;
*ROUTE_ROWS_CACHE.write() = Some(TimedCache {
loaded_at: Instant::now(),
value: rows.clone(),
});
Ok(rows)
}
async fn channel_rows(&self, route_id: i64) -> AppResult<Vec<ChannelRow>> {
if let Some(cache) = CHANNEL_ROWS_CACHE.read().get(&route_id) {
if cache.fresh() {
return Ok(cache.value.clone());
}
}
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,
COALESCE((
SELECT sae.url
FROM site_api_endpoints sae
WHERE sae.site_id = s.id
AND COALESCE(sae.enabled, 1) = 1
AND (sae.cooldown_until IS NULL OR sae.cooldown_until <= datetime('now'))
ORDER BY COALESCE(sae.sort_order, 0), sae.id
LIMIT 1
), 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,
a.unit_cost AS account_unit_cost,
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?;
CHANNEL_ROWS_CACHE.write().insert(
route_id,
TimedCache {
loaded_at: Instant::now(),
value: rows.clone(),
},
);
Ok(rows)
}
pub async fn record_success(&self, channel_id: i64, metrics: SuccessMetrics) -> AppResult<()> {
sqlx::query(
"UPDATE route_channels
@@ -232,6 +303,7 @@ impl TokenRouter {
.bind(channel_id)
.execute(&*self.db)
.await?;
invalidate_channel_cache_for_channel(channel_id);
Ok(())
}
@@ -254,6 +326,7 @@ impl TokenRouter {
.bind(channel_id)
.execute(&*self.db)
.await?;
invalidate_channel_cache_for_channel(channel_id);
Ok(())
}
}
src/routes/admin/compat.rs
+12
View File
@@ -288,6 +288,7 @@ async fn create_route(
.bind(result.last_insert_rowid())
.fetch_one(&*state.db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(json!({ "route": route, "data": route })))
}
@@ -320,6 +321,7 @@ async fn update_route(
.bind(id)
.fetch_one(&*state.db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(json!({ "route": route, "data": route })))
}
@@ -331,6 +333,7 @@ async fn delete_route(
.bind(id)
.execute(&*state.db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(
json!({ "success": true, "data": { "deleted": true, "id": id } }),
))
@@ -375,6 +378,7 @@ async fn add_channel(
.bind(result.last_insert_rowid())
.fetch_one(&*state.db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(json!({ "channel": channel, "data": channel })))
}
@@ -401,6 +405,9 @@ async fn batch_add_channels(
created += 1;
}
}
if created > 0 {
crate::proxy::token_router::router::invalidate_route_cache();
}
Ok(Json(json!({ "success": true, "created": created })))
}
@@ -419,6 +426,7 @@ async fn update_channel(
.bind(id)
.execute(&*state.db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(json!({ "success": true })))
}
@@ -430,6 +438,7 @@ async fn delete_channel(
.bind(id)
.execute(&*state.db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(json!({ "success": true })))
}
@@ -450,6 +459,7 @@ async fn batch_update_channels(
}
}
}
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(json!({ "success": true })))
}
@@ -461,6 +471,7 @@ async fn clear_route_cooldown(
.bind(id)
.execute(&*state.db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(Json(json!({ "success": true })))
}
@@ -502,6 +513,7 @@ async fn batch_routes(
.execute(&*state.db)
.await?;
}
crate::proxy::token_router::router::invalidate_route_cache();
}
Ok(Json(json!({ "success": true })))
}
src/routes/proxy/mod.rs
+470 -381
View File
@@ -7,9 +7,11 @@
use crate::{
auth::api_key::{authorize_downstream_token, DownstreamAuthResult},
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
proxy::conductor::{
AttemptResult, ExecuteInput, ExecuteResult, ProxyConductor, SelectedChannel,
classify_failure, AttemptFailureAction, AttemptResult, ExecuteInput, ExecuteResult,
ProxyConductor, SelectedChannel, SuccessMetrics,
},
proxy::executors::{self, UpstreamRoute},
proxy::token_router::TokenRouter,
@@ -17,15 +19,16 @@ use crate::{
AppState,
};
use axum::{
body::Body,
extract::{Multipart, Query, State},
http::HeaderMap,
response::{IntoResponse, Response, Sse},
http::{header, HeaderMap},
response::{IntoResponse, Response},
routing::{get, post},
Json, Router,
};
use futures::{stream, FutureExt, StreamExt};
use futures::{FutureExt, StreamExt};
use serde_json::{json, Value};
use std::{convert::Infallible, sync::Arc, time::Instant};
use std::{future::Future, sync::Arc, time::Instant};
use tracing::info;
// ── Auth types ────────────────────────────────────────────────────────
@@ -111,14 +114,7 @@ async fn messages_handler(
let log_content = proxy_log_content_from_request(&body);
if canonical.stream {
let chunks =
execute_messages_stream_with_conductor(state, key_id, model, canonical, log_content)
.await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
execute_messages_stream_with_conductor(state, key_id, model, canonical, log_content).await
} else {
let result =
execute_messages_with_conductor(state, key_id, model, canonical, log_content).await?;
@@ -154,12 +150,7 @@ async fn responses_handler(
.unwrap_or(false);
if stream {
let chunks = execute_responses_stream_with_conductor(state, key_id, model, body).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
execute_responses_stream_with_conductor(state, key_id, model, body).await
} else {
let result = execute_responses_with_conductor(state, key_id, model, body).await?;
Ok(Json(result).into_response())
@@ -188,13 +179,7 @@ async fn chat_completions(
.unwrap_or(false);
if stream {
let chunks =
execute_chat_completions_stream_with_conductor(state, key_id, model, body).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
execute_chat_completions_stream_with_conductor(state, key_id, model, body).await
} else {
let result = execute_chat_completions_with_conductor(state, key_id, model, body).await?;
Ok(Json(result).into_response())
@@ -635,16 +620,36 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<Option<i64>
match result {
DownstreamAuthResult::Managed { key } => {
if !key.enabled.unwrap_or(true) {
return Err(AppError::Forbidden("API key is disabled".into()));
}
UsageTracker::new(state.db.clone()).check_rate_limit(key.id).await?;
check_downstream_key_limits(&key)?;
Ok(Some(key.id))
}
DownstreamAuthResult::Global { .. } => Ok(None),
}
}
fn check_downstream_key_limits(key: &DownstreamApiKey) -> AppResult<()> {
if !key.enabled.unwrap_or(true) {
return Err(AppError::Forbidden("API key is disabled".into()));
}
if let Some(max_requests) = key.max_requests {
if max_requests <= 0 || key.used_requests.unwrap_or(0) + 1 > max_requests {
return Err(AppError::RateLimitedWithRetry {
message: "下游密钥请求额度已用完".into(),
retry_after_secs: 60,
});
}
}
if let Some(max_cost) = key.max_cost {
if max_cost <= 0.0 || key.used_cost.unwrap_or(0.0) >= max_cost {
return Err(AppError::RateLimitedWithRetry {
message: "下游密钥成本额度已用完".into(),
retry_after_secs: 60,
});
}
}
Ok(())
}
/// Best-effort usage logging. Only fires when the request was authenticated
/// via a managed downstream key (key_id is Some). Failures are logged but
/// never bubbled — billing must not break a successful proxy response.
@@ -1014,6 +1019,317 @@ fn ensure_stream_usage(mut body: Value) -> Value {
body
}
#[derive(Debug)]
struct OpenedStream {
selected: SelectedChannel,
attempts: usize,
response: reqwest::Response,
}
#[derive(Debug)]
struct StreamOpenFailure {
selected: Option<SelectedChannel>,
attempts: usize,
status: Option<u16>,
raw_error_text: Option<String>,
}
async fn open_stream_with_failover<F, Fut>(
state: &AppState,
requested_model: &str,
mut open_attempt: F,
) -> Result<OpenedStream, StreamOpenFailure>
where
F: FnMut(AppState, SelectedChannel) -> Fut,
Fut: Future<Output = AppResult<reqwest::Response>>,
{
let conductor = build_conductor(state);
let mut selected = (conductor.select_channel)(requested_model.to_string(), None)
.await
.ok_or_else(|| StreamOpenFailure {
selected: None,
attempts: 0,
status: None,
raw_error_text: Some(format!("No route configured for model '{requested_model}'")),
})?;
let mut excluded = Vec::new();
let mut attempts = 0usize;
loop {
attempts += 1;
match open_attempt(state.clone(), selected.clone()).await {
Ok(response) => {
return Ok(OpenedStream {
selected,
attempts,
response,
});
}
Err(err) => {
let status = upstream_status_from_error(&err);
let action = classify_failure(status, Some(&err.to_string()));
if let Err(record_err) = TokenRouter::new(state.db.clone())
.record_failure(
selected.channel_id,
crate::proxy::conductor::FailureInfo {
action: action.as_str().to_string(),
status,
raw_error_text: Some(err.to_string()),
},
)
.await
{
tracing::warn!(%record_err, channel_id = selected.channel_id, "failed to record streaming proxy failure");
}
if matches!(action, AttemptFailureAction::RetrySameChannel) && attempts < 5 {
continue;
}
let raw_error_text = Some(err.to_string());
if matches!(action, AttemptFailureAction::Terminal | AttemptFailureAction::Stop) || attempts >= 5 {
return Err(StreamOpenFailure {
selected: Some(selected),
attempts,
status,
raw_error_text,
});
}
excluded.push(selected.channel_id);
selected = (conductor.select_next_channel)(
requested_model.to_string(),
excluded.clone(),
None,
)
.await
.ok_or_else(|| StreamOpenFailure {
selected: Some(selected.clone()),
attempts,
status,
raw_error_text,
})?;
}
}
}
}
fn sse_passthrough_response(
state: AppState,
opened: OpenedStream,
key_id: Option<i64>,
requested_model: String,
started_at: Instant,
log_content: ProxyLogContent,
usage_extractor: fn(&str) -> Option<(i64, i64)>,
) -> AppResult<Response> {
let selected = opened.selected;
let attempts = opened.attempts;
let billing_state = state.clone();
let stream = async_stream::stream! {
let mut usage = StreamUsageAccumulator::default();
let mut bytes_stream = opened.response.bytes_stream();
let mut stream_error: Option<String> = None;
while let Some(chunk) = bytes_stream.next().await {
let bytes = match chunk {
Ok(bytes) => bytes,
Err(err) => {
let message = format!("upstream stream error: {err}");
stream_error = Some(message.clone());
let event = format!(
"event: error\ndata: {}\n\n",
serde_json::json!({ "error": { "message": message } })
);
yield Ok::<bytes::Bytes, std::io::Error>(bytes::Bytes::from(event));
break;
}
};
usage.ingest_bytes(&bytes, usage_extractor);
yield Ok::<bytes::Bytes, std::io::Error>(bytes);
}
usage.finish(usage_extractor);
let prompt_tokens = usage.prompt_tokens;
let completion_tokens = usage.completion_tokens;
let cost = estimate_request_cost(&selected, prompt_tokens, completion_tokens);
if let Some(error) = stream_error {
if let Err(err) = TokenRouter::new(billing_state.db.clone())
.record_failure(
selected.channel_id,
crate::proxy::conductor::FailureInfo {
action: AttemptFailureAction::Failover.as_str().to_string(),
status: None,
raw_error_text: Some(error.clone()),
},
)
.await
{
tracing::warn!(%err, channel_id = selected.channel_id, "failed to record streaming proxy failure");
}
record_proxy_log_failure(
&billing_state,
key_id,
Some(&selected),
&requested_model,
None,
Some(elapsed_ms_since(started_at)),
Some(&error),
attempts,
true,
log_content,
)
.await;
} else {
if let Err(err) = TokenRouter::new(billing_state.db.clone())
.record_success(
selected.channel_id,
SuccessMetrics {
latency_ms: Some(started_at.elapsed().as_millis().min(u64::MAX as u128) as u64),
cost: Some(cost),
metadata: Default::default(),
},
)
.await
{
tracing::warn!(%err, channel_id = selected.channel_id, "failed to record streaming proxy success");
}
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&requested_model,
Some(elapsed_ms_since(started_at)),
prompt_tokens,
completion_tokens,
cost,
attempts,
true,
log_content,
)
.await;
record_billing(
&billing_state,
key_id,
&requested_model,
prompt_tokens,
completion_tokens,
cost,
)
.await;
}
};
Response::builder()
.header(header::CONTENT_TYPE, "text/event-stream; charset=utf-8")
.header(header::CACHE_CONTROL, "no-cache, no-transform")
.header(header::CONNECTION, "keep-alive")
.header("x-accel-buffering", "no")
.body(Body::from_stream(stream))
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))
}
#[derive(Default)]
struct StreamUsageAccumulator {
buffer: String,
prompt_tokens: i64,
completion_tokens: i64,
}
impl StreamUsageAccumulator {
fn ingest_bytes(
&mut self,
bytes: &bytes::Bytes,
usage_extractor: fn(&str) -> Option<(i64, i64)>,
) {
self.buffer.push_str(&String::from_utf8_lossy(bytes));
while let Some(newline) = self.buffer.find('\n') {
let line = self.buffer[..newline].trim_end_matches('\r').to_string();
self.buffer.drain(..=newline);
self.ingest_line(&line, usage_extractor);
}
if self.buffer.len() > 64 * 1024 {
self.buffer.clear();
}
}
fn finish(&mut self, usage_extractor: fn(&str) -> Option<(i64, i64)>) {
if !self.buffer.trim().is_empty() {
let line = std::mem::take(&mut self.buffer);
self.ingest_line(line.trim_end_matches('\r'), usage_extractor);
}
}
fn ingest_line(&mut self, line: &str, usage_extractor: fn(&str) -> Option<(i64, i64)>) {
if let Some((prompt, completion)) = usage_extractor(line) {
if prompt > 0 {
self.prompt_tokens = prompt;
}
if completion > 0 {
self.completion_tokens = completion;
}
}
}
}
fn openai_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let (prompt, completion) = extract_usage(&value);
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn anthropic_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let mut prompt = 0;
let mut completion = 0;
for holder in [
value.get("message"),
value.get("delta"),
Some(&value),
value.get("response"),
]
.into_iter()
.flatten()
{
let (candidate_prompt, candidate_completion) = extract_usage(holder);
if candidate_prompt > 0 {
prompt = candidate_prompt;
}
if candidate_completion > 0 {
completion = candidate_completion;
}
}
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn responses_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let usage_holder = value.get("response").unwrap_or(&value);
let (prompt, completion) = extract_usage(usage_holder);
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn build_conductor(state: &AppState) -> ProxyConductor {
let db = state.db.clone();
let select_db = db.clone();
@@ -1079,10 +1395,12 @@ async fn execute_messages_with_conductor(
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let mut request = request.clone();
request.model = ctx.selected.actual_model.clone();
async move {
match executors::execute_anthropic_messages(
state,
selected_channel_to_route(ctx.selected),
request,
)
@@ -1174,108 +1492,57 @@ async fn execute_messages_stream_with_conductor(
model: String,
request: crate::transformers::canonical::CanonicalRequest,
log_content: ProxyLogContent,
) -> AppResult<Vec<String>> {
) -> Result<Response, AppError> {
let started_at = Instant::now();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let mut request = request.clone();
request.model = ctx.selected.actual_model.clone();
async move {
match executors::execute_anthropic_messages_stream(
selected_channel_to_route(ctx.selected),
request,
)
.await
{
Ok(chunks) => AttemptResult::Success {
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let request_for_attempt = request.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let mut request = request_for_attempt.clone();
async move {
request.model = selected.actual_model.clone();
executors::execute_anthropic_messages_stream_response(
state,
selected_channel_to_route(selected),
request,
)
.await
}
})
.await;
let (selected, attempts, chunks) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => match response {
Value::Array(chunks) => (
selected,
attempts,
chunks
.into_iter()
.filter_map(|chunk| chunk.as_str().map(str::to_string))
.collect::<Vec<_>>(),
),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
}
},
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
log_content,
anthropic_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&billing_state,
&state,
key_id,
selected.as_ref(),
&billing_model,
status,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
failure.raw_error_text.as_deref(),
failure.attempts,
true,
log_content,
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
};
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
0,
0,
cost,
attempts,
true,
log_content,
)
.await;
record_billing(&billing_state, key_id, &billing_model, 0, 0, cost).await;
Ok(chunks)
}
}
async fn execute_chat_completions_with_conductor(
@@ -1389,148 +1656,57 @@ async fn execute_chat_completions_stream_with_conductor(
key_id: Option<i64>,
model: String,
body: Value,
) -> AppResult<Vec<String>> {
) -> Result<Response, AppError> {
let started_at = Instant::now();
let request_log_body = body.clone();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let body =
ensure_stream_usage(body_with_selected_model(body.clone(), &ctx.selected));
async move {
match executors::execute_chat_completions_stream(
state,
selected_channel_to_route(ctx.selected),
body,
)
.await
{
Ok(chunks) => AttemptResult::Success {
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let body_for_attempt = body.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let body = ensure_stream_usage(body_with_selected_model(body_for_attempt.clone(), &selected));
async move {
executors::execute_chat_completions_stream_response(
state,
selected_channel_to_route(selected),
body,
)
.await
}
})
.await;
let (selected, attempts, chunks) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => match response {
Value::Array(chunks) => (
selected,
attempts,
chunks
.into_iter()
.map(|chunk| {
chunk.as_str().map(str::to_string).ok_or_else(|| {
AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
)
})
})
.collect::<AppResult<Vec<String>>>()?,
),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
}
},
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
proxy_log_content_from_request(&request_log_body),
openai_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&billing_state,
&state,
key_id,
selected.as_ref(),
&billing_model,
status,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
failure.raw_error_text.as_deref(),
failure.attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
};
// Streaming responses don't include a usage block in the aggregate.
// OpenAI emits usage on the final `data: {...}` chunk only when
// `stream_options.include_usage = true` was set; parse it best-effort.
let (prompt, completion) = chunks
.iter()
.rev()
.find_map(|chunk| {
let payload = chunk
.strip_prefix("data: ")
.or_else(|| chunk.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let (p, c) = extract_usage(&value);
if p == 0 && c == 0 {
None
} else {
Some((p, c))
}
})
.unwrap_or((0, 0));
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
cost,
attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
cost,
)
.await;
Ok(chunks)
}
}
async fn execute_responses_with_conductor(
@@ -1644,145 +1820,58 @@ async fn execute_responses_stream_with_conductor(
key_id: Option<i64>,
model: String,
body: Value,
) -> AppResult<Vec<String>> {
) -> Result<Response, AppError> {
let started_at = Instant::now();
let request_log_body = body.clone();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let body = body_with_selected_model(body.clone(), &ctx.selected);
async move {
match executors::execute_streaming_request(
selected_channel_to_route(ctx.selected),
"/responses",
body,
)
.await
{
Ok(chunks) => AttemptResult::Success {
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let body_for_attempt = body.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let body = body_with_selected_model(body_for_attempt.clone(), &selected);
async move {
executors::execute_streaming_response(
state,
selected_channel_to_route(selected),
"/responses",
body,
)
.await
}
})
.await;
let (selected, attempts, chunks) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => match response {
Value::Array(chunks) => (
selected,
attempts,
chunks
.into_iter()
.map(|chunk| {
chunk.as_str().map(str::to_string).ok_or_else(|| {
AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
)
})
})
.collect::<AppResult<Vec<String>>>()?,
),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
}
},
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
proxy_log_content_from_request(&request_log_body),
responses_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&billing_state,
&state,
key_id,
selected.as_ref(),
&billing_model,
status,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
failure.raw_error_text.as_deref(),
failure.attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
};
// Responses API streams `response.completed` events that contain a
// `usage` object with input_tokens / output_tokens. extract_usage already
// handles both naming conventions.
let (prompt, completion) = chunks
.iter()
.rev()
.find_map(|chunk| {
let payload = chunk.strip_prefix("data: ")?;
if payload.trim() == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
// Responses API nests usage under `response.usage`.
let usage_holder = value.get("response").unwrap_or(&value);
let (p, c) = extract_usage(usage_holder);
if p == 0 && c == 0 {
None
} else {
Some((p, c))
}
})
.unwrap_or((0, 0));
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
cost,
attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
cost,
)
.await;
Ok(chunks)
}
}
#[cfg(test)]
src/routes/workspace.rs
+8 -2
View File
@@ -299,7 +299,7 @@ async fn get_asset_row(
AND ((? = 1 AND (? IS NULL OR user_id = ?)) OR (? = 0 AND user_id = ?))
LIMIT 1",
)
.bind(public_id)
.bind(&public_id)
.bind(is_admin)
.bind(scoped_user_id)
.bind(scoped_user_id)
@@ -339,7 +339,7 @@ async fn delete_asset(
AND deleted_at IS NULL
AND ((? = 1 AND (? IS NULL OR user_id = ?)) OR (? = 0 AND user_id = ?))",
)
.bind(public_id)
.bind(&public_id)
.bind(is_admin)
.bind(scoped_user_id)
.bind(scoped_user_id)
@@ -350,6 +350,9 @@ async fn delete_asset(
if result.rows_affected() == 0 {
return Err(AppError::NotFound("Workspace asset not found".into()));
}
if let Err(err) = crate::services::workspace::project_deleted_asset(&state, &public_id).await {
tracing::warn!(%err, %public_id, "failed to update workspace asset delete projection");
}
Ok(StatusCode::NO_CONTENT)
}
@@ -395,6 +398,9 @@ async fn delete_assets_batch(
.execute(&*state.db)
.await?;
if result.rows_affected() > 0 {
if let Err(err) = crate::services::workspace::project_deleted_asset(&state, &public_id).await {
tracing::warn!(%err, %public_id, "failed to update workspace asset batch delete projection");
}
deleted_ids.push(public_id);
} else {
failed_ids.push(public_id);
src/services/cron/mod.rs
+1
View File
@@ -127,6 +127,7 @@ impl CronScheduler {
)
.execute(&*db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(())
}
}
src/services/route_rebuild.rs
+3 -1
View File
@@ -57,7 +57,9 @@ struct AvailRow {
/// single transaction so a partial failure doesn't leave `token_routes` and
/// `route_channels` out of sync.
pub async fn rebuild_routes(state: &AppState) -> AppResult<RebuildSummary> {
rebuild_routes_inner(&state.db).await
let summary = rebuild_routes_inner(&state.db).await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(summary)
}
/// Pool-only entry point. Identical to `rebuild_routes`, just decoupled from
src/services/workspace.rs
+102 -5
View File
@@ -160,6 +160,23 @@ async fn insert_asset(
.execute(&*state.db)
.await?;
if should_project_workspace_asset(status, &metadata) {
if let Err(err) = upsert_workspace_asset_daily(
state,
None,
input.user_id,
&input.source,
&input.kind,
size_bytes,
file_path.is_some(),
1,
)
.await
{
tracing::warn!(%err, "failed to update workspace asset projection");
}
}
Ok(StoredWorkspaceAsset {
public_id: public_id.to_string(),
file_name,
@@ -169,6 +186,51 @@ async fn insert_asset(
})
}
fn should_project_workspace_asset(status: &str, metadata: &Value) -> bool {
status == "completed"
&& (metadata.pointer("/workspaceAsset/role").and_then(Value::as_str) == Some("output")
|| metadata.pointer("/workflowUi/type").and_then(Value::as_str) == Some("output")
|| metadata.get("responseItem").is_some()
|| metadata.get("moneyPrinter").is_some()
|| metadata.get("digitalHuman").is_some())
}
async fn upsert_workspace_asset_daily(
state: &AppState,
local_day: Option<&str>,
user_id: Option<i64>,
source: &str,
kind: &str,
size_bytes: i64,
has_local_file: bool,
delta: i64,
) -> Result<(), AppError> {
let asset_count_delta = delta;
let local_file_delta = if has_local_file { delta } else { 0 };
let size_delta = size_bytes.max(0) * delta;
sqlx::query(
"INSERT INTO workspace_asset_daily
(local_day, user_id, source, kind, asset_count, local_file_count, total_size_bytes, updated_at)
VALUES (COALESCE(?, date('now', 'localtime')), ?, ?, ?, ?, ?, ?, datetime('now'))
ON CONFLICT(local_day, user_id, source, kind) DO UPDATE SET
asset_count = MAX(0, asset_count + excluded.asset_count),
local_file_count = MAX(0, local_file_count + excluded.local_file_count),
total_size_bytes = MAX(0, total_size_bytes + excluded.total_size_bytes),
updated_at = datetime('now')",
)
.bind(local_day)
.bind(user_id.unwrap_or(0))
.bind(source)
.bind(kind)
.bind(asset_count_delta)
.bind(local_file_delta)
.bind(size_delta)
.execute(&*state.db)
.await?;
Ok(())
}
pub async fn record_remote_asset(
state: &AppState,
input: WorkspaceAssetInput,
@@ -188,11 +250,7 @@ pub async fn record_remote_asset(
.await;
};
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(60))
.build()
.unwrap_or_else(|_| reqwest::Client::new());
match client.get(url).send().await {
match state.http.get(url).send().await {
Ok(response) if response.status().is_success() => {
let mime_type = response
.headers()
@@ -307,6 +365,45 @@ pub async fn record_bytes_asset(
.await
}
pub async fn project_deleted_asset(state: &AppState, public_id: &str) -> Result<(), AppError> {
use sqlx::Row;
let row = sqlx::query(
"SELECT user_id, source, kind, status, metadata, size_bytes, file_path, date(created_at, 'localtime') AS local_day
FROM workspace_assets
WHERE public_id = ?
LIMIT 1",
)
.bind(public_id)
.fetch_optional(&*state.db)
.await?;
let Some(row) = row else {
return Ok(());
};
let status: String = row.try_get("status").unwrap_or_default();
let metadata_raw: String = row.try_get("metadata").unwrap_or_else(|_| "{}".to_string());
let metadata = serde_json::from_str::<Value>(&metadata_raw).unwrap_or_else(|_| json!({}));
if !should_project_workspace_asset(&status, &metadata) {
return Ok(());
}
let source: String = row.try_get("source").unwrap_or_default();
let kind: String = row.try_get("kind").unwrap_or_default();
let user_id: Option<i64> = row.try_get("user_id").ok().flatten();
let size_bytes: i64 = row.try_get("size_bytes").unwrap_or_default();
let file_path: Option<String> = row.try_get("file_path").ok().flatten();
let local_day: Option<String> = row.try_get("local_day").ok().flatten();
upsert_workspace_asset_daily(
state,
local_day.as_deref(),
user_id,
&source,
&kind,
size_bytes,
file_path.as_deref().is_some_and(|value| !value.trim().is_empty()),
-1,
)
.await
}
pub async fn record_image_assets_from_response(
state: &AppState,
user_id: Option<i64>,