feat(vertex, management): enhance token handling and OAuth2 integration

Extend `vertexAccessToken` to support proxy-aware HTTP clients and update calls accordingly for better configurability. Add `deleteTokenRecord` to handle token cleanup, improving management of authentication files.
This commit is contained in:
Luis Pater
2025-11-11 23:42:46 +08:00
parent 8ae8a5c296
commit d6bd6f3fb9
2 changed files with 41 additions and 7 deletions

View File

@@ -508,6 +508,10 @@ func (h *Handler) DeleteAuthFile(c *gin.Context) {
} }
} }
if err = os.Remove(full); err == nil { if err = os.Remove(full); err == nil {
if errDel := h.deleteTokenRecord(ctx, full); errDel != nil {
c.JSON(500, gin.H{"error": errDel.Error()})
return
}
deleted++ deleted++
h.disableAuth(ctx, full) h.disableAuth(ctx, full)
} }
@@ -534,6 +538,10 @@ func (h *Handler) DeleteAuthFile(c *gin.Context) {
} }
return return
} }
if err := h.deleteTokenRecord(ctx, full); err != nil {
c.JSON(500, gin.H{"error": err.Error()})
return
}
h.disableAuth(ctx, full) h.disableAuth(ctx, full)
c.JSON(200, gin.H{"status": "ok"}) c.JSON(200, gin.H{"status": "ok"})
} }
@@ -640,9 +648,20 @@ func (h *Handler) disableAuth(ctx context.Context, id string) {
} }
} }
func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (string, error) { func (h *Handler) deleteTokenRecord(ctx context.Context, path string) error {
if record == nil { if strings.TrimSpace(path) == "" {
return "", fmt.Errorf("token record is nil") return fmt.Errorf("auth path is empty")
}
store := h.tokenStoreWithBaseDir()
if store == nil {
return fmt.Errorf("token store unavailable")
}
return store.Delete(ctx, path)
}
func (h *Handler) tokenStoreWithBaseDir() coreauth.Store {
if h == nil {
return nil
} }
store := h.tokenStore store := h.tokenStore
if store == nil { if store == nil {
@@ -654,6 +673,17 @@ func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (s
dirSetter.SetBaseDir(h.cfg.AuthDir) dirSetter.SetBaseDir(h.cfg.AuthDir)
} }
} }
return store
}
func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (string, error) {
if record == nil {
return "", fmt.Errorf("token record is nil")
}
store := h.tokenStoreWithBaseDir()
if store == nil {
return "", fmt.Errorf("token store unavailable")
}
return store.Save(ctx, record) return store.Save(ctx, record)
} }

View File

@@ -22,6 +22,7 @@ import (
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson" "github.com/tidwall/gjson"
"github.com/tidwall/sjson" "github.com/tidwall/sjson"
"golang.org/x/oauth2"
"golang.org/x/oauth2/google" "golang.org/x/oauth2/google"
) )
@@ -89,7 +90,7 @@ func (e *GeminiVertexExecutor) Execute(ctx context.Context, auth *cliproxyauth.A
return resp, errNewReq return resp, errNewReq
} }
httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, saJSON); errTok == nil && token != "" { if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token) httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil { } else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok) log.Errorf("vertex executor: access token error: %v", errTok)
@@ -184,7 +185,7 @@ func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
return nil, errNewReq return nil, errNewReq
} }
httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, saJSON); errTok == nil && token != "" { if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token) httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil { } else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok) log.Errorf("vertex executor: access token error: %v", errTok)
@@ -295,7 +296,7 @@ func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyau
return cliproxyexecutor.Response{}, errNewReq return cliproxyexecutor.Response{}, errNewReq
} }
httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, saJSON); errTok == nil && token != "" { if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token) httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil { } else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok) log.Errorf("vertex executor: access token error: %v", errTok)
@@ -407,7 +408,10 @@ func vertexBaseURL(location string) string {
return fmt.Sprintf("https://%s-aiplatform.googleapis.com", loc) return fmt.Sprintf("https://%s-aiplatform.googleapis.com", loc)
} }
func vertexAccessToken(ctx context.Context, saJSON []byte) (string, error) { func vertexAccessToken(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, saJSON []byte) (string, error) {
if httpClient := newProxyAwareHTTPClient(ctx, cfg, auth, 0); httpClient != nil {
ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient)
}
// Use cloud-platform scope for Vertex AI. // Use cloud-platform scope for Vertex AI.
creds, errCreds := google.CredentialsFromJSON(ctx, saJSON, "https://www.googleapis.com/auth/cloud-platform") creds, errCreds := google.CredentialsFromJSON(ctx, saJSON, "https://www.googleapis.com/auth/cloud-platform")
if errCreds != nil { if errCreds != nil {