i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): harden Gemini endpoint routing

Browse Source
This commit is contained in:
MMEXA
2026-05-18 00:53:34 +00:00
parent 81ff375bfd
commit b004a02e4a
20 changed files with 1348 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/aether-data/src/repository/candidate_selection/memory.rs
View File

@@ -317,11 +317,15 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",

8
crates/aether-data/src/repository/candidate_selection/mysql.rs
View File

@@ -446,11 +446,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",

28
crates/aether-data/src/repository/candidate_selection/postgres.rs
View File

@@ -113,11 +113,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($3) = 'gemini:generate_content'
AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -296,11 +296,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($4) = 'gemini:generate_content'
AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -478,11 +478,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($6) = 'gemini:generate_content'
AND LOWER($6) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -1287,6 +1287,22 @@ mod tests {
}
}
#[test]
fn candidate_selection_sql_allows_vertex_embedding_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'vertex_ai'"));
assert!(sql.contains("gemini:embedding"));
assert!(sql.contains("gemini:generate_content"));
assert!(sql.contains("claude:messages"));
}
}
#[test]
fn requested_model_selection_page_sql_adds_limit_and_offset() {
let sql = requested_model_selection_page_sql();

8
crates/aether-data/src/repository/candidate_selection/sqlite.rs
View File

@@ -446,11 +446,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",