Add Gemini CLI v1internal quota support

This commit is contained in:
Mas0nShi
2026-05-22 16:58:04 +08:00
parent 9533bd7043
commit ce02f1ae8c
17 changed files with 985 additions and 52 deletions
+16 -5
View File
@@ -95,7 +95,7 @@ pub(super) fn classify_ai_public_route(
"openai:video",
true,
))
} else if is_gemini_models_route(normalized_path) {
} else if method == http::Method::POST && is_gemini_models_route(normalized_path) {
if normalized_path.ends_with(":predictLongRunning") {
Some(classified(
"ai_public",
@@ -134,7 +134,9 @@ pub(super) fn classify_ai_public_route(
true,
))
}
} else if is_gemini_operation_route(normalized_path) {
} else if is_gemini_operation_method(method, normalized_path)
&& is_gemini_operation_route(normalized_path)
{
Some(classified(
"ai_public",
"gemini",
@@ -142,9 +144,7 @@ pub(super) fn classify_ai_public_route(
"gemini:video",
true,
))
} else if (method == http::Method::POST && normalized_path == "/upload/v1beta/files")
|| normalized_path.starts_with("/v1beta/files")
{
} else if is_gemini_files_method(method, normalized_path) {
Some(classified(
"ai_public",
"gemini",
@@ -156,3 +156,14 @@ pub(super) fn classify_ai_public_route(
None
}
}
fn is_gemini_operation_method(method: &http::Method, normalized_path: &str) -> bool {
method == http::Method::GET
|| (method == http::Method::POST && normalized_path.ends_with(":cancel"))
}
fn is_gemini_files_method(method: &http::Method, normalized_path: &str) -> bool {
(method == http::Method::POST && normalized_path == "/upload/v1beta/files")
|| ((method == http::Method::GET || method == http::Method::DELETE)
&& normalized_path.starts_with("/v1beta/files"))
}
@@ -219,6 +219,23 @@ fn classifies_gemini_generate_content_api_key_without_cli_marker() {
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn does_not_classify_gemini_options_preflight_as_ai_execution_route() {
let headers = headers(&[
("origin", "http://localhost:3000"),
("access-control-request-method", "POST"),
(
"access-control-request-headers",
"authorization,content-type",
),
]);
let uri: Uri = "/v1beta/models/gemini-3-flash-preview:streamGenerateContent?alt=sse"
.parse()
.expect("uri should parse");
assert!(classify_control_route(&http::Method::OPTIONS, &uri, &headers).is_none());
}
#[test]
fn classifies_gemini_embed_content_as_embedding_route() {
let headers = headers(&[("x-goog-api-key", "gemini-key")]);