mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-03 13:00:52 +08:00
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:
@@ -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)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
Reference in New Issue
Block a user