diff --git a/.github/workflows/sync-release-to-gitcode.yml b/.github/workflows/sync-release-to-gitcode.yml index 96aff722d538..d311f499eece 100644 --- a/.github/workflows/sync-release-to-gitcode.yml +++ b/.github/workflows/sync-release-to-gitcode.yml @@ -5,24 +5,25 @@ permissions: contents: read on: - push: - tags: - - '*' - - '!*-alpha*' workflow_dispatch: inputs: tag_name: description: GitHub release tag to sync required: true type: string + sync_files: + description: Sync GitHub release files to GitCode + required: false + default: false + type: boolean concurrency: - group: gitcode-release-${{ inputs.tag_name || github.ref_name }} + group: gitcode-release-${{ inputs.tag_name }} cancel-in-progress: false jobs: - prepare-release-assets: - name: Prepare GitHub release assets + prepare-release: + name: Create or update GitCode release if: ${{ vars.GITCODE_REPOSITORY != '' }} runs-on: ubuntu-latest timeout-minutes: 90 @@ -30,80 +31,125 @@ jobs: release_tag: ${{ steps.release.outputs.tag }} release_body: ${{ steps.release.outputs.body }} release_prerelease: ${{ steps.release.outputs.prerelease }} - bootstrap_asset: ${{ steps.release.outputs.bootstrap_asset }} - release_asset_matrix: ${{ steps.release.outputs.asset_matrix }} - has_bootstrap_asset: ${{ steps.release.outputs.has_bootstrap_asset }} - has_matrix_assets: ${{ steps.release.outputs.has_matrix_assets }} + release_asset_matrix: ${{ steps.assets.outputs.matrix }} + has_release_assets: ${{ steps.assets.outputs.has_assets }} steps: - - name: Wait for GitHub release workflows - if: ${{ github.event_name == 'push' }} + - name: Create or update GitCode release + id: release env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITCODE_REPOSITORY: ${{ vars.GITCODE_REPOSITORY }} + GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} + RELEASE_TAG: ${{ inputs.tag_name }} run: | set -euo pipefail - wait_for_workflow() { - local workflow="$1" - local label="$2" - local last_state="" - - for attempt in $(seq 1 180); do - run_json="$( - gh run list \ - --repo "$GITHUB_REPOSITORY" \ - --workflow "$workflow" \ - --event push \ - --commit "$GITHUB_SHA" \ - --limit 1 \ - --json conclusion,status,url \ - --jq '.[0] // {}' - )" - status="$(jq -r '.status // empty' <<< "$run_json")" - conclusion="$(jq -r '.conclusion // empty' <<< "$run_json")" - state="${status:-not-found}/${conclusion:-pending}" - - if [[ "$state" != "$last_state" ]]; then - echo "$label: $state" - last_state="$state" - fi - - if [[ "$status" == "completed" ]]; then - if [[ "$conclusion" == "success" ]]; then - return 0 - fi - echo "::error::$label finished with conclusion: $conclusion" - return 1 - fi + release_json="$( + gh release view "$RELEASE_TAG" \ + --repo "$GITHUB_REPOSITORY" \ + --json body,isPrerelease + )" + release_body="$(jq -r '.body // ""' <<< "$release_json" | sed "s/'/’/g")" + release_prerelease="$(jq -r '.isPrerelease' <<< "$release_json")" + if [[ "$release_prerelease" == "true" ]]; then + release_status=pre + else + release_status=latest + fi - sleep 20 - done + gitcode_release_json="$(mktemp)" + gitcode_release_status="$( + curl -L -sS -o "$gitcode_release_json" -w "%{http_code}" \ + -H "PRIVATE-TOKEN: $GITCODE_TOKEN" \ + -H "Accept: application/json" \ + "https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases/tags/$RELEASE_TAG" + )" - echo "::error::Timed out waiting for $label" - return 1 - } + request_json="$(mktemp)" + case "$gitcode_release_status" in + 200) + request_method=PATCH + request_url="https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases/$RELEASE_TAG" + request_action=update + success_action=updated + jq -n \ + --arg tag_name "$RELEASE_TAG" \ + --arg name "$RELEASE_TAG" \ + --arg body "$release_body" \ + --arg release_status "$release_status" \ + '{ + tag_name: $tag_name, + name: $name, + body: $body, + release_status: $release_status + }' > "$request_json" + ;; + 404) + if ! git ls-remote --exit-code --tags \ + "https://gitcode.com/$GITCODE_REPOSITORY.git" \ + "refs/tags/$RELEASE_TAG" > /dev/null; then + echo "::error::Tag $RELEASE_TAG has not been mirrored to GitCode yet. Retry this workflow after the tag appears on GitCode." + exit 1 + fi - wait_for_workflow release.yml "Backend release" + request_method=POST + request_url="https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases" + request_action=create + success_action=created + jq -n \ + --arg tag_name "$RELEASE_TAG" \ + --arg name "$RELEASE_TAG" \ + --arg body "$release_body" \ + --arg release_status "$release_status" \ + '{ + tag_name: $tag_name, + name: $name, + body: $body, + release_status: $release_status + }' > "$request_json" + ;; + *) + echo "::error::Failed to inspect GitCode release. Response code: $gitcode_release_status" + cat "$gitcode_release_json" + exit 1 + ;; + esac - if [[ "$GITHUB_REF_NAME" != *-* ]]; then - wait_for_workflow electron-build.yml "Electron release" + gitcode_response_json="$(mktemp)" + request_status="$( + curl -L -sS -o "$gitcode_response_json" -w "%{http_code}" \ + -X "$request_method" \ + -H "PRIVATE-TOKEN: $GITCODE_TOKEN" \ + -H "Content-Type: application/json" \ + -H "Accept: application/json" \ + --data-binary "@$request_json" \ + "$request_url" + )" + if [[ "$request_status" != "200" ]]; then + echo "::error::Failed to $request_action GitCode release. Response code: $request_status" + cat "$gitcode_response_json" + exit 1 fi + echo "GitCode release $success_action successfully" + + delimiter="release-body-$(openssl rand -hex 16)" + { + echo "tag=$RELEASE_TAG" + echo "body<<$delimiter" + echo "$release_body" + echo "$delimiter" + echo "prerelease=$release_prerelease" + } >> "$GITHUB_OUTPUT" - name: Download GitHub release assets - id: release + if: ${{ inputs.sync_files }} env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GITCODE_REPOSITORY: ${{ vars.GITCODE_REPOSITORY }} - GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} - RELEASE_TAG: ${{ inputs.tag_name || github.ref_name }} + RELEASE_TAG: ${{ inputs.tag_name }} run: | set -euo pipefail mkdir release-assets - release_json="$( - gh release view "$RELEASE_TAG" \ - --repo "$GITHUB_REPOSITORY" \ - --json body,isPrerelease - )" gh release download "$RELEASE_TAG" \ --repo "$GITHUB_REPOSITORY" \ --dir release-assets @@ -128,6 +174,16 @@ jobs: find release-assets -maxdepth 1 -type f -print | sort + - name: Prepare GitCode release assets + id: assets + if: ${{ inputs.sync_files }} + env: + GITCODE_REPOSITORY: ${{ vars.GITCODE_REPOSITORY }} + GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} + RELEASE_TAG: ${{ inputs.tag_name }} + run: | + set -euo pipefail + local_assets="$( find release-assets -maxdepth 1 -type f -printf '%s\t%f\n' \ | sort -n -k1,1 \ @@ -142,23 +198,13 @@ jobs: -H "Accept: application/json" \ "https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases/tags/$RELEASE_TAG" )" + if [[ "$gitcode_release_status" != "200" ]]; then + echo "::error::Failed to inspect GitCode release assets. Response code: $gitcode_release_status" + cat "$gitcode_release_json" + exit 1 + fi - case "$gitcode_release_status" in - 200) - release_exists=true - existing_assets="$(jq -c '[.assets[]?.name]' "$gitcode_release_json")" - ;; - 404) - release_exists=false - existing_assets='[]' - ;; - *) - echo "::error::Failed to inspect GitCode release. Response code: $gitcode_release_status" - cat "$gitcode_release_json" - exit 1 - ;; - esac - + existing_assets="$(jq -c '[.assets[]?.name]' "$gitcode_release_json")" missing_assets="$( jq -cn \ --argjson local_assets "$local_assets" \ @@ -170,37 +216,11 @@ jobs: --argjson missing_assets "$missing_assets" \ '$local_assets - $missing_assets | .[] | "Skipping existing GitCode asset: \(.)"' - if [[ "$release_exists" == "false" ]]; then - bootstrap_asset="$(jq -r '.[0] // ""' <<< "$missing_assets")" - asset_matrix="$(jq -c '.[1:]' <<< "$missing_assets")" - else - bootstrap_asset="" - asset_matrix="$missing_assets" - fi - - if [[ -n "$bootstrap_asset" ]]; then - has_bootstrap_asset=true - else - has_bootstrap_asset=false - fi - has_matrix_assets="$(jq -r 'length > 0' <<< "$asset_matrix")" - - delimiter="release-body-$(openssl rand -hex 16)" - { - echo "tag=$RELEASE_TAG" - echo "bootstrap_asset=$bootstrap_asset" - echo "asset_matrix=$asset_matrix" - echo "has_bootstrap_asset=$has_bootstrap_asset" - echo "has_matrix_assets=$has_matrix_assets" - echo "body<<$delimiter" - # sync_to_gitcode embeds this input in a single-quoted shell string. - # Replace ASCII apostrophes so release notes cannot break its script. - jq -r '.body // ""' <<< "$release_json" | sed "s/'/’/g" - echo "$delimiter" - echo "prerelease=$(jq -r '.isPrerelease' <<< "$release_json")" - } >> "$GITHUB_OUTPUT" + echo "matrix=$missing_assets" >> "$GITHUB_OUTPUT" + echo "has_assets=$(jq -r 'length > 0' <<< "$missing_assets")" >> "$GITHUB_OUTPUT" - name: Upload release assets for GitCode + if: ${{ inputs.sync_files && steps.assets.outputs.has_assets == 'true' }} uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: gitcode-release-assets @@ -208,41 +228,23 @@ jobs: if-no-files-found: error retention-days: 1 - gitcode-release-bootstrap: - name: Create GitCode release - needs: prepare-release-assets - if: ${{ vars.GITCODE_REPOSITORY != '' && needs.prepare-release-assets.outputs.has_bootstrap_asset == 'true' }} - uses: nvdacn/sync_to_gitcode/.github/workflows/CreateReleaseOnGitCode.yaml@18b70112d0e62260bc54085028e5510bbc6323b2 - with: - artifact_name: gitcode-release-assets - gitcode_repository: ${{ vars.GITCODE_REPOSITORY }} - default_branch: main - tag_name: ${{ needs.prepare-release-assets.outputs.release_tag }} - body: ${{ needs.prepare-release-assets.outputs.release_body }} - prerelease: ${{ needs.prepare-release-assets.outputs.release_prerelease == 'true' }} - file_name: ${{ needs.prepare-release-assets.outputs.bootstrap_asset }} - secrets: - GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} - gitcode-release-assets: name: Publish GitCode asset (${{ matrix.file_name }}) - needs: - - prepare-release-assets - - gitcode-release-bootstrap - if: ${{ always() && vars.GITCODE_REPOSITORY != '' && needs.prepare-release-assets.result == 'success' && (needs.gitcode-release-bootstrap.result == 'success' || needs.gitcode-release-bootstrap.result == 'skipped') && needs.prepare-release-assets.outputs.has_matrix_assets == 'true' }} + needs: prepare-release + if: ${{ inputs.sync_files && needs.prepare-release.result == 'success' && needs.prepare-release.outputs.has_release_assets == 'true' }} strategy: fail-fast: false max-parallel: 4 matrix: - file_name: ${{ fromJSON(needs.prepare-release-assets.outputs.release_asset_matrix) }} + file_name: ${{ fromJSON(needs.prepare-release.outputs.release_asset_matrix) }} uses: nvdacn/sync_to_gitcode/.github/workflows/CreateReleaseOnGitCode.yaml@18b70112d0e62260bc54085028e5510bbc6323b2 with: artifact_name: gitcode-release-assets gitcode_repository: ${{ vars.GITCODE_REPOSITORY }} default_branch: main - tag_name: ${{ needs.prepare-release-assets.outputs.release_tag }} - body: ${{ needs.prepare-release-assets.outputs.release_body }} - prerelease: ${{ needs.prepare-release-assets.outputs.release_prerelease == 'true' }} + tag_name: ${{ needs.prepare-release.outputs.release_tag }} + body: ${{ needs.prepare-release.outputs.release_body }} + prerelease: ${{ needs.prepare-release.outputs.release_prerelease == 'true' }} file_name: ${{ matrix.file_name }} secrets: GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} diff --git a/common/body_storage.go b/common/body_storage.go index 094dbda36d3f..515d749e10c9 100644 --- a/common/body_storage.go +++ b/common/body_storage.go @@ -20,6 +20,21 @@ type BodyStorage interface { Size() int64 // IsDisk 是否是磁盘存储 IsDisk() bool + // NewReader returns an independent reader positioned at the start of the + // stored payload. Each call returns a reader with its own cursor, so + // callers (e.g. http.Request.GetBody) can replay the body concurrently + // with, or after, other readers without sharing seek state. Closing the + // returned reader releases only that reader, never the storage itself; + // after the storage has been closed, NewReader returns ErrStorageClosed. + NewReader() (io.ReadCloser, error) +} + +// ReplayableBody is an outbound request body that can report its byte size and +// create independent readers for transport-level retries. +type ReplayableBody interface { + io.Reader + Size() int64 + NewReader() (io.ReadCloser, error) } // ErrStorageClosed 存储已关闭错误 @@ -80,6 +95,18 @@ func (m *memoryStorage) Bytes() ([]byte, error) { return m.data, nil } +func (m *memoryStorage) NewReader() (io.ReadCloser, error) { + m.mu.Lock() + defer m.mu.Unlock() + if atomic.LoadInt32(&m.closed) == 1 { + return nil, ErrStorageClosed + } + // A fresh bytes.Reader over the shared immutable backing array: an + // independent cursor at zero copy cost. NopCloser keeps Close a no-op, so + // the storage lifecycle stays owned by whoever holds the storage itself. + return io.NopCloser(bytes.NewReader(m.data)), nil +} + func (m *memoryStorage) Size() int64 { return m.size } @@ -229,6 +256,24 @@ func (d *diskStorage) Bytes() ([]byte, error) { return data, nil } +func (d *diskStorage) NewReader() (io.ReadCloser, error) { + d.mu.Lock() + defer d.mu.Unlock() + if atomic.LoadInt32(&d.closed) == 1 { + return nil, ErrStorageClosed + } + // A separate file descriptor over the same cache file: an independent + // cursor at zero copy cost. Closing the returned reader closes only that + // descriptor; the storage keeps owning the primary descriptor and the + // file's lifetime. Readers opened before Close stay usable even after the + // file is unlinked, as the descriptor keeps the inode alive. + file, err := os.Open(d.filePath) + if err != nil { + return nil, fmt.Errorf("failed to open body cache file for replay: %w", err) + } + return file, nil +} + func (d *diskStorage) Size() int64 { return d.size } @@ -302,10 +347,27 @@ func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes return storage, nil } -// ReaderOnly wraps an io.Reader to hide io.Closer, preventing http.NewRequest -// from type-asserting io.ReadCloser and closing the underlying BodyStorage. -func ReaderOnly(r io.Reader) io.Reader { - return struct{ io.Reader }{r} +type replayableBodyReader struct { + storage BodyStorage +} + +func (r replayableBodyReader) Read(p []byte) (int, error) { + return r.storage.Read(p) +} + +func (r replayableBodyReader) Size() int64 { + return r.storage.Size() +} + +func (r replayableBodyReader) NewReader() (io.ReadCloser, error) { + return r.storage.NewReader() +} + +// NewReplayableBodyReader exposes the replay capabilities of storage without +// exposing io.Closer. This keeps ownership of the storage lifecycle with the +// caller instead of allowing net/http to close it as the request body. +func NewReplayableBodyReader(storage BodyStorage) ReplayableBody { + return replayableBodyReader{storage: storage} } // CleanupOldCacheFiles 清理旧的缓存文件（用于启动时清理残留） diff --git a/common/body_storage_test.go b/common/body_storage_test.go new file mode 100644 index 000000000000..c3f877a60ff4 --- /dev/null +++ b/common/body_storage_test.go @@ -0,0 +1,37 @@ +package common + +import ( + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewReplayableBodyReaderKeepsStorageLifecycleWithCaller(t *testing.T) { + payload := []byte(`{"model":"test-model","input":"hello"}`) + storage, err := CreateBodyStorage(payload) + require.NoError(t, err) + defer storage.Close() + + body := NewReplayableBodyReader(storage) + assert.EqualValues(t, len(payload), body.Size()) + _, exposesCloser := any(body).(io.Closer) + assert.False(t, exposesCloser, "the request body must not expose the storage closer") + + req, err := http.NewRequest(http.MethodPost, "https://example.com", body) + require.NoError(t, err) + require.NoError(t, req.Body.Close()) + + replayBody, err := body.NewReader() + require.NoError(t, err, "closing the HTTP request body must not close the storage") + replay, err := io.ReadAll(replayBody) + require.NoError(t, err) + require.NoError(t, replayBody.Close()) + assert.Equal(t, payload, replay) + + require.NoError(t, storage.Close()) + _, err = body.NewReader() + require.ErrorIs(t, err, ErrStorageClosed) +} diff --git a/controller/channel-test.go b/controller/channel-test.go index f6e6bd7f1163..a13ed5d607b2 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -151,6 +151,11 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te requestPath = "/v1/responses/compact" } } + // Gemini 原生流式通过 URL action（:streamGenerateContent）表达而非请求体字段， + // GeminiChatRequest.IsStream 依据请求 URL 判定，合成请求路径需与生产入口保持一致 + if isStream && constant.EndpointType(endpointType) == constant.EndpointTypeGemini { + requestPath = strings.Replace(requestPath, ":generateContent", ":streamGenerateContent", 1) + } if strings.HasPrefix(requestPath, "/v1/responses/compact") { testModel = ratio_setting.WithCompactModelSuffix(testModel) } @@ -371,14 +376,18 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te } } default: - // Chat/Completion 等其他请求类型 - if generalReq, ok := request.(*dto.GeneralOpenAIRequest); ok { - convertedRequest, err = adaptor.ConvertOpenAIRequest(c, info, generalReq) - } else { + switch req := request.(type) { + case *dto.GeneralOpenAIRequest: + convertedRequest, err = adaptor.ConvertOpenAIRequest(c, info, req) + case *dto.ClaudeRequest: + convertedRequest, err = adaptor.ConvertClaudeRequest(c, info, req) + case *dto.GeminiChatRequest: + convertedRequest, err = adaptor.ConvertGeminiRequest(c, info, req) + default: return testResult{ context: c, - localErr: errors.New("invalid general request type"), - newAPIError: types.NewError(errors.New("invalid general request type"), types.ErrorCodeConvertRequestFailed), + localErr: errors.New("invalid chat request type"), + newAPIError: types.NewError(errors.New("invalid chat request type"), types.ErrorCodeConvertRequestFailed), } } } @@ -733,12 +742,31 @@ func buildTestRequest(model string, endpointType string, channel *model.Channel, Model: model, Input: testResponsesInput, } - case constant.EndpointTypeAnthropic, constant.EndpointTypeGemini, constant.EndpointTypeOpenAI: - // 返回 GeneralOpenAIRequest - maxTokens := uint(16) - if constant.EndpointType(endpointType) == constant.EndpointTypeGemini { - maxTokens = 3000 + case constant.EndpointTypeAnthropic: + return &dto.ClaudeRequest{ + Model: model, + Stream: lo.ToPtr(isStream), + MaxTokens: lo.ToPtr(uint(16)), + Messages: []dto.ClaudeMessage{ + { + Role: "user", + Content: "hi", + }, + }, + } + case constant.EndpointTypeGemini: + return &dto.GeminiChatRequest{ + Contents: []dto.GeminiChatContent{ + { + Role: "user", + Parts: []dto.GeminiPart{{Text: "hi"}}, + }, + }, + GenerationConfig: dto.GeminiChatGenerationConfig{ + MaxOutputTokens: lo.ToPtr(uint(3000)), + }, } + case constant.EndpointTypeOpenAI: req := &dto.GeneralOpenAIRequest{ Model: model, Stream: lo.ToPtr(isStream), @@ -748,7 +776,7 @@ func buildTestRequest(model string, endpointType string, channel *model.Channel, Content: "hi", }, }, - MaxTokens: lo.ToPtr(maxTokens), + MaxTokens: lo.ToPtr(uint(16)), } if isStream { req.StreamOptions = &dto.StreamOptions{IncludeUsage: true} @@ -1021,6 +1049,9 @@ func selectChannelsForAutomaticTest(channels []*model.Channel, mode string) []*m if channel.Status == common.ChannelStatusManuallyDisabled { continue } + if mode == operation_setting.ChannelTestModeAutoBanOnly && !channel.GetAutoBan() { + continue + } if mode == operation_setting.ChannelTestModePassiveRecovery && channel.Status != common.ChannelStatusAutoDisabled { continue } diff --git a/controller/channel_test_internal_test.go b/controller/channel_test_internal_test.go index abbd2d237070..aa1a6ca7843a 100644 --- a/controller/channel_test_internal_test.go +++ b/controller/channel_test_internal_test.go @@ -265,12 +265,19 @@ func TestBuildTestLogOtherInjectsTieredInfo(t *testing.T) { }, } + requestRules := []billingexpr.RequestRuleTrace{{ + Cond: `param("service_tier") == "fast"`, + Multiplier: 2, + Matched: true, + }} other := buildTestLogOther(ctx, info, priceData, usage, &billingexpr.TieredResult{ - MatchedTier: "base", + MatchedTier: "base", + RequestRules: requestRules, }) require.Equal(t, "tiered_expr", other["billing_mode"]) require.Equal(t, "base", other["matched_tier"]) + require.Equal(t, requestRules, other["request_rules"]) require.NotEmpty(t, other["expr_b64"]) } @@ -312,6 +319,24 @@ func TestSelectChannelsForAutomaticTestScheduledSkipsManualDisabled(t *testing.T require.Equal(t, 2, selected[1].Id) } +func TestSelectChannelsForAutomaticTestAutoBanOnlyUsesEligibleChannels(t *testing.T) { + autoBanEnabled := 1 + autoBanDisabled := 0 + channels := []*model.Channel{ + {Id: 1, Status: common.ChannelStatusEnabled, AutoBan: &autoBanEnabled}, + {Id: 2, Status: common.ChannelStatusEnabled, AutoBan: &autoBanDisabled}, + {Id: 3, Status: common.ChannelStatusAutoDisabled, AutoBan: &autoBanEnabled}, + {Id: 4, Status: common.ChannelStatusManuallyDisabled, AutoBan: &autoBanEnabled}, + {Id: 5, Status: common.ChannelStatusEnabled}, + } + + selected := selectChannelsForAutomaticTest(channels, operation_setting.ChannelTestModeAutoBanOnly) + + require.Len(t, selected, 2) + require.Equal(t, 1, selected[0].Id) + require.Equal(t, 3, selected[1].Id) +} + func TestTestAllChannelsRejectsExistingActiveTask(t *testing.T) { db := setupModelListControllerTestDB(t) require.NoError(t, db.AutoMigrate(&model.SystemTask{}, &model.SystemTaskLock{})) diff --git a/controller/user.go b/controller/user.go index 59aaf5847238..9b8d931ec1f8 100644 --- a/controller/user.go +++ b/controller/user.go @@ -399,11 +399,6 @@ func GetUser(c *gin.Context) { func GenerateAccessToken(c *gin.Context) { id := c.GetInt("id") - user, err := model.GetUserById(id, true) - if err != nil { - common.ApiError(c, err) - return - } // get rand int 28-32 randI := common.GetRandomInt(4) key, err := common.GenerateRandomKey(29 + randI) @@ -412,14 +407,12 @@ func GenerateAccessToken(c *gin.Context) { common.SysLog("failed to generate key: " + err.Error()) return } - user.SetAccessToken(key) - - if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 { + if model.DB.Where("access_token = ?", key).First(&model.User{}).RowsAffected != 0 { common.ApiErrorI18n(c, i18n.MsgUuidDuplicate) return } - if err := user.Update(false); err != nil { + if err := model.UpdateUserAccessToken(id, key); err != nil { common.ApiError(c, err) return } @@ -427,7 +420,7 @@ func GenerateAccessToken(c *gin.Context) { c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", - "data": user.AccessToken, + "data": key, }) return } diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go index 350ee8ec0718..be1820ffee99 100644 --- a/middleware/rate-limit.go +++ b/middleware/rate-limit.go @@ -178,6 +178,17 @@ func CriticalRateLimit() func(c *gin.Context) { return defNext } +func UserCriticalRateLimit(scope string) func(c *gin.Context) { + if !common.CriticalRateLimitEnable { + return defNext + } + return userRateLimitFactory( + common.CriticalRateLimitNum, + common.CriticalRateLimitDuration, + "UC:"+scope, + ) +} + func DownloadRateLimit() func(c *gin.Context) { return rateLimitFactory(common.DownloadRateLimitNum, common.DownloadRateLimitDuration, "DW") } diff --git a/model/user.go b/model/user.go index b25de5e75efa..eb4ea08642d5 100644 --- a/model/user.go +++ b/model/user.go @@ -139,6 +139,22 @@ func (user *User) SetAccessToken(token string) { user.AccessToken = &token } +// UpdateUserAccessToken rotates a dashboard personal access token without +// writing a stale user snapshot back over concurrently updated fields. +func UpdateUserAccessToken(id int, token string) error { + if id == 0 { + return errors.New("id 为空！") + } + result := DB.Model(&User{}).Where("id = ?", id).Update("access_token", token) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return gorm.ErrRecordNotFound + } + return nil +} + func (user *User) GetSetting() dto.UserSetting { setting := dto.UserSetting{} if user.Setting != "" { @@ -489,15 +505,19 @@ func HardDeleteUserById(id int) error { return user.HardDelete() } -func inviteUser(inviterId int) (err error) { - user, err := GetUserById(inviterId, true) - if err != nil { - return err +func inviteUser(inviterId int) error { + result := DB.Model(&User{}).Where("id = ?", inviterId).Updates(map[string]interface{}{ + "aff_count": gorm.Expr("aff_count + ?", 1), + "aff_quota": gorm.Expr("aff_quota + ?", common.QuotaForInviter), + "aff_history": gorm.Expr("aff_history + ?", common.QuotaForInviter), + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return gorm.ErrRecordNotFound } - user.AffCount++ - user.AffQuota += common.QuotaForInviter - user.AffHistoryQuota += common.QuotaForInviter - return DB.Save(user).Error + return nil } func (user *User) TransferAffQuotaToQuota(quota int) error { @@ -514,7 +534,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error { defer tx.Rollback() // 确保在函数退出时事务能回滚 // 加锁查询用户以确保数据一致性 - err := lockForUpdate(tx).First(&user, user.Id).Error + err := lockForUpdate(tx).First(user, user.Id).Error if err != nil { return err } @@ -748,7 +768,16 @@ func (user *User) UpdateWithTx(tx *gorm.DB, updatePassword bool) error { return err } } - if err = tx.Model(¤t).Omit("quota", "used_quota", "request_count", "auth_version").Updates(newUser).Error; err != nil { + if err = tx.Model(¤t).Omit( + "access_token", + "quota", + "used_quota", + "request_count", + "aff_count", + "aff_quota", + "aff_history", + "auth_version", + ).Updates(newUser).Error; err != nil { return err } return tx.First(user, user.Id).Error diff --git a/model/user_update_test.go b/model/user_update_test.go index c68d8f33c85a..8e69ba1ebba2 100644 --- a/model/user_update_test.go +++ b/model/user_update_test.go @@ -27,19 +27,23 @@ func setupUserUpdateTestState(t *testing.T) { }) } -func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { +func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing.T) { setupUserUpdateTestState(t) user := User{ - Id: 1, - Username: "quota-race-user", - Password: "password", - DisplayName: "before", - Status: common.UserStatusEnabled, - Quota: 1000, - UsedQuota: 20, - RequestCount: 3, + Id: 1, + Username: "quota-race-user", + Password: "password", + DisplayName: "before", + Status: common.UserStatusEnabled, + Quota: 1000, + UsedQuota: 20, + RequestCount: 3, + AffCount: 2, + AffQuota: 800, + AffHistoryQuota: 1200, } + user.SetAccessToken("old-token") require.NoError(t, DB.Create(&user).Error) staleUser, err := GetUserById(user.Id, true) @@ -49,6 +53,10 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { "quota": gorm.Expr("quota - ?", 400), "used_quota": gorm.Expr("used_quota + ?", 400), "request_count": gorm.Expr("request_count + ?", 1), + "aff_count": gorm.Expr("aff_count + ?", 1), + "aff_quota": gorm.Expr("aff_quota - ?", 500), + "aff_history": gorm.Expr("aff_history + ?", 500), + "access_token": "rotated-token", }).Error) staleUser.DisplayName = "after" @@ -60,6 +68,63 @@ func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { assert.Equal(t, 600, got.Quota) assert.Equal(t, 420, got.UsedQuota) assert.Equal(t, 4, got.RequestCount) + assert.Equal(t, 3, got.AffCount) + assert.Equal(t, 300, got.AffQuota) + assert.Equal(t, 1700, got.AffHistoryQuota) + assert.Equal(t, "rotated-token", got.GetAccessToken()) +} + +func TestUpdateUserAccessTokenOnlyUpdatesAccessToken(t *testing.T) { + setupUserUpdateTestState(t) + + user := User{ + Id: 2, + Username: "token-rotation-user", + Password: "password", + DisplayName: "before", + Status: common.UserStatusEnabled, + Quota: 1000, + AffQuota: 800, + AffHistoryQuota: 1200, + } + require.NoError(t, DB.Create(&user).Error) + + require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{ + "quota": gorm.Expr("quota + ?", 500), + "aff_quota": gorm.Expr("aff_quota - ?", 500), + "display_name": "concurrent-update", + }).Error) + + require.NoError(t, UpdateUserAccessToken(user.Id, "rotated-token")) + + var got User + require.NoError(t, DB.First(&got, user.Id).Error) + assert.Equal(t, "rotated-token", got.GetAccessToken()) + assert.Equal(t, "concurrent-update", got.DisplayName) + assert.Equal(t, 1500, got.Quota) + assert.Equal(t, 300, got.AffQuota) + assert.Equal(t, 1200, got.AffHistoryQuota) +} + +func TestUpdateUserAccessTokenRejectsSoftDeletedUser(t *testing.T) { + setupUserUpdateTestState(t) + + user := User{ + Id: 3, + Username: "deleted-token-rotation-user", + Password: "password", + Status: common.UserStatusEnabled, + } + user.SetAccessToken("old-token") + require.NoError(t, DB.Create(&user).Error) + require.NoError(t, DB.Delete(&user).Error) + + err := UpdateUserAccessToken(user.Id, "orphaned-token") + require.ErrorIs(t, err, gorm.ErrRecordNotFound) + + var got User + require.NoError(t, DB.Unscoped().First(&got, user.Id).Error) + assert.Equal(t, "old-token", got.GetAccessToken()) } func TestUpdateUserSettingOnlyUpdatesSetting(t *testing.T) { diff --git a/pkg/billingexpr/billingexpr_test.go b/pkg/billingexpr/billingexpr_test.go index 7b59ed258681..90485571a5fb 100644 --- a/pkg/billingexpr/billingexpr_test.go +++ b/pkg/billingexpr/billingexpr_test.go @@ -5,6 +5,8 @@ import ( "testing" "github.com/QuantumNous/new-api/pkg/billingexpr" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) // --------------------------------------------------------------------------- @@ -228,10 +230,11 @@ func TestRequestProbeMissingFieldReturnsNil(t *testing.T) { } } -func TestRequestProbeMultipleRulesMultiply(t *testing.T) { - cost, _, err := billingexpr.RunExprWithRequest( - `(param("service_tier") == "fast" ? 2 : 1) * (has(header("anthropic-beta"), "fast-mode-2026-02-01") ? 2.5 : 1)`, - billingexpr.TokenParams{}, +func TestRequestProbeMultipleRulesTraceAllFactors(t *testing.T) { + exprStr := `(tier("base", p * 2)) * (param("service_tier") == "fast" ? 2 : 1) * (has(header("anthropic-beta"), "fast-mode-2026-02-01") ? 2.5 : 1)` + cost, trace, err := billingexpr.RunExprWithRequest( + exprStr, + billingexpr.TokenParams{P: 10}, billingexpr.RequestInput{ Headers: map[string]string{ "Anthropic-Beta": "fast-mode-2026-02-01", @@ -239,12 +242,62 @@ func TestRequestProbeMultipleRulesMultiply(t *testing.T) { Body: []byte(`{"service_tier":"fast"}`), }, ) - if err != nil { - t.Fatal(err) - } - if math.Abs(cost-5) > 1e-6 { - t.Errorf("cost = %f, want 5", cost) - } + + require.NoError(t, err) + assert.InDelta(t, 100, cost, 1e-6) + assert.Equal(t, "base", trace.MatchedTier) + assert.Equal(t, []billingexpr.RequestRuleTrace{ + {Cond: `param("service_tier") == "fast"`, Multiplier: 2, Matched: true}, + {Cond: `has(header("anthropic-beta"), "fast-mode-2026-02-01")`, Multiplier: 2.5, Matched: true}, + }, trace.RequestRules) +} + +func TestRequestProbeTraceIncludesUnmatchedFactors(t *testing.T) { + exprStr := `(tier("base", p * 2)) * (param("service_tier") == "fast" ? 2 : 1) * (has(header("anthropic-beta"), "fast-mode") ? 2.5 : 1)` + cost, trace, err := billingexpr.RunExprWithRequest( + exprStr, + billingexpr.TokenParams{P: 10}, + billingexpr.RequestInput{Body: []byte(`{"service_tier":"fast"}`)}, + ) + + require.NoError(t, err) + assert.InDelta(t, 40, cost, 1e-6) + assert.Equal(t, []billingexpr.RequestRuleTrace{ + {Cond: `param("service_tier") == "fast"`, Multiplier: 2, Matched: true}, + {Cond: `has(header("anthropic-beta"), "fast-mode")`, Multiplier: 2.5, Matched: false}, + }, trace.RequestRules) +} + +func TestRequestProbeTracePreservesIntegerConditionalType(t *testing.T) { + cost, trace, err := billingexpr.RunExprWithRequest( + `5 % (param("service_tier") == "fast" ? 2 : 1)`, + billingexpr.TokenParams{}, + billingexpr.RequestInput{Body: []byte(`{"service_tier":"fast"}`)}, + ) + + require.NoError(t, err) + assert.Equal(t, float64(1), cost) + assert.Equal(t, []billingexpr.RequestRuleTrace{ + {Cond: `param("service_tier") == "fast"`, Multiplier: 2, Matched: true}, + }, trace.RequestRules) +} + +func TestRequestProbeNonUnitFallbackIsNotTraced(t *testing.T) { + cost, trace, err := billingexpr.RunExprWithRequest( + `10 * (param("service_tier") == "fast" ? 2 : 1.5)`, + billingexpr.TokenParams{}, + billingexpr.RequestInput{Body: []byte(`{"service_tier":"standard"}`)}, + ) + + require.NoError(t, err) + assert.InDelta(t, 15, cost, 1e-6) + assert.Empty(t, trace.RequestRules) +} + +func TestRequestProbeInternalTraceFunctionIsReserved(t *testing.T) { + _, err := billingexpr.CompileFromCache(`_trace(0, true, 5.0)`) + + require.ErrorContains(t, err, `identifier "_trace" is reserved for internal use`) } func TestCeilFloor(t *testing.T) { diff --git a/pkg/billingexpr/compile.go b/pkg/billingexpr/compile.go index a6c7b8f7221a..72a22189f6d9 100644 --- a/pkg/billingexpr/compile.go +++ b/pkg/billingexpr/compile.go @@ -16,6 +16,11 @@ const maxCacheSize = 256 // DefaultExprVersion is used when an expression string has no version prefix. const DefaultExprVersion = 1 +const ( + requestRuleTraceFunction = "_trace" + requestRuleTraceIntFunction = "_trace_int" +) + // ParseExprVersion extracts the version tag and body from an expression string. // Format: "v1:tier(...)" → version=1, body="tier(...)". // No prefix defaults to DefaultExprVersion. @@ -26,10 +31,88 @@ func ParseExprVersion(exprStr string) (version int, body string) { return DefaultExprVersion, exprStr } +// requestRulePatcher adds trace side effects to existing request multipliers +// without changing the stored expression or its numeric result. +type requestRulePatcher struct { + requestRules []RequestRuleTrace + restrictedIdentifier string +} + +func (p *requestRulePatcher) Visit(node *ast.Node) { + if identifier, ok := (*node).(*ast.IdentifierNode); ok { + switch identifier.Value { + case requestRuleTraceFunction, requestRuleTraceIntFunction: + p.restrictedIdentifier = identifier.Value + } + return + } + + conditional, ok := (*node).(*ast.ConditionalNode) + if !ok || !conditional.Ternary || !usesRequestProbe(conditional.Cond) { + return + } + multiplier, ok := requestRuleNumber(conditional.Exp1) + fallback, fallbackOK := requestRuleNumber(conditional.Exp2) + if !ok || !fallbackOK || fallback != 1 { + return + } + + ruleIndex := len(p.requestRules) + p.requestRules = append(p.requestRules, RequestRuleTrace{ + Cond: conditional.Cond.String(), + Multiplier: multiplier, + }) + + traceFunction := requestRuleTraceFunction + var multiplierNode ast.Node = &ast.FloatNode{Value: multiplier} + if _, multiplierIsInt := conditional.Exp1.(*ast.IntegerNode); multiplierIsInt { + if _, fallbackIsInt := conditional.Exp2.(*ast.IntegerNode); fallbackIsInt { + traceFunction = requestRuleTraceIntFunction + multiplierNode = conditional.Exp1 + } + } + + ast.Patch(node, &ast.CallNode{ + Callee: &ast.IdentifierNode{Value: traceFunction}, + Arguments: []ast.Node{ + &ast.IntegerNode{Value: ruleIndex}, + conditional.Cond, + multiplierNode, + }, + }) +} + +func requestRuleNumber(node ast.Node) (float64, bool) { + switch value := node.(type) { + case *ast.IntegerNode: + return float64(value.Value), true + case *ast.FloatNode: + return value.Value, true + default: + return 0, false + } +} + +func usesRequestProbe(node ast.Node) bool { + return ast.Find(node, func(node ast.Node) bool { + identifier, ok := node.(*ast.IdentifierNode) + if !ok { + return false + } + switch identifier.Value { + case "param", "header", "hour", "minute", "weekday", "month", "day": + return true + default: + return false + } + }) != nil +} + type cachedEntry struct { - prog *vm.Program - usedVars map[string]bool - version int + prog *vm.Program + usedVars map[string]bool + requestRules []RequestRuleTrace + version int } var ( @@ -39,30 +122,32 @@ var ( // compileEnvPrototypeV1 is the v1 type-checking prototype used at compile time. var compileEnvPrototypeV1 = map[string]interface{}{ - "p": float64(0), - "c": float64(0), - "len": float64(0), - "cr": float64(0), - "cc": float64(0), - "cc1h": float64(0), - "img": float64(0), - "img_o": float64(0), - "ai": float64(0), - "ao": float64(0), - "tier": func(string, float64) float64 { return 0 }, - "header": func(string) string { return "" }, - "param": func(string) interface{} { return nil }, - "has": func(interface{}, string) bool { return false }, - "hour": func(string) int { return 0 }, - "minute": func(string) int { return 0 }, - "weekday": func(string) int { return 0 }, - "month": func(string) int { return 0 }, - "day": func(string) int { return 0 }, - "max": math.Max, - "min": math.Min, - "abs": math.Abs, - "ceil": math.Ceil, - "floor": math.Floor, + "p": float64(0), + "c": float64(0), + "len": float64(0), + "cr": float64(0), + "cc": float64(0), + "cc1h": float64(0), + "img": float64(0), + "img_o": float64(0), + "ai": float64(0), + "ao": float64(0), + "tier": func(string, float64) float64 { return 0 }, + "_trace": func(int, bool, float64) float64 { return 1 }, + "_trace_int": func(int, bool, int) int { return 1 }, + "header": func(string) string { return "" }, + "param": func(string) interface{} { return nil }, + "has": func(interface{}, string) bool { return false }, + "hour": func(string) int { return 0 }, + "minute": func(string) int { return 0 }, + "weekday": func(string) int { return 0 }, + "month": func(string) int { return 0 }, + "day": func(string) int { return 0 }, + "max": math.Max, + "min": math.Min, + "abs": math.Abs, + "ceil": math.Ceil, + "floor": math.Floor, } func getCompileEnv(version int) map[string]interface{} { @@ -85,29 +170,45 @@ func CompileFromCacheByHash(exprStr, hash string) (*vm.Program, error) { } func compileFromCacheByHash(exprStr, hash string) (*vm.Program, error) { + entry, err := compileEntryFromCacheByHash(exprStr, hash) + if err != nil { + return nil, err + } + return entry.prog, nil +} + +func compileEntryFromCacheByHash(exprStr, hash string) (*cachedEntry, error) { cacheMu.RLock() if entry, ok := cache[hash]; ok { cacheMu.RUnlock() - return entry.prog, nil + return entry, nil } cacheMu.RUnlock() version, body := ParseExprVersion(exprStr) - prog, err := expr.Compile(body, expr.Env(getCompileEnv(version)), expr.AsFloat64()) + patcher := &requestRulePatcher{} + prog, err := expr.Compile(body, expr.Env(getCompileEnv(version)), expr.Patch(patcher), expr.AsFloat64()) + if patcher.restrictedIdentifier != "" { + return nil, fmt.Errorf("expr compile error: identifier %q is reserved for internal use", patcher.restrictedIdentifier) + } if err != nil { return nil, fmt.Errorf("expr compile error: %w", err) } - vars := extractUsedVars(prog) - + entry := &cachedEntry{ + prog: prog, + usedVars: extractUsedVars(prog), + requestRules: patcher.requestRules, + version: version, + } cacheMu.Lock() if len(cache) >= maxCacheSize { cache = make(map[string]*cachedEntry, 64) } - cache[hash] = &cachedEntry{prog: prog, usedVars: vars, version: version} + cache[hash] = entry cacheMu.Unlock() - return prog, nil + return entry, nil } // ExprVersion returns the version of a cached expression. Returns DefaultExprVersion @@ -132,6 +233,10 @@ func extractUsedVars(prog *vm.Program) map[string]bool { node := prog.Node() ast.Find(node, func(n ast.Node) bool { if id, ok := n.(*ast.IdentifierNode); ok { + switch id.Value { + case requestRuleTraceFunction, requestRuleTraceIntFunction: + return false + } vars[id.Value] = true } return false diff --git a/pkg/billingexpr/expr.md b/pkg/billingexpr/expr.md index c52f2e614e5e..f192c74ee6a6 100644 --- a/pkg/billingexpr/expr.md +++ b/pkg/billingexpr/expr.md @@ -116,7 +116,31 @@ Request-conditional multipliers are appended to the expression after a `|||` sep tier("base", p * 5 + c * 25)|||when(header("anthropic-beta") has "fast-mode") * 6 ``` -These are parsed and applied separately by the request rule system. +These factors are stored as ordinary multiplication in the final expression (for example, `(tier(...)) * (condition ? 6 : 1)`) and run in the same billing program. + +### Request Rule Tracing + +At compile time, the engine instruments ternary factors with this exact shape: + +``` + ? : 1 +``` + +The condition must reference at least one request probe (`param`, `header`, `hour`, `minute`, `weekday`, `month`, or `day`). Both branches must be numeric literals and the fallback must equal `1`. Other conditionals, including `(condition ? 2 : 1.5)`, are evaluated normally but are not traced. Integer-only factors use an integer-preserving trace callback, so instrumentation does not change expressions that require an integer operand (for example, `%`). The internal trace callback names are reserved and cannot be used in stored expressions. + +The compiled cache stores the canonical condition and multiplier for every instrumented node. Each run starts with the full detected rule list marked as unmatched; callbacks mark rules that actually evaluate true. Rules skipped by normal expression short-circuiting remain unmatched. This keeps the expression's numeric result unchanged and avoids reparsing it on each request. + +Settlement copies the actual run's traces into the consume log as: + +```json +{ + "request_rules": [ + { "cond": "param(\"service_tier\") == \"fast\"", "multiplier": 2, "matched": true } + ] +} +``` + +The usage-log UI treats `request_rules` as the authoritative rule list and renders directly from it. It parses `cond` only to produce a friendly label and falls back to the canonical condition text when that parser does not recognize the condition. Pricing pages without log context continue to parse the stored expression for display. --- @@ -182,9 +206,9 @@ After the upstream response returns with actual token usage: **Files**: `service/log_info_generate.go`, `web/src/helpers/render.jsx` -Backend: `InjectTieredBillingInfo()` adds `billing_mode`, `expr_b64` (base64 expression), and `matched_tier` to the log's `other` JSON. +Backend: `InjectTieredBillingInfo()` adds `billing_mode`, `expr_b64` (base64 expression), `matched_tier`, and the structured `request_rules` trace list to the log's `other` JSON. -Frontend: Detects `billing_mode === "tiered_expr"`, decodes `expr_b64`, parses tiers via shared `parseTiersFromExpr()`, and renders pricing breakdown. +Frontend: Detects `billing_mode === "tiered_expr"`, decodes `expr_b64`, parses tiers via shared `parseTiersFromExpr()`, and renders request multipliers from `request_rules` when present. Without log traces, it falls back to parsing the stored expression. --- diff --git a/pkg/billingexpr/run.go b/pkg/billingexpr/run.go index 7c0f2ecdd803..397099d557d8 100644 --- a/pkg/billingexpr/run.go +++ b/pkg/billingexpr/run.go @@ -26,11 +26,11 @@ func RunExpr(exprStr string, params TokenParams) (float64, TraceResult, error) { } func RunExprWithRequest(exprStr string, params TokenParams, request RequestInput) (float64, TraceResult, error) { - prog, err := CompileFromCache(exprStr) + entry, err := compileEntryFromCacheByHash(exprStr, ExprHashString(exprStr)) if err != nil { return 0, TraceResult{}, err } - return runProgram(prog, params, request) + return runProgram(entry.prog, entry.requestRules, params, request) } // RunExprByHash is like RunExpr but accepts a pre-computed hash for the cache @@ -41,15 +41,17 @@ func RunExprByHash(exprStr, hash string, params TokenParams) (float64, TraceResu } func RunExprByHashWithRequest(exprStr, hash string, params TokenParams, request RequestInput) (float64, TraceResult, error) { - prog, err := CompileFromCacheByHash(exprStr, hash) + entry, err := compileEntryFromCacheByHash(exprStr, hash) if err != nil { return 0, TraceResult{}, err } - return runProgram(prog, params, request) + return runProgram(entry.prog, entry.requestRules, params, request) } -func runProgram(prog *vm.Program, params TokenParams, request RequestInput) (float64, TraceResult, error) { - trace := TraceResult{} +func runProgram(prog *vm.Program, requestRules []RequestRuleTrace, params TokenParams, request RequestInput) (float64, TraceResult, error) { + trace := TraceResult{ + RequestRules: append([]RequestRuleTrace(nil), requestRules...), + } headers := normalizeHeaders(request.Headers) env := map[string]interface{}{ @@ -68,6 +70,24 @@ func runProgram(prog *vm.Program, params TokenParams, request RequestInput) (flo trace.Cost = value return value }, + requestRuleTraceFunction: func(ruleIndex int, matched bool, multiplier float64) float64 { + if matched && ruleIndex >= 0 && ruleIndex < len(trace.RequestRules) { + trace.RequestRules[ruleIndex].Matched = true + } + if matched { + return multiplier + } + return 1 + }, + requestRuleTraceIntFunction: func(ruleIndex int, matched bool, multiplier int) int { + if matched && ruleIndex >= 0 && ruleIndex < len(trace.RequestRules) { + trace.RequestRules[ruleIndex].Matched = true + } + if matched { + return multiplier + } + return 1 + }, "header": func(key string) string { return headers[strings.ToLower(strings.TrimSpace(key))] }, diff --git a/pkg/billingexpr/settle.go b/pkg/billingexpr/settle.go index f8cf937ae34b..a1c3267686f9 100644 --- a/pkg/billingexpr/settle.go +++ b/pkg/billingexpr/settle.go @@ -32,6 +32,7 @@ func ComputeTieredQuotaWithRequest(snap *BillingSnapshot, params TokenParams, re ActualQuotaBeforeGroup: quotaBeforeGroup, ActualQuotaAfterGroup: afterGroup, MatchedTier: trace.MatchedTier, + RequestRules: trace.RequestRules, CrossedTier: crossed, Clamp: clamp, }, nil diff --git a/pkg/billingexpr/types.go b/pkg/billingexpr/types.go index e95399b4fed0..1bdf6834bd57 100644 --- a/pkg/billingexpr/types.go +++ b/pkg/billingexpr/types.go @@ -28,12 +28,18 @@ type TokenParams struct { AO float64 // audio output tokens } -// TraceResult holds side-channel info captured by the tier() function -// during Expr execution. This replaces the old Breakdown mechanism — -// the Expr itself is the single source of truth for billing logic. +// RequestRuleTrace describes one request-dependent multiplier detected at compile time. +type RequestRuleTrace struct { + Cond string `json:"cond"` + Multiplier float64 `json:"multiplier"` + Matched bool `json:"matched"` +} + +// TraceResult holds side-channel info captured while an expression runs. type TraceResult struct { - MatchedTier string `json:"matched_tier"` - Cost float64 `json:"cost"` + MatchedTier string `json:"matched_tier"` + RequestRules []RequestRuleTrace `json:"request_rules,omitempty"` + Cost float64 `json:"cost"` } // BillingSnapshot captures billing state at pre-consume time. Expression and @@ -57,10 +63,11 @@ type BillingSnapshot struct { // TieredResult holds everything needed after running tiered settlement. type TieredResult struct { - ActualQuotaBeforeGroup float64 `json:"actual_quota_before_group"` - ActualQuotaAfterGroup int `json:"actual_quota_after_group"` - MatchedTier string `json:"matched_tier"` - CrossedTier bool `json:"crossed_tier"` + ActualQuotaBeforeGroup float64 `json:"actual_quota_before_group"` + ActualQuotaAfterGroup int `json:"actual_quota_after_group"` + MatchedTier string `json:"matched_tier"` + RequestRules []RequestRuleTrace `json:"request_rules,omitempty"` + CrossedTier bool `json:"crossed_tier"` // Clamp records an int32 saturation event during quota conversion so the // caller can surface it on the consume log for admin auditing. Nil when no // clamping occurred. Not serialized: the marker is attached separately via diff --git a/relay/alpha_search_handler.go b/relay/alpha_search_handler.go index 23ebafea6883..246eebe864ee 100644 --- a/relay/alpha_search_handler.go +++ b/relay/alpha_search_handler.go @@ -62,12 +62,11 @@ func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError } logger.LogDebug(c, "requestBody: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() - info.UpstreamRequestBodySize = size adaptor := GetAdaptor(info.ApiType) if adaptor == nil { diff --git a/relay/channel/ali/text.go b/relay/channel/ali/text.go index eea13129756e..8e8f4c7cc1f6 100644 --- a/relay/channel/ali/text.go +++ b/relay/channel/ali/text.go @@ -18,11 +18,19 @@ func requestOpenAI2Ali(request dto.GeneralOpenAIRequest, upstreamModelName strin request.ThinkingBudget = nil } - topP := lo.FromPtrOr(request.TopP, 0) - if topP >= 1 { - request.TopP = lo.ToPtr(0.999) - } else if topP <= 0 { - request.TopP = lo.ToPtr(0.001) + // DashScope rejects top_p at the 0 and 1 boundaries, so an explicit value is + // clamped into the open interval. The clamp stays at two decimals because + // some models on the platform reject a third decimal with + // "top_p参数非法：限制小数点[2]位". + // + // A request that omits top_p is left untouched: injecting a value would + // silently replace the model's own default with near-greedy decoding. + if request.TopP != nil { + if *request.TopP >= 1 { + request.TopP = lo.ToPtr(0.99) + } else if *request.TopP <= 0 { + request.TopP = lo.ToPtr(0.01) + } } return &request } diff --git a/relay/channel/ali/text_test.go b/relay/channel/ali/text_test.go new file mode 100644 index 000000000000..dc37e4ed2e26 --- /dev/null +++ b/relay/channel/ali/text_test.go @@ -0,0 +1,59 @@ +package ali + +import ( + "testing" + + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/samber/lo" + "github.com/stretchr/testify/assert" +) + +func TestRequestOpenAI2AliTopP(t *testing.T) { + tests := []struct { + name string + topP *float64 + want *float64 + }{ + { + name: "omitted top_p is not injected", + topP: nil, + want: nil, + }, + { + name: "in-range top_p is preserved", + topP: lo.ToPtr(0.8), + want: lo.ToPtr(0.8), + }, + { + name: "top_p of 1 is clamped to two decimals", + topP: lo.ToPtr(1.0), + want: lo.ToPtr(0.99), + }, + { + name: "top_p above 1 is clamped to two decimals", + topP: lo.ToPtr(1.5), + want: lo.ToPtr(0.99), + }, + { + name: "top_p of 0 is clamped to two decimals", + topP: lo.ToPtr(0.0), + want: lo.ToPtr(0.01), + }, + { + name: "negative top_p is clamped to two decimals", + topP: lo.ToPtr(-0.3), + want: lo.ToPtr(0.01), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := requestOpenAI2Ali(dto.GeneralOpenAIRequest{ + Model: "qwen-plus", + TopP: tt.topP, + }, "qwen-plus") + + assert.Equal(t, tt.want, got.TopP) + }) + } +} diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index fc2855a85bca..48241b14a5e9 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -25,20 +25,26 @@ import ( "github.com/gorilla/websocket" ) -// applyUpstreamContentLength populates req.ContentLength when the upstream -// body is wrapped in a BodyStorage (see relay/common/outbound_body.go). -// -// net/http.NewRequest only auto-detects ContentLength for *bytes.Reader, -// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader -// (which is the case for ReaderOnly(BodyStorage)), the Content-Length header -// would otherwise be omitted, forcing chunked transfer encoding and breaking -// some upstreams that require an explicit Content-Length. -func applyUpstreamContentLength(req *http.Request, info *common.RelayInfo) { - if info == nil { +// ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer from +// a ReplayableBody. Callers must pass the original body because NewRequest +// hides its dynamic type behind req.Body's io.ReadCloser wrapper. +func ApplyUpstreamBodyMetadata(req *http.Request, body io.Reader) { + replayable, ok := body.(common2.ReplayableBody) + if !ok { return } - if info.UpstreamRequestBodySize > 0 && req.ContentLength <= 0 { - req.ContentLength = info.UpstreamRequestBodySize + + // BodyStorage structurally satisfies ReplayableBody, but it also exposes + // io.Closer. If a caller passes the storage directly instead of using + // NewReplayableBodyReader, hide Close before the transport takes ownership + // of req.Body so the shared replay source remains available to GetBody. + if _, rawStorage := body.(common2.BodyStorage); rawStorage { + req.Body = io.NopCloser(body) + } + + req.ContentLength = replayable.Size() + if req.GetBody == nil { + req.GetBody = replayable.NewReader } } @@ -314,7 +320,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) + ApplyUpstreamBodyMetadata(req, requestBody) headers := req.Header err = a.SetupRequestHeader(c, &headers, info) if err != nil { @@ -344,7 +350,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) + ApplyUpstreamBodyMetadata(req, requestBody) // set form data req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) headers := req.Header @@ -474,11 +480,24 @@ func sendPingData(c *gin.Context, mutex *sync.Mutex) error { func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { return doRequest(c, req, info) } + +// keepUpstreamRedirectResponse stops net/http from following redirects while +// returning the upstream 3xx response to the relay without an extra error. +func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse +} + func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting) if err != nil { return nil, fmt.Errorf("new proxy http client failed: %w", err) } + // Clients are cached and shared across channels, so override redirect + // behavior on a shallow copy instead of mutating the cached client. This + // still reuses its transport and connection pools, including HTTP/2's + // transparent stream retries. + relayClient := *client + relayClient.CheckRedirect = keepUpstreamRedirectResponse if common2.DebugEnabled && req != nil && req.URL != nil { policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting) logger.LogDebug(c, fmt.Sprintf( @@ -510,7 +529,7 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http } } - resp, err := client.Do(req) + resp, err := relayClient.Do(req) if err != nil { logger.LogError(c, "do request failed: "+err.Error()) return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed")) @@ -548,10 +567,15 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, req if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) - req.GetBody = func() (io.ReadCloser, error) { - return io.NopCloser(requestBody), nil - } + ApplyUpstreamBodyMetadata(req, requestBody) + // Do NOT wrap requestBody in a GetBody closure here: returning the same + // (already consumed) reader would make any transport-level retry silently + // replay an empty body. http.NewRequest already derives a correct, + // snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies + // (which most task adaptors pass in); ApplyUpstreamBodyMetadata wires the + // same contract for bodies that explicitly implement ReplayableBody. + // Otherwise GetBody stays nil so the transport fails the retry instead of + // sending a corrupted request. err = a.BuildRequestHeader(c, req, info) if err != nil { diff --git a/relay/channel/api_request_getbody_test.go b/relay/channel/api_request_getbody_test.go new file mode 100644 index 000000000000..9a2de73337cd --- /dev/null +++ b/relay/channel/api_request_getbody_test.go @@ -0,0 +1,594 @@ +package channel + +import ( + "bytes" + "context" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/net/http2" + "golang.org/x/net/http2/hpack" +) + +func TestApplyUpstreamBodyMetadataSetsReplayableMetadata(t *testing.T) { + t.Parallel() + + payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`) + + body, closer, err := relaycommon.NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + // NewRequest hides the body's dynamic type behind req.Body, so metadata + // extraction must use the original body passed to NewRequest. + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body) + require.NoError(t, err) + assert.Nil(t, req.GetBody) + assert.Zero(t, req.ContentLength) + _, requestBodyIsReplayable := req.Body.(common.ReplayableBody) + assert.False(t, requestBodyIsReplayable) + + ApplyUpstreamBodyMetadata(req, req.Body) + assert.Nil(t, req.GetBody) + assert.Zero(t, req.ContentLength) + + ApplyUpstreamBodyMetadata(req, body) + + assert.EqualValues(t, len(payload), req.ContentLength) + require.NotNil(t, req.GetBody) + + // Drain the primary body as the transport does on the first attempt, then + // make sure GetBody can replay the complete payload repeatedly. + sent, err := io.ReadAll(req.Body) + require.NoError(t, err) + assert.Equal(t, payload, sent) + + for i := 0; i < 2; i++ { + rc, err := req.GetBody() + require.NoError(t, err) + replay, err := io.ReadAll(rc) + require.NoError(t, err) + require.NoError(t, rc.Close()) + assert.Equal(t, payload, replay, "replay %d must equal the original payload", i+1) + } +} + +func TestApplyUpstreamBodyMetadataHidesRawBodyStorageCloser(t *testing.T) { + t.Parallel() + + payload := []byte(`{"model":"test-model","input":"raw storage"}`) + storage, err := common.CreateBodyStorage(payload) + require.NoError(t, err) + defer storage.Close() + + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", storage) + require.NoError(t, err) + _, exposesStorageBeforeApply := req.Body.(common.BodyStorage) + require.True(t, exposesStorageBeforeApply) + + ApplyUpstreamBodyMetadata(req, storage) + + _, exposesStorageAfterApply := req.Body.(common.BodyStorage) + assert.False(t, exposesStorageAfterApply) + assert.EqualValues(t, len(payload), req.ContentLength) + require.NotNil(t, req.GetBody) + + sent, err := io.ReadAll(req.Body) + require.NoError(t, err) + assert.Equal(t, payload, sent) + require.NoError(t, req.Body.Close()) + + replayBody, err := req.GetBody() + require.NoError(t, err, "closing the HTTP request body must not close the shared storage") + replay, err := io.ReadAll(replayBody) + require.NoError(t, err) + require.NoError(t, replayBody.Close()) + assert.Equal(t, payload, replay) +} + +func TestApplyUpstreamBodyMetadataKeepsNativeMetadataForNonReplayableBody(t *testing.T) { + tests := []struct { + name string + body func() io.Reader + }{ + {name: "bytes reader", body: func() io.Reader { return bytes.NewReader([]byte("original")) }}, + {name: "bytes buffer", body: func() io.Reader { return bytes.NewBufferString("original") }}, + {name: "strings reader", body: func() io.Reader { return strings.NewReader("original") }}, + } + + for _, test := range tests { + test := test + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + body := test.body() + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body) + require.NoError(t, err) + require.NotNil(t, req.GetBody, "net/http must derive GetBody for the concrete reader") + + ApplyUpstreamBodyMetadata(req, body) + + rc, err := req.GetBody() + require.NoError(t, err) + got, err := io.ReadAll(rc) + require.NoError(t, err) + require.NoError(t, rc.Close()) + assert.Equal(t, "original", string(got), "an already correct GetBody must not be overwritten") + assert.EqualValues(t, len("original"), req.ContentLength, "native content length must not be overwritten") + }) + } +} + +func TestApplyUpstreamBodyMetadataKeepsExistingGetBody(t *testing.T) { + t.Parallel() + + payload := []byte(`{"model":"test-model"}`) + body, closer, err := relaycommon.NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body) + require.NoError(t, err) + req.ContentLength = 99 + req.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader([]byte("existing"))), nil + } + + ApplyUpstreamBodyMetadata(req, body) + + assert.EqualValues(t, len(payload), req.ContentLength) + rc, err := req.GetBody() + require.NoError(t, err) + got, err := io.ReadAll(rc) + require.NoError(t, err) + require.NoError(t, rc.Close()) + assert.Equal(t, "existing", string(got)) +} + +func TestApplyUpstreamBodyMetadataEmptyStorageRemainsReplayable(t *testing.T) { + t.Parallel() + + storage, err := common.CreateBodyStorage(nil) + require.NoError(t, err) + defer storage.Close() + body := common.NewReplayableBodyReader(storage) + + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body) + require.NoError(t, err) + ApplyUpstreamBodyMetadata(req, body) + + assert.Zero(t, req.ContentLength) + require.NotNil(t, req.GetBody) + rc, err := req.GetBody() + require.NoError(t, err) + replay, err := io.ReadAll(rc) + require.NoError(t, err) + require.NoError(t, rc.Close()) + assert.Empty(t, replay) +} + +// stubTaskAdaptor implements just enough of TaskAdaptor for DoTaskApiRequest. +type stubTaskAdaptor struct { + TaskAdaptor + baseURL string + capturedReq *http.Request +} + +func (s *stubTaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) { + return s.baseURL + "/v1/video/generations", nil +} + +func (s *stubTaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error { + s.capturedReq = req + return nil +} + +// TestDoTaskApiRequest_KeepsReplayableGetBody guards against reintroducing the +// hand-rolled GetBody override that wrapped the already consumed request +// reader: any transport-level retry would then have silently replayed an empty +// body. net/http derives a correct snapshot-based GetBody from the +// *bytes.Reader bodies the task adaptors pass in, and it must be left intact. +func TestDoTaskApiRequest_KeepsReplayableGetBody(t *testing.T) { + service.InitHttpClient() + + payload := []byte(`{"model":"test-model","prompt":"hello"}`) + + type receivedBody struct { + body []byte + err error + } + receivedCh := make(chan receivedBody, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + receivedCh <- receivedBody{body: body, err: err} + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/video/generations", bytes.NewReader(payload)) + + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{}, + } + + adaptor := &stubTaskAdaptor{baseURL: server.URL} + resp, err := DoTaskApiRequest(adaptor, ctx, info, bytes.NewReader(payload)) + require.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + received := <-receivedCh + require.NoError(t, received.err) + assert.Equal(t, payload, received.body) + + req := adaptor.capturedReq + require.NotNil(t, req) + require.NotNil(t, req.GetBody) + // Even after the request body has been fully written, GetBody must still + // return the complete payload, repeatedly. + for i := 0; i < 2; i++ { + rc, err := req.GetBody() + require.NoError(t, err) + replay, err := io.ReadAll(rc) + require.NoError(t, err) + require.NoError(t, rc.Close()) + assert.Equal(t, payload, replay, "replay %d must equal the original payload", i+1) + } +} + +type h2ServerResult struct { + err error + streamCount int + attemptBodies [][]byte +} + +func acceptH2TestConnection(ln net.Listener) (net.Conn, *http2.Framer, error) { + conn, err := ln.Accept() + if err != nil { + return nil, nil, err + } + _ = conn.SetDeadline(time.Now().Add(15 * time.Second)) + + preface := make([]byte, len(http2.ClientPreface)) + if _, err := io.ReadFull(conn, preface); err != nil { + conn.Close() + return nil, nil, fmt.Errorf("read client preface: %w", err) + } + if !bytes.Equal(preface, []byte(http2.ClientPreface)) { + conn.Close() + return nil, nil, fmt.Errorf("unexpected client preface") + } + + framer := http2.NewFramer(conn, conn) + framer.ReadMetaHeaders = hpack.NewDecoder(4096, nil) + if err := framer.WriteSettings(); err != nil { + conn.Close() + return nil, nil, err + } + return conn, framer, nil +} + +func readH2TestRequest(framer *http2.Framer) (uint32, []byte, error) { + var streamID uint32 + var body []byte + for { + frame, err := framer.ReadFrame() + if err != nil { + return 0, nil, fmt.Errorf("read frame: %w", err) + } + switch f := frame.(type) { + case *http2.SettingsFrame: + if !f.IsAck() { + if err := framer.WriteSettingsAck(); err != nil { + return 0, nil, err + } + } + case *http2.MetaHeadersFrame: + streamID = f.Header().StreamID + if f.StreamEnded() { + return streamID, body, nil + } + case *http2.DataFrame: + if streamID == 0 { + streamID = f.Header().StreamID + } + if f.Header().StreamID != streamID { + continue + } + body = append(body, f.Data()...) + if f.StreamEnded() { + return streamID, body, nil + } + } + } +} + +func writeH2TestResponse(framer *http2.Framer, streamID uint32) error { + var hpackBuf bytes.Buffer + henc := hpack.NewEncoder(&hpackBuf) + if err := henc.WriteField(hpack.HeaderField{Name: ":status", Value: "200"}); err != nil { + return err + } + if err := framer.WriteHeaders(http2.HeadersFrameParam{ + StreamID: streamID, + BlockFragment: hpackBuf.Bytes(), + EndHeaders: true, + }); err != nil { + return err + } + return framer.WriteData(streamID, true, []byte(`{}`)) +} + +func awaitH2ServerResult(t *testing.T, resultCh <-chan h2ServerResult) h2ServerResult { + t.Helper() + select { + case result := <-resultCh: + return result + case <-time.After(20 * time.Second): + t.Fatal("timed out waiting for HTTP/2 test server") + return h2ServerResult{} + } +} + +// runResetOnFirstStreamServer speaks just enough raw HTTP/2 to emulate an +// upstream that accepts the first request, waits until the request body has +// been fully written, and then resets the stream with REFUSED_STREAM (the +// retry-safe reset some proxy/CDN-fronted upstreams send under load or during +// graceful shutdown, see RFC 9113 section 8.7). When expectRetry is true it +// serves the retried stream a 200 response; otherwise it stops after the reset. +func runResetOnFirstStreamServer(ln net.Listener, expectRetry bool) <-chan h2ServerResult { + resCh := make(chan h2ServerResult, 1) + go func() { + res := h2ServerResult{} + defer func() { resCh <- res }() + + conn, framer, err := acceptH2TestConnection(ln) + if err != nil { + res.err = err + return + } + defer conn.Close() + + attempts: + for attempt := 0; ; attempt++ { + streamID, body, err := readH2TestRequest(framer) + if err != nil { + res.err = err + return + } + res.streamCount++ + res.attemptBodies = append(res.attemptBodies, body) + + if attempt == 0 { + if err := framer.WriteRSTStream(streamID, http2.ErrCodeRefusedStream); err != nil { + res.err = err + return + } + if !expectRetry { + break attempts + } + continue + } + + if err := writeH2TestResponse(framer, streamID); err != nil { + res.err = err + } + return + } + }() + return resCh +} + +func runGoAwayAfterFirstRequestServer(ln net.Listener) <-chan h2ServerResult { + resCh := make(chan h2ServerResult, 1) + go func() { + res := h2ServerResult{} + defer func() { resCh <- res }() + + for attempt := 0; attempt < 2; attempt++ { + conn, framer, err := acceptH2TestConnection(ln) + if err != nil { + res.err = err + return + } + streamID, body, err := readH2TestRequest(framer) + if err != nil { + conn.Close() + res.err = err + return + } + res.streamCount++ + res.attemptBodies = append(res.attemptBodies, body) + + if attempt == 0 { + err = framer.WriteGoAway(0, http2.ErrCodeNo, nil) + conn.Close() + if err != nil { + res.err = err + return + } + continue + } + + err = writeH2TestResponse(framer, streamID) + conn.Close() + if err != nil { + res.err = err + } + return + } + }() + return resCh +} + +func newH2PriorKnowledgeClient(ln net.Listener) (*http.Client, *http2.Transport) { + transport := &http2.Transport{ + AllowHTTP: true, + DialTLSContext: func(ctx context.Context, network, _ string, _ *tls.Config) (net.Conn, error) { + var dialer net.Dialer + return dialer.DialContext(ctx, network, ln.Addr().String()) + }, + } + return &http.Client{Transport: transport, Timeout: 15 * time.Second}, transport +} + +func newPassThroughBody(t *testing.T, payload []byte) (common.ReplayableBody, common.BodyStorage) { + t.Helper() + storage, err := common.CreateBodyStorage(payload) + require.NoError(t, err) + return common.NewReplayableBodyReader(storage), storage +} + +// TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset exercises the actual +// failure this change fixes: an HTTP/2 upstream resets the stream with a +// retryable error after the request body has been written. With GetBody wired +// up the transport must transparently retry, and the retried request must +// carry the complete body. +func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset(t *testing.T) { + payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"retry me"}]}`) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + resCh := runResetOnFirstStreamServer(ln, true) + + client, transport := newH2PriorKnowledgeClient(ln) + defer transport.CloseIdleConnections() + + // Build the upstream request exactly the way DoApiRequest does: pass the + // original replayable body to the metadata helper after NewRequest. + body, closer, err := relaycommon.NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) + require.NoError(t, err) + ApplyUpstreamBodyMetadata(req, body) + require.NotNil(t, req.GetBody) + + resp, err := client.Do(req) + require.NoError(t, err, "the transport must transparently retry after RST_STREAM") + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + srv := awaitH2ServerResult(t, resCh) + require.NoError(t, srv.err) + assert.Equal(t, 2, srv.streamCount, "the request must have been attempted twice") + require.Len(t, srv.attemptBodies, 2) + assert.Equal(t, payload, srv.attemptBodies[0], "first attempt must carry the full body") + assert.Equal(t, payload, srv.attemptBodies[1], "the retried request must carry the complete body") +} + +func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset_PassThrough(t *testing.T) { + payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"pass through"}]}`) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + resCh := runResetOnFirstStreamServer(ln, true) + + client, transport := newH2PriorKnowledgeClient(ln) + defer transport.CloseIdleConnections() + + body, storage := newPassThroughBody(t, payload) + defer storage.Close() + req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) + require.NoError(t, err) + ApplyUpstreamBodyMetadata(req, body) + require.NotNil(t, req.GetBody) + assert.EqualValues(t, len(payload), req.ContentLength) + + resp, err := client.Do(req) + require.NoError(t, err, "the transport must transparently retry a pass-through body after RST_STREAM") + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + srv := awaitH2ServerResult(t, resCh) + require.NoError(t, srv.err) + assert.Equal(t, 2, srv.streamCount) + require.Len(t, srv.attemptBodies, 2) + assert.Equal(t, payload, srv.attemptBodies[0]) + assert.Equal(t, payload, srv.attemptBodies[1]) +} + +func TestUpstreamGetBody_HTTP2RetryAfterGracefulGoAway_PassThrough(t *testing.T) { + payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"go away"}]}`) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + resCh := runGoAwayAfterFirstRequestServer(ln) + + client, transport := newH2PriorKnowledgeClient(ln) + defer transport.CloseIdleConnections() + + body, storage := newPassThroughBody(t, payload) + defer storage.Close() + req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) + require.NoError(t, err) + ApplyUpstreamBodyMetadata(req, body) + require.NotNil(t, req.GetBody) + + resp, err := client.Do(req) + require.NoError(t, err, "the transport must retry on a new connection after graceful GOAWAY") + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + srv := awaitH2ServerResult(t, resCh) + require.NoError(t, srv.err) + assert.Equal(t, 2, srv.streamCount) + require.Len(t, srv.attemptBodies, 2) + assert.Equal(t, payload, srv.attemptBodies[0]) + assert.Equal(t, payload, srv.attemptBodies[1]) +} + +// TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody documents the pre-fix +// behavior: without GetBody the transport cannot safely retry once the body +// has been written, and the whole relay request fails. +func TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody(t *testing.T) { + payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"retry me"}]}`) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + resCh := runResetOnFirstStreamServer(ln, false) + + client, transport := newH2PriorKnowledgeClient(ln) + defer transport.CloseIdleConnections() + + body, closer, err := relaycommon.NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) + require.NoError(t, err) + req.ContentLength = body.Size() + assert.Nil(t, req.GetBody) + + resp, err := client.Do(req) //nolint:bodyclose // Do fails, no body to close + require.Error(t, err) + assert.Nil(t, resp) + require.ErrorContains(t, err, "cannot retry err") + require.ErrorContains(t, err, "Request.Body was written") + + srv := awaitH2ServerResult(t, resCh) + require.NoError(t, srv.err) + assert.Equal(t, 1, srv.streamCount) + require.Len(t, srv.attemptBodies, 1) + assert.Equal(t, payload, srv.attemptBodies[0]) +} diff --git a/relay/channel/api_request_redirect_test.go b/relay/channel/api_request_redirect_test.go new file mode 100644 index 000000000000..e24685ad56c4 --- /dev/null +++ b/relay/channel/api_request_redirect_test.go @@ -0,0 +1,88 @@ +package channel + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "reflect" + "sync/atomic" + "testing" + + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDoRequestReturnsUpstreamRedirectWithoutFollowing(t *testing.T) { + service.InitHttpClient() + gin.SetMode(gin.TestMode) + sharedClient := service.GetHttpClient() + require.NotNil(t, sharedClient) + require.NotNil(t, sharedClient.CheckRedirect) + originalRedirectPolicy := reflect.ValueOf(sharedClient.CheckRedirect).Pointer() + + var targetRequests atomic.Int32 + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + targetRequests.Add(1) + w.WriteHeader(http.StatusTeapot) + })) + defer target.Close() + + const responseBody = "redirect response" + tests := []int{ + http.StatusMovedPermanently, + http.StatusFound, + http.StatusSeeOther, + http.StatusTemporaryRedirect, + http.StatusPermanentRedirect, + } + + for _, statusCode := range tests { + t.Run(http.StatusText(statusCode), func(t *testing.T) { + targetRequests.Store(0) + var sourceRequests atomic.Int32 + type sourceResult struct { + body []byte + err error + } + sourceResultCh := make(chan sourceResult, 1) + source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sourceRequests.Add(1) + body, err := io.ReadAll(r.Body) + sourceResultCh <- sourceResult{body: body, err: err} + w.Header().Set("Location", target.URL+"/redirect-target") + w.WriteHeader(statusCode) + _, _ = io.WriteString(w, responseBody) + })) + defer source.Close() + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/relay", nil) + + req, err := http.NewRequest(http.MethodPost, source.URL, bytes.NewReader([]byte("request body"))) + require.NoError(t, err) + info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}} + + resp, err := doRequest(ctx, req, info) + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + gotSource := <-sourceResultCh + require.NoError(t, gotSource.err) + + assert.Equal(t, statusCode, resp.StatusCode) + assert.Equal(t, target.URL+"/redirect-target", resp.Header.Get("Location")) + assert.Equal(t, responseBody, string(body)) + assert.Equal(t, []byte("request body"), gotSource.body) + assert.EqualValues(t, 1, sourceRequests.Load()) + assert.Zero(t, targetRequests.Load()) + }) + } + + assert.Equal(t, originalRedirectPolicy, reflect.ValueOf(sharedClient.CheckRedirect).Pointer(), "the cached client must not be mutated") +} diff --git a/relay/channel/deepseek/adaptor.go b/relay/channel/deepseek/adaptor.go index ff13efa44bf4..72805c8c8654 100644 --- a/relay/channel/deepseek/adaptor.go +++ b/relay/channel/deepseek/adaptor.go @@ -115,7 +115,7 @@ func applyDeepSeekV4OpenAIThinkingSuffix(info *relaycommon.RelayInfo, request *d if info.ChannelMeta != nil { info.UpstreamModelName = baseModel } - info.ReasoningEffort = effort + info.SetReasoningEffort(effort) } return nil } @@ -146,7 +146,7 @@ func applyDeepSeekV4ClaudeThinkingSuffix(info *relaycommon.RelayInfo, request *d if info.ChannelMeta != nil { info.UpstreamModelName = baseModel } - info.ReasoningEffort = effort + info.SetReasoningEffort(effort) } return nil } @@ -185,7 +185,7 @@ func applyDeepSeekV4ResponsesThinkingSuffix(info *relaycommon.RelayInfo, request } } if info != nil && request.Reasoning != nil { - info.ReasoningEffort = request.Reasoning.Effort + info.SetReasoningEffort(request.Reasoning.Effort) } } diff --git a/relay/channel/jimeng/adaptor.go b/relay/channel/jimeng/adaptor.go index 658d9236420d..776c1829a03e 100644 --- a/relay/channel/jimeng/adaptor.go +++ b/relay/channel/jimeng/adaptor.go @@ -112,6 +112,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } + channel.ApplyUpstreamBodyMetadata(req, requestBody) err = Sign(c, req, info.ApiKey) if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) diff --git a/relay/channel/ollama/dto.go b/relay/channel/ollama/dto.go index 07aeb17a75c6..80e1ecfbc97a 100644 --- a/relay/channel/ollama/dto.go +++ b/relay/channel/ollama/dto.go @@ -5,12 +5,13 @@ import ( ) type OllamaChatMessage struct { - Role string `json:"role"` - Content string `json:"content,omitempty"` - Images []string `json:"images,omitempty"` - ToolCalls []OllamaToolCall `json:"tool_calls,omitempty"` - ToolName string `json:"tool_name,omitempty"` - Thinking json.RawMessage `json:"thinking,omitempty"` + Role string `json:"role"` + Content string `json:"content,omitempty"` + Images []string `json:"images,omitempty"` + ToolCalls []OllamaToolCall `json:"tool_calls,omitempty"` + ToolName string `json:"tool_name,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + Thinking json.RawMessage `json:"thinking,omitempty"` } type OllamaToolFunction struct { @@ -25,6 +26,7 @@ type OllamaTool struct { } type OllamaToolCall struct { + ID string `json:"id,omitempty"` Function struct { Name string `json:"name"` Arguments interface{} `json:"arguments"` @@ -36,7 +38,7 @@ type OllamaChatRequest struct { Messages []OllamaChatMessage `json:"messages"` Tools interface{} `json:"tools,omitempty"` Format interface{} `json:"format,omitempty"` - Stream bool `json:"stream,omitempty"` + Stream bool `json:"stream"` Options map[string]any `json:"options,omitempty"` KeepAlive interface{} `json:"keep_alive,omitempty"` Think json.RawMessage `json:"think,omitempty"` @@ -48,7 +50,7 @@ type OllamaGenerateRequest struct { Suffix string `json:"suffix,omitempty"` Images []string `json:"images,omitempty"` Format interface{} `json:"format,omitempty"` - Stream bool `json:"stream,omitempty"` + Stream bool `json:"stream"` Options map[string]any `json:"options,omitempty"` KeepAlive interface{} `json:"keep_alive,omitempty"` Think json.RawMessage `json:"think,omitempty"` diff --git a/relay/channel/ollama/relay-ollama.go b/relay/channel/ollama/relay-ollama.go index 874d2e9b3240..e517a1e6aa03 100644 --- a/relay/channel/ollama/relay-ollama.go +++ b/relay/channel/ollama/relay-ollama.go @@ -1,7 +1,6 @@ package ollama import ( - "encoding/json" "fmt" "io" "net/http" @@ -19,24 +18,67 @@ import ( "github.com/samber/lo" ) +func toOllamaResponseFormat(responseFormat *dto.ResponseFormat) (any, error) { + if responseFormat == nil { + return nil, nil + } + switch responseFormat.Type { + case "json", "json_object": + return "json", nil + case "json_schema": + if len(responseFormat.JsonSchema) == 0 { + return nil, nil + } + var jsonSchema dto.FormatJsonSchema + if err := common.Unmarshal(responseFormat.JsonSchema, &jsonSchema); err != nil { + return nil, fmt.Errorf("invalid ollama response format: %w", err) + } + return jsonSchema.Schema, nil + default: + return nil, nil + } +} + func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaChatRequest, error) { + think := r.Think + if len(think) == 0 { + effort := r.ReasoningEffort + if len(r.Reasoning) > 0 { + var reasoning dto.Reasoning + if err := common.Unmarshal(r.Reasoning, &reasoning); err != nil { + return nil, fmt.Errorf("invalid ollama reasoning: %w", err) + } + effort = lo.CoalesceOrEmpty(reasoning.Effort, effort) + } + if effort != "" { + var thinkValue any + switch effort { + case "none": + thinkValue = false + case "low", "medium", "high", "max": + thinkValue = effort + default: + return nil, fmt.Errorf("unsupported ollama reasoning effort %q", effort) + } + var err error + think, err = common.Marshal(thinkValue) + if err != nil { + return nil, fmt.Errorf("marshal ollama think: %w", err) + } + } + } + chatReq := &OllamaChatRequest{ Model: r.Model, Stream: lo.FromPtrOr(r.Stream, false), Options: map[string]any{}, - Think: r.Think, + Think: think, } - if r.ResponseFormat != nil { - if r.ResponseFormat.Type == "json" { - chatReq.Format = "json" - } else if r.ResponseFormat.Type == "json_schema" { - if len(r.ResponseFormat.JsonSchema) > 0 { - var schema any - _ = json.Unmarshal(r.ResponseFormat.JsonSchema, &schema) - chatReq.Format = schema - } - } + format, err := toOllamaResponseFormat(r.ResponseFormat) + if err != nil { + return nil, err } + chatReq.Format = format // options mapping if r.Temperature != nil { @@ -68,12 +110,10 @@ func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*Ollam case []string: chatReq.Options["stop"] = v case []any: - arr := make([]string, 0, len(v)) - for _, i := range v { - if s, ok := i.(string); ok { - arr = append(arr, s) - } - } + arr := lo.FilterMap(v, func(item any, _ int) (string, bool) { + value, ok := item.(string) + return value, ok + }) if len(arr) > 0 { chatReq.Options["stop"] = arr } @@ -81,14 +121,20 @@ func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*Ollam } if len(r.Tools) > 0 { - tools := make([]OllamaTool, 0, len(r.Tools)) - for _, t := range r.Tools { - tools = append(tools, OllamaTool{Type: "function", Function: OllamaToolFunction{Name: t.Function.Name, Description: t.Function.Description, Parameters: t.Function.Parameters}}) - } - chatReq.Tools = tools + chatReq.Tools = lo.Map(r.Tools, func(tool dto.ToolCallRequest, _ int) OllamaTool { + return OllamaTool{ + Type: "function", + Function: OllamaToolFunction{ + Name: tool.Function.Name, + Description: tool.Function.Description, + Parameters: tool.Function.Parameters, + }, + } + }) } chatReq.Messages = make([]OllamaChatMessage, 0, len(r.Messages)) + toolNamesByCallID := make(map[string]string) for _, m := range r.Messages { var textBuilder strings.Builder var images []string @@ -117,8 +163,18 @@ func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*Ollam if len(images) > 0 { cm.Images = images } - if m.Role == "tool" && m.Name != nil { - cm.ToolName = *m.Name + if m.Role == "assistant" { + if reasoning, ok := lo.Coalesce(m.ReasoningContent, m.Reasoning); ok { + thinking, err := common.Marshal(*reasoning) + if err != nil { + return nil, fmt.Errorf("marshal ollama thinking: %w", err) + } + cm.Thinking = thinking + } + } + if m.Role == "tool" { + cm.ToolCallID = m.ToolCallId + cm.ToolName = lo.CoalesceOrEmpty(lo.FromPtr(m.Name), toolNamesByCallID[m.ToolCallId]) } if m.ToolCalls != nil && len(m.ToolCalls) > 0 { parsed := m.ParseToolCalls() @@ -127,15 +183,18 @@ func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*Ollam for _, tc := range parsed { var args interface{} if tc.Function.Arguments != "" { - _ = json.Unmarshal([]byte(tc.Function.Arguments), &args) + _ = common.Unmarshal([]byte(tc.Function.Arguments), &args) } if args == nil { args = map[string]any{} } - oc := OllamaToolCall{} + oc := OllamaToolCall{ID: tc.ID} oc.Function.Name = tc.Function.Name oc.Function.Arguments = args calls = append(calls, oc) + if tc.ID != "" { + toolNamesByCallID[tc.ID] = tc.Function.Name + } } cm.ToolCalls = calls } @@ -175,15 +234,11 @@ func openAIToGenerate(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaGener gen.Suffix = s } } - if r.ResponseFormat != nil { - if r.ResponseFormat.Type == "json" { - gen.Format = "json" - } else if r.ResponseFormat.Type == "json_schema" { - var schema any - _ = json.Unmarshal(r.ResponseFormat.JsonSchema, &schema) - gen.Format = schema - } + format, err := toOllamaResponseFormat(r.ResponseFormat) + if err != nil { + return nil, err } + gen.Format = format if r.Temperature != nil { gen.Options["temperature"] = r.Temperature } @@ -212,12 +267,10 @@ func openAIToGenerate(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaGener case []string: gen.Options["stop"] = v case []any: - arr := make([]string, 0, len(v)) - for _, i := range v { - if s, ok := i.(string); ok { - arr = append(arr, s) - } - } + arr := lo.FilterMap(v, func(item any, _ int) (string, bool) { + value, ok := item.(string) + return value, ok + }) if len(arr) > 0 { gen.Options["stop"] = arr } @@ -510,7 +563,7 @@ func FetchOllamaVersion(baseURL, apiKey string) (string, error) { Version string `json:"version"` } - if err := json.Unmarshal(body, &versionResp); err != nil { + if err := common.Unmarshal(body, &versionResp); err != nil { return "", fmt.Errorf("解析响应失败: %v", err) } diff --git a/relay/channel/ollama/stream.go b/relay/channel/ollama/stream.go index 20e36015ef20..a0d7839f9f6d 100644 --- a/relay/channel/ollama/stream.go +++ b/relay/channel/ollama/stream.go @@ -58,8 +58,12 @@ func ollamaToolCallsToOpenAI(toolCalls []OllamaToolCall, startIndex int, include argBytes = []byte("{}") } } + toolCallID := tc.ID + if toolCallID == "" { + toolCallID = fmt.Sprintf("call_%d", startIndex) + } tr := dto.ToolCallResponse{ - ID: fmt.Sprintf("call_%d", startIndex), + ID: toolCallID, Type: "function", Function: dto.FunctionResponse{ Name: tc.Function.Name, diff --git a/relay/channel/ollama/stream_test.go b/relay/channel/ollama/stream_test.go index 8ba58b191857..69aff1ce8f0f 100644 --- a/relay/channel/ollama/stream_test.go +++ b/relay/channel/ollama/stream_test.go @@ -21,12 +21,14 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) { gin.SetMode(gin.TestMode) tests := []struct { - name string - raw string + name string + raw string + wantID string }{ { - name: "compact json per-line parse path", - raw: `{"model":"llama3.1","created_at":"2026-05-27T12:00:00Z","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"get_weather","arguments":{"city":"Paris","days":0}}}]},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":7}`, + name: "compact json per-line parse path", + raw: `{"model":"llama3.1","created_at":"2026-05-27T12:00:00Z","message":{"role":"assistant","content":"","tool_calls":[{"id":"call_upstream","function":{"name":"get_weather","arguments":{"city":"Paris","days":0}}}]},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":7}`, + wantID: "call_upstream", }, { name: "pretty json fallback parse path", @@ -53,6 +55,7 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) { "prompt_eval_count": 5, "eval_count": 7 }`, + wantID: "call_0", }, } @@ -82,7 +85,7 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) { var toolCalls []dto.ToolCallResponse require.NoError(t, common.Unmarshal(out.Choices[0].Message.ToolCalls, &toolCalls)) require.Len(t, toolCalls, 1) - assert.NotEmpty(t, toolCalls[0].ID) + assert.Equal(t, tt.wantID, toolCalls[0].ID) assert.Equal(t, "function", toolCalls[0].Type) assert.Equal(t, "get_weather", toolCalls[0].Function.Name) assert.Nil(t, toolCalls[0].Index) diff --git a/relay/channel/openai/adaptor.go b/relay/channel/openai/adaptor.go index 4f1c42863dba..64ae3102b3c2 100644 --- a/relay/channel/openai/adaptor.go +++ b/relay/channel/openai/adaptor.go @@ -352,7 +352,7 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn request.Model = originModel } - info.ReasoningEffort = request.ReasoningEffort + info.SetReasoningEffort(request.ReasoningEffort) // o系列模型developer适配（o1-mini除外） if !strings.HasPrefix(info.UpstreamModelName, "o1-mini") && !strings.HasPrefix(info.UpstreamModelName, "o1-preview") { @@ -615,7 +615,7 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo request.Model = originModel } if info != nil && request.Reasoning != nil && request.Reasoning.Effort != "" { - info.ReasoningEffort = request.Reasoning.Effort + info.SetReasoningEffort(request.Reasoning.Effort) } return request, nil } diff --git a/relay/channel/task/sora/adaptor.go b/relay/channel/task/sora/adaptor.go index e9029aa20d46..7f81e5335ebb 100644 --- a/relay/channel/task/sora/adaptor.go +++ b/relay/channel/task/sora/adaptor.go @@ -216,7 +216,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn return &buf, nil } - return common.ReaderOnly(storage), nil + return common.NewReplayableBodyReader(storage), nil } // DoRequest delegates to common helper. diff --git a/relay/channel/task/sora/adaptor_test.go b/relay/channel/task/sora/adaptor_test.go new file mode 100644 index 000000000000..7021f2a9d6dc --- /dev/null +++ b/relay/channel/task/sora/adaptor_test.go @@ -0,0 +1,41 @@ +package sora + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSoraBuildRequestBodyReturnsReplayablePassThroughBody(t *testing.T) { + payload := []byte("opaque-sora-request-body") + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(payload)) + c.Request.Header.Set("Content-Type", "application/octet-stream") + defer common.CleanupBodyStorage(c) + + info := &relaycommon.RelayInfo{} + body, err := (&TaskAdaptor{}).BuildRequestBody(c, info) + require.NoError(t, err) + replayable, ok := body.(common.ReplayableBody) + require.True(t, ok) + + sent, err := io.ReadAll(body) + require.NoError(t, err) + assert.Equal(t, payload, sent) + assert.EqualValues(t, len(payload), replayable.Size()) + + replayBody, err := replayable.NewReader() + require.NoError(t, err) + replay, err := io.ReadAll(replayBody) + require.NoError(t, err) + require.NoError(t, replayBody.Close()) + assert.Equal(t, payload, replay) +} diff --git a/relay/channel/xai/adaptor.go b/relay/channel/xai/adaptor.go index 0db042f06ea5..62b41b33987c 100644 --- a/relay/channel/xai/adaptor.go +++ b/relay/channel/xai/adaptor.go @@ -85,7 +85,7 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn request.ReasoningEffort = "low" request.Model = strings.TrimSuffix(request.Model, "-low") } - info.ReasoningEffort = request.ReasoningEffort + info.SetReasoningEffort(request.ReasoningEffort) info.UpstreamModelName = request.Model } return request, nil diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index ccc7fb281220..b8a6fc875872 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -128,13 +128,12 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size var requestBody io.Reader = body var httpResp *http.Response diff --git a/relay/claude_handler.go b/relay/claude_handler.go index 63bc11f5889c..ff7854469d7f 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -106,6 +106,11 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ } info.UpstreamModelName = request.Model } + if !model_setting.GetGlobalSettings().PassThroughRequestEnabled && !info.ChannelSetting.PassThroughBodyEnabled { + if effort := request.GetEfforts(); effort != "" { + info.SetReasoningEffort(effort) + } + } if info.ChannelSetting.SystemPrompt != "" { if request.System == nil { @@ -159,8 +164,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - info.UpstreamRequestBodySize = storage.Size() - requestBody = common.ReaderOnly(storage) + requestBody = common.NewReplayableBodyReader(storage) } else { convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request) if err != nil { @@ -187,13 +191,12 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ } logger.LogDebug(c, "requestBody: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size requestBody = body } diff --git a/relay/common/outbound_body.go b/relay/common/outbound_body.go index 94ef8dde1da6..ac9e110e0eba 100644 --- a/relay/common/outbound_body.go +++ b/relay/common/outbound_body.go @@ -18,14 +18,14 @@ import ( // The caller MUST invoke closer.Close() once the upstream call has finished // (typically via defer) to release the disk file / memory accounting. // -// The returned reader is wrapped with common.ReaderOnly to prevent the HTTP -// transport from prematurely closing the underlying BodyStorage. The returned -// size is meant to be propagated to http.Request.ContentLength because the -// type-erased io.Reader prevents net/http from auto-detecting it. -func NewOutboundJSONBody(data []byte) (body io.Reader, size int64, closer io.Closer, err error) { +// The returned body exposes its size and replay capability without exposing +// io.Closer. Request construction uses that metadata to populate ContentLength +// and GetBody, while the caller retains ownership of the underlying storage +// through the separately returned closer. +func NewOutboundJSONBody(data []byte) (body common.ReplayableBody, closer io.Closer, err error) { storage, err := common.CreateBodyStorage(data) if err != nil { - return nil, 0, nil, err + return nil, nil, err } - return common.ReaderOnly(storage), storage.Size(), storage, nil + return common.NewReplayableBodyReader(storage), storage, nil } diff --git a/relay/common/outbound_body_test.go b/relay/common/outbound_body_test.go new file mode 100644 index 000000000000..28c3aa699c8f --- /dev/null +++ b/relay/common/outbound_body_test.go @@ -0,0 +1,162 @@ +package common + +import ( + "io" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewOutboundJSONBody_GetBodyReplaysFullBody(t *testing.T) { + t.Parallel() + + payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hello"}]}`) + + body, closer, err := NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + assert.EqualValues(t, len(payload), body.Size()) + + // Consume the primary body, as the HTTP transport does on the first attempt. + first, err := io.ReadAll(body) + require.NoError(t, err) + assert.Equal(t, payload, first) + + // GetBody must hand out the complete body again — and repeatedly, since the + // transport may need more than one retry. + for i := 0; i < 2; i++ { + rc, err := body.NewReader() + require.NoError(t, err) + replay, err := io.ReadAll(rc) + require.NoError(t, err) + require.NoError(t, rc.Close()) + assert.Equal(t, payload, replay, "replay %d must equal the original payload", i+1) + } +} + +func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) { + t.Parallel() + + payload := []byte(`{"model":"test-model","input":"0123456789"}`) + + body, closer, err := NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + // Simulate an aborted first attempt that only wrote part of the body. + partial := make([]byte, 10) + _, err = io.ReadFull(body, partial) + require.NoError(t, err) + + rc, err := body.NewReader() + require.NoError(t, err) + replay, err := io.ReadAll(rc) + require.NoError(t, err) + assert.Equal(t, payload, replay) + + // Closing the replayed body must not close the underlying storage: the + // handler owns the storage lifetime via the returned closer. + require.NoError(t, rc.Close()) + rc2, err := body.NewReader() + require.NoError(t, err) + replay2, err := io.ReadAll(rc2) + require.NoError(t, err) + require.NoError(t, rc2.Close()) + assert.Equal(t, payload, replay2) +} + +// assertIndependentReplayReaders proves that readers handed out by getBody own +// independent cursors, per the http.Request.GetBody contract of returning a +// new copy of the body: interleaved reads across two replay readers and the +// primary body each observe exactly their own byte stream. +func assertIndependentReplayReaders(t *testing.T, payload []byte, body common.ReplayableBody) { + t.Helper() + + half := len(payload) / 2 + + // Partially drain the primary body first, as if attempt N's body write + // were still in flight when the transport builds attempt N+1 via GetBody. + primaryHead := make([]byte, half) + _, err := io.ReadFull(body, primaryHead) + require.NoError(t, err) + assert.Equal(t, payload[:half], primaryHead) + + // Interleave two replay readers: A reads half, B reads everything, then A + // reads the rest. + a, err := body.NewReader() + require.NoError(t, err) + b, err := body.NewReader() + require.NoError(t, err) + + aHead := make([]byte, half) + _, err = io.ReadFull(a, aHead) + require.NoError(t, err) + assert.Equal(t, payload[:half], aHead) + + bAll, err := io.ReadAll(b) + require.NoError(t, err) + require.NoError(t, b.Close()) + assert.Equal(t, payload, bAll, "reader B must see the complete body even while A is mid-read") + + aRest, err := io.ReadAll(a) + require.NoError(t, err) + require.NoError(t, a.Close()) + assert.Equal(t, payload[half:], aRest, "reader A must resume from its own cursor, unaffected by B") + + // The replays must not have disturbed the primary body's cursor either. + primaryRest, err := io.ReadAll(body) + require.NoError(t, err) + assert.Equal(t, payload[half:], primaryRest, "the primary body must be unaffected by replay readers") +} + +func TestNewOutboundJSONBody_GetBodyReadersAreIndependent(t *testing.T) { + t.Parallel() + + payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`) + + body, closer, err := NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + assertIndependentReplayReaders(t, payload, body) + + // Once the handler releases the storage, GetBody must fail loudly instead + // of replaying stale data. + require.NoError(t, closer.Close()) + _, err = body.NewReader() + require.ErrorIs(t, err, common.ErrStorageClosed) +} + +// TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage runs the +// same independence assertions against the disk-backed storage. Deliberately +// not parallel: it temporarily lowers the global disk-cache threshold so the +// payload takes the diskStorage path. +func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing.T) { + prev := common.GetDiskCacheConfig() + common.SetDiskCacheConfig(common.DiskCacheConfig{ + Enabled: true, + ThresholdMB: 0, + MaxSizeMB: 64, + Path: t.TempDir(), + }) + defer common.SetDiskCacheConfig(prev) + + payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`) + + body, closer, err := NewOutboundJSONBody(payload) + require.NoError(t, err) + defer closer.Close() + + storage, ok := closer.(common.BodyStorage) + require.True(t, ok) + assert.True(t, storage.IsDisk(), "the payload must have taken the diskStorage path") + + assertIndependentReplayReaders(t, payload, body) + + require.NoError(t, closer.Close()) + _, err = body.NewReader() + require.ErrorIs(t, err, common.ErrStorageClosed) +} diff --git a/relay/common/override.go b/relay/common/override.go index f92daf77552f..b1a7d17744fa 100644 --- a/relay/common/override.go +++ b/relay/common/override.go @@ -30,6 +30,11 @@ var paramOverrideSensitivePathPrefixes = []string{ "model", "original_model", "upstream_model", + "reasoning", + "reasoning_effort", + "output_config", + "generationConfig.thinkingConfig", + "generation_config.thinking_config", "service_tier", "inference_geo", "speed", @@ -191,6 +196,7 @@ func ApplyParamOverrideWithRelayInfo(jsonData []byte, info *RelayInfo) ([]byte, if err != nil { return nil, err } + syncReasoningEffortAfterParamOverride(info, jsonData, result) syncRuntimeHeaderOverrideFromContext(info, overrideCtx) if info != nil { if recorder != nil { @@ -202,6 +208,51 @@ func ApplyParamOverrideWithRelayInfo(jsonData []byte, info *RelayInfo) ([]byte, return result, nil } +func syncReasoningEffortAfterParamOverride(info *RelayInfo, before, after []byte) { + if info == nil { + return + } + _, existedBefore := extractReasoningEffortFromJSON(info.GetFinalRequestRelayFormat(), before) + effort, existsAfter := extractReasoningEffortFromJSON(info.GetFinalRequestRelayFormat(), after) + if existsAfter { + info.SetReasoningEffort(effort) + return + } + if existedBefore { + info.SetReasoningEffort("") + } +} + +func extractReasoningEffortFromJSON(format types.RelayFormat, data []byte) (string, bool) { + var paths []string + switch format { + case types.RelayFormatOpenAI: + paths = []string{"reasoning_effort", "reasoning.effort"} + case types.RelayFormatOpenAIResponses: + paths = []string{"reasoning.effort"} + case types.RelayFormatClaude: + paths = []string{"output_config.effort"} + case types.RelayFormatGemini: + paths = []string{ + "generationConfig.thinkingConfig.thinkingLevel", + "generation_config.thinking_config.thinking_level", + } + default: + return "", false + } + for _, path := range paths { + value := gjson.GetBytes(data, path) + if !value.Exists() { + continue + } + if value.Type != gjson.String { + return "", true + } + return strings.TrimSpace(value.String()), true + } + return "", false +} + func shouldEnableParamOverrideAudit(paramOverride map[string]interface{}) bool { if common.DebugEnabled { return true @@ -2043,6 +2094,10 @@ func mergeObjects(data []byte, path string, value interface{}, keepOrigin bool) // 目前内置以下字段： // - upstream_model/model：始终为通道映射后的上游模型名。 // - original_model：请求最初指定的模型名。 +// - user_id：已认证用户 ID。 +// - user_group：用户所属分组。 +// - token_group：令牌指定的分组；未指定时回退为用户分组。 +// - using_group：当前实际使用的分组，自动跨分组重试时可能变化。 // - request_path：请求路径 // - is_channel_test：是否为渠道测试请求（同 is_test）。 func BuildParamOverrideContext(info *RelayInfo) map[string]interface{} { @@ -2051,6 +2106,10 @@ func BuildParamOverrideContext(info *RelayInfo) map[string]interface{} { } ctx := make(map[string]interface{}) + ctx["user_id"] = info.UserId + ctx["user_group"] = info.UserGroup + ctx["token_group"] = info.TokenGroup + ctx["using_group"] = info.UsingGroup if info.ChannelMeta != nil && info.ChannelMeta.UpstreamModelName != "" { ctx["model"] = info.ChannelMeta.UpstreamModelName ctx["upstream_model"] = info.ChannelMeta.UpstreamModelName diff --git a/relay/common/override_test.go b/relay/common/override_test.go index c5fc7f5d5f04..19af348df99e 100644 --- a/relay/common/override_test.go +++ b/relay/common/override_test.go @@ -12,6 +12,7 @@ import ( "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/setting/model_setting" "github.com/samber/lo" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -1259,6 +1260,116 @@ func TestApplyParamOverrideConditionFromRetryAndLastErrorContext(t *testing.T) { assertJSONEqual(t, `{"temperature":0.1}`, string(out)) } +func TestApplyParamOverrideConditionByUserAndGPTModel(t *testing.T) { + paramOverride := map[string]interface{}{ + "operations": []interface{}{ + map[string]interface{}{ + "path": "service_tier", + "mode": "set", + "value": "priority", + "logic": "AND", + "conditions": []interface{}{ + map[string]interface{}{ + "path": "user_id", + "mode": "full", + "value": 1, + }, + map[string]interface{}{ + "path": "upstream_model", + "mode": "contains", + "value": "gpt", + }, + }, + }, + }, + } + + tests := []struct { + name string + userID int + model string + expected string + }{ + { + name: "target user and GPT model", + userID: 1, + model: "gpt-5.2", + expected: `{"model":"gpt-5.2","service_tier":"priority"}`, + }, + { + name: "other user", + userID: 2, + model: "gpt-5.2", + expected: `{"model":"gpt-5.2"}`, + }, + { + name: "non-GPT model", + userID: 1, + model: "claude-sonnet-4-5", + expected: `{"model":"claude-sonnet-4-5"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + info := &RelayInfo{ + UserId: tt.userID, + ChannelMeta: &ChannelMeta{ + ParamOverride: paramOverride, + UpstreamModelName: tt.model, + }, + } + input := []byte(fmt.Sprintf(`{"model":%q}`, tt.model)) + + out, err := ApplyParamOverrideWithRelayInfo(input, info) + + require.NoError(t, err) + require.JSONEq(t, tt.expected, string(out)) + }) + } +} + +func TestApplyParamOverrideConditionByGroupContext(t *testing.T) { + info := &RelayInfo{ + UserGroup: "vip", + TokenGroup: "premium", + UsingGroup: "priority-route", + } + ctx := BuildParamOverrideContext(info) + paramOverride := map[string]interface{}{ + "operations": []interface{}{ + map[string]interface{}{ + "path": "service_tier", + "mode": "set", + "value": "priority", + "logic": "AND", + "conditions": []interface{}{ + map[string]interface{}{ + "path": "user_group", + "mode": "full", + "value": "vip", + }, + map[string]interface{}{ + "path": "token_group", + "mode": "full", + "value": "premium", + }, + map[string]interface{}{ + "path": "using_group", + "mode": "full", + "value": "priority-route", + }, + }, + }, + }, + } + + out, err := ApplyParamOverride([]byte(`{"model":"gpt-5.2"}`), paramOverride, ctx) + + require.NoError(t, err) + require.JSONEq(t, `{"model":"gpt-5.2","service_tier":"priority"}`, string(out)) +} + func TestApplyParamOverrideConditionFromRequestHeaders(t *testing.T) { input := []byte(`{"temperature":0.7}`) override := map[string]interface{}{ @@ -2302,3 +2413,113 @@ func assertJSONEqual(t *testing.T, want, got string) { t.Fatalf("json not equal\nwant: %s\ngot: %s", want, got) } } + +func TestApplyParamOverrideWithRelayInfoSynchronizesReasoningEffort(t *testing.T) { + originalDebugEnabled := common2.DebugEnabled + common2.DebugEnabled = false + t.Cleanup(func() { + common2.DebugEnabled = originalDebugEnabled + }) + + tests := []struct { + name string + relayFormat types.RelayFormat + initialEffort string + input string + operation map[string]interface{} + expected string + }{ + { + name: "Responses set", + relayFormat: types.RelayFormatOpenAIResponses, + initialEffort: "high", + input: `{"reasoning":{"effort":"high"}}`, + operation: map[string]interface{}{"mode": "set", "path": "reasoning.effort", "value": "max"}, + expected: "max", + }, + { + name: "chat delete", + relayFormat: types.RelayFormatOpenAI, + initialEffort: "high", + input: `{"reasoning_effort":"high"}`, + operation: map[string]interface{}{"mode": "delete", "path": "reasoning_effort"}, + expected: "", + }, + { + name: "OpenRouter nested set", + relayFormat: types.RelayFormatOpenAI, + initialEffort: "medium", + input: `{"reasoning":{"effort":"medium"}}`, + operation: map[string]interface{}{"mode": "set", "path": "reasoning.effort", "value": "xhigh"}, + expected: "xhigh", + }, + { + name: "Claude output config set", + relayFormat: types.RelayFormatClaude, + initialEffort: "high", + input: `{"output_config":{"effort":"high"}}`, + operation: map[string]interface{}{"mode": "set", "path": "output_config.effort", "value": "max"}, + expected: "max", + }, + { + name: "Gemini thinking level set", + relayFormat: types.RelayFormatGemini, + initialEffort: "medium", + input: `{"generationConfig":{"thinkingConfig":{"thinkingLevel":"medium"}}}`, + operation: map[string]interface{}{"mode": "set", "path": "generationConfig.thinkingConfig.thinkingLevel", "value": "high"}, + expected: "high", + }, + { + name: "non-string value clears effort", + relayFormat: types.RelayFormatOpenAIResponses, + initialEffort: "high", + input: `{"reasoning":{"effort":"high"}}`, + operation: map[string]interface{}{"mode": "set", "path": "reasoning.effort", "value": 42}, + expected: "", + }, + { + name: "unrelated override preserves converter-derived effort", + relayFormat: types.RelayFormatClaude, + initialEffort: "high", + input: `{"thinking":{"type":"adaptive"},"max_tokens":4096}`, + operation: map[string]interface{}{"mode": "set", "path": "max_tokens", "value": 8192}, + expected: "high", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + info := &RelayInfo{ + RelayFormat: tt.relayFormat, + ReasoningEffort: tt.initialEffort, + ChannelMeta: &ChannelMeta{ParamOverride: map[string]interface{}{ + "operations": []interface{}{tt.operation}, + }}, + } + + _, err := ApplyParamOverrideWithRelayInfo([]byte(tt.input), info) + require.NoError(t, err) + assert.Equal(t, tt.expected, info.ReasoningEffort) + }) + } +} + +func TestReasoningEffortOverrideIsAuditedWithoutDebugMode(t *testing.T) { + originalDebugEnabled := common2.DebugEnabled + common2.DebugEnabled = false + t.Cleanup(func() { + common2.DebugEnabled = originalDebugEnabled + }) + info := &RelayInfo{ + RelayFormat: types.RelayFormatOpenAIResponses, + ChannelMeta: &ChannelMeta{ParamOverride: map[string]interface{}{ + "operations": []interface{}{ + map[string]interface{}{"mode": "set", "path": "reasoning.effort", "value": "max"}, + }, + }}, + } + + _, err := ApplyParamOverrideWithRelayInfo([]byte(`{"reasoning":{"effort":"high"}}`), info) + require.NoError(t, err) + assert.Equal(t, []string{"set reasoning.effort = max"}, info.ParamOverrideAudit) +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 45ae30bf9fe8..b0bb19bdca3b 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -149,13 +149,6 @@ type RelayInfo struct { UseRuntimeHeadersOverride bool ParamOverrideAudit []string - // UpstreamRequestBodySize is the byte size of the marshaled upstream request - // body. It is set when the body is wrapped in a BodyStorage (see - // relay/common/outbound_body.go), so that DoApiRequest can populate - // http.Request.ContentLength manually (net/http only auto-detects it for - // *bytes.Reader/Buffer/strings.Reader). 0 means "let net/http decide". - UpstreamRequestBodySize int64 - PriceData hosttypes.PriceData // QuotaClamp is set (non-nil) when a quota conversion saturated at the @@ -241,6 +234,11 @@ func (info *RelayInfo) InitChannelMeta(c *gin.Context) { // Channel identity feeds the converter options snapshot (e.g. // OpenRouterDialect); drop the cache so a cross-channel retry rebuilds it. info.convOptions = nil + if model_setting.GetGlobalSettings().PassThroughRequestEnabled || channelMeta.ChannelSetting.PassThroughBodyEnabled { + info.ReasoningEffort = "" + } else { + info.ReasoningEffort = reasoningEffortFromRequest(info.Request) + } // reset some fields based on channel meta // 重置某些字段，例如模型名称等 @@ -442,6 +440,36 @@ func GenRelayInfoOpenAI(c *gin.Context, request dto.Request) *RelayInfo { return info } +func reasoningEffortFromRequest(request dto.Request) string { + var effort string + switch req := request.(type) { + case *dto.GeneralOpenAIRequest: + if req == nil { + return "" + } + effort = req.ReasoningEffort + if strings.TrimSpace(effort) == "" && len(req.Reasoning) > 0 { + value := gjson.GetBytes(req.Reasoning, "effort") + if value.Type == gjson.String { + effort = value.String() + } + } + case *dto.OpenAIResponsesRequest: + if req != nil && req.Reasoning != nil { + effort = req.Reasoning.Effort + } + case *dto.ClaudeRequest: + if req != nil { + effort = req.GetEfforts() + } + case *dto.GeminiChatRequest: + if req != nil && req.GenerationConfig.ThinkingConfig != nil { + effort = req.GenerationConfig.ThinkingConfig.ThinkingLevel + } + } + return strings.TrimSpace(effort) +} + func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { //channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType) @@ -472,8 +500,10 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { if reqId == "" { reqId = common.NewRequestId() } + reasoningEffort := reasoningEffortFromRequest(request) info := &RelayInfo{ - Request: request, + Request: request, + ReasoningEffort: reasoningEffort, RequestId: reqId, UserId: common.GetContextKeyInt(c, constant.ContextKeyUserId), @@ -747,7 +777,7 @@ func (info *RelayInfo) SetReasoningEffort(effort string) { if info == nil { return } - info.ReasoningEffort = effort + info.ReasoningEffort = strings.TrimSpace(effort) } func (info *RelayInfo) EnsureClaudeConvertInfo() *convmeta.ClaudeConvertInfo { diff --git a/relay/common/relay_info_test.go b/relay/common/relay_info_test.go index 9e58f3f92d8b..42a0f8567bfe 100644 --- a/relay/common/relay_info_test.go +++ b/relay/common/relay_info_test.go @@ -1,10 +1,14 @@ package common import ( + "encoding/json" + "net/http/httptest" "testing" + "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta" "github.com/QuantumNous/new-api/relaykit/types" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -78,3 +82,97 @@ func TestRelayInfoMetaTypedNilReceiver(t *testing.T) { assert.NotNil(t, firstOptions.Gemini.SafetySetting) assert.NotNil(t, firstOptions.PreserveThinkingSuffix) } + +func TestGenRelayInfoCapturesRequestReasoningEffort(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + path string + relayFormat types.RelayFormat + request dto.Request + expected string + }{ + { + name: "OpenAI chat top-level effort", + path: "/v1/chat/completions", + relayFormat: types.RelayFormatOpenAI, + request: &dto.GeneralOpenAIRequest{Model: "gpt-5.6-sol", ReasoningEffort: " high "}, + expected: "high", + }, + { + name: "OpenRouter nested chat effort", + path: "/v1/chat/completions", + relayFormat: types.RelayFormatOpenAI, + request: &dto.GeneralOpenAIRequest{Model: "anthropic/claude", Reasoning: json.RawMessage(`{"effort":"xhigh"}`)}, + expected: "xhigh", + }, + { + name: "OpenAI Responses effort", + path: "/v1/responses", + relayFormat: types.RelayFormatOpenAIResponses, + request: &dto.OpenAIResponsesRequest{Model: "gpt-5.6-sol", Reasoning: &dto.Reasoning{Effort: "max"}}, + expected: "max", + }, + { + name: "explicit none is preserved", + path: "/v1/responses", + relayFormat: types.RelayFormatOpenAIResponses, + request: &dto.OpenAIResponsesRequest{Model: "gpt-5.6-sol", Reasoning: &dto.Reasoning{Effort: "none"}}, + expected: "none", + }, + { + name: "non-string nested effort is ignored", + path: "/v1/chat/completions", + relayFormat: types.RelayFormatOpenAI, + request: &dto.GeneralOpenAIRequest{Model: "anthropic/claude", Reasoning: json.RawMessage(`{"effort":42}`)}, + expected: "", + }, + { + name: "Claude output config effort", + path: "/v1/messages", + relayFormat: types.RelayFormatClaude, + request: &dto.ClaudeRequest{Model: "claude-opus-4-7", OutputConfig: json.RawMessage(`{"effort":"medium"}`)}, + expected: "medium", + }, + { + name: "Gemini thinking level", + path: "/v1beta/models/gemini-3-pro:generateContent", + relayFormat: types.RelayFormatGemini, + request: &dto.GeminiChatRequest{GenerationConfig: dto.GeminiChatGenerationConfig{ + ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingLevel: "low"}, + }}, + expected: "low", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = httptest.NewRequest("POST", tt.path, nil) + + info, err := GenRelayInfo(ctx, tt.relayFormat, tt.request, nil) + require.NoError(t, err) + assert.Equal(t, tt.expected, info.ReasoningEffort) + }) + } +} + +func TestInitChannelMetaRestoresRequestReasoningEffortForRetry(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = httptest.NewRequest("POST", "/v1/responses", nil) + request := &dto.OpenAIResponsesRequest{ + Model: "gpt-5.6-sol", + Reasoning: &dto.Reasoning{Effort: "max"}, + } + info, err := GenRelayInfo(ctx, types.RelayFormatOpenAIResponses, request, nil) + require.NoError(t, err) + + info.SetReasoningEffort("high") + info.InitChannelMeta(ctx) + assert.Equal(t, "max", info.ReasoningEffort) + + info.SetReasoningEffort("low") + info.InitChannelMeta(ctx) + assert.Equal(t, "max", info.ReasoningEffort) +} diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index b1e25e036f8d..8edb3362b6e0 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -104,7 +104,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types logger.LogDebug(c, "requestBody: %s", debugBytes) } } - requestBody = common.ReaderOnly(storage) + requestBody = common.NewReplayableBodyReader(storage) } else { convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request) if err != nil { @@ -175,13 +175,12 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types logger.LogDebug(c, "text request body: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size requestBody = body } diff --git a/relay/embedding_handler.go b/relay/embedding_handler.go index 3be543bd91a1..44d3de486dc4 100644 --- a/relay/embedding_handler.go +++ b/relay/embedding_handler.go @@ -58,13 +58,12 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * } logger.LogDebug(c, "converted embedding request body: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size var requestBody io.Reader = body statusCodeMappingStr := c.GetString("status_code_mapping") resp, err := adaptor.DoRequest(c, info, requestBody) diff --git a/relay/gemini_handler.go b/relay/gemini_handler.go index 786feb8ba345..57010d87c380 100644 --- a/relay/gemini_handler.go +++ b/relay/gemini_handler.go @@ -141,7 +141,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.ReaderOnly(storage) + requestBody = common.NewReplayableBodyReader(storage) } else { // 使用 ConvertGeminiRequest 转换请求格式 convertedRequest, err := adaptor.ConvertGeminiRequest(c, info, request) @@ -164,13 +164,12 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ logger.LogDebug(c, "Gemini request body: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size requestBody = body } @@ -269,13 +268,12 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI } } logger.LogDebug(c, "Gemini embedding request body: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size requestBody = body resp, err := adaptor.DoRequest(c, info, requestBody) diff --git a/relay/image_handler.go b/relay/image_handler.go index f970a43e8e20..690f229f3ff0 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -51,7 +51,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.ReaderOnly(storage) + requestBody = common.NewReplayableBodyReader(storage) } else { convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request) if err != nil { @@ -77,13 +77,12 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type } logger.LogDebug(c, "image request body: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size requestBody = body } } diff --git a/relay/rerank_handler.go b/relay/rerank_handler.go index 769feb227415..460fe3056801 100644 --- a/relay/rerank_handler.go +++ b/relay/rerank_handler.go @@ -47,7 +47,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.ReaderOnly(storage) + requestBody = common.NewReplayableBodyReader(storage) } else { convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request) if err != nil { @@ -68,13 +68,12 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ } logger.LogDebug(c, "Rerank request body: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size requestBody = body } diff --git a/relay/responses_handler.go b/relay/responses_handler.go index c736c6a546b5..f8b4dec2be7b 100644 --- a/relay/responses_handler.go +++ b/relay/responses_handler.go @@ -82,7 +82,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * if err != nil { return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry()) } - requestBody = common.ReaderOnly(storage) + requestBody = common.NewReplayableBodyReader(storage) } else { convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request) if err != nil { @@ -109,13 +109,12 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * } logger.LogDebug(c, "requestBody: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil - info.UpstreamRequestBodySize = size requestBody = body } diff --git a/router/api-router.go b/router/api-router.go index 907cf1ed2885..31c595e00db2 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -90,7 +90,7 @@ func SetApiRouter(router *gin.Engine) { selfRoute.GET("/models", controller.GetUserModels) selfRoute.PUT("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.UpdateSelf) selfRoute.DELETE("/self", controller.DeleteSelf) - selfRoute.GET("/token", middleware.DisableCache(), controller.GenerateAccessToken) + selfRoute.GET("/token", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit("access-token"), middleware.DisableCache(), controller.GenerateAccessToken) selfRoute.GET("/passkey", controller.PasskeyStatus) selfRoute.POST("/passkey/register/begin", middleware.DisableCache(), controller.PasskeyRegisterBegin) selfRoute.POST("/passkey/register/finish", middleware.DisableCache(), controller.PasskeyRegisterFinish) @@ -110,7 +110,7 @@ func SetApiRouter(router *gin.Engine) { selfRoute.POST("/waffo/pay", middleware.CriticalRateLimit(), controller.RequestWaffoPay) selfRoute.POST("/waffo-pancake/amount", controller.RequestWaffoPancakeAmount) selfRoute.POST("/waffo-pancake/pay", middleware.CriticalRateLimit(), controller.RequestWaffoPancakePay) - selfRoute.POST("/aff_transfer", controller.TransferAffQuota) + selfRoute.POST("/aff_transfer", middleware.UserCriticalRateLimit("aff-transfer"), controller.TransferAffQuota) selfRoute.PUT("/setting", controller.UpdateUserSetting) // 2FA routes diff --git a/service/log_info_generate.go b/service/log_info_generate.go index e510ec02606c..353f7098f7a1 100644 --- a/service/log_info_generate.go +++ b/service/log_info_generate.go @@ -316,5 +316,8 @@ func InjectTieredBillingInfo(other map[string]interface{}, relayInfo *relaycommo other["expr_b64"] = base64.StdEncoding.EncodeToString([]byte(snap.ExprString)) if result != nil { other["matched_tier"] = result.MatchedTier + if len(result.RequestRules) > 0 { + other["request_rules"] = result.RequestRules + } } } diff --git a/setting/console_setting/validation.go b/setting/console_setting/validation.go index d6e4342c3d8f..5c664a1feadd 100644 --- a/setting/console_setting/validation.go +++ b/setting/console_setting/validation.go @@ -1,13 +1,15 @@ package console_setting import ( - "encoding/json" "fmt" "net/url" "regexp" "sort" "strings" "time" + "unicode/utf16" + + "github.com/QuantumNous/new-api/common" ) var ( @@ -24,12 +26,16 @@ var ( func parseJSONArray(jsonStr string, typeName string) ([]map[string]interface{}, error) { var list []map[string]interface{} - if err := json.Unmarshal([]byte(jsonStr), &list); err != nil { + if err := common.UnmarshalJsonStr(jsonStr, &list); err != nil { return nil, fmt.Errorf("%s格式错误：%s", typeName, err.Error()) } return list, nil } +func exceedsMaxCharacters(s string, max int) bool { + return len(utf16.Encode([]rune(s))) > max +} + func validateURL(urlStr string, index int, itemType string) error { if !urlRegex.MatchString(urlStr) { return fmt.Errorf("第%d个%s的URL格式不正确", index, itemType) @@ -55,7 +61,7 @@ func getJSONList(jsonStr string) []map[string]interface{} { return []map[string]interface{}{} } var list []map[string]interface{} - json.Unmarshal([]byte(jsonStr), &list) + _ = common.UnmarshalJsonStr(jsonStr, &list) return list } @@ -110,13 +116,13 @@ func validateApiInfo(apiInfoStr string) error { return err } - if len(urlStr) > 500 { + if exceedsMaxCharacters(urlStr, 500) { return fmt.Errorf("第%d个API信息的URL长度不能超过500字符", i+1) } - if len(route) > 100 { + if exceedsMaxCharacters(route, 100) { return fmt.Errorf("第%d个API信息的线路描述长度不能超过100字符", i+1) } - if len(description) > 200 { + if exceedsMaxCharacters(description, 200) { return fmt.Errorf("第%d个API信息的说明长度不能超过200字符", i+1) } @@ -172,12 +178,12 @@ func validateAnnouncements(announcementsStr string) error { } } } - if len(content) > 500 { + if exceedsMaxCharacters(content, 500) { return fmt.Errorf("第%d个公告的内容长度不能超过500字符", i+1) } if extra, exists := ann["extra"]; exists { - if extraStr, ok := extra.(string); ok && len(extraStr) > 200 { - return fmt.Errorf("第%d个公告的说明长度不能超过200字符", i+1) + if extraStr, ok := extra.(string); ok && exceedsMaxCharacters(extraStr, 100) { + return fmt.Errorf("第%d个公告的说明长度不能超过100字符", i+1) } } } @@ -201,10 +207,10 @@ func validateFAQ(faqStr string) error { if !ok || answer == "" { return fmt.Errorf("第%d个FAQ缺少答案字段", i+1) } - if len(question) > 200 { + if exceedsMaxCharacters(question, 200) { return fmt.Errorf("第%d个FAQ的问题长度不能超过200字符", i+1) } - if len(answer) > 1000 { + if exceedsMaxCharacters(answer, 1000) { return fmt.Errorf("第%d个FAQ的答案长度不能超过1000字符", i+1) } } @@ -272,16 +278,16 @@ func validateUptimeKumaGroups(groupsStr string) error { return err } - if len(categoryName) > 50 { + if exceedsMaxCharacters(categoryName, 50) { return fmt.Errorf("第%d个分组的分类名称长度不能超过50字符", i+1) } - if len(urlStr) > 500 { + if exceedsMaxCharacters(urlStr, 500) { return fmt.Errorf("第%d个分组的URL长度不能超过500字符", i+1) } - if len(slug) > 100 { + if exceedsMaxCharacters(slug, 100) { return fmt.Errorf("第%d个分组的Slug长度不能超过100字符", i+1) } - if len(description) > 200 { + if exceedsMaxCharacters(description, 200) { return fmt.Errorf("第%d个分组的描述长度不能超过200字符", i+1) } diff --git a/setting/operation_setting/monitor_setting.go b/setting/operation_setting/monitor_setting.go index 8593d8349a4c..a88087f21569 100644 --- a/setting/operation_setting/monitor_setting.go +++ b/setting/operation_setting/monitor_setting.go @@ -15,6 +15,7 @@ type MonitorSetting struct { const ( ChannelTestModeScheduledAll = "scheduled_all" + ChannelTestModeAutoBanOnly = "auto_ban_only" ChannelTestModePassiveRecovery = "passive_recovery" ) @@ -45,7 +46,9 @@ func GetMonitorSetting() *MonitorSetting { monitorSetting.AutoTestChannelEnabled = parsed } } - if monitorSetting.ChannelTestMode != ChannelTestModePassiveRecovery { + switch monitorSetting.ChannelTestMode { + case ChannelTestModeAutoBanOnly, ChannelTestModePassiveRecovery: + default: monitorSetting.ChannelTestMode = ChannelTestModeScheduledAll } return &monitorSetting diff --git a/setting/operation_setting/monitor_setting_test.go b/setting/operation_setting/monitor_setting_test.go index 7aef7eaaa9cc..c31023e84b8b 100644 --- a/setting/operation_setting/monitor_setting_test.go +++ b/setting/operation_setting/monitor_setting_test.go @@ -41,3 +41,17 @@ func TestGetMonitorSetting_ChannelTestEnabledEnvCanEnableDisabledConfig(t *testi assert.True(t, setting.AutoTestChannelEnabled) assert.Equal(t, float64(12), setting.AutoTestChannelMinutes) } + +func TestGetMonitorSettingPreservesAutoBanOnlyMode(t *testing.T) { + orig := monitorSetting + t.Cleanup(func() { monitorSetting = orig }) + + t.Setenv("CHANNEL_TEST_ENABLED", "") + t.Setenv("CHANNEL_TEST_FREQUENCY", "") + monitorSetting = MonitorSetting{ChannelTestMode: ChannelTestModeAutoBanOnly} + + setting := GetMonitorSetting() + + require.NotNil(t, setting) + assert.Equal(t, ChannelTestModeAutoBanOnly, setting.ChannelTestMode) +} diff --git a/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx b/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx index cca720d97b5d..6d7ce554eefa 100644 --- a/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx +++ b/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx @@ -18,7 +18,7 @@ For commercial licensing, please contact support@quantumnous.com */ import { useQueryClient } from '@tanstack/react-query' import { Loader2, Search, Info, ChevronDown } from 'lucide-react' -import { useState, useEffect, useMemo } from 'react' +import { useState, useEffect, useMemo, type ReactNode } from 'react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' @@ -41,17 +41,16 @@ import { import { fetchUpstreamModels, updateChannel } from '../../api' import { - channelsQueryKeys, + categorizeModels, categorizeModelsWithRedirect, + channelsQueryKeys, normalizeModelName, parseModelsString, } from '../../lib' import { useChannels } from '../channels-provider' function normalizeModelNameList(models: readonly string[]): string[] { - return Array.from( - new Set(models.map((m) => normalizeModelName(m)).filter(Boolean)) - ) + return [...new Set(models.map((m) => normalizeModelName(m)).filter(Boolean))] } type FetchModelsDialogProps = { @@ -140,8 +139,8 @@ export function FetchModelsDialog({ setFetchedModels(list) setSelectedModels(existingModels) toast.success(t('Fetched {{count}} models', { count: list.length })) - } else { - const response = await fetchUpstreamModels(activeChannel!.id) + } else if (activeChannel) { + const response = await fetchUpstreamModels(activeChannel.id) if (response.success) { const list = Array.isArray(response.data) ? response.data : [] setFetchedModels(list) @@ -202,45 +201,6 @@ export function FetchModelsDialog({ onOpenChange(false) } - // Categorize models by common prefixes - const categorizeModels = (models: string[]) => { - const categories: Record = {} - - models.forEach((model) => { - let category = 'Other' - - // Determine category based on model name - if ( - model.toLowerCase().includes('gpt') || - model.toLowerCase().includes('o1') || - model.toLowerCase().includes('o3') - ) { - category = 'OpenAI' - } else if (model.toLowerCase().includes('claude')) { - category = 'Anthropic' - } else if (model.toLowerCase().includes('gemini')) { - category = 'Gemini' - } else if (model.toLowerCase().includes('qwen')) { - category = 'Qwen' - } else if (model.toLowerCase().includes('deepseek')) { - category = 'DeepSeek' - } else if (model.toLowerCase().includes('glm')) { - category = 'Zhipu' - } else if (model.toLowerCase().includes('llama')) { - category = 'Meta' - } else if (model.toLowerCase().includes('mistral')) { - category = 'Mistral' - } - - if (!categories[category]) { - categories[category] = [] - } - categories[category].push(model) - }) - - return categories - } - // Filter models by search const filteredModels = useMemo(() => { if (!searchKeyword) return fetchedModels @@ -249,18 +209,30 @@ export function FetchModelsDialog({ ) }, [fetchedModels, searchKeyword]) - // Helper to check if a model is considered "existing" (in selected or redirect) - const isExistingModel = (model: string) => - classificationSet.has(normalizeModelName(model)) - - // Separate new and existing models - const newModels = filteredModels.filter((m) => !isExistingModel(m)) - const existingFilteredModels = filteredModels.filter((m) => - isExistingModel(m) - ) + const { + newModels, + existingFilteredModels, + newModelsByCategory, + existingModelsByCategory, + } = useMemo(() => { + const newModels: string[] = [] + const existingFilteredModels: string[] = [] + + for (const model of filteredModels) { + if (classificationSet.has(normalizeModelName(model))) { + existingFilteredModels.push(model) + } else { + newModels.push(model) + } + } - const newModelsByCategory = categorizeModels(newModels) - const existingModelsByCategory = categorizeModels(existingFilteredModels) + return { + newModels, + existingFilteredModels, + newModelsByCategory: categorizeModels(newModels), + existingModelsByCategory: categorizeModels(existingFilteredModels), + } + }, [classificationSet, filteredModels]) // 厂商分类按 a-z 排序，Other 放最后，便于查找 const getSortedCategoryEntries = ( @@ -345,7 +317,7 @@ export function FetchModelsDialog({ } - > + /> {t('From model redirect, not yet added to models list')} @@ -365,24 +337,143 @@ export function FetchModelsDialog({ !isFetching && (fetchedModels.length > 0 || removedModels.length > 0) + let dialogDescription: ReactNode = t('Fetch available models from upstream') + if (activeChannel) { + dialogDescription = ( + <> + {t('Channel:')} {activeChannel.name} + + ) + } else if (channelName) { + dialogDescription = ( + <> + {t('Channel:')} {channelName} + + ) + } + + let defaultTab = 'existing' + if (newModels.length > 0) { + defaultTab = 'new' + } else if (removedModels.length > 0) { + defaultTab = 'removed' + } + + let dialogBody: ReactNode + if (!activeChannel && !customFetcher) { + dialogBody = ( +
+ {t('No channel selected')} +
+ ) + } else if (isFetching) { + dialogBody = ( +
+ +
+ ) + } else if (fetchedModels.length === 0 && removedModels.length === 0) { + dialogBody = ( +
+

{t('No models fetched yet.')}

+ +
+ ) + } else { + dialogBody = ( +
+ {/* Search Bar */} +
+ + setSearchKeyword(e.target.value)} + className='pl-9' + /> +
+ + {/* Tabs for New vs Existing vs Removed */} + + 0 ? 'grid-cols-3' : 'grid-cols-2'}`} + > + + {t('New Models ({{count}})', { count: newModels.length })} + + + {t('Existing Models ({{count}})', { + count: existingFilteredModels.length, + })} + + {removedModels.length > 0 && ( + + {t('Removed Models ({{count}})', { + count: removedModels.length, + })} + + )} + + + + {getSortedCategoryEntries(newModelsByCategory).map( + ([category, models]) => renderModelCategory(category, models) + )} + + + + {getSortedCategoryEntries(existingModelsByCategory).map( + ([category, models]) => renderModelCategory(category, models) + )} + + + {removedModels.length > 0 && ( + +

+ {t( + 'These models are still in your selection but were not returned by the upstream listing. Entries that are only model_mapping source aliases are omitted. Toggle to adjust before saving.' + )} +

+ {renderModelCategory(t('Removed'), removedModels)} + + )} + + + {/* Selection Summary */} +
+ {t('{{n}} model(s) selected', { n: selectedModels.length })} +
+
+ ) + } + return ( - {t('Channel:')} {activeChannel.name} - - ) : channelName ? ( - <> - {t('Channel:')} {channelName} - - ) : ( - t('Fetch available models from upstream') - ) - } + description={dialogDescription} contentClassName='max-w-3xl' contentHeight='auto' bodyClassName='space-y-4' @@ -400,113 +491,7 @@ export function FetchModelsDialog({ ) : null } > - {!activeChannel && !customFetcher ? ( -
- {t('No channel selected')} -
- ) : isFetching ? ( -
- -
- ) : fetchedModels.length === 0 && removedModels.length === 0 ? ( -
-

{t('No models fetched yet.')}

- -
- ) : ( - <> -
- {/* Search Bar */} -
- - setSearchKeyword(e.target.value)} - className='pl-9' - /> -
- - {/* Tabs for New vs Existing vs Removed */} - 0 - ? 'new' - : removedModels.length > 0 - ? 'removed' - : 'existing' - } - > - 0 ? 'grid-cols-3' : 'grid-cols-2'}`} - > - - {t('New Models ({{count}})', { count: newModels.length })} - - - {t('Existing Models ({{count}})', { - count: existingFilteredModels.length, - })} - - {removedModels.length > 0 && ( - - {t('Removed Models ({{count}})', { - count: removedModels.length, - })} - - )} - - - - {getSortedCategoryEntries(newModelsByCategory).map( - ([category, models]) => renderModelCategory(category, models) - )} - - - - {getSortedCategoryEntries(existingModelsByCategory).map( - ([category, models]) => renderModelCategory(category, models) - )} - - - {removedModels.length > 0 && ( - -

- {t( - 'These models are still in your selection but were not returned by the upstream listing. Entries that are only model_mapping source aliases are omitted. Toggle to adjust before saving.' - )} -

- {renderModelCategory(t('Removed'), removedModels)} - - )} - - - {/* Selection Summary */} -
- {t('{{n}} model(s) selected', { n: selectedModels.length })} -
-
- - )} + {dialogBody} ) } diff --git a/web/src/features/channels/lib/index.ts b/web/src/features/channels/lib/index.ts index 43eb7773b4cf..8c18151ceb5f 100644 --- a/web/src/features/channels/lib/index.ts +++ b/web/src/features/channels/lib/index.ts @@ -26,3 +26,4 @@ export * from './channel-type-config' export * from './channel-utils' export * from './multi-key-utils' export * from './model-mapping-validation' +export * from './model-categories' diff --git a/web/src/features/channels/lib/model-categories.ts b/web/src/features/channels/lib/model-categories.ts new file mode 100644 index 000000000000..5db0dd738219 --- /dev/null +++ b/web/src/features/channels/lib/model-categories.ts @@ -0,0 +1,174 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +type ModelCategoryRule = { + name: string + keywords?: readonly string[] + pattern?: RegExp +} + +// Rules are ordered so platform-specific IDs such as Perplexity's Sonar and +// NVIDIA's Nemotron take precedence over the base Llama/Mixtral family name. +const MODEL_CATEGORY_RULES: readonly ModelCategoryRule[] = [ + { name: 'Perplexity', keywords: ['perplexity', 'sonar-'] }, + { name: 'NVIDIA', keywords: ['nvidia/', 'nvidia.', 'nemotron'] }, + { + name: 'OpenAI', + keywords: [ + 'openai/', + 'openai.', + 'gpt-', + 'chatgpt-', + 'codex-', + 'dall-e-', + 'whisper-', + 'omni-moderation-', + 'text-moderation-', + 'text-embedding-ada-', + 'text-embedding-3-', + 'text-ada-', + 'text-babbage-', + 'text-curie-', + 'davinci-', + 'babbage-', + 'computer-use-preview', + 'sora', + ], + pattern: /(?:^|[/.:])(?:o(?:1|3|4)(?=$|[-.:])|tts-)/, + }, + { name: 'Anthropic', keywords: ['anthropic', 'claude'] }, + { + name: 'Gemini', + keywords: [ + 'gemini', + 'gemma', + 'learnlm', + 'imagen', + 'veo', + 'nano-banana', + 'palm-', + ], + pattern: /(?:^|[/.:])aqa$/, + }, + { name: 'xAI', keywords: ['x-ai/', 'xai/', 'xai-', 'grok'] }, + { name: 'DeepSeek', keywords: ['deepseek'] }, + { + name: 'Qwen', + keywords: ['qwen', 'qwq-', 'qvq-', 'tongyi', 'gte-'], + pattern: /(?:^|[/.:])(?:text-embedding-v\d+|gui-plus|z-image)(?:$|[-_.:])/, + }, + { name: 'Wan', pattern: /(?:^|[/.:])wan(?:x?\d|[-_])/ }, + { name: 'Moonshot', keywords: ['moonshot', 'kimi-'] }, + { + name: 'MiniMax', + keywords: ['minimax', 'abab', 'hailuo'], + pattern: /^(?:t2v|i2v|s2v)-01(?:-|$)/, + }, + { + name: 'Doubao', + keywords: ['doubao', 'volcengine', 'seedance', 'seedream', 'seed-1-'], + }, + { + name: 'Zhipu', + keywords: ['zhipu', 'zai-org', 'thudm', 'chatglm', 'cogview', 'cogvideo'], + pattern: /(?:^|[/._-])glm(?=$|[-._])/, + }, + { name: 'Baidu', keywords: ['baidu', 'wenxin', 'ernie'] }, + { name: 'Yi', keywords: ['01-ai/'], pattern: /(?:^|[/.:])yi(?=$|[-_])/ }, + { name: 'iFlytek', keywords: ['iflytek', 'sparkdesk'] }, + { + name: 'Tencent', + keywords: ['tencent', 'hunyuan'], + pattern: /(?:^|[/.:])hy\d*(?=$|[-_.:])/, + }, + { name: 'Baichuan', keywords: ['baichuan'] }, + { name: 'InternLM', keywords: ['internlm'] }, + { name: 'StepFun', keywords: ['stepfun', 'step-'] }, + { name: 'MiMo', keywords: ['xiaomi', 'mimo-'] }, + { + name: 'Mistral', + keywords: [ + 'mistral', + 'mixtral', + 'codestral', + 'ministral', + 'pixtral', + 'magistral', + ], + }, + { name: 'Meta', keywords: ['meta-llama', 'llama-', 'llama2', 'llama3'] }, + { + name: 'Cohere', + keywords: ['cohere', 'command-', 'c4ai-aya', 'aya-'], + pattern: /(?:^|[/.:])command$/, + }, + { name: 'Jina', keywords: ['jinaai', 'jina-'] }, + { name: 'BAAI', keywords: ['baai/', 'bge-'] }, + { name: 'Black Forest Labs', keywords: ['black-forest-labs', 'flux.'] }, + { + name: 'Microsoft', + keywords: ['microsoft/'], + pattern: /(?:^|[/.:])phi(?=$|[-._])/, + }, + { + name: 'Amazon', + keywords: ['amazon/', 'amazon.', 'nova-', 'titan-'], + }, + { name: 'AI21 Labs', keywords: ['ai21', 'jamba'] }, + { + name: 'Stability AI', + keywords: ['stabilityai', 'stable-diffusion', 'stable-image', 'sdxl-'], + }, + { name: 'Nous Research', keywords: ['nousresearch', 'hermes-'] }, + { name: '360 AI', keywords: ['360gpt', '360zhinao'] }, + { name: 'Midjourney', keywords: ['midjourney', 'mj_', 'mj-', 'swap_face'] }, + { name: 'Kling', keywords: ['kling'] }, + { name: 'Vidu', keywords: ['vidu'] }, + { name: 'Suno', keywords: ['suno'] }, + { name: 'Jimeng', keywords: ['jimeng'] }, +] + +export function getModelCategory(modelName: string): string { + const normalizedName = modelName.trim().toLowerCase() + + for (const rule of MODEL_CATEGORY_RULES) { + if ( + rule.keywords?.some((keyword) => normalizedName.includes(keyword)) || + rule.pattern?.test(normalizedName) + ) { + return rule.name + } + } + + return 'Other' +} + +export function categorizeModels( + models: readonly string[] +): Record { + const categories: Record = {} + + for (const model of models) { + const category = getModelCategory(model) + categories[category] ??= [] + categories[category].push(model) + } + + return categories +} diff --git a/web/src/features/keys/components/api-keys-table.tsx b/web/src/features/keys/components/api-keys-table.tsx index 3561df5254c9..2ee497924944 100644 --- a/web/src/features/keys/components/api-keys-table.tsx +++ b/web/src/features/keys/components/api-keys-table.tsx @@ -307,6 +307,7 @@ export function ApiKeysTable() { applyHeaderSize toolbarProps={{ searchPlaceholder: t('Filter by name...'), + searchDebounceMs: 500, additionalSearch: ( . For commercial licensing, please contact support@quantumnous.com */ import { useQueryClient } from '@tanstack/react-query' -import { type ColumnDef, type RowSelectionState } from '@tanstack/react-table' +import type { ColumnDef, RowSelectionState } from '@tanstack/react-table' import { Search, Info, @@ -48,6 +48,7 @@ import { SelectTrigger, SelectValue, } from '@/components/ui/select' +import { useDebounce } from '@/hooks/use-debounce' import { useIsMobile } from '@/hooks/use-mobile' import { applyUpstreamOverwrite } from '../../api' @@ -117,6 +118,7 @@ export function UpstreamConflictDialog({ } = useModels() const isMobile = useIsMobile() const [search, setSearch] = useState('') + const debouncedSearch = useDebounce(search, 200) const [isSubmitting, setIsSubmitting] = useState(false) const [rowSelection, setRowSelection] = useState({}) const [pageSize, setPageSize] = useState(10) @@ -149,7 +151,7 @@ export function UpstreamConflictDialog({ const totalModels = upstreamConflicts.length const totalFields = conflictRows.length - const normalizedSearch = search.trim().toLowerCase() + const normalizedSearch = debouncedSearch.trim().toLowerCase() const { matchingModelNames, visibleRowIds } = useMemo(() => { if (!normalizedSearch) { @@ -396,7 +398,7 @@ export function UpstreamConflictDialog({ const payload: SyncOverwritePayload[] = Object.entries(groupedSelections) .map(([modelName, fields]) => ({ model_name: modelName, - fields: Array.from(fields), + fields: [...fields], })) .filter((item) => item.fields.length > 0) diff --git a/web/src/features/models/components/models-table.tsx b/web/src/features/models/components/models-table.tsx index 91640ddbf522..4a9b7b6ecd07 100644 --- a/web/src/features/models/components/models-table.tsx +++ b/web/src/features/models/components/models-table.tsx @@ -202,6 +202,7 @@ export function ModelsTable() { applyHeaderSize toolbarProps={{ searchPlaceholder: t('Filter by model name...'), + searchDebounceMs: 500, filters: [ { columnId: 'status', diff --git a/web/src/features/pricing/components/dynamic-pricing-breakdown.tsx b/web/src/features/pricing/components/dynamic-pricing-breakdown.tsx index 79be84ddea70..848454cc3ec6 100644 --- a/web/src/features/pricing/components/dynamic-pricing-breakdown.tsx +++ b/web/src/features/pricing/components/dynamic-pricing-breakdown.tsx @@ -36,11 +36,13 @@ import { SOURCE_TIME, normalizeTierLabel, parseTiersFromExpr, + requestRuleGroupsFromTrace, splitBillingExprAndRequestRules, tryParseRequestRuleExpr, type ParsedTier, type RequestCondition, type RequestRuleGroup, + type RequestRuleTrace, type TierCondition, } from '../lib/billing-expr' @@ -52,6 +54,8 @@ type DynamicPricingBreakdownProps = { * the usage-log details dialog to show which tier the engine selected. */ matchedTierLabel?: string | null + /** Request-rule traces emitted by the settlement run. */ + requestRules?: RequestRuleTrace[] | null /** * Hide cache-pricing columns regardless of the per-tier values. The log * details dialog passes this when the actual request did not consume any @@ -148,14 +152,25 @@ function describeGroup( group: RequestRuleGroup, t: (key: string) => string ): string { - return (group.conditions || []) - .map((c) => describeCondition(c, t)) + const description = (group.conditions || []) + .map((condition) => describeCondition(condition, t)) .join(' && ') + return description || group.conditionText || '' +} + +function nextOccurrenceKey( + baseKey: string, + occurrences: Map +): string { + const occurrence = occurrences.get(baseKey) || 0 + occurrences.set(baseKey, occurrence + 1) + return `${baseKey}:${occurrence}` } export function DynamicPricingBreakdown({ billingExpr, matchedTierLabel, + requestRules, hideCacheColumns = false, compact = false, }: DynamicPricingBreakdownProps) { @@ -179,12 +194,15 @@ export function DynamicPricingBreakdown({ const { tiers, ruleGroups } = useMemo(() => { const split = splitBillingExprAndRequestRules(expr) const parsedTiers = parseTiersFromExpr(split.billingExpr) - const parsedRules = tryParseRequestRuleExpr(split.requestRuleExpr || '') + const parsedRules = + requestRules != null + ? requestRuleGroupsFromTrace(requestRules) + : tryParseRequestRuleExpr(split.requestRuleExpr || '') return { tiers: parsedTiers, ruleGroups: parsedRules || [], } - }, [expr]) + }, [expr, requestRules]) const hasTiers = tiers.length > 0 const hasRules = ruleGroups.length > 0 @@ -229,6 +247,8 @@ export function DynamicPricingBreakdown({ (tier) => Number(tier[v.field as string as keyof ParsedTier] || 0) > 0 ) }) + const mobileTierKeyOccurrences = new Map() + const requestRuleKeyOccurrences = new Map() return (
@@ -260,15 +280,19 @@ export function DynamicPricingBreakdown({ {t('Tiered price table')}
- {tiers.map((tier, i) => { + {tiers.map((tier) => { const condSummary = formatConditionSummary(tier.conditions, t) const isMatched = matchedTierLabel != null && matchedTierLabel !== '' && tier.label === matchedTierLabel + const rowKey = nextOccurrenceKey( + JSON.stringify(tier), + mobileTierKeyOccurrences + ) return (
    - {ruleGroups.map((group, gi) => ( -
  • - { + const isMatched = group.matched === true + const rowKey = nextOccurrenceKey( + `${group.conditionText || JSON.stringify(group.conditions)}:${group.multiplier}`, + requestRuleKeyOccurrences + ) + return ( +
  • - {describeGroup(group, t)} - - - {group.multiplier}x - -
    • - ))} + + {describeGroup(group, t)} + + + {group.multiplier}x{isMatched && ` · ${t('Matched')}`} + + + ) + })}
)} diff --git a/web/src/features/pricing/hooks/use-filters.ts b/web/src/features/pricing/hooks/use-filters.ts index 18415b8a434d..c055153f612a 100644 --- a/web/src/features/pricing/hooks/use-filters.ts +++ b/web/src/features/pricing/hooks/use-filters.ts @@ -19,6 +19,8 @@ For commercial licensing, please contact support@quantumnous.com import { useSearch } from '@tanstack/react-router' import { useMemo, useCallback, useState } from 'react' +import { useDebounce } from '@/hooks/use-debounce' + import { FILTER_ALL, SORT_OPTIONS, @@ -67,6 +69,7 @@ export function useFilters(models: PricingModel[]) { })) const searchInput = filterState.search || '' + const debouncedSearchInput = useDebounce(searchInput, 200) const sortBy = filterState.sort || SORT_OPTIONS.NAME const vendorFilter = filterState.vendor || FILTER_ALL const groupFilter = filterState.group || FILTER_ALL @@ -147,7 +150,7 @@ export function useFilters(models: PricingModel[]) { if (!models || models.length === 0) return [] return filterAndSortModels(models, { - search: searchInput, + search: debouncedSearchInput, vendor: vendorFilter, group: groupFilter, quotaType: quotaTypeFilter, @@ -157,7 +160,7 @@ export function useFilters(models: PricingModel[]) { }) }, [ models, - searchInput, + debouncedSearchInput, vendorFilter, groupFilter, quotaTypeFilter, diff --git a/web/src/features/pricing/lib/billing-expr.ts b/web/src/features/pricing/lib/billing-expr.ts index 59f64b7a3b99..f1428f674904 100644 --- a/web/src/features/pricing/lib/billing-expr.ts +++ b/web/src/features/pricing/lib/billing-expr.ts @@ -226,6 +226,14 @@ export type RequestCondition = TimeCondition | ParamHeaderCondition export type RequestRuleGroup = { conditions: RequestCondition[] multiplier: string + conditionText?: string + matched?: boolean +} + +export type RequestRuleTrace = { + cond: string + multiplier: number + matched: boolean } export type TierCondition = { @@ -307,9 +315,9 @@ export function parseTiersFromExpr(exprStr: string): ParsedTier[] { export function normalizeTierLabel(label: string | undefined): string { if (!label) return '' return label - .replace(/<[=＝]?|≤|＜[=＝]?/g, '<') - .replace(/>[=＝]?|≥|＞[=＝]?/g, '>') - .replace(/\s+/g, '') + .replaceAll(/<[=＝]?|≤|＜[=＝]?/g, '<') + .replaceAll(/>[=＝]?|≥|＞[=＝]?/g, '>') + .replaceAll(/\s+/g, '') .toLowerCase() } @@ -426,24 +434,26 @@ function tryParseRequestCondition(expr: string): RequestCondition | null { if (m) return { source: 'param', path: m[1], mode: MATCH_EXISTS, value: '' } m = expr.match(/^has\(header\("([^"]+)"\), ((?:"(?:[^"\\]|\\.)*"))\)$/) - if (m) + if (m) { return { source: 'header', path: m[1], mode: MATCH_CONTAINS, value: JSON.parse(m[2]) as string, } + } m = expr.match( /^param\("([^"]+)"\) != nil && has\(param\("([^"]+)"\), ((?:"(?:[^"\\]|\\.)*"))\)$/ ) - if (m && m[1] === m[2]) + if (m && m[1] === m[2]) { return { source: 'param', path: m[1], mode: MATCH_CONTAINS, value: JSON.parse(m[3]) as string, } + } m = expr.match( /^param\("([^"]+)"\) != nil && param\("([^"]+)"\) (>|>=|<|<=) ([\d.eE+-]+)$/ @@ -473,22 +483,40 @@ function tryParseRequestCondition(expr: string): RequestCondition | null { return null } +function tryParseRequestConditions( + conditionStr: string +): RequestCondition[] | null { + const andParts = splitTopLevelAnd(conditionStr) + const conditions: RequestCondition[] = [] + for (const part of andParts) { + const condition = tryParseRequestCondition(part.trim()) + if (!condition) return null + conditions.push(condition) + } + return conditions.length > 0 ? conditions : null +} + function tryParseRuleGroupFactor(part: string): RequestRuleGroup | null { const m = part.match(/^\((.+) \? ([\d.eE+-]+) : 1\)$/s) if (!m) return null - const conditionStr = m[1] - const multiplier = m[2] + const conditions = tryParseRequestConditions(m[1]) + if (!conditions) return null + return { conditions, multiplier: m[2] } +} - const andParts = splitTopLevelAnd(conditionStr) - const conditions: RequestCondition[] = [] - for (const ap of andParts) { - const cond = tryParseRequestCondition(ap.trim()) - if (!cond) return null - conditions.push(cond) - } - if (conditions.length === 0) return null - return { conditions, multiplier } +export function requestRuleGroupsFromTrace( + requestRules: RequestRuleTrace[] +): RequestRuleGroup[] { + return requestRules.map((rule) => { + const conditionText = rule.cond.trim() + return { + conditions: tryParseRequestConditions(conditionText) || [], + multiplier: String(rule.multiplier), + conditionText, + matched: rule.matched, + } + }) } export function tryParseRequestRuleExpr( @@ -642,12 +670,12 @@ function isTimeFunc(value: unknown): value is TimeFunc { export function normalizeCondition( cond: Partial | null | undefined ): RequestCondition { - const source = - cond?.source === 'time' - ? 'time' - : cond?.source === 'header' - ? 'header' - : 'param' + let source: RequestCondition['source'] = 'param' + if (cond?.source === 'time') { + source = 'time' + } else if (cond?.source === 'header') { + source = 'header' + } if (source === 'time') { const timeCond = cond as Partial | null | undefined diff --git a/web/src/features/profile/components/dialogs/access-token-dialog.tsx b/web/src/features/profile/components/dialogs/access-token-dialog.tsx index 550166caa158..c506cc95f1fb 100644 --- a/web/src/features/profile/components/dialogs/access-token-dialog.tsx +++ b/web/src/features/profile/components/dialogs/access-token-dialog.tsx @@ -16,13 +16,21 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ -import { RefreshCw, Loader2 } from 'lucide-react' -import { useEffect } from 'react' +import { KeyRound, Loader2, RefreshCw } from 'lucide-react' +import { useState } from 'react' import { useTranslation } from 'react-i18next' +import { ConfirmDialog } from '@/components/confirm-dialog' import { CopyButton } from '@/components/copy-button' import { Dialog } from '@/components/dialog' import { Button } from '@/components/ui/button' +import { + Empty, + EmptyDescription, + EmptyHeader, + EmptyMedia, + EmptyTitle, +} from '@/components/ui/empty' import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' @@ -42,77 +50,149 @@ export function AccessTokenDialog({ onOpenChange, }: AccessTokenDialogProps) { const { t } = useTranslation() - const { token, generating, generate } = useAccessToken() + const { token, generating, generate, clearToken } = useAccessToken() + const [confirmOpen, setConfirmOpen] = useState(false) - // Auto-generate token when dialog opens if no token exists - useEffect(() => { - if (open && !token) { - generate() + const handleOpenChange = (nextOpen: boolean) => { + if (generating) return + + if (!nextOpen) { + setConfirmOpen(false) + clearToken() + } + onOpenChange(nextOpen) + } + + const handleGenerate = async () => { + if (await generate()) { + setConfirmOpen(false) } - }, [open, token, generate]) + } return ( - - - - - } - > -
-
- -
- - + +
-

- {t('Use this token for API authentication')} -

+ onClick={() => handleOpenChange(false)} + disabled={generating} + > + {t('Close')} + + + + } + > +
+ {token ? ( +
+ +
+ + +
+

+ {t( + "Save this token now. You won't be able to view it again after closing this dialog." + )} +

+
+ ) : ( + + + + + + {t('Access tokens are shown only once')} + + + {t( + 'For security, existing access tokens cannot be displayed. Regenerate only when you need a new token.' + )} + + + {t( + 'Regenerating immediately invalidates any existing token.' + )} + + + + )}
-
-
+ + + { + if (!generating) setConfirmOpen(nextOpen) + }} + title={t('Regenerate access token?')} + desc={ +
+

+ {t( + 'This will immediately invalidate your existing access token. Any applications or scripts using it will stop working.' + )} +

+

+ {t( + 'The new token will only be shown once. Copy it and store it securely.' + )} +

+
+ } + confirmText={ + generating ? ( + <> +