mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-14 08:06:40 +00:00
315 lines
12 KiB
Rust
315 lines
12 KiB
Rust
use std::time::Instant;
|
|
|
|
use axum::{
|
|
body::Body,
|
|
http::{
|
|
header::{CACHE_CONTROL, CONTENT_TYPE},
|
|
StatusCode,
|
|
},
|
|
response::{IntoResponse, Response},
|
|
Json,
|
|
};
|
|
use serde_json::json;
|
|
|
|
use crate::{
|
|
codex::events::event::EventLevel,
|
|
codex::gateway::protocol::{
|
|
codex_to_openai::{
|
|
chat_completion_from_codex_sse, chat_completion_stream_from_codex_sse, openai_error,
|
|
},
|
|
openai_to_codex::{translate_chat_to_codex, ChatCompletionRequest},
|
|
},
|
|
codex::models::catalog::ModelCatalog,
|
|
codex::models::service::ModelService,
|
|
codex::serving::dispatch::{
|
|
classify_upstream_account_retry, no_available_accounts_response,
|
|
normalize_service_tier_for_upstream, CodexRequestLogContext, CodexUpstreamService,
|
|
},
|
|
codex::serving::http::errors::{
|
|
codex_client_error_message, codex_client_error_response,
|
|
codex_client_error_response_with_status_and_message, model_not_found_response,
|
|
},
|
|
config::ModelConfig,
|
|
};
|
|
|
|
#[derive(Clone)]
|
|
pub struct ChatService {
|
|
model_config: ModelConfig,
|
|
models: ModelService,
|
|
upstream: CodexUpstreamService,
|
|
}
|
|
|
|
impl ChatService {
|
|
pub(crate) fn new(
|
|
model_config: ModelConfig,
|
|
models: ModelService,
|
|
upstream: CodexUpstreamService,
|
|
) -> Self {
|
|
Self {
|
|
model_config,
|
|
models,
|
|
upstream,
|
|
}
|
|
}
|
|
|
|
pub async fn handle(
|
|
&self,
|
|
request_id: &str,
|
|
chat_request: ChatCompletionRequest,
|
|
started_at: Instant,
|
|
) -> Response {
|
|
let client_stream = chat_request.stream;
|
|
let requested_model = chat_request.model.clone();
|
|
let catalog = self.models.catalog().await;
|
|
if !catalog.is_recognized_model_name(&requested_model) {
|
|
return model_not_found_response().into_response();
|
|
}
|
|
let parsed_model = catalog.parse_model_name(&requested_model);
|
|
let display_model = ModelCatalog::build_display_model_name(&parsed_model);
|
|
let mut codex_request = match translate_chat_to_codex(chat_request) {
|
|
Ok(request) => request,
|
|
Err(_) => {
|
|
return (
|
|
StatusCode::BAD_REQUEST,
|
|
Json(openai_error(
|
|
"Invalid chat completion request",
|
|
"invalid_request",
|
|
)),
|
|
)
|
|
.into_response();
|
|
}
|
|
};
|
|
codex_request.model = parsed_model.model_id.clone();
|
|
if codex_request.reasoning.is_none() {
|
|
let effort = parsed_model
|
|
.reasoning_effort
|
|
.clone()
|
|
.or_else(|| self.model_config.default_reasoning_effort.clone());
|
|
if let Some(effort) = effort {
|
|
codex_request.reasoning = Some(json!({"effort": effort, "summary": "auto"}));
|
|
}
|
|
}
|
|
if codex_request.service_tier.is_none() {
|
|
codex_request.service_tier = parsed_model
|
|
.service_tier
|
|
.clone()
|
|
.or_else(|| self.model_config.service_tier.clone());
|
|
}
|
|
codex_request.service_tier = codex_request
|
|
.service_tier
|
|
.map(normalize_service_tier_for_upstream);
|
|
let include_reasoning = codex_request.reasoning.is_some();
|
|
|
|
let Some(mut acquired) = self.upstream.acquire_account(&codex_request.model).await else {
|
|
return no_available_accounts_response().into_response();
|
|
};
|
|
let mut log_context = CodexRequestLogContext::new(
|
|
request_id,
|
|
&acquired.account.id,
|
|
&codex_request.model,
|
|
client_stream,
|
|
started_at,
|
|
);
|
|
|
|
let mut excluded_account_ids = Vec::new();
|
|
let mut model_unsupported_retry_used = false;
|
|
let response = loop {
|
|
self.upstream
|
|
.stagger_request(acquired.previous_slot_at)
|
|
.await;
|
|
let response = self
|
|
.upstream
|
|
.send_codex_request_with_upstream_retries(
|
|
&codex_request,
|
|
&acquired.account,
|
|
request_id,
|
|
)
|
|
.await;
|
|
self.upstream.release_account(&acquired.account.id).await;
|
|
|
|
match response {
|
|
Ok(response) => break response,
|
|
Err(error) => {
|
|
if let Some(retry) =
|
|
classify_upstream_account_retry(&error, model_unsupported_retry_used)
|
|
{
|
|
if retry.is_model_unsupported() {
|
|
model_unsupported_retry_used = true;
|
|
}
|
|
let fallback = self
|
|
.upstream
|
|
.apply_retry_and_acquire_fallback(
|
|
&acquired.account,
|
|
retry,
|
|
&codex_request.model,
|
|
&mut excluded_account_ids,
|
|
false,
|
|
)
|
|
.await;
|
|
self.upstream
|
|
.log_response(
|
|
&log_context,
|
|
retry.status(),
|
|
EventLevel::Warn,
|
|
"v1 chat completions 上游请求将使用备用账户重试",
|
|
retry.metadata(client_stream),
|
|
)
|
|
.await;
|
|
if let Some(fallback) = fallback {
|
|
log_context = log_context.with_account(&fallback.account.id);
|
|
acquired = fallback;
|
|
continue;
|
|
}
|
|
let retry_message =
|
|
retry.fallback_response_message(codex_client_error_message(&error));
|
|
let message = self
|
|
.upstream
|
|
.fallback_exhausted_message(&retry_message)
|
|
.await;
|
|
let error_response = codex_client_error_response_with_status_and_message(
|
|
error,
|
|
retry.status(),
|
|
&message,
|
|
);
|
|
self.upstream
|
|
.log_response(
|
|
&log_context,
|
|
error_response.0,
|
|
EventLevel::Error,
|
|
"v1 chat completions fallback 已耗尽",
|
|
json!({"stream": client_stream}),
|
|
)
|
|
.await;
|
|
return error_response.into_response();
|
|
}
|
|
if self
|
|
.upstream
|
|
.record_request_attempt(&acquired.account.id, false)
|
|
.await
|
|
.is_err()
|
|
{
|
|
self.upstream
|
|
.log_response(
|
|
&log_context,
|
|
StatusCode::OK,
|
|
EventLevel::Warn,
|
|
"v1 chat completions 记录失败请求计数失败",
|
|
json!({"stream": client_stream, "requestAttemptStoreError": true}),
|
|
)
|
|
.await;
|
|
}
|
|
let error_response = codex_client_error_response(error);
|
|
self.upstream
|
|
.log_response(
|
|
&log_context,
|
|
error_response.0,
|
|
EventLevel::Error,
|
|
"v1 chat completions 上游请求失败",
|
|
json!({"stream": client_stream}),
|
|
)
|
|
.await;
|
|
return error_response.into_response();
|
|
}
|
|
}
|
|
};
|
|
if self
|
|
.upstream
|
|
.persist_cookies(&acquired.account.id, &response.set_cookie_headers)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return (
|
|
StatusCode::INTERNAL_SERVER_ERROR,
|
|
Json(openai_error(
|
|
"Failed to persist upstream cookies",
|
|
"cookie_store_error",
|
|
)),
|
|
)
|
|
.into_response();
|
|
}
|
|
if let Some(usage) = response.usage {
|
|
if self
|
|
.upstream
|
|
.record_usage(&acquired.account.id, usage, false)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return (
|
|
StatusCode::INTERNAL_SERVER_ERROR,
|
|
Json(openai_error(
|
|
"Failed to record account usage",
|
|
"usage_store_error",
|
|
)),
|
|
)
|
|
.into_response();
|
|
}
|
|
}
|
|
|
|
if client_stream {
|
|
match chat_completion_stream_from_codex_sse(
|
|
&response.body,
|
|
&display_model,
|
|
include_reasoning,
|
|
codex_request.tuple_schema.as_ref(),
|
|
) {
|
|
Ok(Some(body)) => Response::builder()
|
|
.status(StatusCode::OK)
|
|
.header(CONTENT_TYPE, "text/event-stream")
|
|
.header(CACHE_CONTROL, "no-cache")
|
|
.body(Body::from(body))
|
|
.unwrap_or_else(|_| {
|
|
(
|
|
StatusCode::INTERNAL_SERVER_ERROR,
|
|
Json(openai_error(
|
|
"Failed to build stream response",
|
|
"stream_response_error",
|
|
)),
|
|
)
|
|
.into_response()
|
|
}),
|
|
Ok(None) => (
|
|
StatusCode::BAD_GATEWAY,
|
|
Json(openai_error(
|
|
"Codex response did not include response.completed",
|
|
"empty_upstream_response",
|
|
)),
|
|
)
|
|
.into_response(),
|
|
Err(_) => (
|
|
StatusCode::BAD_GATEWAY,
|
|
Json(openai_error(
|
|
"Invalid Codex SSE response",
|
|
"invalid_upstream_sse",
|
|
)),
|
|
)
|
|
.into_response(),
|
|
}
|
|
} else {
|
|
match chat_completion_from_codex_sse(
|
|
&response.body,
|
|
&display_model,
|
|
include_reasoning,
|
|
codex_request.tuple_schema.as_ref(),
|
|
) {
|
|
Ok(Some(body)) => (StatusCode::OK, Json(body)).into_response(),
|
|
Ok(None) => (
|
|
StatusCode::BAD_GATEWAY,
|
|
Json(openai_error(
|
|
"Codex response did not include response.completed",
|
|
"empty_upstream_response",
|
|
)),
|
|
)
|
|
.into_response(),
|
|
Err(_) => (
|
|
StatusCode::BAD_GATEWAY,
|
|
Json(openai_error(
|
|
"Invalid Codex SSE response",
|
|
"invalid_upstream_sse",
|
|
)),
|
|
)
|
|
.into_response(),
|
|
}
|
|
}
|
|
}
|
|
}
|