Support per-format provider key auth

This commit is contained in:
fawney19
2026-04-29 15:46:50 +08:00
parent 07a319259b
commit e751289dfb
67 changed files with 1244 additions and 420 deletions

View File

@@ -212,6 +212,8 @@ pub struct AdminSystemConfigProviderKey {
#[serde(default)]
pub global_priority_by_format: Option<Value>,
#[serde(default)]
pub auth_type_by_format: Option<Value>,
#[serde(default)]
pub rpm_limit: Option<u32>,
#[serde(default)]
pub allowed_models: Option<Vec<String>>,

View File

@@ -555,6 +555,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
@@ -621,6 +623,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["openai:responses".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
@@ -694,6 +698,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["openai:responses".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
@@ -754,6 +760,8 @@ mod tests {
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
@@ -817,6 +825,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:messages".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -248,7 +248,8 @@ pub struct StoredProviderCatalogKey {
pub capabilities: Option<serde_json::Value>,
pub is_active: bool,
pub api_formats: Option<serde_json::Value>,
pub encrypted_api_key: String,
pub auth_type_by_format: Option<serde_json::Value>,
pub encrypted_api_key: Option<String>,
pub encrypted_auth_config: Option<String>,
pub note: Option<String>,
pub internal_priority: i32,
@@ -321,7 +322,8 @@ impl StoredProviderCatalogKey {
capabilities,
is_active,
api_formats: None,
encrypted_api_key: String::new(),
auth_type_by_format: None,
encrypted_api_key: None,
encrypted_auth_config: None,
note: None,
internal_priority: 50,
@@ -371,7 +373,7 @@ impl StoredProviderCatalogKey {
pub fn with_transport_fields(
mut self,
api_formats: Option<serde_json::Value>,
encrypted_api_key: String,
encrypted_api_key: impl Into<Option<String>>,
encrypted_auth_config: Option<String>,
rate_multipliers: Option<serde_json::Value>,
global_priority_by_format: Option<serde_json::Value>,
@@ -380,7 +382,11 @@ impl StoredProviderCatalogKey {
proxy: Option<serde_json::Value>,
fingerprint: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
if encrypted_api_key.trim().is_empty() {
let encrypted_api_key = encrypted_api_key.into();
if encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
{
return Err(crate::DataLayerError::UnexpectedValue(
"provider_api_keys.api_key is empty".to_string(),
));
@@ -629,3 +635,60 @@ pub trait ProviderCatalogWriteRepository: Send + Sync {
circuit_breaker_by_format: Option<&serde_json::Value>,
) -> Result<bool, crate::DataLayerError>;
}
#[cfg(test)]
mod tests {
use super::StoredProviderCatalogKey;
fn sample_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key".to_string(),
"service_account".to_string(),
None,
true,
)
.expect("key should build")
}
#[test]
fn transport_fields_allow_null_encrypted_api_key() {
let key = sample_key()
.with_transport_fields(
None,
None::<String>,
None,
None,
None,
None,
None,
None,
None,
)
.expect("null api key should be accepted");
assert_eq!(key.encrypted_api_key, None);
}
#[test]
fn transport_fields_reject_empty_encrypted_api_key_string() {
let err = sample_key()
.with_transport_fields(
None,
Some(" ".to_string()),
None,
None,
None,
None,
None,
None,
None,
)
.expect_err("empty api key string should be rejected");
assert!(err
.to_string()
.contains("provider_api_keys.api_key is empty"));
}
}

View File

@@ -465,7 +465,7 @@ CREATE TABLE IF NOT EXISTS public.payment_orders (
CREATE TABLE IF NOT EXISTS public.provider_api_keys (
id character varying(36) NOT NULL,
api_key text NOT NULL,
api_key text,
name character varying(100) NOT NULL,
note character varying(500),
internal_priority integer DEFAULT 50,
@@ -496,6 +496,7 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
updated_at timestamp with time zone DEFAULT now() NOT NULL,
provider_id character varying(36) NOT NULL,
api_formats json,
auth_type_by_format json,
rate_multipliers json,
health_by_format jsonb,
circuit_breaker_by_format jsonb,
@@ -1992,21 +1993,6 @@ END $mig$;
--
-- Name: provider_endpoints uq_provider_api_format; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.provider_endpoints
ADD CONSTRAINT uq_provider_api_format UNIQUE (provider_id, api_format);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: models uq_provider_model; Type: CONSTRAINT; Schema: public; Owner: -
--
@@ -2503,6 +2489,14 @@ CREATE INDEX IF NOT EXISTS idx_endpoint_format_active ON public.provider_endpoin
--
-- Name: idx_provider_endpoints_provider_api_format; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_endpoints_provider_api_format ON public.provider_endpoints USING btree (provider_id, api_format);
--
-- Name: idx_gemini_file_mappings_expires; Type: INDEX; Schema: public; Owner: -
--

View File

@@ -464,7 +464,7 @@ CREATE TABLE IF NOT EXISTS public.payment_orders (
CREATE TABLE IF NOT EXISTS public.provider_api_keys (
id character varying(36) NOT NULL,
api_key text NOT NULL,
api_key text,
name character varying(100) NOT NULL,
note character varying(500),
internal_priority integer DEFAULT 50,
@@ -495,6 +495,7 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
updated_at timestamp with time zone DEFAULT now() NOT NULL,
provider_id character varying(36) NOT NULL,
api_formats json,
auth_type_by_format json,
rate_multipliers json,
health_by_format jsonb,
circuit_breaker_by_format jsonb,

View File

@@ -14,89 +14,38 @@ AS $$
END
$$;
DO $$
DECLARE
conflict_summary text;
BEGIN
SELECT string_agg(
DISTINCT provider_id::text || ':' || canonical_api_format,
', ' ORDER BY provider_id::text || ':' || canonical_api_format
)
INTO conflict_summary
FROM (
SELECT
left_endpoint.provider_id,
public.aether_canonical_api_format_alias(left_endpoint.api_format) AS canonical_api_format
FROM public.provider_endpoints AS left_endpoint
INNER JOIN public.provider_endpoints AS right_endpoint
ON right_endpoint.provider_id = left_endpoint.provider_id
AND right_endpoint.id > left_endpoint.id
AND public.aether_canonical_api_format_alias(right_endpoint.api_format)
= public.aether_canonical_api_format_alias(left_endpoint.api_format)
WHERE left_endpoint.api_format IN ('openai:responses', 'openai:cli', 'openai:responses:compact', 'openai:compact', 'claude:messages', 'claude:chat', 'claude:cli', 'gemini:generate_content', 'gemini:chat', 'gemini:cli')
AND right_endpoint.api_format IN ('openai:responses', 'openai:cli', 'openai:responses:compact', 'openai:compact', 'claude:messages', 'claude:chat', 'claude:cli', 'gemini:generate_content', 'gemini:chat', 'gemini:cli')
AND (
left_endpoint.base_url IS DISTINCT FROM right_endpoint.base_url
OR left_endpoint.custom_path IS DISTINCT FROM right_endpoint.custom_path
OR left_endpoint.max_retries IS DISTINCT FROM right_endpoint.max_retries
OR left_endpoint.header_rules::jsonb IS DISTINCT FROM right_endpoint.header_rules::jsonb
OR left_endpoint.body_rules::jsonb IS DISTINCT FROM right_endpoint.body_rules::jsonb
OR left_endpoint.config::jsonb IS DISTINCT FROM right_endpoint.config::jsonb
OR left_endpoint.proxy IS DISTINCT FROM right_endpoint.proxy
OR left_endpoint.format_acceptance_config::jsonb IS DISTINCT FROM right_endpoint.format_acceptance_config::jsonb
)
) AS conflicts;
ALTER TABLE IF EXISTS public.provider_endpoints
DROP CONSTRAINT IF EXISTS uq_provider_api_format;
IF conflict_summary IS NOT NULL THEN
RAISE EXCEPTION
'Cannot normalize OpenAI/Claude/Gemini provider_endpoints because transport fields differ for: %',
conflict_summary;
END IF;
END $$;
CREATE INDEX IF NOT EXISTS idx_provider_endpoints_provider_api_format
ON public.provider_endpoints USING btree (provider_id, api_format);
WITH grouped AS (
ALTER TABLE IF EXISTS public.provider_api_keys
ADD COLUMN IF NOT EXISTS auth_type_by_format json;
ALTER TABLE IF EXISTS public.provider_api_keys
ALTER COLUMN api_key DROP NOT NULL;
WITH normalized AS (
SELECT
id,
provider_id,
api_format,
public.aether_canonical_api_format_alias(api_format) AS canonical_api_format,
ROW_NUMBER() OVER (
PARTITION BY provider_id, public.aether_canonical_api_format_alias(api_format)
ORDER BY
CASE
WHEN api_format = public.aether_canonical_api_format_alias(api_format) THEN 0
ELSE 1
END,
created_at ASC,
id ASC
) AS rank
public.aether_canonical_api_format_alias(api_format) AS canonical_api_format
FROM public.provider_endpoints
WHERE api_format IN ('openai:responses', 'openai:cli', 'openai:responses:compact', 'openai:compact', 'claude:messages', 'claude:chat', 'claude:cli', 'gemini:generate_content', 'gemini:chat', 'gemini:cli')
),
survivors AS (
SELECT *
FROM grouped
WHERE rank = 1
),
retired AS (
UPDATE public.provider_endpoints AS endpoint
SET
is_active = FALSE,
updated_at = NOW()
FROM grouped
WHERE endpoint.id = grouped.id
AND grouped.rank > 1
RETURNING endpoint.id
)
UPDATE public.provider_endpoints AS endpoint
SET
api_format = survivors.canonical_api_format,
api_family = SPLIT_PART(survivors.canonical_api_format, ':', 1),
endpoint_kind = SUBSTRING(survivors.canonical_api_format FROM POSITION(':' IN survivors.canonical_api_format) + 1),
api_format = normalized.canonical_api_format,
api_family = SPLIT_PART(normalized.canonical_api_format, ':', 1),
endpoint_kind = SUBSTRING(normalized.canonical_api_format FROM POSITION(':' IN normalized.canonical_api_format) + 1),
updated_at = NOW()
FROM survivors
WHERE endpoint.id = survivors.id
AND endpoint.api_format IS DISTINCT FROM survivors.canonical_api_format;
FROM normalized
WHERE endpoint.id = normalized.id
AND (
endpoint.api_format IS DISTINCT FROM normalized.canonical_api_format
OR endpoint.api_family IS DISTINCT FROM SPLIT_PART(normalized.canonical_api_format, ':', 1)
OR endpoint.endpoint_kind IS DISTINCT FROM SUBSTRING(normalized.canonical_api_format FROM POSITION(':' IN normalized.canonical_api_format) + 1)
);
WITH expanded AS (
SELECT

View File

@@ -735,6 +735,266 @@ SELECT EXISTS (
assert!(!BASELINE_V2_SQL.contains("api_formats json DEFAULT '[]'::json NOT NULL"));
}
#[test]
fn provider_api_keys_api_key_is_nullable() {
let baseline_migration = MIGRATOR
.iter()
.find(|migration| migration.version == 20260403000000)
.expect("baseline migration should be embedded");
let normalization_migration = MIGRATOR
.iter()
.find(|migration| migration.version == 20260428000000)
.expect("api format normalization migration should be embedded");
assert!(baseline_migration.sql.contains("api_key text,"));
assert!(!baseline_migration.sql.contains("api_key text NOT NULL"));
assert!(BASELINE_V2_SQL.contains("api_key text,"));
assert!(!BASELINE_V2_SQL.contains("api_key text NOT NULL"));
assert!(normalization_migration
.sql
.contains("ALTER COLUMN api_key DROP NOT NULL"));
}
#[test]
fn normalized_endpoint_formats_do_not_require_unique_provider_format_pairs() {
let normalization_migration = MIGRATOR
.iter()
.find(|migration| migration.version == 20260428000000)
.expect("api format normalization migration should be embedded");
assert!(normalization_migration
.sql
.contains("DROP CONSTRAINT IF EXISTS uq_provider_api_format"));
assert!(normalization_migration
.sql
.contains("idx_provider_endpoints_provider_api_format"));
assert!(!BASELINE_V2_SQL.contains("uq_provider_api_format"));
assert!(BASELINE_V2_SQL.contains("idx_provider_endpoints_provider_api_format"));
}
#[tokio::test]
async fn api_format_normalization_migration_preserves_duplicate_endpoint_transports() {
let Some(server) = ManagedPostgresServer::try_start()
.await
.expect("postgres migration test should start or skip")
else {
return;
};
let pool = PgPool::connect(server.database_url())
.await
.expect("pool should connect");
let normalization_migration = MIGRATOR
.iter()
.find(|migration| migration.version == 20260428000000)
.expect("api format normalization migration should be embedded");
sqlx::raw_sql(
r#"
CREATE TABLE public.provider_endpoints (
id text PRIMARY KEY,
provider_id text NOT NULL,
api_format text NOT NULL,
api_family text,
endpoint_kind text,
base_url text NOT NULL,
max_retries integer,
is_active boolean DEFAULT true NOT NULL,
custom_path text,
config json,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,
proxy jsonb,
header_rules json,
format_acceptance_config json,
body_rules json
);
ALTER TABLE ONLY public.provider_endpoints
ADD CONSTRAINT uq_provider_api_format UNIQUE (provider_id, api_format);
CREATE TABLE public.provider_api_keys (
id text PRIMARY KEY,
api_formats json,
updated_at timestamp with time zone DEFAULT now() NOT NULL,
rate_multipliers json,
global_priority_by_format json,
health_by_format jsonb,
circuit_breaker_by_format jsonb
);
CREATE TABLE public.api_keys (
id text PRIMARY KEY,
allowed_api_formats json,
updated_at timestamp with time zone DEFAULT now() NOT NULL
);
CREATE TABLE public.users (
id text PRIMARY KEY,
allowed_api_formats json,
updated_at timestamp with time zone DEFAULT now() NOT NULL
);
CREATE TABLE public.models (
id text PRIMARY KEY,
provider_model_mappings jsonb,
updated_at timestamp with time zone DEFAULT now() NOT NULL
);
INSERT INTO public.provider_endpoints (
id,
provider_id,
api_format,
api_family,
endpoint_kind,
base_url,
custom_path,
max_retries,
header_rules,
body_rules,
config,
proxy,
format_acceptance_config
) VALUES
(
'endpoint-claude-chat',
'provider-conflict',
'claude:chat',
'claude',
'chat',
'https://claude-chat.example',
'/v1/messages',
2,
'{"x-channel":"chat"}'::json,
'{"mode":"chat"}'::json,
'{"transport":"chat"}'::json,
'{"url":"http://proxy-chat"}'::jsonb,
'{"accept":"chat"}'::json
),
(
'endpoint-claude-cli',
'provider-conflict',
'claude:cli',
'claude',
'cli',
'https://claude-cli.example',
'/v1/messages',
3,
'{"x-channel":"cli"}'::json,
'{"mode":"cli"}'::json,
'{"transport":"cli"}'::json,
'{"url":"http://proxy-cli"}'::jsonb,
'{"accept":"cli"}'::json
);
INSERT INTO public.provider_api_keys (
id,
api_formats,
rate_multipliers,
global_priority_by_format,
health_by_format,
circuit_breaker_by_format
) VALUES (
'provider-key',
'["claude:chat","claude:cli","openai:cli","openai:responses"]'::json,
'{"claude:chat":1,"openai:compact":2}'::json,
'{"gemini:cli":3}'::json,
'{"openai:cli":{"health_score":0.9}}'::jsonb,
'{"openai:compact":{"open":false}}'::jsonb
);
INSERT INTO public.api_keys (id, allowed_api_formats)
VALUES ('api-key', '["gemini:chat","gemini:cli"]'::json);
INSERT INTO public.users (id, allowed_api_formats)
VALUES ('user', '["openai:compact","openai:responses:compact"]'::json);
INSERT INTO public.models (id, provider_model_mappings)
VALUES (
'model',
'[{"api_formats":["claude:chat","claude:cli","gemini:chat"]}]'::jsonb
);
"#,
)
.execute(&pool)
.await
.expect("fixture schema should be created");
sqlx::raw_sql(&normalization_migration.sql)
.execute(&pool)
.await
.expect("api format normalization migration should preserve duplicate endpoints");
let endpoint_rows = sqlx::query_as::<_, (String, String, Option<String>, Option<String>)>(
r#"
SELECT id, api_format, api_family, endpoint_kind
FROM public.provider_endpoints
WHERE provider_id = 'provider-conflict'
ORDER BY id
"#,
)
.fetch_all(&pool)
.await
.expect("endpoint rows should be readable");
assert_eq!(
endpoint_rows,
vec![
(
"endpoint-claude-chat".to_string(),
"claude:messages".to_string(),
Some("claude".to_string()),
Some("messages".to_string())
),
(
"endpoint-claude-cli".to_string(),
"claude:messages".to_string(),
Some("claude".to_string()),
Some("messages".to_string())
),
]
);
let base_urls = sqlx::query_as::<_, (String,)>(
r#"
SELECT base_url
FROM public.provider_endpoints
WHERE provider_id = 'provider-conflict'
ORDER BY id
"#,
)
.fetch_all(&pool)
.await
.expect("endpoint transport rows should be readable")
.into_iter()
.map(|(base_url,)| base_url)
.collect::<Vec<_>>();
assert_eq!(
base_urls,
vec![
"https://claude-chat.example".to_string(),
"https://claude-cli.example".to_string()
]
);
let provider_key_formats: serde_json::Value = query_scalar(
"SELECT api_formats::jsonb FROM public.provider_api_keys WHERE id = 'provider-key'",
)
.fetch_one(&pool)
.await
.expect("provider key formats should be readable");
assert_eq!(
provider_key_formats,
serde_json::json!(["claude:messages", "openai:responses"])
);
let provider_format_constraint_count: i64 = query_scalar(
"SELECT COUNT(*)::BIGINT FROM pg_constraint WHERE conname = 'uq_provider_api_format'",
)
.fetch_one(&pool)
.await
.expect("constraint count should be readable");
assert_eq!(provider_format_constraint_count, 0);
}
#[test]
fn deprecation_migration_and_baseline_mark_legacy_usage_columns() {
let settlement_migration = MIGRATOR

View File

@@ -537,7 +537,7 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
return Ok(false);
};
key.encrypted_api_key = encrypted_api_key.to_string();
key.encrypted_api_key = Some(encrypted_api_key.to_string());
key.encrypted_auth_config = encrypted_auth_config.map(ToOwned::to_owned);
key.expires_at_unix_secs = expires_at_unix_secs;
Ok(true)
@@ -705,7 +705,10 @@ mod tests {
.await
.expect("keys should read");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].encrypted_api_key, "ciphertext-updated-token");
assert_eq!(
stored[0].encrypted_api_key.as_deref(),
Some("ciphertext-updated-token")
);
assert_eq!(
stored[0].encrypted_auth_config.as_deref(),
Some("ciphertext-auth-2")

View File

@@ -140,6 +140,7 @@ SELECT
capabilities,
is_active,
api_formats,
auth_type_by_format,
api_key,
auth_config,
note,
@@ -196,6 +197,7 @@ SELECT
capabilities,
is_active,
api_formats,
auth_type_by_format,
api_key,
auth_config,
note,
@@ -252,6 +254,7 @@ SELECT
NULL::jsonb AS capabilities,
is_active,
api_formats,
NULL::jsonb AS auth_type_by_format,
'summary' AS api_key,
CASE
WHEN auth_config IS NULL THEN NULL
@@ -592,6 +595,7 @@ SELECT
capabilities,
is_active,
api_formats,
auth_type_by_format,
api_key,
auth_config,
note,
@@ -1183,6 +1187,7 @@ INSERT INTO provider_api_keys (
id,
provider_id,
api_formats,
auth_type_by_format,
auth_type,
api_key,
auth_config,
@@ -1255,55 +1260,56 @@ INSERT INTO provider_api_keys (
$22,
$23,
$24,
CASE
WHEN $25::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($25::double precision)
END,
$25,
CASE
WHEN $26::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($26::double precision)
END,
$27,
$28,
COALESCE($29, 0),
COALESCE($30, 0),
CASE
WHEN $31::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($31::double precision)
WHEN $27::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($27::double precision)
END,
$28,
$29,
COALESCE($30, 0),
COALESCE($31, 0),
CASE
WHEN $32::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($32::double precision)
END,
$32,
$33,
$34,
$35,
CASE
WHEN $35::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($35::double precision)
WHEN $36::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($36::double precision)
END,
$36,
COALESCE($37, 0),
$37,
COALESCE($38, 0),
COALESCE($39, 0),
COALESCE($40, 0),
COALESCE($41, 0),
COALESCE($42, 0),
CASE
WHEN $43::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($43::double precision)
END,
COALESCE($43, 0),
CASE
WHEN $44::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($44::double precision)
END,
$45,
CASE
WHEN $45::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($45::double precision)
END,
$46,
$47,
$48,
CASE
WHEN $49::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($49::double precision)
END,
$49,
CASE
WHEN $50::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($50::double precision)
END,
CASE
WHEN $51::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($51::double precision)
END
)
"#,
@@ -1311,6 +1317,7 @@ INSERT INTO provider_api_keys (
.bind(&key.id)
.bind(&key.provider_id)
.bind(&key.api_formats)
.bind(&key.auth_type_by_format)
.bind(&key.auth_type)
.bind(&key.encrypted_api_key)
.bind(&key.encrypted_auth_config)
@@ -1715,6 +1722,7 @@ UPDATE provider_api_keys
SET
provider_id = $2,
api_formats = $3,
auth_type_by_format = $39,
auth_type = $4,
api_key = $5,
auth_config = $6,
@@ -1809,6 +1817,7 @@ WHERE id = $1
.bind(key.is_active)
.bind(key.updated_at_unix_secs.map(|value| value as f64))
.bind(key.expires_at_unix_secs.map(|value| value as f64))
.bind(&key.auth_type_by_format)
.execute(&self.pool)
.await
.map_postgres_err()?
@@ -2437,7 +2446,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
)?
.with_transport_fields(
row_get(row, "api_formats")?,
row_get(row, "api_key")?,
row_get::<Option<String>>(row, "api_key")?,
row_get(row, "auth_config")?,
row_get(row, "rate_multipliers")?,
row_get(row, "global_priority_by_format")?,
@@ -2469,6 +2478,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
row.try_get("circuit_breaker_by_format").ok(),
);
key.note = row.try_get("note").ok();
key.auth_type_by_format = row.try_get("auth_type_by_format").ok();
key.internal_priority = row.try_get("internal_priority").unwrap_or(50);
key.cache_ttl_minutes = row.try_get("cache_ttl_minutes").unwrap_or(5);
key.max_probe_interval_minutes = row.try_get("max_probe_interval_minutes").unwrap_or(32);

View File

@@ -1172,6 +1172,8 @@ mod tests {
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -562,6 +562,8 @@ mod tests {
auth_type: auth_type.to_string(),
is_active: true,
api_formats: Some(vec![api_format.to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -230,7 +230,7 @@ fn merge_comma_header_values(left: Option<&str>, right: Option<&str>) -> Option<
pub fn resolve_local_openai_bearer_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let auth_type = resolve_local_auth_type_for_transport_format(transport);
if !matches!(auth_type.as_str(), "api_key" | "bearer") {
return None;
}
@@ -242,7 +242,7 @@ pub fn resolve_local_openai_bearer_auth(
pub fn resolve_local_standard_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let auth_type = resolve_local_auth_type_for_transport_format(transport);
let secret = resolved_local_secret(transport)?;
match auth_type.as_str() {
@@ -255,7 +255,7 @@ pub fn resolve_local_standard_auth(
pub fn resolve_local_gemini_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let auth_type = resolve_local_auth_type_for_transport_format(transport);
let secret = resolved_local_secret(transport)?;
match auth_type.as_str() {
@@ -265,6 +265,30 @@ pub fn resolve_local_gemini_auth(
}
}
pub(crate) fn resolve_local_auth_type_for_transport_format(
transport: &GatewayProviderTransportSnapshot,
) -> String {
let default_auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let api_format = aether_ai_formats::normalize_api_format_alias(&transport.endpoint.api_format);
let Some(overrides) = transport
.key
.auth_type_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
else {
return default_auth_type;
};
overrides
.get(&api_format)
.or_else(|| overrides.get(transport.endpoint.api_format.trim()))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.filter(|value| matches!(value.as_str(), "api_key" | "bearer"))
.unwrap_or(default_auth_type)
}
fn resolved_local_secret(transport: &GatewayProviderTransportSnapshot) -> Option<&str> {
let secret = transport.key.decrypted_api_key.trim();
(!secret.is_empty() && secret != PLACEHOLDER_API_KEY).then_some(secret)
@@ -322,6 +346,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
@@ -443,6 +469,15 @@ mod tests {
assert!(resolve_local_standard_auth(&sample_transport()).is_none());
}
#[test]
fn local_standard_auth_rejects_empty_secret() {
let mut transport = sample_transport();
transport.key.auth_type = "api_key".to_string();
transport.key.decrypted_api_key = String::new();
assert!(resolve_local_standard_auth(&transport).is_none());
}
#[test]
fn local_openai_bearer_auth_maps_api_key_to_bearer_authorization() {
let mut transport = sample_transport();
@@ -466,4 +501,35 @@ mod tests {
Some(("authorization".to_string(), "Bearer sk-openai".to_string(),))
);
}
#[test]
fn local_standard_auth_uses_format_auth_type_override() {
let mut transport = sample_transport();
transport.key.auth_type = "api_key".to_string();
transport.key.auth_type_by_format = Some(serde_json::json!({
"claude:messages": "bearer"
}));
transport.key.decrypted_api_key = "sk-claude".to_string();
assert_eq!(
resolve_local_standard_auth(&transport),
Some(("authorization".to_string(), "Bearer sk-claude".to_string(),))
);
}
#[test]
fn local_gemini_auth_falls_back_to_default_when_other_format_is_overridden() {
let mut transport = sample_transport();
transport.endpoint.api_format = "gemini:generate_content".to_string();
transport.key.auth_type = "api_key".to_string();
transport.key.auth_type_by_format = Some(serde_json::json!({
"claude:messages": "bearer"
}));
transport.key.decrypted_api_key = "sk-gemini".to_string();
assert_eq!(
super::resolve_local_gemini_auth(&transport),
Some(("x-goog-api-key".to_string(), "sk-gemini".to_string(),))
);
}
}

View File

@@ -80,6 +80,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -119,6 +119,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:messages".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -196,6 +196,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:messages".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -136,6 +136,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:messages".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -232,6 +232,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:messages".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -298,6 +298,8 @@ mod tests {
auth_type: "api_key".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -671,6 +671,8 @@ mod tests {
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -306,6 +306,8 @@ mod tests {
auth_type: "api_key".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -61,6 +61,7 @@ pub struct GatewayProviderTransportKey {
pub auth_type: String,
pub is_active: bool,
pub api_formats: Option<Vec<String>>,
pub auth_type_by_format: Option<serde_json::Value>,
pub allowed_models: Option<Vec<String>>,
pub capabilities: Option<serde_json::Value>,
pub rate_multipliers: Option<serde_json::Value>,
@@ -386,6 +387,8 @@ mod tests {
"openai:chat".to_string(),
"openai:responses".to_string(),
]),
auth_type_by_format: None,
allowed_models: Some(vec!["gpt-4.1".to_string(), "gpt-4.1-mini".to_string(),]),
capabilities: Some(serde_json::json!({"cache_1h": true})),
rate_multipliers: Some(serde_json::json!({"openai:chat": 0.8})),
@@ -402,6 +405,31 @@ mod tests {
);
}
#[tokio::test]
async fn reads_snapshot_when_provider_key_api_key_is_null() {
let mut key = sample_key();
key.auth_type = "service_account".to_string();
key.encrypted_api_key = None;
let state = TestSnapshotSource::new(
vec![sample_provider()],
vec![sample_endpoint()],
vec![key],
Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()),
);
let snapshot =
read_provider_transport_snapshot(&state, "provider-1", "endpoint-1", "key-1")
.await
.expect("snapshot should read")
.expect("snapshot should exist");
assert_eq!(snapshot.key.decrypted_api_key, "");
assert_eq!(
snapshot.key.decrypted_auth_config.as_deref(),
Some("{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}")
);
}
#[tokio::test]
async fn returns_none_when_encryption_key_is_not_configured() {
let state = TestSnapshotSource::new(

View File

@@ -54,12 +54,21 @@ pub(super) fn map_key(
encryption_key: &str,
fallback_encryption_keys: &[String],
) -> Result<GatewayProviderTransportKey, DataLayerError> {
let decrypted_api_key = decrypt_secret(
encryption_key,
fallback_encryption_keys,
&key.encrypted_api_key,
"provider_api_keys.api_key",
)?;
let decrypted_api_key = key
.encrypted_api_key
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|ciphertext| {
decrypt_secret(
encryption_key,
fallback_encryption_keys,
ciphertext,
"provider_api_keys.api_key",
)
})
.transpose()?
.unwrap_or_default();
let decrypted_auth_config = key
.encrypted_auth_config
.as_deref()
@@ -85,6 +94,7 @@ pub(super) fn map_key(
normalize_optional_json(key.api_formats),
"provider_api_keys.api_formats",
)?,
auth_type_by_format: normalize_optional_json(key.auth_type_by_format),
allowed_models: normalize_string_list(
normalize_optional_json(key.allowed_models),
"provider_api_keys.allowed_models",

View File

@@ -88,6 +88,8 @@ mod tests {
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -1,5 +1,6 @@
use url::Url;
use super::super::auth::resolve_local_auth_type_for_transport_format;
use super::super::snapshot::GatewayProviderTransportSnapshot;
const VERTEX_AI_HOST: &str = "aiplatform.googleapis.com";
@@ -32,10 +33,7 @@ pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportS
.trim()
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
{
return transport
.key
.auth_type
.trim()
return resolve_local_auth_type_for_transport_format(transport)
.eq_ignore_ascii_case("api_key");
}
@@ -48,11 +46,7 @@ pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportS
return false;
}
transport
.key
.auth_type
.trim()
.eq_ignore_ascii_case("api_key")
resolve_local_auth_type_for_transport_format(transport).eq_ignore_ascii_case("api_key")
}
pub fn uses_vertex_api_key_query_auth(
@@ -117,6 +111,8 @@ mod tests {
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -203,6 +203,8 @@ mod tests {
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,

View File

@@ -157,6 +157,8 @@ mod tests {
auth_type: auth_type.to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,