diff --git a/common/codex_credential.go b/common/codex_credential.go new file mode 100644 index 0000000..a5938ea --- /dev/null +++ b/common/codex_credential.go @@ -0,0 +1,58 @@ +package common + +import ( + "errors" + "strings" +) + +type CodexCredentialMode string + +const ( + CodexCredentialModeAPIKey CodexCredentialMode = "api_key" + CodexCredentialModeOAuth CodexCredentialMode = "oauth" +) + +var ( + ErrCodexOAuthCredentialRequired = errors.New("codex channel: oauth credential required") + ErrCodexOAuthCredentialInvalidJSON = errors.New("codex channel: invalid oauth key json") + ErrCodexOAuthAccessTokenRequired = errors.New("codex channel: access_token is required") + ErrCodexOAuthAccountIDRequired = errors.New("codex channel: account_id is required") +) + +type CodexOAuthCredential struct { + IDToken string `json:"id_token,omitempty"` + AccessToken string `json:"access_token,omitempty"` + RefreshToken string `json:"refresh_token,omitempty"` + AccountID string `json:"account_id,omitempty"` + LastRefresh string `json:"last_refresh,omitempty"` + Email string `json:"email,omitempty"` + Type string `json:"type,omitempty"` + Expired string `json:"expired,omitempty"` +} + +func ParseCodexOAuthCredential(raw string) (*CodexOAuthCredential, error) { + trimmed := strings.TrimSpace(raw) + if trimmed == "" || !strings.HasPrefix(trimmed, "{") { + return nil, ErrCodexOAuthCredentialRequired + } + + var credential CodexOAuthCredential + if err := Unmarshal([]byte(trimmed), &credential); err != nil { + return nil, ErrCodexOAuthCredentialInvalidJSON + } + if strings.TrimSpace(credential.AccessToken) == "" { + return nil, ErrCodexOAuthAccessTokenRequired + } + if strings.TrimSpace(credential.AccountID) == "" { + return nil, ErrCodexOAuthAccountIDRequired + } + + return &credential, nil +} + +func DetectCodexCredentialMode(raw string) CodexCredentialMode { + if _, err := ParseCodexOAuthCredential(raw); err == nil { + return CodexCredentialModeOAuth + } + return CodexCredentialModeAPIKey +} diff --git a/controller/channel.go b/controller/channel.go index f9a0d0c..5d4b939 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -3,6 +3,7 @@ package controller import ( "context" "encoding/json" + "errors" "fmt" "net/http" "strconv" @@ -615,20 +616,22 @@ func validateChannel(channel *model.Channel, isAdd bool) error { // Codex OAuth key validation (optional, only when JSON object is provided) if channel.Type == constant.ChannelTypeCodex { trimmedKey := strings.TrimSpace(channel.Key) - if isAdd || trimmedKey != "" { - if !strings.HasPrefix(trimmedKey, "{") { - return fmt.Errorf("Codex key must be a valid JSON object") - } - var keyMap map[string]any - if err := common.Unmarshal([]byte(trimmedKey), &keyMap); err != nil { + if isAdd && trimmedKey == "" { + return fmt.Errorf("Codex key cannot be empty") + } + if strings.HasPrefix(trimmedKey, "{") { + if _, err := common.ParseCodexOAuthCredential(trimmedKey); err != nil { + if errors.Is(err, common.ErrCodexOAuthCredentialInvalidJSON) { + return fmt.Errorf("Codex key must be a valid JSON object") + } + if errors.Is(err, common.ErrCodexOAuthAccessTokenRequired) { + return fmt.Errorf("Codex key JSON must include access_token") + } + if errors.Is(err, common.ErrCodexOAuthAccountIDRequired) { + return fmt.Errorf("Codex key JSON must include account_id") + } return fmt.Errorf("Codex key must be a valid JSON object") } - if v, ok := keyMap["access_token"]; !ok || v == nil || strings.TrimSpace(fmt.Sprintf("%v", v)) == "" { - return fmt.Errorf("Codex key JSON must include access_token") - } - if v, ok := keyMap["account_id"]; !ok || v == nil || strings.TrimSpace(fmt.Sprintf("%v", v)) == "" { - return fmt.Errorf("Codex key JSON must include account_id") - } } } @@ -647,6 +650,10 @@ func RefreshCodexChannelCredential(c *gin.Context) { oauthKey, ch, err := service.RefreshCodexChannelCredential(ctx, channelId, service.CodexCredentialRefreshOptions{ResetCaches: true}) if err != nil { + if errors.Is(err, common.ErrCodexOAuthCredentialRequired) { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "当前凭证方式不支持刷新凭证"}) + return + } common.SysError("failed to refresh codex channel credential: " + err.Error()) c.JSON(http.StatusOK, gin.H{"success": false, "message": "刷新凭证失败,请稍后重试"}) return @@ -2112,11 +2119,11 @@ func GetUserChannelsForBinding(c *gin.Context) { result := make([]gin.H, 0, len(channels)) for _, ch := range channels { result = append(result, gin.H{ - "id": ch.Id, - "name": ch.Name, + "id": ch.Id, + "name": ch.Name, "public_name": ch.PublicName, - "type": ch.Type, - "remark": ch.Remark, + "type": ch.Type, + "remark": ch.Remark, }) } diff --git a/controller/codex_channel_test.go b/controller/codex_channel_test.go new file mode 100644 index 0000000..2071c5b --- /dev/null +++ b/controller/codex_channel_test.go @@ -0,0 +1,111 @@ +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupCodexChannelDB(t *testing.T, key string) *gorm.DB { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + origLogDB := model.LOG_DB + model.DB = db + model.LOG_DB = db + common.UsingSQLite = true + common.RedisEnabled = false + + require.NoError(t, db.AutoMigrate(&model.Channel{})) + require.NoError(t, db.Create(&model.Channel{ + Id: 1, + Name: "codex-channel", + PublicName: "codex-channel", + Type: constant.ChannelTypeCodex, + Key: key, + Status: common.ChannelStatusEnabled, + }).Error) + + t.Cleanup(func() { + model.DB = origDB + model.LOG_DB = origLogDB + _ = sqlDB.Close() + }) + + return db +} + +func setupCodexChannelRouter(t *testing.T, key string) *gin.Engine { + t.Helper() + setupCodexChannelDB(t, key) + + gin.SetMode(gin.TestMode) + r := gin.New() + g := r.Group("/api/channel") + g.POST("/:id/codex/refresh", RefreshCodexChannelCredential) + g.GET("/:id/codex/usage", GetCodexChannelUsage) + return r +} + +func TestValidateChannelAcceptsCodexAPIKey(t *testing.T) { + channel := &model.Channel{ + Name: "codex-api-key", + PublicName: "codex-api-key", + Type: constant.ChannelTypeCodex, + Key: "sk-codex-api-key", + } + + require.NoError(t, validateChannel(channel, true)) +} + +func TestValidateChannelRejectsCodexOAuthWithoutAccountID(t *testing.T) { + channel := &model.Channel{ + Name: "codex-oauth", + PublicName: "codex-oauth", + Type: constant.ChannelTypeCodex, + Key: `{"access_token":"token-only"}`, + } + + err := validateChannel(channel, true) + require.Error(t, err) + assert.Contains(t, err.Error(), "account_id") +} + +func TestRefreshCodexChannelCredentialRejectsAPIKeyMode(t *testing.T) { + router := setupCodexChannelRouter(t, "sk-codex-api-key") + + req := httptest.NewRequest(http.MethodPost, "/api/channel/1/codex/refresh", bytes.NewReader([]byte(`{}`))) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Contains(t, w.Body.String(), "当前凭证方式不支持刷新凭证") +} + +func TestGetCodexChannelUsageRejectsAPIKeyMode(t *testing.T) { + router := setupCodexChannelRouter(t, "sk-codex-api-key") + + req := httptest.NewRequest(http.MethodGet, "/api/channel/1/codex/usage", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Contains(t, w.Body.String(), "当前凭证方式不支持查看用量") +} diff --git a/controller/codex_usage.go b/controller/codex_usage.go index 52fdbdf..98dc465 100644 --- a/controller/codex_usage.go +++ b/controller/codex_usage.go @@ -2,6 +2,7 @@ package controller import ( "context" + "errors" "fmt" "net/http" "strconv" @@ -11,7 +12,6 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" - "github.com/QuantumNous/new-api/relay/channel/codex" "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" @@ -42,22 +42,19 @@ func GetCodexChannelUsage(c *gin.Context) { return } - oauthKey, err := codex.ParseOAuthKey(strings.TrimSpace(ch.Key)) + oauthKey, err := common.ParseCodexOAuthCredential(strings.TrimSpace(ch.Key)) if err != nil { + if errors.Is(err, common.ErrCodexOAuthCredentialRequired) { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "当前凭证方式不支持查看用量"}) + return + } common.SysError("failed to parse oauth key: " + err.Error()) c.JSON(http.StatusOK, gin.H{"success": false, "message": "解析凭证失败,请检查渠道配置"}) return } + accessToken := strings.TrimSpace(oauthKey.AccessToken) accountID := strings.TrimSpace(oauthKey.AccountID) - if accessToken == "" { - c.JSON(http.StatusOK, gin.H{"success": false, "message": "codex channel: access_token is required"}) - return - } - if accountID == "" { - c.JSON(http.StatusOK, gin.H{"success": false, "message": "codex channel: account_id is required"}) - return - } client, err := service.NewProxyHttpClient(ch.GetSetting().Proxy) if err != nil { @@ -98,6 +95,7 @@ func GetCodexChannelUsage(c *gin.Context) { ctx2, cancel2 := context.WithTimeout(c.Request.Context(), 15*time.Second) defer cancel2() + statusCode, body, err = service.FetchCodexWhamUsage(ctx2, client, ch.GetBaseURL(), oauthKey.AccessToken, accountID) if err != nil { common.SysError("failed to fetch codex usage after refresh: " + err.Error()) diff --git a/service/codex_credential_refresh.go b/service/codex_credential_refresh.go index 2e681ee..8193eb2 100644 --- a/service/codex_credential_refresh.go +++ b/service/codex_credential_refresh.go @@ -2,7 +2,6 @@ package service import ( "context" - "errors" "fmt" "strings" "time" @@ -16,28 +15,7 @@ type CodexCredentialRefreshOptions struct { ResetCaches bool } -type CodexOAuthKey struct { - IDToken string `json:"id_token,omitempty"` - AccessToken string `json:"access_token,omitempty"` - RefreshToken string `json:"refresh_token,omitempty"` - - AccountID string `json:"account_id,omitempty"` - LastRefresh string `json:"last_refresh,omitempty"` - Email string `json:"email,omitempty"` - Type string `json:"type,omitempty"` - Expired string `json:"expired,omitempty"` -} - -func parseCodexOAuthKey(raw string) (*CodexOAuthKey, error) { - if strings.TrimSpace(raw) == "" { - return nil, errors.New("codex channel: empty oauth key") - } - var key CodexOAuthKey - if err := common.Unmarshal([]byte(raw), &key); err != nil { - return nil, errors.New("codex channel: invalid oauth key json") - } - return &key, nil -} +type CodexOAuthKey = common.CodexOAuthCredential func RefreshCodexChannelCredential(ctx context.Context, channelID int, opts CodexCredentialRefreshOptions) (*CodexOAuthKey, *model.Channel, error) { ch, err := model.GetChannelById(channelID, true) @@ -51,7 +29,7 @@ func RefreshCodexChannelCredential(ctx context.Context, channelID int, opts Code return nil, nil, fmt.Errorf("channel type is not Codex") } - oauthKey, err := parseCodexOAuthKey(strings.TrimSpace(ch.Key)) + oauthKey, err := common.ParseCodexOAuthCredential(strings.TrimSpace(ch.Key)) if err != nil { return nil, nil, err } diff --git a/service/codex_credential_refresh_task.go b/service/codex_credential_refresh_task.go index 627ab92..aa1806a 100644 --- a/service/codex_credential_refresh_task.go +++ b/service/codex_credential_refresh_task.go @@ -93,7 +93,7 @@ func runCodexCredentialAutoRefreshOnce() { continue } - oauthKey, err := parseCodexOAuthKey(rawKey) + oauthKey, err := common.ParseCodexOAuthCredential(rawKey) if err != nil { continue } diff --git a/web/src/components/table/channels/modals/EditChannelModal.jsx b/web/src/components/table/channels/modals/EditChannelModal.jsx index 1b7a0ab..de16b09 100644 --- a/web/src/components/table/channels/modals/EditChannelModal.jsx +++ b/web/src/components/table/channels/modals/EditChannelModal.jsx @@ -26,6 +26,10 @@ import { showSuccess, verifyJSON, } from '../../../../helpers'; +import { + CODEX_CREDENTIAL_MODE, + detectCodexCredentialMode, +} from '../../../../helpers/codex'; import { useIsMobile } from '../../../../hooks/common/useIsMobile'; import { CHANNEL_OPTIONS } from '../../../../constants'; import { @@ -235,8 +239,14 @@ const EditChannelModal = (props) => { const [isIonetChannel, setIsIonetChannel] = useState(false); const [ionetMetadata, setIonetMetadata] = useState(null); const [codexOAuthModalVisible, setCodexOAuthModalVisible] = useState(false); + const [codexCredentialMode, setCodexCredentialMode] = useState( + CODEX_CREDENTIAL_MODE.API_KEY, + ); const [codexCredentialRefreshing, setCodexCredentialRefreshing] = useState(false); + const isCodexOAuthMode = + inputs.type === 57 && + codexCredentialMode === CODEX_CREDENTIAL_MODE.OAUTH; // 密钥显示状态 const [keyDisplayState, setKeyDisplayState] = useState({ @@ -537,6 +547,7 @@ const EditChannelModal = (props) => { setUseManualInput(false); if (value === 57) { + setCodexCredentialMode(CODEX_CREDENTIAL_MODE.API_KEY); setBatch(false); setMultiToSingle(false); setMultiKeyMode('random'); @@ -546,6 +557,8 @@ const EditChannelModal = (props) => { formApiRef.current.setValue('vertex_files', []); } setInputs((prev) => ({ ...prev, vertex_files: [] })); + } else { + setCodexCredentialMode(CODEX_CREDENTIAL_MODE.API_KEY); } } //setAutoBan @@ -689,6 +702,12 @@ const EditChannelModal = (props) => { data.base_url = 'https://ark.cn-beijing.volces.com'; } + if (data.type === 57) { + setCodexCredentialMode(detectCodexCredentialMode(data.key)); + } else { + setCodexCredentialMode(CODEX_CREDENTIAL_MODE.API_KEY); + } + setInputs(data); if (formApiRef.current) { formApiRef.current.setValues(data); @@ -922,6 +941,7 @@ const EditChannelModal = (props) => { }; const handleCodexOAuthGenerated = (key) => { + setCodexCredentialMode(CODEX_CREDENTIAL_MODE.OAUTH); handleInputChange('key', key); formatJsonField('key'); }; @@ -1073,6 +1093,7 @@ const EditChannelModal = (props) => { setIsEnterpriseAccount(false); // 重置豆包隐藏入口状态 setDoubaoApiEditUnlocked(false); + setCodexCredentialMode(CODEX_CREDENTIAL_MODE.API_KEY); doubaoApiClickCountRef.current = 0; // 清空表单中的key_mode字段 if (formApiRef.current) { @@ -1235,6 +1256,9 @@ const EditChannelModal = (props) => { } if (rawKey !== '') { + if (codexCredentialMode === CODEX_CREDENTIAL_MODE.API_KEY) { + localInputs.key = rawKey; + } else { if (!verifyJSON(rawKey)) { showInfo(t('密钥必须是合法的 JSON 格式!')); return; @@ -1260,6 +1284,7 @@ const EditChannelModal = (props) => { showInfo(t('密钥必须是合法的 JSON 格式!')); return; } + } } } @@ -2137,6 +2162,32 @@ const EditChannelModal = (props) => { <> {inputs.type === 57 ? ( <> +