From c7eed0c0c194dbc310bc17accf916d49eefb5845 Mon Sep 17 00:00:00 2001 From: woody Date: Wed, 13 May 2026 10:18:29 +0800 Subject: [PATCH] fix(vertex): inject api key for express raw endpoints (#3777) Signed-off-by: wydream Co-authored-by: EndlessSeeker <153817598+EndlessSeeker@users.noreply.github.com> --- .../extensions/ai-proxy/provider/vertex.go | 6 +- .../extensions/ai-proxy/test/vertex.go | 58 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go b/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go index 561e139d..b1c21378 100644 --- a/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go +++ b/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go @@ -56,6 +56,10 @@ const ( // 允许任意 basePath 前缀，兼容 basePathHandling 配置 var vertexRawPathRegex = regexp.MustCompile(`^.*/([^/]+)/projects/([^/]+)/locations/([^/]+)/publishers/([^/]+)/models/([^/:]+):([^/?]+)`) +// vertexExpressRawPathRegex 匹配 Vertex AI Express Mode 专用 REST API 路径 +// 格式: [任意前缀]/{api-version}/publishers/{publisher}/models/{model}:{action} +var vertexExpressRawPathRegex = regexp.MustCompile(`^.*/(v[^/]+)/publishers/([^/]+)/models/([^/:]+):([^/?]+)`) + type vertexProviderInitializer struct{} func (v *vertexProviderInitializer) ValidateConfig(config *ProviderConfig) error { @@ -158,7 +162,7 @@ func (v *vertexProvider) GetApiName(path string) ApiName { // 优先匹配原生 Vertex AI REST API 路径，支持任意 basePath 前缀 // 格式: [任意前缀]/{api-version}/projects/{project}/locations/{location}/publishers/{publisher}/models/{model}:{action} // 必须在其他 action 检查之前，因为 :predict、:generateContent 等 action 会被其他规则匹配 - if vertexRawPathRegex.MatchString(path) { + if vertexRawPathRegex.MatchString(path) || (v.isExpressMode() && vertexExpressRawPathRegex.MatchString(path)) { return ApiNameVertexRaw } if strings.HasSuffix(path, vertexChatCompletionAction) || strings.HasSuffix(path, vertexChatCompletionStreamAction) { diff --git a/plugins/wasm-go/extensions/ai-proxy/test/vertex.go b/plugins/wasm-go/extensions/ai-proxy/test/vertex.go index 8ab18a68..4b411f84 100644 --- a/plugins/wasm-go/extensions/ai-proxy/test/vertex.go +++ b/plugins/wasm-go/extensions/ai-proxy/test/vertex.go @@ -2280,6 +2280,27 @@ func RunVertexRawModeOnHttpRequestHeadersTests(t *testing.T) { "Host header should be changed to vertex domain without region prefix") }) + // 测试 Vertex Raw 模式请求头处理（Express Mode 专用路径，无 project/location） + t.Run("vertex raw mode express - request headers with express endpoint path", func(t *testing.T) { + host, status := test.NewTestHost(vertexRawModeExpressConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/publishers/google/models/gemini-3.1-flash-lite-preview:streamGenerateContent"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + require.Equal(t, types.HeaderStopIteration, action) + + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + require.True(t, test.HasHeaderWithValue(requestHeaders, ":authority", "aiplatform.googleapis.com"), + "Host header should be changed to vertex domain without region prefix") + }) + // 测试 Vertex Raw 模式请求头处理（标准模式 + 原生 Vertex API 路径） t.Run("vertex raw mode standard - request headers with native vertex path", func(t *testing.T) { host, status := test.NewTestHost(vertexRawModeStandardConfig) @@ -2416,6 +2437,43 @@ func RunVertexRawModeOnHttpRequestBodyTests(t *testing.T) { "Authorization header should be removed in Express Mode") }) + // 测试 Vertex Raw 模式请求体处理（Express Mode 专用路径 + API Key 认证） + t.Run("vertex raw mode express - express endpoint request body with api key", func(t *testing.T) { + host, status := test.NewTestHost(vertexRawModeExpressConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/publishers/google/models/gemini-3.1-flash-lite-preview:streamGenerateContent"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + {"Authorization", "Bearer some-token"}, + }) + + requestBody := `{"contents":[{"role":"user","parts":[{"text":"Tell me a story"}]}]}` + action := host.CallOnHttpRequestBody([]byte(requestBody)) + + require.Equal(t, types.ActionContinue, action) + + processedBody := host.GetRequestBody() + require.NotNil(t, processedBody) + require.Equal(t, requestBody, string(processedBody), "Request body should be passed through unchanged") + + requestHeaders := host.GetRequestHeaders() + var pathHeader string + for _, header := range requestHeaders { + if header[0] == ":path" { + pathHeader = header[1] + break + } + } + require.Equal(t, "/v1/publishers/google/models/gemini-3.1-flash-lite-preview:streamGenerateContent?key=test-api-key-for-raw-mode", pathHeader, + "API key should be appended to express endpoint path as query parameter") + require.False(t, test.HasHeaderWithValue(requestHeaders, "Authorization", "Bearer some-token"), + "Authorization header should be removed in Express Mode") + }) + // 测试 Vertex Raw 模式请求体处理（标准模式 - 需要 OAuth token） // 注意：使用 countTokens action，因为 generateContent/predict 等会被识别为其他 API 类型 // 注意：在单元测试环境中，由于测试配置使用的是无效的私钥，JWT 创建会失败，