1package provider
2
3// Google Vertex AI client.
4//
5// Wire format is identical to the public Generative Language API
6// (generateContent / streamGenerateContent), so we reuse the existing
7// geminiClient and rewrite the outgoing URL + auth via a RoundTripper.
8//
9// URL shape:
10//
11// POST https://{location}-aiplatform.googleapis.com/v1/projects/{project}/
12// locations/{location}/publishers/google/models/{model}:streamGenerateContent?alt=sse
13//
14// Auth: two paths supported.
15//
16// - **API key** (GOOGLE_CLOUD_API_KEY): `x-goog-api-key` header. Simplest
17// setup; works for users who have created a Vertex AI API key in the
18// GCP console. We use this directly, no token exchange.
19// - **Service-account JSON** pointed to by GOOGLE_APPLICATION_CREDENTIALS:
20// we read the JSON, build a JWT signed with RS256 using the private
21// key, exchange it at `https://oauth2.googleapis.com/token` for a
22// short-lived (1h) access token, cache it in memory, and refresh on
23// demand. Sent as `Authorization: Bearer <token>`.
24// - **gcloud user OAuth** (~/.config/gcloud/application_default_credentials.json
25// with `type: "authorized_user"`): we read the JSON, pull the
26// client_id/client_secret/refresh_token, exchange the refresh token
27// at `https://oauth2.googleapis.com/token` for an access token,
28// same cache + refresh logic as service-account.
29//
30// Config env vars (read at construction time):
31// GOOGLE_CLOUD_PROJECT required
32// GOOGLE_CLOUD_LOCATION required (e.g. us-central1)
33// GOOGLE_CLOUD_API_KEY optional; preferred over service-account
34// GOOGLE_APPLICATION_CREDENTIALS optional; service-account JSON path
35
36import (
37 "context"
38 "crypto"
39 "crypto/rand"
40 "crypto/rsa"
41 "crypto/sha256"
42 "crypto/x509"
43 "encoding/base64"
44 "encoding/json"
45 "encoding/pem"
46 "fmt"
47 "io"
48 "net/http"
49 "net/url"
50 "os"
51 "strings"
52 "sync"
53 "time"
54
55 providerauth "github.com/bnema/zut/packages/provider/auth"
56)
57
58type vertexConfig struct {
59 project string
60 location string
61 apiKey string // empty if using OAuth-style auth
62
63 // Auth mode. Exactly one of (saEmail+saPrivateKey) or
64 // (userClientID+userClientSecret+userRefreshToken) is populated when
65 // apiKey is empty.
66 saEmail string
67 saPrivateKey *rsa.PrivateKey
68 saTokenURI string
69
70 userClientID string
71 userClientSecret string
72 userRefreshToken string
73 userTokenURI string
74}
75
76// cacheKey identifies the credential set for the token cache. We key on
77// the auth principal (SA email or user client_id) so multiple configs
78// in the same process don't trample each other.
79func (c *vertexConfig) cacheKey() string {
80 if c.saEmail != "" {
81 return "sa:" + c.saEmail
82 }
83 return "user:" + c.userClientID
84}
85
86func loadVertexConfig() (*vertexConfig, error) {
87 cfg := &vertexConfig{
88 project: os.Getenv("GOOGLE_CLOUD_PROJECT"),
89 location: os.Getenv("GOOGLE_CLOUD_LOCATION"),
90 apiKey: os.Getenv("GOOGLE_CLOUD_API_KEY"),
91 }
92 if cfg.location == "" {
93 cfg.location = "us-central1"
94 }
95 if cfg.project == "" {
96 return nil, fmt.Errorf("vertex: GOOGLE_CLOUD_PROJECT not set")
97 }
98 if cfg.apiKey != "" {
99 return cfg, nil
100 }
101 // Try service-account JSON.
102 credPath := os.Getenv("GOOGLE_APPLICATION_CREDENTIALS")
103 if credPath == "" {
104 // ADC default path. Mirrors `gcloud auth application-default login`.
105 if candidate, err := providerauth.GoogleApplicationDefaultCredentialsPath(); err == nil {
106 if _, err := os.Stat(candidate); err == nil {
107 credPath = candidate
108 }
109 }
110 }
111 if credPath == "" {
112 return nil, fmt.Errorf("vertex: no auth — set GOOGLE_CLOUD_API_KEY or GOOGLE_APPLICATION_CREDENTIALS")
113 }
114 b, err := os.ReadFile(credPath)
115 if err != nil {
116 return nil, fmt.Errorf("vertex: read credentials %q: %w", credPath, err)
117 }
118 var raw struct {
119 Type string `json:"type"`
120 ClientEmail string `json:"client_email"`
121 PrivateKey string `json:"private_key"`
122 TokenURI string `json:"token_uri"`
123 ClientID string `json:"client_id"`
124 ClientSecret string `json:"client_secret"`
125 RefreshToken string `json:"refresh_token"`
126 }
127 if err := json.Unmarshal(b, &raw); err != nil {
128 return nil, fmt.Errorf("vertex: parse credentials: %w", err)
129 }
130 switch raw.Type {
131 case "service_account":
132 if raw.ClientEmail == "" || raw.PrivateKey == "" {
133 return nil, fmt.Errorf("vertex: service_account JSON missing client_email or private_key")
134 }
135 block, _ := pem.Decode([]byte(raw.PrivateKey))
136 if block == nil {
137 return nil, fmt.Errorf("vertex: private_key is not PEM-encoded")
138 }
139 key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
140 if err != nil {
141 // Try PKCS#1 as a fallback (older keys).
142 k1, err2 := x509.ParsePKCS1PrivateKey(block.Bytes)
143 if err2 != nil {
144 return nil, fmt.Errorf("vertex: parse private_key: %w", err)
145 }
146 cfg.saPrivateKey = k1
147 } else {
148 rsaKey, ok := key.(*rsa.PrivateKey)
149 if !ok {
150 return nil, fmt.Errorf("vertex: private_key is not RSA")
151 }
152 cfg.saPrivateKey = rsaKey
153 }
154 cfg.saEmail = raw.ClientEmail
155 cfg.saTokenURI = raw.TokenURI
156 if cfg.saTokenURI == "" {
157 cfg.saTokenURI = "https://oauth2.googleapis.com/token"
158 }
159 case "authorized_user":
160 // gcloud auth application-default login produces this shape.
161 if raw.ClientID == "" || raw.ClientSecret == "" || raw.RefreshToken == "" {
162 return nil, fmt.Errorf("vertex: authorized_user JSON missing client_id/client_secret/refresh_token")
163 }
164 cfg.userClientID = raw.ClientID
165 cfg.userClientSecret = raw.ClientSecret
166 cfg.userRefreshToken = raw.RefreshToken
167 cfg.userTokenURI = raw.TokenURI
168 if cfg.userTokenURI == "" {
169 cfg.userTokenURI = "https://oauth2.googleapis.com/token"
170 }
171 default:
172 return nil, fmt.Errorf("vertex: unsupported credentials type %q (expected service_account or authorized_user)", raw.Type)
173 }
174 return cfg, nil
175}
176
177// vertexTokenCache caches one access token per service-account email.
178type vertexTokenCache struct {
179 mu sync.Mutex
180 tokens map[string]struct {
181 value string
182 expiresAt time.Time
183 }
184 http *http.Client
185}
186
187var vertexCache = &vertexTokenCache{
188 tokens: map[string]struct {
189 value string
190 expiresAt time.Time
191 }{},
192 http: &http.Client{Timeout: 30 * time.Second},
193}
194
195func b64url(b []byte) string { return base64.RawURLEncoding.EncodeToString(b) }
196
197// signServiceAccountJWT builds and signs a JWT for the
198// `urn:ietf:params:oauth:grant-type:jwt-bearer` token exchange.
199func signServiceAccountJWT(cfg *vertexConfig) (string, error) {
200 header := map[string]string{"alg": "RS256", "typ": "JWT"}
201 now := time.Now().Unix()
202 claims := map[string]any{
203 "iss": cfg.saEmail,
204 "scope": "https://www.googleapis.com/auth/cloud-platform",
205 "aud": cfg.saTokenURI,
206 "exp": now + 3600,
207 "iat": now,
208 }
209 hb, _ := json.Marshal(header)
210 cb, _ := json.Marshal(claims)
211 signingInput := b64url(hb) + "." + b64url(cb)
212 hashed := sha256.Sum256([]byte(signingInput))
213 sig, err := rsa.SignPKCS1v15(rand.Reader, cfg.saPrivateKey, crypto.SHA256, hashed[:])
214 if err != nil {
215 return "", err
216 }
217 return signingInput + "." + b64url(sig), nil
218}
219
220func (c *vertexTokenCache) get(ctx context.Context, cfg *vertexConfig) (string, error) {
221 key := cfg.cacheKey()
222 c.mu.Lock()
223 tok, ok := c.tokens[key]
224 c.mu.Unlock()
225 if ok && time.Now().Before(tok.expiresAt) {
226 return tok.value, nil
227 }
228
229 var form url.Values
230 var endpoint string
231 switch {
232 case cfg.saPrivateKey != nil:
233 jwt, err := signServiceAccountJWT(cfg)
234 if err != nil {
235 return "", fmt.Errorf("vertex: sign jwt: %w", err)
236 }
237 form = url.Values{
238 "grant_type": {"urn:ietf:params:oauth:grant-type:jwt-bearer"},
239 "assertion": {jwt},
240 }
241 endpoint = cfg.saTokenURI
242 case cfg.userRefreshToken != "":
243 form = url.Values{
244 "grant_type": {"refresh_token"},
245 "client_id": {cfg.userClientID},
246 "client_secret": {cfg.userClientSecret},
247 "refresh_token": {cfg.userRefreshToken},
248 }
249 endpoint = cfg.userTokenURI
250 default:
251 return "", fmt.Errorf("vertex: no usable auth (need service_account or authorized_user creds)")
252 }
253
254 req, err := http.NewRequestWithContext(ctx, "POST", endpoint, strings.NewReader(form.Encode()))
255 if err != nil {
256 return "", err
257 }
258 req.Header.Set("content-type", "application/x-www-form-urlencoded")
259 resp, err := c.http.Do(req)
260 if err != nil {
261 return "", fmt.Errorf("vertex: token exchange: %w", err)
262 }
263 defer resp.Body.Close()
264 body, _ := io.ReadAll(resp.Body)
265 if resp.StatusCode != http.StatusOK {
266 return "", fmt.Errorf("vertex: token exchange http %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
267 }
268 var out struct {
269 AccessToken string `json:"access_token"`
270 ExpiresIn int `json:"expires_in"`
271 }
272 if err := json.Unmarshal(body, &out); err != nil {
273 return "", fmt.Errorf("vertex: parse token: %w", err)
274 }
275 if out.AccessToken == "" {
276 return "", fmt.Errorf("vertex: empty access token")
277 }
278 exp := time.Now().Add(time.Duration(out.ExpiresIn-60) * time.Second)
279 c.mu.Lock()
280 c.tokens[key] = struct {
281 value string
282 expiresAt time.Time
283 }{out.AccessToken, exp}
284 c.mu.Unlock()
285 return out.AccessToken, nil
286}
287
288// vertexTransport rewrites the outgoing Gemini URL into the Vertex shape
289// and replaces the x-goog-api-key header with either another
290// x-goog-api-key (for the API-key flow) or Authorization: Bearer (for
291// the service-account flow).
292type vertexTransport struct {
293 inner http.RoundTripper
294 cfg *vertexConfig
295}
296
297func (t *vertexTransport) RoundTrip(req *http.Request) (*http.Response, error) {
298 // Path looks like /v1beta/models/{model}:streamGenerateContent. Pull
299 // out the model id and rewrite.
300 const prefix = "/v1beta/models/"
301 if !strings.HasPrefix(req.URL.Path, prefix) {
302 return nil, fmt.Errorf("vertex: unexpected request path %q", req.URL.Path)
303 }
304 rest := strings.TrimPrefix(req.URL.Path, prefix)
305 // rest = "<model>:streamGenerateContent"
306 sep := strings.IndexByte(rest, ':')
307 if sep < 0 {
308 return nil, fmt.Errorf("vertex: malformed request path %q", req.URL.Path)
309 }
310 modelID := rest[:sep]
311 verb := rest[sep:] // ":streamGenerateContent"
312 newPath := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/google/models/%s%s",
313 t.cfg.project, t.cfg.location, modelID, verb)
314 clone := req.Clone(req.Context())
315 if t.cfg.location == "global" {
316 clone.URL.Host = "aiplatform.googleapis.com"
317 } else {
318 clone.URL.Host = t.cfg.location + "-aiplatform.googleapis.com"
319 }
320 clone.URL.Path = newPath
321 // Drop existing auth header and pick the right one.
322 clone.Header.Del("x-goog-api-key")
323 if t.cfg.apiKey != "" {
324 clone.Header.Set("x-goog-api-key", t.cfg.apiKey)
325 } else {
326 tok, err := vertexCache.get(req.Context(), t.cfg)
327 if err != nil {
328 return nil, err
329 }
330 clone.Header.Set("Authorization", "Bearer "+tok)
331 }
332 return t.inner.RoundTrip(clone)
333}
334
335// NewVertex returns a Vertex AI client. The apiKey argument is ignored
336// in favor of env-based config (GOOGLE_CLOUD_API_KEY or
337// GOOGLE_APPLICATION_CREDENTIALS), since Vertex's auth model doesn't fit
338// the "just paste a key" interface other providers use.
339func NewVertex(_ string, _ string) Client {
340 cfg, err := loadVertexConfig()
341 if err != nil {
342 return &unimplementedClient{name: "google-vertex", hint: err.Error()}
343 }
344 baseHost := "aiplatform.googleapis.com"
345 if cfg.location != "global" {
346 baseHost = cfg.location + "-aiplatform.googleapis.com"
347 }
348 inner := &geminiClient{
349 providerName: "google-vertex",
350 apiKey: "vertex-placeholder", // overwritten by transport
351 baseURL: "https://" + baseHost,
352 http: &http.Client{
353 Transport: &vertexTransport{inner: http.DefaultTransport, cfg: cfg},
354 Timeout: 0,
355 },
356 }
357 // Keep the public wrapper and wire client on the same provider identity.
358 return &renamedClient{inner: inner, name: "google-vertex"}
359}
360
361// renamedClient wraps a Client and overrides its Name(). Used for
362// provider-id remapping (Vertex reuses gemini's wire client but should
363// report as google-vertex for cost / cache / log purposes).
364type renamedClient struct {
365 inner Client
366 name string
367}
368
369func (r *renamedClient) Name() string { return r.name }
370func (r *renamedClient) SetModelMetadata(model Model) {
371 setModelMetadata(r.inner, model)
372}
373func (r *renamedClient) Stream(ctx context.Context, req Request) (<-chan Event, error) {
374 if err := ValidateFastMode(r.Name(), req.FastMode); err != nil {
375 return nil, err
376 }
377 return r.inner.Stream(ctx, req)
378}
379
380func (r *renamedClient) Close() error {
381 if closer, ok := r.inner.(ClientCloser); ok {
382 return closer.Close()
383 }
384 return nil
385}