From 6609315131017ea789dbc48296d66cf1d759d3e8 Mon Sep 17 00:00:00 2001 From: Charlie Gildawie Date: Wed, 22 Jul 2026 11:34:17 +0100 Subject: [PATCH] fix(vertex): support API-key/context-header auth in cached content methods MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Vertex cached content methods (create/list/retrieve/update/delete) all authenticate via the shared vertexAuthHeaders helper, which unconditionally fetched an OAuth token from the key credentials and overwrote the Authorization header. This mirrors the pre-fix Embedding behaviour and prevents callers from supplying their own bearer token via context extra headers (e.g. a proxy that holds short-lived credentials out of band). Make vertexAuthHeaders take the API-key query-parameter path when the key carries a value — the same escape hatch the Gemini generation endpoints already use — leaving any Authorization header set from context extra headers intact. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Charlie Gildawie --- core/providers/vertex/cachedcontents.go | 10 +++++- .../vertex/cachedcontents_auth_test.go | 32 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 core/providers/vertex/cachedcontents_auth_test.go diff --git a/core/providers/vertex/cachedcontents.go b/core/providers/vertex/cachedcontents.go index 4cd87352005..dac48cae328 100644 --- a/core/providers/vertex/cachedcontents.go +++ b/core/providers/vertex/cachedcontents.go @@ -85,8 +85,16 @@ func expandVertexModelPath(model, projectID, region string) string { return fmt.Sprintf("projects/%s/locations/%s/publishers/google/models/%s", projectID, region, model) } -// vertexAuthHeaders pulls an OAuth bearer token from the key and applies it. +// vertexAuthHeaders applies Vertex AI authentication to the request. When the key +// carries an API key value, it is passed as the "key" query parameter (mirroring +// the Gemini generation endpoints) and any Authorization header already set from +// context extra headers is left intact. Otherwise an OAuth bearer token is fetched +// from the key credentials and set on the Authorization header. func vertexAuthHeaders(req *fasthttp.Request, key schemas.Key) *schemas.BifrostError { + if key.Value.GetValue() != "" { + req.URI().QueryArgs().Set("key", key.Value.GetValue()) + return nil + } tokenSource, err := getAuthTokenSource(key) if err != nil { return providerUtils.NewBifrostOperationError("error creating auth token source", err) diff --git a/core/providers/vertex/cachedcontents_auth_test.go b/core/providers/vertex/cachedcontents_auth_test.go new file mode 100644 index 00000000000..72a9166b6fc --- /dev/null +++ b/core/providers/vertex/cachedcontents_auth_test.go @@ -0,0 +1,32 @@ +package vertex + +import ( + "testing" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/valyala/fasthttp" +) + +// TestVertexAuthHeaders_APIKeyPreservesInjectedAuthHeader verifies that when the +// key carries an API key value, vertexAuthHeaders passes it as the "key" query +// parameter and leaves an Authorization header (set upstream from context extra +// headers) untouched. This mirrors the Gemini generation endpoints and lets a +// caller inject its own bearer token via context extra headers. +func TestVertexAuthHeaders_APIKeyPreservesInjectedAuthHeader(t *testing.T) { + req := fasthttp.AcquireRequest() + defer fasthttp.ReleaseRequest(req) + req.SetRequestURI("https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/cachedContents") + req.Header.Set("Authorization", "Bearer injected-token") + + key := schemas.Key{Value: *schemas.NewSecretVar("api-key-123")} + if err := vertexAuthHeaders(req, key); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if got := string(req.Header.Peek("Authorization")); got != "Bearer injected-token" { + t.Errorf("Authorization header was overwritten: got %q, want the injected token preserved", got) + } + if got := string(req.URI().QueryArgs().Peek("key")); got != "api-key-123" { + t.Errorf("key query parameter: got %q, want %q", got, "api-key-123") + } +}