diff --git a/backend/crates/gateway-api/src/openai/endpoint.rs b/backend/crates/gateway-api/src/openai/endpoint.rs index b36a43fe2..a1b17a9fb 100644 --- a/backend/crates/gateway-api/src/openai/endpoint.rs +++ b/backend/crates/gateway-api/src/openai/endpoint.rs @@ -62,7 +62,7 @@ where let started = match service .start_prepared_provider_endpoint( prepared, operation, - // 图像/搜索等原生端点没有对应的上游模型权限。 + // 原生端点的模型语义由 Provider 从正文提取并检查账号政策 None, client_ip, user_agent, endpoint, ) .await diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index 5f85bd933..d98accdfc 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -68,7 +68,9 @@ pub use selector::{ CodexAccountFailure, CodexCredentialLease, CodexCredentialSelector, CredentialSelectionError, SelectCodexCredential, }; -pub(crate) use selector::{CodexCyberPolicyScope, SelectCodexProviderEndpointCredential}; +pub(crate) use selector::{ + CodexCyberPolicyScope, CodexSelectionModel, SelectCodexProviderEndpointCredential, +}; pub use types::{ CODEX_AUTHENTICATION_KIND_OAUTH, CodexAccountProfile, CodexCookie, CodexCookieCaptureOutcome, CodexCredentialData, CodexCredentialPrincipal, CodexOAuthCredentialData, CodexOAuthSecret, diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index 7ff035086..f70de56f0 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -7,11 +7,13 @@ use std::time::{Duration, SystemTime}; use gateway_core::error::ErrorSource; +use gateway_core::account::scope::FrozenAccountScope; use gateway_core::account::{ AccountAffinity, AccountCandidate, AccountCapacitySnapshot, AccountEligibilityPolicy, - AccountErrorReason, AccountFeedbackStats, AccountRuntimeSignals, AccountSchedulingBlocker, - AccountSelectionContext, AccountSelectionPolicy, AccountSelector, AccountStatus, - CredentialState, PreferredAccountSelection, ProviderAccount, ProviderAccountId, QuotaEvidence, + AccountErrorReason, AccountFeedbackStats, AccountModelAccessMode, AccountRuntimeSignals, + AccountSchedulingBlocker, AccountSelectionContext, AccountSelectionPolicy, AccountSelector, + AccountStatus, CredentialState, PreferredAccountSelection, ProviderAccount, ProviderAccountId, + QuotaEvidence, }; use gateway_core::concurrency::{CapacityWait, ConcurrencyWaitQueue, QueueRejection, WaitPriority}; use gateway_core::engine::{AttemptContext, ContinuationAttempt, policy::AccountPolicyError}; @@ -101,12 +103,53 @@ pub(crate) struct SelectCodexProviderEndpointCredential<'a> { pub request_url: &'a Url, pub attempt: &'a AttemptContext, pub session_affinity: Option<&'a CodexSessionAffinity>, - /// 本端点的上游模型；提供后按账号模型权限过滤候选（如 live 语音）。 - pub upstream_model: Option<&'a str>, + pub model: CodexSelectionModel<'a>, /// 本端点只接受 OAuth 凭据时排除 API Key 账号，避免混合池选中后必然失败。 pub requires_oauth: bool, } +/// 选号使用的模型类别，区分图片资格与不涉及模型的端点 +#[derive(Clone, Copy)] +pub(crate) enum CodexSelectionModel<'a> { + NotApplicable, + Requested(&'a str), + Image(Option<&'a str>), +} + +impl<'a> CodexSelectionModel<'a> { + pub(crate) const fn model(self) -> Option<&'a str> { + match self { + Self::NotApplicable => None, + Self::Requested(model) => Some(model), + Self::Image(model) => model, + } + } + + pub(crate) fn allows_account(self, account: &ProviderAccount) -> bool { + // 官方 Codex 的 image_generation_available 排除 Free，API Key 不套用 ChatGPT 套餐 + // https://github.com/openai/codex/commit/0a0a9b6c8f + !(matches!(self, Self::Image(_)) + && account.authentication_kind() == CODEX_AUTHENTICATION_KIND_OAUTH + && account + .plan_type() + .is_some_and(|plan| plan.trim().eq_ignore_ascii_case("free"))) + } + + pub(crate) fn allows(self, scope: &FrozenAccountScope, account: &ProviderAccountId) -> bool { + match self { + Self::NotApplicable => scope.allows(account), + Self::Requested(model) | Self::Image(Some(model)) => scope.allows_model(account, model), + // 缺失、重复或不可解析的图片模型不能绕过黑白名单；不猜上游默认值 + Self::Image(None) => { + scope.allows(account) + && scope.directory().account(account).is_some_and(|account| { + account.model_access().mode() == AccountModelAccessMode::All + }) + } + } + } +} + struct CredentialSelectionInput<'a> { requires_websocket: bool, requires_oauth: bool, @@ -304,8 +347,12 @@ impl CodexCredentialSelector { session_affinity_observation: None, guardian: false, }; - self.select_inner(&input, None, Some(request.upstream_model)) - .await + self.select_inner( + &input, + None, + CodexSelectionModel::Requested(request.upstream_model), + ) + .await } pub(crate) async fn select_with_cyber_policy( @@ -328,7 +375,7 @@ impl CodexCredentialSelector { self.select_inner( &input, cyber_policy_session_key, - Some(request.upstream_model), + CodexSelectionModel::Requested(request.upstream_model), ) .await } @@ -385,9 +432,8 @@ impl CodexCredentialSelector { /// 为不属于 Responses 文本模型目录的 Provider 原生端点选择账号 /// /// 账号范围、健康度、配额、并发租约、cookie 与认证准备仍走同一套选择链路； - /// 原生端点默认没有 Responses 模型，不套用管理员配置的文本模型权限。 - /// 端点有明确上游模型（如 live 语音）时通过 `upstream_model` 让账号 - /// 模型权限参与候选过滤；`requires_oauth` 限定本端点支持的认证类型。 + /// Images 与 live 语音按实际请求模型执行账号权限，不依赖文本模型目录 + /// `requires_oauth` 限定本端点支持的认证类型 pub(crate) async fn select_for_provider_endpoint( &self, request: &SelectCodexProviderEndpointCredential<'_>, @@ -401,15 +447,14 @@ impl CodexCredentialSelector { session_affinity_observation: request.session_affinity, guardian: false, }; - self.select_inner(&input, None, request.upstream_model) - .await + self.select_inner(&input, None, request.model).await } async fn select_inner( &self, request: &CredentialSelectionInput<'_>, cyber_policy_session_key: Option<&ProviderSessionAffinityKey>, - upstream_model: Option<&str>, + model: CodexSelectionModel<'_>, ) -> Result { let follow_only = !request.attempt.is_diagnostic_required_account() && request @@ -423,7 +468,7 @@ impl CodexCredentialSelector { let result = tokio::select! { biased; () = request.attempt.cancellation().cancelled() => Err(CredentialSelectionError::Cancelled), - result = self.select_inner_loop(request, cyber_policy_session_key, upstream_model, follow_only) => result, + result = self.select_inner_loop(request, cyber_policy_session_key, model, follow_only) => result, }; result.map_err(|error| { if follow_only @@ -446,7 +491,7 @@ impl CodexCredentialSelector { &self, request: &CredentialSelectionInput<'_>, cyber_policy_session_key: Option<&ProviderSessionAffinityKey>, - upstream_model: Option<&str>, + model: CodexSelectionModel<'_>, follow_only: bool, ) -> Result { let diagnostic = request.attempt.is_diagnostic_required_account(); @@ -561,17 +606,17 @@ impl CodexCredentialSelector { .is_some_and(|scope| scope.allows(account.id()))) && (!request.requires_oauth || account.authentication_kind() == CODEX_AUTHENTICATION_KIND_OAUTH) - && (diagnostic - || upstream_model.is_none_or(|upstream_model| { - let allowed = - request.attempt.account_scope().is_some_and(|scope| { - scope.allows_model(account.id(), upstream_model) - }); - if !allowed { - model_access_rejected += 1; - } - allowed - })) + && model.allows_account(account) + && (diagnostic || { + let allowed = request + .attempt + .account_scope() + .is_some_and(|scope| model.allows(scope, account.id())); + if !allowed { + model_access_rejected += 1; + } + allowed + }) }) .collect::>(); let mut eligible = Vec::with_capacity(accounts.len()); @@ -742,7 +787,7 @@ impl CodexCredentialSelector { } let selection = match request .attempt - .select_account(&self.provider_kind, upstream_model, &candidates, &context) + .select_account(&self.provider_kind, model.model(), &candidates, &context) .await { Ok(selection) => selection, diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index cc2807bf3..42414604a 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -3,6 +3,7 @@ use gateway_core::metering::{CalculatedCost, Usage}; use super::*; +use crate::credential::CodexSelectionModel; use gateway_core::operation::RawJsonPayload; impl CodexProvider { @@ -150,6 +151,15 @@ impl CodexProvider { context: AttemptContext, request: RawJsonEndpointRequest, ) -> Result { + let requested_model = match &request.operation { + Operation::GenerateImage(image) => endpoint_requested_model(image.payload()), + _ => None, + }; + let model = if matches!(&request.operation, Operation::GenerateImage(_)) { + CodexSelectionModel::Image(requested_model.as_ref().map(|model| model.as_str())) + } else { + CodexSelectionModel::NotApplicable + }; let selection_started_at = Instant::now(); let lease = self .selector @@ -157,8 +167,7 @@ impl CodexProvider { request_url: &request.response_origin, attempt: &context, session_affinity: request.session_affinity.as_ref(), - // 图像/搜索端点没有对应的上游模型权限，也不限定认证类型。 - upstream_model: None, + model, requires_oauth: false, }) .await @@ -175,7 +184,7 @@ impl CodexProvider { .execute_middleware( operation, provider_kind, - None, + requested_model.map(|model| model.as_str().to_owned()), account_id, Box::new(move |operation, middleware_headers| { Box::pin(async move { @@ -206,6 +215,23 @@ impl CodexProvider { ) -> Result { let affinity = match &operation { Operation::GenerateImage(image) => { + let model = endpoint_requested_model(image.payload()); + let selection = + CodexSelectionModel::Image(model.as_ref().map(|model| model.as_str())); + // 中间件可改写正文，最终图片模型必须仍被当前租约账号允许 + // 校验先于亲和变更和发送，不能在持有租约时重新选号 + if !selection.allows_account(lease.account()) + || (!context.is_diagnostic_required_account() + && !context + .account_scope() + .is_some_and(|scope| selection.allows(scope, lease.account_id()))) + { + return Err(provider_error( + ProviderErrorKind::NoEligibleAccount, + UpstreamSendState::NotSent, + ) + .with_retry_prohibited()); + } self.image_session_affinity(image.payload(), &middleware_headers, &context) .await? } diff --git a/backend/crates/providers/openai/src/provider/live.rs b/backend/crates/providers/openai/src/provider/live.rs index f8b2b57ea..b04f40710 100644 --- a/backend/crates/providers/openai/src/provider/live.rs +++ b/backend/crates/providers/openai/src/provider/live.rs @@ -32,7 +32,7 @@ use uuid::Uuid; use super::*; use crate::credential::{ - CODEX_AUTHENTICATION_KIND_OAUTH, CodexCredentialRepository, + CODEX_AUTHENTICATION_KIND_OAUTH, CodexCredentialRepository, CodexSelectionModel, SelectCodexProviderEndpointCredential, }; use crate::transport::headers::websocket_header_pairs; @@ -252,7 +252,7 @@ impl CodexProvider { request_url: &self.live_calls_url, attempt: &context, session_affinity: session_affinity.as_ref(), - upstream_model: Some(upstream_model.as_str()), + model: CodexSelectionModel::Requested(upstream_model.as_str()), // realtime calls 端点绑定 ChatGPT OAuth 身份；在候选阶段就排除 // API Key 账号，避免混合账号池选中不支持语音的账号后必然失败。 requires_oauth: true, diff --git a/backend/crates/providers/openai/src/provider/observation.rs b/backend/crates/providers/openai/src/provider/observation.rs index bdd400179..3cd16040b 100644 --- a/backend/crates/providers/openai/src/provider/observation.rs +++ b/backend/crates/providers/openai/src/provider/observation.rs @@ -8,7 +8,7 @@ pub(super) fn endpoint_requested_model( if payload.protocol() != PROVIDER_NAME { return None; } - // 只读取 model，避免把编辑请求中的整张图片复制到观测数据 + // 选号与观测共用模型提取，跳过编辑请求中的图片且不改写原始正文 #[derive(Deserialize)] struct RequestModel { model: Option, diff --git a/backend/crates/providers/openai/tests/provider/contract/image_account_eligibility.rs b/backend/crates/providers/openai/tests/provider/contract/image_account_eligibility.rs new file mode 100644 index 000000000..db9ab3221 --- /dev/null +++ b/backend/crates/providers/openai/tests/provider/contract/image_account_eligibility.rs @@ -0,0 +1,317 @@ +//! Images 的 OAuth 套餐资格、混合池、亲和和失败换号回归 + +use super::*; +use gateway_core::account::{AccountAffinity, ProviderAccountUpdate}; + +const PLUS: &str = "acct_provider_contract"; +const FREE: [&str; 3] = ["acct_scope_new", "acct_scope_old", "acct_scope_same"]; + +async fn oauth_account(store: &MemoryAccountStore, id: &str, plan: Option<&str>) { + let mut verified_account = profile(&format!("chatgpt-{id}")); + verified_account.plan_type = plan.map(str::to_owned); + store + .seed_oauth_credential(ImportCodexOAuthCredential { + account_id: id.to_owned(), + name: id.to_owned(), + secret: secret(&format!("at-{id}")), + verified_account, + next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await; +} + +fn image(kind: ImageRequestKind, body: Value) -> Operation { + Operation::GenerateImage(ImageRequest::from_raw_json( + kind, + RawJsonPayload::new("openai", Bytes::from(serde_json::to_vec(&body).unwrap())).unwrap(), + )) +} + +fn selection_attempt( + accounts: AccountAttemptContext, + affinity: AccountAffinity, + index: u32, +) -> AttemptContext { + AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_image_eligibility").unwrap(), + ClientApiKeyId::new("key_openai_contract").unwrap(), + ), + NonZeroU32::new(index).unwrap(), + SystemTime::now() + Duration::from_secs(30), + account_policy().with_openai_account_affinity(affinity), + accounts.with_account_scope(contract_account_scope()), + None, + CancellationToken::new(), + ) +} + +#[tokio::test] +async fn images_use_plus_in_a_pool_with_three_higher_weight_free_accounts() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(header("chatgpt-account-id", format!("chatgpt-{PLUS}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[]}))) + .expect(2) + .mount(&server) + .await; + let store = Arc::new(MemoryAccountStore::default()); + oauth_account(&store, PLUS, Some("plus")).await; + for id in FREE { + oauth_account(&store, id, Some("free")).await; + store.set_scheduling(id, None, AccountWeight::new(100).unwrap()); + } + let leases = Arc::new(TestLeaseCoordinator::default()); + let provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::default(), + server.uri(), + leases.clone(), + ); + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + let mut stream = provider + .clone() + .execute( + planned_provider_endpoint_request( + "openai", + image(kind, json!({"model":"gpt-image-2","prompt":"test"})), + ), + context("req_plus_images", CancellationToken::new()), + ) + .await + .unwrap(); + assert_eq!(stream.metadata().provider_account_id().as_str(), PLUS); + while let Some(event) = stream.next().await { + event.unwrap(); + } + } + assert!( + leases + .requests + .lock() + .unwrap() + .iter() + .all(|request| request.account_id().as_str() == PLUS) + ); +} + +#[tokio::test] +async fn images_never_fall_back_to_free_when_plus_is_excluded_or_free_is_pinned() { + let server = MockServer::start().await; + let store = Arc::new(MemoryAccountStore::default()); + oauth_account(&store, PLUS, Some("plus")).await; + oauth_account(&store, FREE[0], Some("free")).await; + let leases = Arc::new(TestLeaseCoordinator::default()); + let provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::default(), + server.uri(), + leases.clone(), + ); + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + for body in [json!({"model":"gpt-image-2"}), json!({"prompt":"test"})] { + for (required, excluded) in [ + ( + Some(ProviderAccountId::new(FREE[0]).unwrap()), + BTreeSet::new(), + ), + ( + None, + BTreeSet::from([ProviderAccountId::new(PLUS).unwrap()]), + ), + ] { + for index in [1, 2] { + let error = provider + .clone() + .execute( + planned_provider_endpoint_request("openai", image(kind, body.clone())), + selection_attempt( + AccountAttemptContext::new( + excluded.clone(), + required.clone(), + None, + ), + AccountAffinity::Strict, + index, + ), + ) + .await + .err() + .expect("Free must remain ineligible for Images"); + assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + } + } + } + } + assert!(leases.requests.lock().unwrap().is_empty()); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn images_escape_free_session_bindings_while_search_remains_available() { + for mode in [ + AccountAffinity::Strict, + AccountAffinity::Preferred, + AccountAffinity::Relaxed, + ] { + let server = MockServer::start().await; + let store = Arc::new(MemoryAccountStore::default()); + oauth_account(&store, PLUS, Some("plus")).await; + oauth_account(&store, FREE[0], Some("free")).await; + store.set_scheduling(FREE[0], None, AccountWeight::new(100).unwrap()); + let affinity = Arc::new(MemorySessionAffinity::default()); + let provider = provider_with_affinity_and_base_url(&store, affinity.clone(), server.uri()); + let search = Operation::Search(StandaloneSearchRequest::from_raw_json( + RawJsonPayload::new( + "openai", + Bytes::from_static(br#"{"id":"image-plan-session","commands":{}}"#), + ) + .unwrap(), + )); + let stream = provider + .clone() + .execute( + planned_provider_endpoint_request("openai", search), + selection_attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None), + mode, + 1, + ), + ) + .await + .unwrap(); + assert_eq!(stream.metadata().provider_account_id().as_str(), FREE[0]); + drop(stream); + assert_eq!(affinity.binding_count(), 1); + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + let key = affinity.lookup_keys()[0].clone(); + affinity.seed_binding( + &ProviderKind::new("openai").unwrap(), + &key, + ProviderAccountId::new(FREE[0]).unwrap(), + ); + let stream = provider + .clone() + .execute( + planned_provider_endpoint_request( + "openai", + image( + kind, + json!({ + "model":"gpt-image-2", "session_id":"image-plan-session" + }), + ), + ), + selection_attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None), + mode, + 1, + ), + ) + .await + .unwrap(); + assert_eq!(stream.metadata().provider_account_id().as_str(), PLUS); + drop(stream); + assert_eq!(affinity.binding_count(), 1); + } + assert!(server.received_requests().await.unwrap().is_empty()); + } +} + +#[tokio::test] +async fn image_plan_eligibility_tracks_account_updates_without_affecting_text() { + let server = MockServer::start().await; + let store = Arc::new(MemoryAccountStore::default()); + oauth_account(&store, PLUS, Some("free")).await; + let provider = provider_with_base_url(&store, server.uri()); + drop( + provider + .clone() + .execute( + planned_request("openai", generate_operation()), + context("req_free_text", CancellationToken::new()), + ) + .await + .expect("Free can still use Responses"), + ); + for plan in [ + Some("free"), + Some("plus"), + Some("free"), + None, + Some("future-plan"), + ] { + store + .update_account(ProviderAccountUpdate { + account_id: ProviderAccountId::new(PLUS).unwrap(), + name: PLUS.to_owned(), + email: None, + plan_type: plan.map(str::to_owned), + }) + .await + .unwrap(); + let result = provider + .clone() + .execute( + planned_provider_endpoint_request( + "openai", + image(ImageRequestKind::Generation, json!({"model":"gpt-image-2"})), + ), + context("req_plan_change", CancellationToken::new()), + ) + .await; + if plan == Some("free") { + let error = result.err().expect("Free must be excluded"); + assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + } else { + drop(result.expect("unknown plans are not inferred to be Free")); + } + } +} + +#[tokio::test] +async fn api_key_image_accounts_do_not_inherit_chatgpt_free_plan_restrictions() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[]}))) + .expect(2) + .mount(&server) + .await; + let store = Arc::new(MemoryAccountStore::default()); + store + .seed_api_key( + PLUS, + server.uri(), + provider_openai::credential::ResponsesTransport::Http, + ) + .await; + store + .update_account(ProviderAccountUpdate { + account_id: ProviderAccountId::new(PLUS).unwrap(), + name: PLUS.to_owned(), + email: None, + plan_type: Some("free".to_owned()), + }) + .await + .unwrap(); + let provider = provider_with_base_url(&store, server.uri()); + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + let mut stream = provider + .clone() + .execute( + planned_provider_endpoint_request( + "openai", + image(kind, json!({"model":"gpt-image-2"})), + ), + context("req_api_key_images", CancellationToken::new()), + ) + .await + .unwrap(); + while let Some(event) = stream.next().await { + event.unwrap(); + } + } +} diff --git a/backend/crates/providers/openai/tests/provider/contract/image_model_access.rs b/backend/crates/providers/openai/tests/provider/contract/image_model_access.rs new file mode 100644 index 000000000..6441b0abe --- /dev/null +++ b/backend/crates/providers/openai/tests/provider/contract/image_model_access.rs @@ -0,0 +1,461 @@ +//! Images 混合账号池的模型权限、亲和、重试与中间件改写回归 + +use super::*; +use gateway_core::account::{AccountAffinity, AccountModelAccess, AccountModelAccessMode}; + +const CPA_ACCOUNT: &str = "acct_provider_contract"; +const IMAGE_ACCOUNT: &str = "acct_scope_new"; +const IMAGE_BODY: &[u8] = + br#"{ "model":"gpt-image-2", "prompt":"first", "prompt":"last", "future":9007199254740993 }"#; + +fn policy(mode: AccountModelAccessMode, model: &str) -> AccountModelAccess { + AccountModelAccess::new(mode, vec![model.to_owned()]).unwrap() +} + +fn scope(first: AccountModelAccess, second: AccountModelAccess) -> Arc { + Arc::new(FrozenAccountScope::new( + Arc::new(RuntimeAccountDirectory::new( + [(CPA_ACCOUNT, first), (IMAGE_ACCOUNT, second)] + .into_iter() + .map(|(id, policy)| { + ( + ProviderAccountId::new(id).unwrap(), + RuntimeAccount::new(ProviderKind::new("openai").unwrap(), BTreeSet::new()) + .with_model_access(policy), + ) + }) + .collect(), + )), + ClientRoutingScope::all_accounts(), + )) +} + +fn attempt(accounts: AccountAttemptContext) -> AttemptContext { + configured_attempt(accounts, AccountAffinity::Strict, 1, None) +} + +fn configured_attempt( + accounts: AccountAttemptContext, + affinity: AccountAffinity, + index: u32, + middleware: Option>, +) -> AttemptContext { + let middleware = middleware.map(|plan| { + FrozenMiddlewarePlan::new( + plan, + ExtensionSetReference::new( + ExtensionSetId::new("image-access".to_owned()).unwrap(), + Arc::new(TestExtensionLease), + ), + ) + }); + AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_image_access").unwrap(), + ClientApiKeyId::new("key_openai_contract").unwrap(), + ) + .with_middleware( + middleware, + Arc::from([]), + "/v1/images/generations".to_owned(), + ClientTransport::HttpJson, + ), + NonZeroU32::new(index).unwrap(), + SystemTime::now() + Duration::from_secs(30), + account_policy().with_openai_account_affinity(affinity), + accounts, + None, + CancellationToken::new(), + ) +} + +#[tokio::test] +async fn images_do_not_use_forbidden_accounts_when_pinned_or_after_failover() { + let server = MockServer::start().await; + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, CPA_ACCOUNT).await; + create_account(&store, IMAGE_ACCOUNT).await; + let leases = Arc::new(TestLeaseCoordinator::default()); + let provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::default(), + server.uri(), + leases.clone(), + ); + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + for (required, excluded) in [ + ( + Some(ProviderAccountId::new(CPA_ACCOUNT).unwrap()), + BTreeSet::new(), + ), + ( + None, + BTreeSet::from([ProviderAccountId::new(IMAGE_ACCOUNT).unwrap()]), + ), + ] { + for index in [1, 2] { + let scope = scope( + policy(AccountModelAccessMode::Allowlist, "kimi-k3"), + policy(AccountModelAccessMode::Allowlist, "gpt-image-2"), + ); + let error = provider + .clone() + .execute( + planned_provider_endpoint_request( + "openai", + image_operation(kind, IMAGE_BODY), + ), + configured_attempt( + AccountAttemptContext::new(excluded.clone(), required.clone(), None) + .with_account_scope(scope), + AccountAffinity::Strict, + index, + None, + ), + ) + .await + .err() + .expect("forbidden account must not be used"); + assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + } + } + } + assert!(leases.requests.lock().unwrap().is_empty()); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn unresolved_image_models_only_use_unrestricted_accounts_without_rewriting_bodies() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[]}))) + .mount(&server) + .await; + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, CPA_ACCOUNT).await; + create_account(&store, IMAGE_ACCOUNT).await; + store.set_scheduling(CPA_ACCOUNT, None, AccountWeight::new(100).unwrap()); + let provider = provider_with_base_url(&store, server.uri()); + let bodies: &[&[u8]] = &[ + br#"{"prompt":"default model"}"#, + br#"{"model":null}"#, + br#"{"model":{ "future":"unknown" }}"#, + br#"{"model":""}"#, + br#"{"model":"gpt-image-2","model":"forbidden"}"#, + br#"{"model":"forbidden","model":"gpt-image-2"}"#, + br#"{"model":"gpt-image-2", broken"#, + b"--image-boundary\r\nopaque multipart body\r\n--image-boundary--\r\n", + ]; + for mode in [ + AccountModelAccessMode::Allowlist, + AccountModelAccessMode::Denylist, + ] { + for body in bodies { + for unrestricted in [true, false] { + let scope = scope( + policy(mode, "gpt-image-2"), + if unrestricted { + AccountModelAccess::all() + } else { + policy(mode, "gpt-image-2") + }, + ); + let result = provider + .clone() + .execute( + planned_provider_endpoint_request( + "openai", + image_operation(ImageRequestKind::Edit, body), + ), + attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None) + .with_account_scope(scope), + ), + ) + .await; + if unrestricted { + let mut stream = result.expect("unrestricted account accepts opaque payloads"); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + IMAGE_ACCOUNT + ); + while let Some(event) = stream.next().await { + event.unwrap(); + } + assert_eq!( + server + .received_requests() + .await + .unwrap() + .last() + .unwrap() + .body, + *body + ); + } else { + let error = result + .err() + .expect("unknown model cannot bypass either model policy"); + assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + } + } + } + } + assert_eq!( + server.received_requests().await.unwrap().len(), + bodies.len() * 2 + ); +} + +#[tokio::test] +async fn image_model_access_overrides_existing_session_affinity() { + for mode in [ + AccountAffinity::Strict, + AccountAffinity::Preferred, + AccountAffinity::Relaxed, + ] { + let server = MockServer::start().await; + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, CPA_ACCOUNT).await; + create_account(&store, IMAGE_ACCOUNT).await; + store.set_scheduling(CPA_ACCOUNT, None, AccountWeight::new(100).unwrap()); + let affinity = Arc::new(MemorySessionAffinity::default()); + let provider = provider_with_affinity_and_base_url(&store, affinity.clone(), server.uri()); + let operation = image_operation( + ImageRequestKind::Generation, + br#"{"model":"gpt-image-2","session_id":"image-model-session","prompt":"test"}"#, + ); + let stream = provider + .clone() + .execute( + planned_provider_endpoint_request("openai", operation.clone()), + configured_attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None).with_account_scope( + scope(AccountModelAccess::all(), AccountModelAccess::all()), + ), + mode, + 1, + None, + ), + ) + .await + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + CPA_ACCOUNT + ); + drop(stream); + let stream = provider + .clone() + .execute( + planned_provider_endpoint_request("openai", operation), + configured_attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None).with_account_scope( + scope( + policy(AccountModelAccessMode::Allowlist, "kimi-k3"), + policy(AccountModelAccessMode::Allowlist, "gpt-image-2"), + ), + ), + mode, + 1, + None, + ), + ) + .await + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + IMAGE_ACCOUNT + ); + assert!(server.received_requests().await.unwrap().is_empty()); + } +} + +#[tokio::test] +async fn final_image_model_after_middleware_must_be_allowed_by_the_selected_account() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[]}))) + .mount(&server) + .await; + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, CPA_ACCOUNT).await; + let provider = provider_with_base_url(&store, server.uri()); + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + for (replacement, allowed) in [ + (json!("forbidden"), false), + (Value::Null, false), + (json!("gpt-image-future"), true), + ] { + let scope = scope( + AccountModelAccess::new( + AccountModelAccessMode::Allowlist, + vec!["gpt-image-2".to_owned(), "gpt-image-future".to_owned()], + ) + .unwrap(), + AccountModelAccess::all(), + ); + let middleware = Arc::new(RecordingMiddleware { + observed: Arc::default(), + replacement: ("model".to_owned(), replacement.clone()), + request_headers: Vec::new(), + }); + let result = provider + .clone() + .execute( + planned_provider_endpoint_request("openai", image_operation(kind, IMAGE_BODY)), + configured_attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None) + .with_account_scope(scope), + AccountAffinity::Strict, + 1, + Some(middleware), + ), + ) + .await; + if allowed { + let mut stream = result.unwrap(); + while let Some(event) = stream.next().await { + event.unwrap(); + } + assert_eq!( + captured_request_body( + server.received_requests().await.unwrap().last().unwrap() + )["model"], + replacement + ); + } else { + let error = result + .err() + .expect("rewriting cannot escape account policy"); + assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + assert!(error.retry_is_prohibited()); + } + } + } + assert_eq!(server.received_requests().await.unwrap().len(), 2); +} + +#[tokio::test] +async fn standalone_search_does_not_inherit_image_model_restrictions() { + let server = MockServer::start().await; + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, CPA_ACCOUNT).await; + let search = Operation::Search(StandaloneSearchRequest::from_raw_json( + RawJsonPayload::new( + "openai", + Bytes::from_static(br#"{"id":"search","commands":{}}"#), + ) + .unwrap(), + )); + let stream = provider_with_base_url(&store, server.uri()) + .execute( + planned_provider_endpoint_request("openai", search), + attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None).with_account_scope(scope( + policy(AccountModelAccessMode::Allowlist, "kimi-k3"), + AccountModelAccess::all(), + )), + ), + ) + .await + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + CPA_ACCOUNT + ); +} + +fn image_operation(kind: ImageRequestKind, body: &[u8]) -> Operation { + Operation::GenerateImage(ImageRequest::from_raw_json( + kind, + RawJsonPayload::new("openai", Bytes::copy_from_slice(body)).unwrap(), + )) +} + +#[tokio::test] +async fn images_filter_mixed_accounts_by_model_before_weight_and_preserve_wire() { + let cpa = MockServer::start().await; + let oauth = MockServer::start().await; + let store = Arc::new(MemoryAccountStore::default()); + store + .seed_api_key( + CPA_ACCOUNT, + cpa.uri(), + provider_openai::credential::ResponsesTransport::Http, + ) + .await; + store.set_scheduling(CPA_ACCOUNT, None, AccountWeight::new(100).unwrap()); + create_account(&store, IMAGE_ACCOUNT).await; + Mock::given(method("POST")) + .and(body_bytes(IMAGE_BODY.to_vec())) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data":[{"b64_json":"AAEC"}]})), + ) + .expect(4) + .mount(&oauth) + .await; + let leases = Arc::new(TestLeaseCoordinator::default()); + let provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::default(), + oauth.uri(), + leases.clone(), + ); + for cpa_policy in [ + policy(AccountModelAccessMode::Allowlist, "kimi-k3"), + policy(AccountModelAccessMode::Denylist, "gpt-image-2"), + ] { + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + let scope = scope( + cpa_policy.clone(), + policy(AccountModelAccessMode::Allowlist, "gpt-image-2"), + ); + let mut stream = provider + .clone() + .execute( + planned_provider_endpoint_request("openai", image_operation(kind, IMAGE_BODY)), + attempt( + AccountAttemptContext::new(BTreeSet::new(), None, None) + .with_account_scope(scope), + ), + ) + .await + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + IMAGE_ACCOUNT + ); + while let Some(event) = stream.next().await { + event.unwrap(); + } + } + } + assert!(cpa.received_requests().await.unwrap().is_empty()); + assert!( + leases + .requests + .lock() + .unwrap() + .iter() + .all(|request| request.account_id().as_str() == IMAGE_ACCOUNT) + ); + assert_eq!( + oauth + .received_requests() + .await + .unwrap() + .iter() + .map(|request| request.url.path().to_owned()) + .collect::>(), + [ + "/codex/images/generations", + "/codex/images/edits", + "/codex/images/generations", + "/codex/images/edits" + ] + ); +} diff --git a/backend/crates/providers/openai/tests/provider/contract/mod.rs b/backend/crates/providers/openai/tests/provider/contract/mod.rs index 5ab0c07b7..5c18742be 100644 --- a/backend/crates/providers/openai/tests/provider/contract/mod.rs +++ b/backend/crates/providers/openai/tests/provider/contract/mod.rs @@ -3,6 +3,8 @@ mod account_isolation; mod capacity; mod error_details; +mod image_account_eligibility; +mod image_model_access; mod passthrough; mod precommit; mod response_interrupt; diff --git a/docs/api.md b/docs/api.md index 313dd40ed..8eecbf6c9 100644 --- a/docs/api.md +++ b/docs/api.md @@ -408,9 +408,13 @@ OpenAI 上游错误 envelope 和允许下发的 opaque header 值也不由 canon `response.metadata` 的业务内容继续交付，仅其中的响应头沿用凭据、账号身份和逐跳头隔离规则。 未改写的 WebSocket 文本不依赖 JSON 旁路解析成功；插件改写后的响应仍复核必要关联，原始 SSE 心跳 和无法提取事实的未改写内容不构成协议失败。 -Images 请求不读取或重建 JSON，也不要求或映射模型字段； -它固定使用 OpenAI Provider， -只在原始字节之外完成账号选择、鉴权头替换和端点路由，成功与非容量失败响应正文保持原始字节。 +Images 固定使用 OpenAI Provider，不重建正文或映射模型字段；Provider 从原始 JSON 旁路提取图片模型， +按[账号模型限制](#账号模型限制)选号，并复核插件改写后的模型。账号选择、鉴权头替换和端点路由 +不改变请求正文，成功与非容量失败响应正文保持原始字节。 + +Images 排除已知为 Free 套餐的 OAuth 账号；套餐缺失或未知时不推断为 Free，API Key 账号不套用该套餐规则。 +其余账号仍需满足模型限制、额度等调度条件，实际生图权限由上游决定 + `/v1/alpha/search` 使用相同的 OpenAI Provider 原生端点边界：body（包括 `model`）不解析、不映射， `x-codex-turn-metadata` 在移除客户端账号身份并按当前 lease 重写 installation ID 后转发；上游账号 Authorization、Cookie、account ID、originator 和 User-Agent 均由代理安全重建。 @@ -717,9 +721,11 @@ Pro 使用 `all` 时也能参与 luna 调度。套餐名称不自动生成或修 批量接口的调度、分组、模型与代理字段均可省略；省略的字段保持各账号原值。 提供 `groupIds` 时替换完整分组集合，提供 `concurrencyLimit: null` 时恢复继承运行参数 -限制适用于 Responses HTTP、WebSocket 及其带压缩触发的请求选号，包括重试、亲和和换号；没有合规账号时沿用 +限制适用于 Responses HTTP、WebSocket 及其带压缩触发的请求，以及 Images 生图与编辑选号，包括重试、亲和和换号；没有合规账号时沿用 无可用账号错误，不会回退到被禁止的账号。已开始请求使用冻结的政策，新请求使用已发布的新配置。 -Images、独立 Search 及管理员连接测试不受该文本模型限制；连接测试成功只证明指定账号的上游能力。 +Images 按请求正文的图片模型 ID 检查，不依赖文本模型目录；可手动添加目录未列出的图片模型 ID。 +图片模型缺失、重复或无法识别时，仅使用模型政策为 `all` 的账号，不推测上游默认模型。 +独立 Search 及管理员连接测试不受该模型限制；连接测试成功只证明指定账号的上游能力。 普通 `/v1/models` 和单模型查询将已发现的 OpenAI 模型关联到来源账号，至少一个来源账号在当前 Key 范围内且政策允许时才展示； 同 Provider 中未发现该模型的 `all` 账号不会使它进入列表。原生目录保留已有来源选择和完整模型对象 diff --git a/docs/architecture.md b/docs/architecture.md index db1e0a4c1..b1c1b6372 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -384,8 +384,9 @@ Client Key 鉴权完成后，API adapter 从有界请求头识别客户端，Cor Responses 按模型目录编译候选；全局模型映射是精确映射，未命中时模型名原样交给候选 Provider。 Images 与 standalone Search 是 OpenAI Provider 自有端点：两者都不参与文本模型映射，只在 Client Key -的账号范围确实包含 OpenAI 账号时生成单一 OpenAI 候选。Images 不要求模型字段；Search body 中的模型 -及其他字段保持原始 bytes 并由上游解释 +的账号范围确实包含 OpenAI 账号时生成单一 OpenAI 候选。Images 的图片模型由 Provider 从正文旁路 +提取，用于账号模型政策检查，并在中间件改写后复验当前租约账号；不要求图片模型进入文本目录。 +Search body 中的模型及其他字段保持原始 bytes 并由上游解释 ## 5. Provider 与协议边界 @@ -487,6 +488,9 @@ client，OIDC 的 JWKS 缓存与单飞归属对应出口状态。自动刷新提 客户端据此开放原生生图工具，图片生成和编辑由 Images 路由处理。 这只解决客户端能力识别；模型能力、客户端限制和上游账号的实际生图权限仍分别检查 +OpenAI Provider 将 OAuth 套餐资格作为 Images 选号条件，复用账号持久事实， +在获取租约前与冻结的模型黑白名单共同过滤候选；亲和与重试沿用同一资格检查 + `X-OpenAI-Actor-Authorization` 是客户端能力标记，不是账号凭据。API 解码和 OpenAI Provider 都过滤该 header，网关继续校验 Client Key，上游认证由服务端账号产生。 客户端配置示例维护在 [部署文档](../deploy/README.md#客户端配置)