fix(vertex): add API key auth support to Embedding method

[TransactCharlie]

Summary

The Vertex provider's Embedding() method unconditionally calls getAuthTokenSource(key), which attempts google.FindDefaultCredentials(). This fails in environments where GCP auth is provided externally via context headers (e.g. Workload Identity Federation) rather than Application Default Credentials.

Other methods (ChatCompletion, Responses, ResponsesStream) already handle this correctly by checking key.Value first and skipping the internal auth when set — allowing auth provided via BifrostContextKeyExtraHeaders to take effect. This PR applies the same pattern to Embedding().

Changes

• Added authQuery check to Embedding() in core/providers/vertex/vertex.go, matching the pattern used by ChatCompletion(), Responses(), and ResponsesStream()
• When key.Value is set, the value is appended as a ?key=... query parameter and getAuthTokenSource() is skipped
• When key.Value is empty, the existing OAuth2 flow via getAuthTokenSource() is preserved unchanged
• Renamed local url variable to completeURL to avoid shadowing the net/url import (needed for url.QueryEscape)

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

go version
go test ./...

Additionally, to validate the fix end-to-end:

1. Configure a Vertex key with a non-empty Value and set Authorization header via BifrostContextKeyExtraHeaders
2. Send an embedding request to the Vertex provider
3. Verify the request succeeds (previously failed with "error creating auth token source")
4. Verify that embedding requests with empty key.Value and valid AuthCredentials still work via the OAuth2 path

Screenshots/Recordings

Thats difficult to provide. We are using bifrost as the interal layer of a ai-gateway service we are building. With this change in play we can successfully route to vertex using our auth header in the same way we do for Responses etc.

here is an example of a successful call:

rpc POST /service.ai-gateway/service/embeddings {"model":"text-embedding-005","input_text":"Hello world","provider":"vertex","virtual_key":"service.ai-gateway:image-test", "dimensions":4}
{
    "data": [
        {
            "embedding": [
                -0.03997058793902397,
                0.043299295008182526,
                -0.00944140087813139,
                -0.04618649184703827
            ],
            "object": "embedding"
        }
    ],
    "object": "list",
    "provider": "vertex",
    "usage": {
        "prompt_tokens": 2,
        "total_tokens": 2
    }
}

Without this patch we always receive a key error instead:

Breaking changes

• Yes
• No

Related issues

N/A

Security considerations

No new auth mechanisms introduced. The change reuses the existing authQuery pattern already present in ChatCompletion(), Responses(), and ResponsesStream(). The API key is URL-encoded via url.QueryEscape consistent with all other call sites.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

Summary by CodeRabbit

• New Features

  • Vertex provider now supports API key authentication, giving users an alternative method to authenticate with embedding services. This provides greater flexibility for different use cases and deployment scenarios.

• Improvements

  • Refined request handling and configuration for Vertex provider embeddings to ensure proper setup of authentication and network parameters, improving reliability.

[CLAassistant]

All committers have signed the CLA.

[greptile-apps]

Confidence Score: 4/5

The change is a one-method, pattern-matched fix consistent with every other auth branch in the same file; the OAuth fallback is unchanged and the API-key path is straightforward.

The logic is correct and well-precedented within the file. The only gap is the complete absence of unit tests for Embedding() — neither the new API-key branch nor the existing OAuth path has any test coverage, so a future regression here would not be caught automatically.

core/providers/vertex/vertex.go — specifically the new authQuery block in Embedding() and the missing test coverage for both auth branches

Important Files Changed

Filename	Overview
core/providers/vertex/vertex.go	Adds API key auth (key= query param) to Embedding(), matching the existing pattern in ChatCompletion/Responses/ResponsesStream; no tests cover the new Embedding auth branch

_{Reviews (1): Last reviewed commit: "fix(vertex): add API key auth support to..." | Re-trigger Greptile}

[greptile-apps]

Missing test coverage for the new auth branch

vertex_test.go has no Embedding test cases at all, so neither the new API-key path nor the fallback OAuth path is exercised by the test suite. The PR author's checklist also notes tests were not added. If a regression is introduced in the authQuery != "" branch (e.g., a bad URL construction) it will only surface at runtime. Consider adding at minimum a table-driven unit test that covers (1) key.Value set → ?key=… appended, OAuth skipped, and (2) key.Value empty → OAuth path invoked, no query param appended.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 27e806ba-2955-4878-8173-9ef05fa6c73c

📥 Commits

Reviewing files that changed from the base of the PR and between fed162e and 9139499.

📒 Files selected for processing (1)

• core/providers/vertex/vertex.go

---

📝 Walkthrough

Walkthrough

Updated VertexProvider.Embedding method to support API key authentication via query string parameter. When an API key is configured, the request appends it as a query parameter instead of using OAuth2 bearer token authentication. Request setup was reordered to configure headers and content-type before the authentication conditional logic.

Changes

Vertex Embedding API Key Authentication

Layer / File(s)	Summary
Embedding request authentication routing core/providers/vertex/vertex.go	Embedding method builds authQuery from key.Value when present, constructs completeURL via Gemini endpoint, and conditionally appends ?key=... to the URL or fetches OAuth2 token. Request setup order was adjusted to configure content-type and headers before authentication, then set request URI from completeURL.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Suggested reviewers

• akshaydeo

Poem

🐰 A rabbit hops through auth pathways bright,
Now API keys bypass the OAuth night—
Query strings dance where tokens once stood,
Embedding requests flow as they should! 🔑

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: adding API key authentication support to the Vertex Embedding method.
Description check	✅ Passed	The description comprehensively covers all key template sections including summary, changes, type of change, affected areas, testing instructions, security considerations, and checklist items.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}