diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml new file mode 100644 index 000000000000..7eb7780ddb3a --- /dev/null +++ b/.github/workflows/docker-build.yml @@ -0,0 +1,235 @@ +name: Publish fork release images + +on: + push: + tags: + - 'v*' + workflow_dispatch: + inputs: + tag: + description: 'Existing fork release tag, for example v1.0.0-rc.23-0' + required: true + type: string + +env: + DOCKERHUB_IMAGE: karlorz/new-api + GHCR_IMAGE: ghcr.io/karlorz/new-api + +jobs: + prepare: + name: Validate fork release tag + runs-on: ubuntu-latest + outputs: + tag: ${{ steps.release.outputs.tag }} + publish_latest: ${{ steps.release.outputs.publish_latest }} + sha: ${{ steps.release.outputs.sha }} + permissions: + contents: read + steps: + - name: Check out release tag + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 0 + ref: ${{ github.event.inputs.tag || github.ref }} + + - name: Validate tag and resolve release metadata + id: release + env: + REQUESTED_TAG: ${{ github.event.inputs.tag || github.ref_name }} + run: | + scripts/validate-fork-release-tag.sh "$REQUESTED_TAG" + if ! git rev-parse --verify "refs/tags/$REQUESTED_TAG^{commit}" >/dev/null 2>&1; then + echo "::error::Tag '$REQUESTED_TAG' does not exist" + exit 1 + fi + TAG_SHA=$(git rev-list -n 1 "$REQUESTED_TAG") + HEAD_SHA=$(git rev-parse HEAD) + if [ "$TAG_SHA" != "$HEAD_SHA" ]; then + echo "::error::Checked out commit $HEAD_SHA does not match tag $REQUESTED_TAG at $TAG_SHA" + exit 1 + fi + PUBLISH_LATEST=false + case "$REQUESTED_TAG" in + v1.*) PUBLISH_LATEST=true ;; + esac + { + echo "tag=$REQUESTED_TAG" + echo "publish_latest=$PUBLISH_LATEST" + echo "sha=$HEAD_SHA" + } >> "$GITHUB_OUTPUT" + + build_single_arch: + name: Build and push ${{ matrix.arch }} + needs: prepare + strategy: + fail-fast: false + matrix: + include: + - arch: amd64 + platform: linux/amd64 + runner: ubuntu-latest + - arch: arm64 + platform: linux/arm64 + runner: ubuntu-24.04-arm + runs-on: ${{ matrix.runner }} + permissions: + contents: read + packages: write + id-token: write + steps: + - name: Check out release commit + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 1 + ref: ${{ needs.prepare.outputs.sha }} + + - name: Write release version + env: + RELEASE_TAG: ${{ needs.prepare.outputs.tag }} + run: echo "$RELEASE_TAG" > VERSION + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 + + - name: Log in to Docker Hub + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Log in to GHCR + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Resolve image tags + id: tags + env: + ARCH: ${{ matrix.arch }} + RELEASE_TAG: ${{ needs.prepare.outputs.tag }} + PUBLISH_LATEST: ${{ needs.prepare.outputs.publish_latest }} + run: | + { + echo 'value<> "$GITHUB_OUTPUT" + + - name: Extract OCI labels + id: meta + uses: docker/metadata-action@dc802804100637a589fabce1cb79ff13a1411302 # v6.2.0 + with: + images: | + ${{ env.DOCKERHUB_IMAGE }} + ${{ env.GHCR_IMAGE }} + labels: | + org.opencontainers.image.version=${{ needs.prepare.outputs.tag }} + org.opencontainers.image.revision=${{ needs.prepare.outputs.sha }} + + - name: Build and push architecture image + id: build + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0 + with: + context: . + platforms: ${{ matrix.platform }} + push: true + tags: ${{ steps.tags.outputs.value }} + labels: ${{ steps.meta.outputs.labels }} + cache-from: type=gha,scope=release-${{ matrix.arch }} + cache-to: type=gha,mode=max,scope=release-${{ matrix.arch }} + provenance: mode=max + sbom: true + + - name: Install cosign + uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 + + - name: Sign architecture images + env: + DIGEST: ${{ steps.build.outputs.digest }} + run: | + cosign sign --yes "$DOCKERHUB_IMAGE@$DIGEST" + cosign sign --yes "$GHCR_IMAGE@$DIGEST" + + create_manifests: + name: Create multi-architecture manifests + needs: [prepare, build_single_arch] + runs-on: ubuntu-latest + permissions: + contents: read + packages: write + id-token: write + steps: + - name: Log in to Docker Hub + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Log in to GHCR + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Create immutable manifests + env: + RELEASE_TAG: ${{ needs.prepare.outputs.tag }} + run: | + docker buildx imagetools create \ + -t "$DOCKERHUB_IMAGE:$RELEASE_TAG" \ + "$DOCKERHUB_IMAGE:$RELEASE_TAG-amd64" \ + "$DOCKERHUB_IMAGE:$RELEASE_TAG-arm64" + docker buildx imagetools create \ + -t "$GHCR_IMAGE:$RELEASE_TAG" \ + "$GHCR_IMAGE:$RELEASE_TAG-amd64" \ + "$GHCR_IMAGE:$RELEASE_TAG-arm64" + + - name: Create current-line latest manifests + if: needs.prepare.outputs.publish_latest == 'true' + run: | + docker buildx imagetools create \ + -t "$DOCKERHUB_IMAGE:latest" \ + "$DOCKERHUB_IMAGE:latest-amd64" \ + "$DOCKERHUB_IMAGE:latest-arm64" + docker buildx imagetools create \ + -t "$GHCR_IMAGE:latest" \ + "$GHCR_IMAGE:latest-amd64" \ + "$GHCR_IMAGE:latest-arm64" + + - name: Install cosign + uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 + + - name: Sign manifests + env: + RELEASE_TAG: ${{ needs.prepare.outputs.tag }} + PUBLISH_LATEST: ${{ needs.prepare.outputs.publish_latest }} + run: | + cosign sign --yes "$DOCKERHUB_IMAGE:$RELEASE_TAG" + cosign sign --yes "$GHCR_IMAGE:$RELEASE_TAG" + if [ "$PUBLISH_LATEST" = true ]; then + cosign sign --yes "$DOCKERHUB_IMAGE:latest" + cosign sign --yes "$GHCR_IMAGE:latest" + fi + + - name: Release summary + env: + RELEASE_TAG: ${{ needs.prepare.outputs.tag }} + run: | + { + echo '### Published fork images' + echo + echo "- $DOCKERHUB_IMAGE:$RELEASE_TAG" + echo "- $GHCR_IMAGE:$RELEASE_TAG" + if [ '${{ needs.prepare.outputs.publish_latest }}' = true ]; then + echo "- $DOCKERHUB_IMAGE:latest" + echo "- $GHCR_IMAGE:latest" + fi + } >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/docker-image-alpha.yml b/.github/workflows/docker-image-alpha.yml index 116dd1452152..13603e8bfbda 100644 --- a/.github/workflows/docker-image-alpha.yml +++ b/.github/workflows/docker-image-alpha.yml @@ -13,6 +13,8 @@ on: jobs: build_single_arch: name: Build & push (${{ matrix.arch }}) [native] + # Fork branch images are published explicitly by docker-image-branch.yml. + if: github.repository == 'Calcium-Ion/new-api' strategy: fail-fast: false matrix: diff --git a/.github/workflows/docker-image-arm64.yml b/.github/workflows/docker-image-arm64.yml index 83303ee30612..be19e4883d34 100644 --- a/.github/workflows/docker-image-arm64.yml +++ b/.github/workflows/docker-image-arm64.yml @@ -15,6 +15,10 @@ on: jobs: build_single_arch: name: Build & push (${{ matrix.arch }}) [native] + # Preserve the upstream legacy publisher for its original repository, but + # never publish calciumion/new-api from this fork. Fork releases use + # docker-build.yml and the karlorz namespaces instead. + if: github.repository == 'Calcium-Ion/new-api' strategy: fail-fast: false matrix: diff --git a/.github/workflows/docker-image-branch.yml b/.github/workflows/docker-image-branch.yml new file mode 100644 index 000000000000..3183a69571ef --- /dev/null +++ b/.github/workflows/docker-image-branch.yml @@ -0,0 +1,217 @@ +name: Publish Docker image (manual branch) + +on: + workflow_dispatch: + inputs: + branch: + description: "Branch name to build (e.g. alpha, nightly)" + required: true + type: string + +env: + DOCKERHUB_IMAGE: karlorz/new-api + GHCR_IMAGE: ghcr.io/karlorz/new-api + +jobs: + prepare: + name: Prepare Docker tags + runs-on: ubuntu-latest + outputs: + branch: ${{ steps.version.outputs.branch }} + sha: ${{ steps.version.outputs.sha }} + tag_prefix: ${{ steps.version.outputs.tag_prefix }} + version: ${{ steps.version.outputs.version }} + permissions: + contents: read + steps: + - name: Check out branch + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 1 + ref: ${{ inputs.branch }} + + - name: Resolve Docker tags + id: version + env: + BRANCH_NAME: ${{ inputs.branch }} + run: | + TAG_PREFIX=$(printf '%s' "$BRANCH_NAME" | tr '[:upper:]' '[:lower:]' | sed -E 's/[^a-z0-9_.-]+/-/g; s/^[.-]+//; s/[.-]+$//') + TAG_PREFIX=${TAG_PREFIX:0:105} + TAG_PREFIX=$(printf '%s' "$TAG_PREFIX" | sed -E 's/[.-]+$//') + if [ -z "$TAG_PREFIX" ]; then + echo "::error::Branch '$BRANCH_NAME' cannot be converted to a valid Docker tag prefix" + exit 1 + fi + + SHA=$(git rev-parse HEAD) + SHORT_SHA=$(git rev-parse --short HEAD) + VERSION="${TAG_PREFIX}-$(date +'%Y%m%d')-${SHORT_SHA}" + + { + echo "branch=$BRANCH_NAME" + echo "sha=$SHA" + echo "tag_prefix=$TAG_PREFIX" + echo "version=$VERSION" + } >> "$GITHUB_OUTPUT" + echo "Prepared Docker tags for $BRANCH_NAME at $SHORT_SHA" + + build_single_arch: + name: Build & push (${{ matrix.arch }}) [native] + needs: [prepare] + strategy: + fail-fast: false + matrix: + include: + - arch: amd64 + platform: linux/amd64 + runner: ubuntu-latest + - arch: arm64 + platform: linux/arm64 + runner: ubuntu-24.04-arm + runs-on: ${{ matrix.runner }} + permissions: + contents: read + packages: write + id-token: write + steps: + - name: Check out branch + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 1 + ref: ${{ needs.prepare.outputs.sha }} + + - name: Write VERSION + run: | + echo "${{ needs.prepare.outputs.version }}" > VERSION + echo "Publishing version: ${{ needs.prepare.outputs.version }} for ${{ matrix.arch }}" + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 + + - name: Log in to Docker Hub + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Log in to GHCR + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Extract metadata (labels) + id: meta + uses: docker/metadata-action@dc802804100637a589fabce1cb79ff13a1411302 # v6.2.0 + with: + images: | + ${{ env.DOCKERHUB_IMAGE }} + ${{ env.GHCR_IMAGE }} + + - name: Build & push single-arch + id: build + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0 + with: + context: . + platforms: ${{ matrix.platform }} + push: true + tags: | + ${{ env.DOCKERHUB_IMAGE }}:${{ needs.prepare.outputs.tag_prefix }}-${{ matrix.arch }} + ${{ env.DOCKERHUB_IMAGE }}:${{ needs.prepare.outputs.version }}-${{ matrix.arch }} + ${{ env.GHCR_IMAGE }}:${{ needs.prepare.outputs.tag_prefix }}-${{ matrix.arch }} + ${{ env.GHCR_IMAGE }}:${{ needs.prepare.outputs.version }}-${{ matrix.arch }} + labels: ${{ steps.meta.outputs.labels }} + cache-from: type=gha + cache-to: type=gha,mode=max + provenance: mode=max + sbom: true + + - name: Install cosign + uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 + + - name: Sign image with cosign + run: | + cosign sign --yes "$DOCKERHUB_IMAGE@${{ steps.build.outputs.digest }}" + cosign sign --yes "$GHCR_IMAGE@${{ steps.build.outputs.digest }}" + + - name: Output digest + run: | + { + echo "### Docker Image Digest (${{ matrix.arch }})" + echo '```' + echo "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.tag_prefix }}-${{ matrix.arch }}" + echo "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.version }}-${{ matrix.arch }}" + echo "$GHCR_IMAGE:${{ needs.prepare.outputs.tag_prefix }}-${{ matrix.arch }}" + echo "$GHCR_IMAGE:${{ needs.prepare.outputs.version }}-${{ matrix.arch }}" + echo "${{ steps.build.outputs.digest }}" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + create_manifests: + name: Create multi-arch manifests + needs: [prepare, build_single_arch] + runs-on: ubuntu-latest + permissions: + packages: write + id-token: write + steps: + - name: Log in to Docker Hub + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Log in to GHCR + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Create and push branch manifests + run: | + docker buildx imagetools create \ + -t "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.tag_prefix }}" \ + "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.tag_prefix }}-amd64" \ + "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.tag_prefix }}-arm64" + docker buildx imagetools create \ + -t "$GHCR_IMAGE:${{ needs.prepare.outputs.tag_prefix }}" \ + "$GHCR_IMAGE:${{ needs.prepare.outputs.tag_prefix }}-amd64" \ + "$GHCR_IMAGE:${{ needs.prepare.outputs.tag_prefix }}-arm64" + + - name: Create and push versioned manifests + run: | + docker buildx imagetools create \ + -t "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.version }}" \ + "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.version }}-amd64" \ + "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.version }}-arm64" + docker buildx imagetools create \ + -t "$GHCR_IMAGE:${{ needs.prepare.outputs.version }}" \ + "$GHCR_IMAGE:${{ needs.prepare.outputs.version }}-amd64" \ + "$GHCR_IMAGE:${{ needs.prepare.outputs.version }}-arm64" + + - name: Install cosign + uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 + + - name: Sign manifests with cosign + run: | + cosign sign --yes "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.tag_prefix }}" + cosign sign --yes "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.version }}" + cosign sign --yes "$GHCR_IMAGE:${{ needs.prepare.outputs.tag_prefix }}" + cosign sign --yes "$GHCR_IMAGE:${{ needs.prepare.outputs.version }}" + + - name: Output manifest digest + run: | + { + echo "### Multi-arch Manifest Digests" + echo '```' + docker buildx imagetools inspect "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.tag_prefix }}" + echo "---" + docker buildx imagetools inspect "$DOCKERHUB_IMAGE:${{ needs.prepare.outputs.version }}" + echo "---" + docker buildx imagetools inspect "$GHCR_IMAGE:${{ needs.prepare.outputs.tag_prefix }}" + echo "---" + docker buildx imagetools inspect "$GHCR_IMAGE:${{ needs.prepare.outputs.version }}" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/docker-image-nightly.yml b/.github/workflows/docker-image-nightly.yml index 2125fa9dd925..9fc1176f42f6 100644 --- a/.github/workflows/docker-image-nightly.yml +++ b/.github/workflows/docker-image-nightly.yml @@ -13,6 +13,8 @@ on: jobs: build_single_arch: name: Build & push (${{ matrix.arch }}) [native] + # Fork branch images are published explicitly by docker-image-branch.yml. + if: github.repository == 'Calcium-Ion/new-api' strategy: fail-fast: false matrix: diff --git a/README.md b/README.md index 8f23d5dcd380..fd04c77193aa 100644 --- a/README.md +++ b/README.md @@ -170,6 +170,7 @@ docker run --name new-api -d --restart always \ | 🚀 Deployment Guide | [Installation Documentation](https://docs.newapi.pro/en/docs/installation) | | ⚙️ Environment Configuration | [Environment Variables](https://docs.newapi.pro/en/docs/installation/config-maintenance/environment-variables) | | 📡 API Documentation | [API Documentation](https://docs.newapi.pro/en/docs/api) | +| 🔌 Fork Compatibility | [Codex Responses WebSocket fork status and upstream migration](./docs/responses-websocket-tracking.md) | | ❓ FAQ | [FAQ](https://docs.newapi.pro/en/docs/support/faq) | | 💬 Community Interaction | [Communication Channels](https://docs.newapi.pro/en/docs/support/community-interaction) | diff --git a/common/json.go b/common/json.go index 1625be6d51f7..68bbe7ea7c2c 100644 --- a/common/json.go +++ b/common/json.go @@ -6,6 +6,8 @@ import ( "io" ) +type RawMessage = json.RawMessage + func Unmarshal(data []byte, v any) error { return json.Unmarshal(data, v) } @@ -22,7 +24,7 @@ func Marshal(v any) ([]byte, error) { return json.Marshal(v) } -func GetJsonType(data json.RawMessage) string { +func GetJsonType(data RawMessage) string { trimmed := bytes.TrimSpace(data) if len(trimmed) == 0 { return "unknown" @@ -45,7 +47,7 @@ func GetJsonType(data json.RawMessage) string { } // JsonRawMessageToString returns JSON strings as their decoded value and other JSON values as raw text. -func JsonRawMessageToString(data json.RawMessage) string { +func JsonRawMessageToString(data RawMessage) string { trimmed := bytes.TrimSpace(data) if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { return "" diff --git a/common/origin.go b/common/origin.go new file mode 100644 index 000000000000..ba70bfc80805 --- /dev/null +++ b/common/origin.go @@ -0,0 +1,62 @@ +package common + +import ( + "fmt" + "net" + "net/url" + "os" + "strings" +) + +// SessionCookieTrustedURLs is also used as the trusted browser-origin list for +// authenticated WebSocket handshakes. Legacy deployments may configure the +// equivalent list with WEBSOCKET_TRUSTED_ORIGINS. +var SessionCookieTrustedURLs = loadWebSocketTrustedOrigins() + +// NormalizeOrigin validates and canonicalizes a browser origin. Only an exact +// scheme, host and effective-port match is accepted; paths and wildcards are +// not valid origins for credential-bearing WebSocket handshakes. +func NormalizeOrigin(raw string) (string, error) { + raw = strings.TrimSpace(raw) + if raw == "" || raw == "null" || strings.ContainsAny(raw, "\r\n") { + return "", fmt.Errorf("origin is empty or invalid") + } + parsedURL, err := url.Parse(raw) + if err != nil { + return "", fmt.Errorf("invalid origin: %w", err) + } + if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" { + return "", fmt.Errorf("origin scheme must be http or https") + } + if parsedURL.Host == "" || parsedURL.User != nil || parsedURL.RawQuery != "" || parsedURL.Fragment != "" || (parsedURL.Path != "" && parsedURL.Path != "/") { + return "", fmt.Errorf("origin must contain only scheme and host") + } + hostname := strings.ToLower(parsedURL.Hostname()) + if hostname == "" || strings.Contains(hostname, "*") { + return "", fmt.Errorf("origin host is empty") + } + port := parsedURL.Port() + normalizedHost := hostname + if strings.Contains(hostname, ":") { + normalizedHost = "[" + hostname + "]" + } + if port == "" || (parsedURL.Scheme == "http" && port == "80") || (parsedURL.Scheme == "https" && port == "443") { + return parsedURL.Scheme + "://" + normalizedHost, nil + } + return parsedURL.Scheme + "://" + net.JoinHostPort(hostname, port), nil +} + +func loadWebSocketTrustedOrigins() []string { + raw := strings.TrimSpace(os.Getenv("WEBSOCKET_TRUSTED_ORIGINS")) + if raw == "" { + return nil + } + trusted := make([]string, 0) + for _, candidate := range strings.Split(raw, ",") { + normalized, err := NormalizeOrigin(candidate) + if err == nil { + trusted = append(trusted, normalized) + } + } + return trusted +} diff --git a/common/rate-limit.go b/common/rate-limit.go index 301c101c9748..73fea0835a37 100644 --- a/common/rate-limit.go +++ b/common/rate-limit.go @@ -45,6 +45,9 @@ func (l *InMemoryRateLimiter) clearExpiredItems() { func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration int64) bool { l.mutex.Lock() defer l.mutex.Unlock() + if maxRequestNum == 0 { + return true + } // [old <-- new] queue, ok := l.store[key] now := time.Now().Unix() @@ -68,3 +71,19 @@ func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration in } return true } + +// Check reports whether a request would be allowed without recording it. +// The duration parameter's unit is seconds. +func (l *InMemoryRateLimiter) Check(key string, maxRequestNum int, duration int64) bool { + l.mutex.Lock() + defer l.mutex.Unlock() + if maxRequestNum == 0 { + return true + } + queue, ok := l.store[key] + if !ok || len(*queue) < maxRequestNum { + return true + } + now := time.Now().Unix() + return now-(*queue)[0] >= duration +} diff --git a/common/str.go b/common/str.go index 71391f722acf..aeb39c03b479 100644 --- a/common/str.go +++ b/common/str.go @@ -3,6 +3,7 @@ package common import ( "encoding/base64" "encoding/json" + "fmt" "net/url" "regexp" "strconv" @@ -12,6 +13,16 @@ import ( "github.com/samber/lo" ) +const LocalLogContentLimit = 2048 + +// LocalLogPreview limits log-only content unless debug logging is enabled. +func LocalLogPreview(content string) string { + if DebugEnabled || len(content) <= LocalLogContentLimit { + return content + } + return fmt.Sprintf("%s... [truncated, original_length=%d, limit=%d]", content[:LocalLogContentLimit], len(content), LocalLogContentLimit) +} + var ( maskURLPattern = regexp.MustCompile(`(http|https)://[^\s/$.?#].[^\s]*`) maskDomainPattern = regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`) diff --git a/controller/channel.go b/controller/channel.go index b0dd22861507..1851fb6fc746 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -68,6 +68,52 @@ func clearChannelInfo(channel *model.Channel) { } } +func channelIDsFromChannels(channels []*model.Channel) []int { + ids := make([]int, 0, len(channels)) + for _, channel := range channels { + if channel != nil && channel.Id > 0 { + ids = append(ids, channel.Id) + } + } + return ids +} + +func closeActiveChannelWebSockets(channelIDs []int) { + service.CloseActiveWebSocketsForChannels(channelIDs, service.ChannelDisabledCloseReason) +} + +func hasEnabledMultiKey(channel *model.Channel) bool { + if channel == nil || !channel.ChannelInfo.IsMultiKey { + return true + } + keys := channel.GetKeys() + if len(keys) == 0 { + return false + } + for i := range keys { + if channel.ChannelInfo.MultiKeyStatusList == nil { + return true + } + if status, ok := channel.ChannelInfo.MultiKeyStatusList[i]; !ok || status == common.ChannelStatusEnabled { + return true + } + } + return false +} + +func disableMultiKeyChannelIfUnavailable(channel *model.Channel) bool { + if channel == nil || !channel.ChannelInfo.IsMultiKey || hasEnabledMultiKey(channel) { + return false + } + if channel.Status != common.ChannelStatusEnabled { + return true + } + if !model.UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "All keys are disabled") { + return false + } + channel.Status = common.ChannelStatusManuallyDisabled + return true +} func GetAllChannels(c *gin.Context) { pageInfo := common.GetPageQuery(c) channelData := make([]*model.Channel, 0) @@ -672,6 +718,7 @@ func DeleteChannel(c *gin.Context) { return } model.InitChannelCache() + closeActiveChannelWebSockets([]int{id}) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -680,12 +727,20 @@ func DeleteChannel(c *gin.Context) { } func DeleteDisabledChannel(c *gin.Context) { + var ids []int + if err := model.DB.Model(&model.Channel{}). + Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled). + Pluck("id", &ids).Error; err != nil { + common.ApiError(c, err) + return + } rows, err := model.DeleteDisabledChannel() if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + closeActiveChannelWebSockets(ids) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -716,12 +771,19 @@ func DisableTagChannels(c *gin.Context) { }) return } + channels, err := model.GetChannelsByTag(channelTag.Tag, false, false) + if err != nil { + common.ApiError(c, err) + return + } + ids := channelIDsFromChannels(channels) err = model.DisableChannelByTag(channelTag.Tag) if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + closeActiveChannelWebSockets(ids) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -825,6 +887,7 @@ func DeleteChannelBatch(c *gin.Context) { return } model.InitChannelCache() + closeActiveChannelWebSockets(channelBatch.Ids) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -959,6 +1022,9 @@ func UpdateChannel(c *gin.Context) { return } model.InitChannelCache() + if channel.Status != common.ChannelStatusEnabled { + closeActiveChannelWebSockets([]int{channel.Id}) + } service.ResetProxyClientCache() channel.Key = "" clearChannelInfo(&channel.Channel) @@ -1419,7 +1485,11 @@ func ManageMultiKeys(c *gin.Context) { return } + shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel) model.InitChannelCache() + if shouldCloseWebSocket { + closeActiveChannelWebSockets([]int{channel.Id}) + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "密钥已禁用", @@ -1532,7 +1602,11 @@ func ManageMultiKeys(c *gin.Context) { return } + shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel) model.InitChannelCache() + if shouldCloseWebSocket { + closeActiveChannelWebSockets([]int{channel.Id}) + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": fmt.Sprintf("已禁用 %d 个密钥", disabledCount), @@ -1612,7 +1686,11 @@ func ManageMultiKeys(c *gin.Context) { return } + shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel) model.InitChannelCache() + if shouldCloseWebSocket { + closeActiveChannelWebSockets([]int{channel.Id}) + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "密钥已删除", @@ -1680,7 +1758,11 @@ func ManageMultiKeys(c *gin.Context) { return } + shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel) model.InitChannelCache() + if shouldCloseWebSocket { + closeActiveChannelWebSockets([]int{channel.Id}) + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount), diff --git a/controller/relay.go b/controller/relay.go index c97ab45b4ac4..024deb73fb7b 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -1,13 +1,13 @@ package controller import ( + "crypto/subtle" "errors" "fmt" "io" "log" "net/http" "strings" - "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -22,9 +22,9 @@ import ( "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/system_setting" "github.com/QuantumNous/new-api/types" - "github.com/bytedance/gopkg/util/gopool" "github.com/samber/lo" "github.com/gin-gonic/gin" @@ -64,6 +64,22 @@ func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewA return err } +func ResponsesWebSocket(c *gin.Context) { + requestId := c.GetString(common.RequestIdKey) + ws, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + return + } + defer ws.Close() + + if newAPIError := relay.ResponsesWebSocketHelper(c, ws); newAPIError != nil { + errorPreview := common.LocalLogPreview(newAPIError.Error()) + logger.LogError(c, fmt.Sprintf("responses websocket relay error: %s", errorPreview)) + newAPIError.SetMessage(common.MessageWithRequestId(newAPIError.Error(), requestId)) + helper.WssError(c, ws, newAPIError.ToOpenAIError()) + } +} + func Relay(c *gin.Context, relayFormat types.RelayFormat) { requestId := c.GetString(common.RequestIdKey) @@ -241,11 +257,56 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } } +// EnableCompression is deliberately left off. gorilla's permessage-deflate is +// experimental and has no context takeover, so each message compresses in +// isolation — the streamed delta events are too small to gain from that, while +// every frame costs a deflate round trip. More importantly, SetReadLimit is +// enforced against the compressed wire length, so enabling compression would +// turn the WebSocket read limits into a compressed bound and reopen the zip +// bomb hole that MAX_REQUEST_BODY_MB exists to close on the HTTP side. var upgrader = websocket.Upgrader{ - Subprotocols: []string{"realtime"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol TODO add other protocol - CheckOrigin: func(r *http.Request) bool { - return true // 允许跨域 - }, + Subprotocols: []string{"realtime", "responses"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol + CheckOrigin: isAllowedWebSocketOrigin, +} + +// isAllowedWebSocketOrigin keeps authenticated CLI clients compatible while +// preventing a browser on an unrelated site from opening a credential-bearing +// relay socket. CLI clients such as Codex normally omit Origin. When a browser +// supplies it, require an exact normalized match against the request host, the +// configured server address, or the existing trusted browser-origin list. +func isAllowedWebSocketOrigin(request *http.Request) bool { + if request == nil { + return false + } + originValues := request.Header.Values("Origin") + if len(originValues) == 0 { + return true + } + if len(originValues) != 1 || strings.Contains(originValues[0], ",") { + return false + } + origin, err := common.NormalizeOrigin(originValues[0]) + if err != nil { + return false + } + + allowedOrigins := make([]string, 0, 3+len(common.SessionCookieTrustedURLs)) + for _, candidate := range []string{ + "http://" + request.Host, + "https://" + request.Host, + system_setting.ServerAddress, + } { + if normalized, normalizeErr := common.NormalizeOrigin(candidate); normalizeErr == nil { + allowedOrigins = append(allowedOrigins, normalized) + } + } + allowedOrigins = append(allowedOrigins, common.SessionCookieTrustedURLs...) + for _, allowedOrigin := range allowedOrigins { + if subtle.ConstantTimeCompare([]byte(origin), []byte(allowedOrigin)) == 1 { + return true + } + } + return false } func addUsedChannel(c *gin.Context, channelId int) { @@ -316,82 +377,11 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service } func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { - if openaiErr == nil { - return false - } - if service.ShouldSkipRetryAfterChannelAffinityFailure(c) { - return false - } - if types.IsChannelError(openaiErr) { - return true - } - if types.IsSkipRetryError(openaiErr) { - return false - } - if retryTimes <= 0 { - return false - } - if _, ok := c.Get("specific_channel_id"); ok { - return false - } - code := openaiErr.StatusCode - if code >= 200 && code < 300 { - return false - } - if code < 100 || code > 599 { - return true - } - if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) { - return false - } - return operation_setting.ShouldRetryByStatusCode(code) + return service.ShouldRetryRelayError(c, openaiErr, retryTimes) } func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) { - logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, err.Error())) - // 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况 - // do not use context to get channel info, there may be inconsistent channel info when processing asynchronously - if service.ShouldDisableChannel(err) && channelError.AutoBan { - gopool.Go(func() { - service.DisableChannel(channelError, err.ErrorWithStatusCode()) - }) - } - - if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) { - // 保存错误日志到mysql中 - userId := c.GetInt("id") - tokenName := c.GetString("token_name") - modelName := c.GetString("original_model") - tokenId := c.GetInt("token_id") - userGroup := c.GetString("group") - channelId := c.GetInt("channel_id") - other := make(map[string]interface{}) - if c.Request != nil && c.Request.URL != nil { - other["request_path"] = c.Request.URL.Path - } - other["error_type"] = err.GetErrorType() - other["error_code"] = err.GetErrorCode() - other["status_code"] = err.StatusCode - other["channel_id"] = channelId - other["channel_name"] = c.GetString("channel_name") - other["channel_type"] = c.GetInt("channel_type") - adminInfo := make(map[string]interface{}) - adminInfo["use_channel"] = c.GetStringSlice("use_channel") - isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey) - if isMultiKey { - adminInfo["is_multi_key"] = true - adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex) - } - service.AppendChannelAffinityAdminInfo(c, adminInfo) - other["admin_info"] = adminInfo - startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime) - if startTime.IsZero() { - startTime = time.Now() - } - useTimeSeconds := int(time.Since(startTime).Seconds()) - model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other) - } - + service.ProcessChannelError(c, channelError, err) } func RelayMidjourney(c *gin.Context) { diff --git a/controller/relay_websocket_origin_test.go b/controller/relay_websocket_origin_test.go new file mode 100644 index 000000000000..875f7564af5d --- /dev/null +++ b/controller/relay_websocket_origin_test.go @@ -0,0 +1,54 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/stretchr/testify/assert" +) + +func TestAllowedWebSocketOrigin(t *testing.T) { + previousServerAddress := system_setting.ServerAddress + previousTrustedURLs := common.SessionCookieTrustedURLs + system_setting.ServerAddress = "https://panel.example.com" + common.SessionCookieTrustedURLs = []string{"https://trusted.example.com"} + t.Cleanup(func() { + system_setting.ServerAddress = previousServerAddress + common.SessionCookieTrustedURLs = previousTrustedURLs + }) + + tests := []struct { + name string + origin []string + allowed bool + }{ + {name: "codex cli without origin", allowed: true}, + {name: "same https host", origin: []string{"https://api.example.com"}, allowed: true}, + {name: "same http host", origin: []string{"http://api.example.com"}, allowed: true}, + {name: "configured server address", origin: []string{"https://panel.example.com"}, allowed: true}, + {name: "trusted browser origin", origin: []string{"https://trusted.example.com"}, allowed: true}, + {name: "untrusted host", origin: []string{"https://evil.example.com"}}, + {name: "trusted suffix attack", origin: []string{"https://trusted.example.com.evil.test"}}, + {name: "null origin", origin: []string{"null"}}, + {name: "origin with path", origin: []string{"https://api.example.com/path"}}, + {name: "comma joined origins", origin: []string{"https://api.example.com, https://evil.example.com"}}, + {name: "multiple origin headers", origin: []string{"https://api.example.com", "https://evil.example.com"}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + request := httptest.NewRequest("GET", "http://api.example.com/v1/responses", nil) + request.Host = "api.example.com" + for _, origin := range test.origin { + request.Header.Add("Origin", origin) + } + assert.Equal(t, test.allowed, isAllowedWebSocketOrigin(request)) + }) + } +} + +func TestAllowedWebSocketOriginRejectsNilRequest(t *testing.T) { + assert.False(t, isAllowedWebSocketOrigin(nil)) +} diff --git a/docs/responses-websocket-tracking.md b/docs/responses-websocket-tracking.md new file mode 100644 index 000000000000..55639a42d98c --- /dev/null +++ b/docs/responses-websocket-tracking.md @@ -0,0 +1,324 @@ +# Codex Responses WebSocket fork tracking + +This page tracks the fork-maintained Codex Responses WebSocket relay and the +conditions required to replace it with official upstream New API support. + +> [!IMPORTANT] +> The relay described here is available on the maintained `karlorz/new-api` +> fork release branches. It is not yet part of an official tagged +> `QuantumNous/new-api` release. Do not retire the fork implementation merely +> because an upstream pull request is opened or merged. + +## Status at a glance + +Last verified: **2026-08-02** + + + +| Item | Status | +| --- | --- | +| Official upstream `GET /v1/responses` WebSocket relay | Not available in an official tagged release | +| Upstream implementation candidate | [QuantumNous/new-api PR #5062](https://github.com/QuantumNous/new-api/pull/5062), open and conflicting as of the last verification | +| Current fork branch | `release/responses-websocket-v1` | +| Current fork release | [`v1.0.0-rc.23-0`](https://github.com/karlorz/new-api/releases/tag/v1.0.0-rc.23-0) | +| Legacy fork branch | `release/responses-websocket-v0.13` | +| Legacy fork release | [`v0.13.2-0`](https://github.com/karlorz/new-api/releases/tag/v0.13.2-0) | +| Fork state | Maintained while waiting for a verified official upstream tag | + + + +## Why the fork exists + +Codex can attempt to establish a WebSocket connection at: + +```text +wss:///v1/responses +``` + +Without an authenticated `GET /v1/responses` route, the request returns 404 +and Codex falls back to an HTTPS transport. The fork adds a long-lived +Responses WebSocket relay while preserving the gateway responsibilities that +cannot safely be bypassed: authentication, model and channel selection, rate +limits, quota pre-consumption, observed-usage settlement, streaming logs, +channel lifecycle, and retry behavior. + +The implementation lineage includes the protocol work represented by upstream +PR #5062, hardening and accounting work from the Yorick-Ryu fork line, and +additional integration, security, release, and lifecycle corrections made in +this fork. The complete release deltas are available here: + + + +- [Current release delta](https://github.com/karlorz/new-api/compare/v1.0.0-rc.23...v1.0.0-rc.23-0) +- [Legacy release delta: `v0.13.2...v0.13.2-0`](https://github.com/karlorz/new-api/compare/v0.13.2...v0.13.2-0) + + + +## Supported release lines + + + +| Line | Upstream base | Fork tag | Architectures | Docker `latest` | +| --- | --- | --- | --- | --- | +| Current | `v1.0.0-rc.23` | `v1.0.0-rc.23-0` | `linux/amd64`, `linux/arm64` | Yes | +| Legacy | `v0.13.2` | `v0.13.2-0` | `linux/amd64`, `linux/arm64` | Never | + + + +Published images: + +```text +docker.io/karlorz/new-api:v1.0.0-rc.23-0 +ghcr.io/karlorz/new-api:v1.0.0-rc.23-0 +sha256:d1bfdb1200e50ad17a9a5e59e0c206566f87cf694eb10a5ab3222440f1027076 + +docker.io/karlorz/new-api:v0.13.2-0 +ghcr.io/karlorz/new-api:v0.13.2-0 +sha256:07335c8511794e69f82e393043e4304daa625b23ba0aa3e78f3f43bcc6c22284 +``` + +The Docker Hub and GHCR `latest` tags are current-line aliases and must resolve +to the same digest as `v1.0.0-rc.23-0`. A v0.13 release must never update +`latest`. + +The upstream tags remain preserved. Fork releases use an additional numeric +suffix and do not move, replace, or overwrite the corresponding upstream tag. + +## Public protocol contract + +The maintained fork contract is: + +- Route: authenticated `GET /v1/responses` with a WebSocket upgrade. +- Supported subprotocols: `responses` and `realtime`. +- The first data event must be `response.create`. +- A connection may carry multiple sequential Responses calls. +- The first logical call performs channel selection and establishes the locked + model/channel target. Each call independently performs model validation, + rate-limit evaluation, quota pre-consumption, and settlement against that + established target. +- Responses WebSocket upstream relay is limited to compatible OpenAI/Codex + channel paths. +- Terminal response events complete the current call without forcing the + client socket to close when the protocol permits another call. +- Invalid events, relay failures, channel closure, and idle timeout produce + deterministic errors or close behavior. + +The route and controller entry points are visible in the current release at: + +- [`router/relay-router.go`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/router/relay-router.go) +- [`controller/relay.go`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/controller/relay.go) +- [`relay/responses_websocket.go`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/relay/responses_websocket.go) + +The corresponding legacy implementation is branch-specific: + +- [`router/relay-router.go`](https://github.com/karlorz/new-api/blob/9cde69dfcc3d622f3dea600ad77fb55554de6aef/router/relay-router.go) +- [`controller/relay.go`](https://github.com/karlorz/new-api/blob/9cde69dfcc3d622f3dea600ad77fb55554de6aef/controller/relay.go) +- [`relay/responses_websocket.go`](https://github.com/karlorz/new-api/blob/9cde69dfcc3d622f3dea600ad77fb55554de6aef/relay/responses_websocket.go) + +## Safety and accounting contract + +The fork implementation includes the following controls: + +- WebSocket Origin validation when an Origin header is supplied. +- Per-message read limits. +- WebSocket compression disabled for the Responses relay. +- Idle read deadlines refreshed by data messages. +- Per-user active Responses WebSocket limits. +- Per-call model rate-limit checks and commits rather than a handshake-only + allowance. +- Model/channel lock behavior that prevents unsafe switching after a session + has established its supported target. +- Quota pre-consumption before upstream work. +- Observed-usage settlement and refund behavior after terminal events or + failures. +- Streaming usage logs, including supported tool and image accounting on the + current RelayKit line. +- Channel affinity and retry behavior. +- Local and Redis-backed connection registration so channel disable/delete can + close affected sockets across nodes. +- Explicit process-shutdown closure is not currently implemented for hijacked + Responses WebSockets. Channel lifecycle and idle-policy closes are covered; + shutdown behavior remains a documented comparison point for future fork and + upstream releases. + +Relevant current-line sources: + +- [`middleware/model-rate-limit.go`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/middleware/model-rate-limit.go) +- [`relay/common/websocket_idle.go`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/relay/common/websocket_idle.go) +- [`pkg/wsmanager/wsmanager.go`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/pkg/wsmanager/wsmanager.go) +- [`service/ws_close.go`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/service/ws_close.go) + +## Configuration + + + +| Variable | Default | Purpose | +| --- | --- | --- | +| `WEBSOCKET_MAX_MESSAGE_MB` | Falls back to `MAX_REQUEST_BODY_MB` | Maximum accepted client WebSocket data message size | +| `MAX_REQUEST_BODY_MB` | 128 MiB | Standard effective fallback for the WebSocket message limit | +| `WEBSOCKET_IDLE_TIMEOUT_MINUTES` | 10 minutes | Idle read timeout, refreshed by data messages | +| `RESPONSES_WEBSOCKET_MAX_PER_USER` | 8 | Maximum concurrent Responses WebSockets per authenticated user | + + + +The implementation contains a final defensive 32 MiB fallback only if the +initialized `MAX_REQUEST_BODY_MB` value is also non-positive. Normal startup +initializes `MAX_REQUEST_BODY_MB` to 128 MiB. + +Deployments behind Cloudflare or another reverse proxy must allow WebSocket +upgrade forwarding. Compose-based deployments should prefer their private +service network for application, database, and Redis communication; joining a +shared global Docker network can introduce ambiguous service-name resolution. + +## Branch implementation distinction + +The two release lines expose the same intended external contract, but they are +not byte-identical implementations: + +- The v1 line follows the current architecture and RelayKit request, usage, + performance, and tool/image accounting paths. +- The v0.13 line retains the legacy root DTO and relay architecture required by + deployments pinned to the v0.13 database/application line. + +Fixes must be reviewed and tested on both branches. A current-line RelayKit +change must not be assumed to apply directly to v0.13, and the legacy branch +must not receive a wholesale merge from current `main`. + +## Release and registry policy + +Fork release tags must: + +- preserve the upstream base tag; +- add a numeric fork suffix, such as `-0`, `-1`, or a later integer; +- point to the exact intended release-branch commit; +- publish immutable multi-architecture images to Docker Hub and GHCR; +- include provenance/SBOM attestations and signatures where supported; +- update `latest` only for the current v1 line. + +The release policy is implemented in: + +- [`.github/workflows/docker-build.yml`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/.github/workflows/docker-build.yml) +- [`scripts/validate-fork-release-tag.sh`](https://github.com/karlorz/new-api/blob/749eb8c605c01f7c368f3ff8540d89a3424161fc/scripts/validate-fork-release-tag.sh) + +## Deployment verification + +Use all of the following checks before calling a deployment WebSocket-ready: + +1. `GET /api/status` returns HTTP 200 and the expected fork version. +2. An unauthenticated WebSocket-upgrade probe to `GET /v1/responses` reaches + New API and returns 401 rather than 404. This proves routing, not a complete + authenticated WebSocket exchange. +3. An authenticated client receives `101 Switching Protocols` with a supported + subprotocol. +4. A valid `response.create` exchange reaches terminal response events. +5. Multiple sequential calls on the same connection complete without stale + events or cross-call leakage. +6. Usage logs, rate-limit state, quota settlement, and refunds match the + completed or failed call. +7. Channel disable/delete closes an associated connection. +8. Idle timeout, message limit, per-user connection limit, invalid Origin, and + malformed first-event behavior are exercised. +9. The test passes through the actual reverse proxy and deployment platform, + not only through a local direct port. +10. The previously published image remains available for rollback. + +During the recorded operator verification on 2026-08-02, the fork staging +deployment passed health and proxy routing: `/api/status` returned 200 and an +unauthenticated WebSocket-upgrade request reached the relay authentication +layer with 401 rather than the previous 404. This observation is deployment +evidence rather than a repository test. An authenticated 101 and complete +Responses exchange remain required whenever a new fork or official upstream +image is evaluated. + +## Upstream watch list + +Monitor: + +- [QuantumNous/new-api PR #5062](https://github.com/QuantumNous/new-api/pull/5062), + or any official successor implementation; +- upstream changes that register `GET /v1/responses`; +- controller, relay, billing, rate-limit, Origin, frame-limit, idle-timeout, + connection-limit, and channel-lifecycle behavior; +- official New API release notes and tags; +- official tagged multi-architecture container images and digests; +- Codex transport behavior and fallback telemetry when relevant. + +The tracker state should move through these stages: + +```text +fork-maintained + ↓ +upstream-merged + ↓ +upstream-tagged + ↓ +staging-verified + ↓ +production-approved + ↓ +fork-relay-retired +``` + +## Official-support acceptance gate + +Official upstream support can replace the fork only when a tagged upstream +artifact satisfies every applicable condition: + +- Authenticated `GET /v1/responses` is present. +- The handshake returns 101 through the real Cloudflare/deployment path. +- Supported Responses subprotocol negotiation works. +- First-event validation and deterministic error/close behavior work. +- Multiple sequential Responses calls work on one connection. +- Channel/model selection, retries, affinity, and failure behavior are safe. +- Rate limits are evaluated per logical Responses call. +- Quota pre-consumption, observed usage, tool/image usage, settlement, refunds, + and logs are correct. +- Origin validation, compression policy, message limits, idle timeout, and + per-user connection limits are present. +- Channel disable/delete closes associated sockets. Process-shutdown behavior + is explicitly tested and documented rather than assumed from HTTP server + shutdown behavior. +- Official release notes and source identify the included implementation. +- A tagged, digest-pinnable multi-architecture official image is available. +- Fork-versus-upstream staging comparison and rollback tests pass. + +An open PR, merged PR, untagged main-branch commit, unauthenticated 401 probe, +or mutable container tag is not sufficient by itself. + +## Retirement and rollback procedure + +When the acceptance gate passes: + +1. Pin the official upstream image by version and digest in staging. +2. Run the complete deployment-verification checklist. +3. Compare protocol events, usage logs, rate limits, quota settlement, and + channel lifecycle with the fork release. +4. Record the result and obtain production approval. +5. Deploy the official image while retaining the last verified fork image as + rollback. +6. Observe the agreed production window. +7. Remove fork-specific relay code only in a separate reviewed change. +8. Archive this tracker only after the fork retirement and rollback decisions + are recorded. + +If verification fails, return to the last verified fork digest and record the +failed upstream version and failure reason. Do not move or overwrite existing +release tags. + +## Maintenance record + + + +| Date | Event | +| --- | --- | +| 2026-08-01 | Published current fork release `v1.0.0-rc.23-0` to GitHub, Docker Hub, and GHCR. | +| 2026-08-01 | Published legacy fork release `v0.13.2-0` to GitHub, Docker Hub, and GHCR without changing `latest`. | +| 2026-08-02 | Verified both release manifests as `linux/amd64` and `linux/arm64`; verified `latest` points to the current-line digest. | +| 2026-08-02 | Operator verification observed staging health and a WebSocket-upgrade probe reaching `/v1/responses` authentication with 401 instead of 404. | +| 2026-08-02 | Confirmed upstream PR #5062 remains open and conflicting; fork remains the maintained implementation. | + + + +Update this table when an upstream implementation merges, an official tag is +published, staging verification is attempted, or the migration/rollback state +changes. diff --git a/dto/claude.go b/dto/claude.go index d7fed412aaa9..3ef91610f3ec 100644 --- a/dto/claude.go +++ b/dto/claude.go @@ -414,7 +414,7 @@ func (c *ClaudeRequest) GetTools() []any { func (c *ClaudeRequest) GetEfforts() string { var OutputConfig OutputConfigForEffort - if err := json.Unmarshal(c.OutputConfig, &OutputConfig); err == nil { + if err := common.Unmarshal(c.OutputConfig, &OutputConfig); err == nil { effort := OutputConfig.Effort return effort } diff --git a/dto/gemini.go b/dto/gemini.go index 489ebea534b4..673eea92c353 100644 --- a/dto/gemini.go +++ b/dto/gemini.go @@ -44,9 +44,9 @@ func (r *GeminiChatRequest) UnmarshalJSON(data []byte) error { } type ToolConfig struct { - FunctionCallingConfig *FunctionCallingConfig `json:"functionCallingConfig,omitempty"` - RetrievalConfig *RetrievalConfig `json:"retrievalConfig,omitempty"` - IncludeServerSideToolInvocations *bool `json:"includeServerSideToolInvocations,omitempty"` + FunctionCallingConfig *FunctionCallingConfig `json:"functionCallingConfig,omitempty"` + RetrievalConfig *RetrievalConfig `json:"retrievalConfig,omitempty"` + IncludeServerSideToolInvocations *bool `json:"includeServerSideToolInvocations,omitempty"` } type FunctionCallingConfig struct { @@ -160,10 +160,9 @@ func (r *GeminiChatRequest) SetTools(tools []GeminiChatTool) { } type GeminiThinkingConfig struct { - IncludeThoughts bool `json:"includeThoughts,omitempty"` - ThinkingBudget *int `json:"thinkingBudget,omitempty"` - // TODO Conflict with thinkingbudget. - ThinkingLevel string `json:"thinkingLevel,omitempty"` + IncludeThoughts *bool `json:"includeThoughts,omitempty"` + ThinkingBudget *int `json:"thinkingBudget,omitempty"` + ThinkingLevel *string `json:"thinkingLevel,omitempty"` } // UnmarshalJSON allows GeminiThinkingConfig to accept both snake_case and camelCase fields. @@ -171,9 +170,9 @@ func (c *GeminiThinkingConfig) UnmarshalJSON(data []byte) error { type Alias GeminiThinkingConfig var aux struct { Alias - IncludeThoughtsSnake *bool `json:"include_thoughts,omitempty"` - ThinkingBudgetSnake *int `json:"thinking_budget,omitempty"` - ThinkingLevelSnake string `json:"thinking_level,omitempty"` + IncludeThoughtsSnake *bool `json:"include_thoughts,omitempty"` + ThinkingBudgetSnake *int `json:"thinking_budget,omitempty"` + ThinkingLevelSnake *string `json:"thinking_level,omitempty"` } if err := common.Unmarshal(data, &aux); err != nil { @@ -183,14 +182,14 @@ func (c *GeminiThinkingConfig) UnmarshalJSON(data []byte) error { *c = GeminiThinkingConfig(aux.Alias) if aux.IncludeThoughtsSnake != nil { - c.IncludeThoughts = *aux.IncludeThoughtsSnake + c.IncludeThoughts = aux.IncludeThoughtsSnake } if aux.ThinkingBudgetSnake != nil { c.ThinkingBudget = aux.ThinkingBudgetSnake } - if aux.ThinkingLevelSnake != "" { + if aux.ThinkingLevelSnake != nil { c.ThinkingLevel = aux.ThinkingLevelSnake } diff --git a/dto/gemini_generation_config_test.go b/dto/gemini_generation_config_test.go index ed4beb301943..3a2b02c9b112 100644 --- a/dto/gemini_generation_config_test.go +++ b/dto/gemini_generation_config_test.go @@ -87,3 +87,30 @@ func TestGeminiChatGenerationConfigPreservesExplicitZeroValuesSnakeCase(t *testi assert.Equal(t, float64(0), generationConfig["seed"]) assert.Equal(t, false, generationConfig["responseLogprobs"]) } + +func TestGeminiThinkingConfigPreservesExplicitFalseAndEmptyLevel(t *testing.T) { + raw := []byte(`{ + "contents":[{"role":"user","parts":[{"text":"hello"}]}], + "generationConfig":{"thinkingConfig":{"includeThoughts":false,"thinkingLevel":""}} + }`) + + var req GeminiChatRequest + require.NoError(t, common.Unmarshal(raw, &req)) + require.NotNil(t, req.GenerationConfig.ThinkingConfig) + require.NotNil(t, req.GenerationConfig.ThinkingConfig.IncludeThoughts) + require.NotNil(t, req.GenerationConfig.ThinkingConfig.ThinkingLevel) + assert.False(t, *req.GenerationConfig.ThinkingConfig.IncludeThoughts) + assert.Equal(t, "", *req.GenerationConfig.ThinkingConfig.ThinkingLevel) + + encoded, err := common.Marshal(req) + require.NoError(t, err) + + var out map[string]any + require.NoError(t, common.Unmarshal(encoded, &out)) + generationConfig, ok := out["generationConfig"].(map[string]any) + require.True(t, ok) + thinkingConfig, ok := generationConfig["thinkingConfig"].(map[string]any) + require.True(t, ok) + assert.Equal(t, false, thinkingConfig["includeThoughts"]) + assert.Equal(t, "", thinkingConfig["thinkingLevel"]) +} diff --git a/dto/message_reasoning_test.go b/dto/message_reasoning_test.go new file mode 100644 index 000000000000..630aa94395cd --- /dev/null +++ b/dto/message_reasoning_test.go @@ -0,0 +1,104 @@ +package dto + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +// TestMessageReasoningContentPreservesEmptyString verifies that an explicitly +// set empty reasoning_content string survives the JSON round-trip. +// +// This is critical for the request-forwarding path (non-passThrough mode): +// the gateway unmarshals the client request into GeneralOpenAIRequest, then +// re-marshals it before sending upstream. If Message.ReasoningContent were +// `string` + `omitempty` (the old type), the empty string would be silently +// dropped, causing the upstream to never receive the field. +// +// With the fix (`*string` + `omitempty`), nil = absent, &"" = explicit empty. +func TestMessageReasoningContentPreservesEmptyString(t *testing.T) { + raw := []byte(`{ + "role": "assistant", + "content": "Hello", + "reasoning_content": "", + "reasoning": "" + }`) + + var msg Message + err := common.Unmarshal(raw, &msg) + require.NoError(t, err) + + // Pointers must be non-nil: the field was explicitly set to "" + require.NotNil(t, msg.ReasoningContent, "reasoning_content should be non-nil when explicitly set to empty string") + require.NotNil(t, msg.Reasoning, "reasoning should be non-nil when explicitly set to empty string") + require.Equal(t, "", *msg.ReasoningContent) + require.Equal(t, "", *msg.Reasoning) + + // Re-marshal — the fields must still be present in the output JSON + encoded, err := common.Marshal(msg) + require.NoError(t, err) + + require.True(t, gjson.GetBytes(encoded, "reasoning_content").Exists(), + "reasoning_content should exist in re-marshaled JSON when explicitly set to empty string") + require.True(t, gjson.GetBytes(encoded, "reasoning").Exists(), + "reasoning should exist in re-marshaled JSON when explicitly set to empty string") + require.Equal(t, "", gjson.GetBytes(encoded, "reasoning_content").String()) + require.Equal(t, "", gjson.GetBytes(encoded, "reasoning").String()) +} + +// TestMessageReasoningContentOmitsAbsentField verifies that when +// reasoning_content / reasoning are absent from the input JSON, they remain +// absent after a round-trip (nil pointer → omitted by omitempty). +func TestMessageReasoningContentOmitsAbsentField(t *testing.T) { + raw := []byte(`{ + "role": "assistant", + "content": "Hello" + }`) + + var msg Message + err := common.Unmarshal(raw, &msg) + require.NoError(t, err) + + // Pointers must be nil: the fields were not present in the input + require.Nil(t, msg.ReasoningContent) + require.Nil(t, msg.Reasoning) + + // Re-marshal — the fields must NOT appear in the output JSON + encoded, err := common.Marshal(msg) + require.NoError(t, err) + + require.False(t, gjson.GetBytes(encoded, "reasoning_content").Exists(), + "reasoning_content should not exist in re-marshaled JSON when absent from input") + require.False(t, gjson.GetBytes(encoded, "reasoning").Exists(), + "reasoning should not exist in re-marshaled JSON when absent from input") +} + +// TestMessageGetReasoningContent verifies the GetReasoningContent helper +// method that is used in token-counting code paths. +func TestMessageGetReasoningContent(t *testing.T) { + t.Run("both nil returns empty", func(t *testing.T) { + msg := Message{Role: "assistant"} + require.Equal(t, "", msg.GetReasoningContent()) + }) + + t.Run("ReasoningContent takes priority", func(t *testing.T) { + rc := "thinking..." + r := "should be ignored" + msg := Message{ReasoningContent: &rc, Reasoning: &r} + require.Equal(t, "thinking...", msg.GetReasoningContent()) + }) + + t.Run("falls back to Reasoning when ReasoningContent is nil", func(t *testing.T) { + r := "fallback reasoning" + msg := Message{Reasoning: &r} + require.Equal(t, "fallback reasoning", msg.GetReasoningContent()) + }) + + t.Run("empty string values returned correctly", func(t *testing.T) { + empty := "" + msg := Message{ReasoningContent: &empty} + require.Equal(t, "", msg.GetReasoningContent()) + }) +} diff --git a/dto/openai_request.go b/dto/openai_request.go index 25ef3a21aa51..487340add42e 100644 --- a/dto/openai_request.go +++ b/dto/openai_request.go @@ -279,8 +279,8 @@ type Message struct { Content any `json:"content"` Name *string `json:"name,omitempty"` Prefix *bool `json:"prefix,omitempty"` - ReasoningContent string `json:"reasoning_content,omitempty"` - Reasoning string `json:"reasoning,omitempty"` + ReasoningContent *string `json:"reasoning_content,omitempty"` + Reasoning *string `json:"reasoning,omitempty"` ToolCalls json.RawMessage `json:"tool_calls,omitempty"` ToolCallId string `json:"tool_call_id,omitempty"` parsedContent []MediaContent @@ -431,6 +431,16 @@ const ( //ContentTypeAudioUrl = "audio_url" ) +func (m *Message) GetReasoningContent() string { + if m.ReasoningContent == nil && m.Reasoning == nil { + return "" + } + if m.ReasoningContent != nil { + return *m.ReasoningContent + } + return *m.Reasoning +} + func (m *Message) GetPrefix() bool { if m.Prefix == nil { return false @@ -946,6 +956,11 @@ type Input struct { Type string `json:"type,omitempty"` Role string `json:"role,omitempty"` Content json.RawMessage `json:"content,omitempty"` + // Output carries function_call_output payloads. It must be counted: an + // incremental turn (notably over the Responses WebSocket transport) can + // consist of nothing but a large tool result, which would otherwise be + // estimated at zero tokens and under-reserve quota. + Output json.RawMessage `json:"output,omitempty"` } type MediaInput struct { @@ -985,6 +1000,14 @@ func (r *OpenAIResponsesRequest) ParseInput() []MediaInput { var inputs []Input _ = common.Unmarshal(r.Input, &inputs) for _, input := range inputs { + if len(input.Output) > 0 { + text := string(input.Output) + if common.GetJsonType(input.Output) == "string" { + _ = common.Unmarshal(input.Output, &text) + } + mediaInputs = append(mediaInputs, MediaInput{Type: "input_text", Text: text}) + } + if common.GetJsonType(input.Content) == "string" { var str string _ = common.Unmarshal(input.Content, &str) diff --git a/dto/openai_response.go b/dto/openai_response.go index 0e6b818dbd8b..5256151654b4 100644 --- a/dto/openai_response.go +++ b/dto/openai_response.go @@ -3,6 +3,7 @@ package dto import ( "encoding/json" "fmt" + "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/types" @@ -228,11 +229,12 @@ type Usage struct { UsageSemantic string `json:"usage_semantic,omitempty"` UsageSource string `json:"usage_source,omitempty"` - PromptTokensDetails InputTokenDetails `json:"prompt_tokens_details"` - CompletionTokenDetails OutputTokenDetails `json:"completion_tokens_details"` - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - InputTokensDetails *InputTokenDetails `json:"input_tokens_details"` + PromptTokensDetails InputTokenDetails `json:"prompt_tokens_details"` + CompletionTokenDetails OutputTokenDetails `json:"completion_tokens_details"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + InputTokensDetails *InputTokenDetails `json:"input_tokens_details"` + OutputTokensDetails *OutputTokenDetails `json:"output_tokens_details,omitempty"` // claude cache 1h ClaudeCacheCreation5mTokens int `json:"claude_cache_creation_5_m_tokens"` @@ -403,6 +405,19 @@ type ResponsesStreamResponse struct { Part *ResponsesReasoningSummaryPart `json:"part,omitempty"` } +// IsResponsesTransportEventType reports non-semantic frames that some upstream +// proxies or CLI gateways inject to keep long Responses streams/sockets open. +// These are not part of the OpenAI Responses streaming event model and break +// clients that deserialize event types into a closed enum (e.g. serde). +func IsResponsesTransportEventType(eventType string) bool { + switch strings.ToLower(strings.TrimSpace(eventType)) { + case "keepalive", "keep_alive", "heartbeat": + return true + default: + return false + } +} + // GetOpenAIError 从动态错误类型中提取OpenAIError结构 func GetOpenAIError(errorField any) *types.OpenAIError { if errorField == nil { diff --git a/dto/openai_response_keepalive_test.go b/dto/openai_response_keepalive_test.go new file mode 100644 index 000000000000..65a198286974 --- /dev/null +++ b/dto/openai_response_keepalive_test.go @@ -0,0 +1,33 @@ +package dto + +import "testing" + +func TestIsResponsesTransportEventType(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + eventType string + want bool + }{ + {name: "keepalive", eventType: "keepalive", want: true}, + {name: "keep_alive", eventType: "keep_alive", want: true}, + {name: "heartbeat", eventType: "heartbeat", want: true}, + {name: "case insensitive", eventType: "KeepAlive", want: true}, + {name: "trimmed", eventType: " keepalive ", want: true}, + {name: "response created", eventType: "response.created", want: false}, + {name: "response completed", eventType: "response.completed", want: false}, + {name: "error", eventType: "error", want: false}, + {name: "empty", eventType: "", want: false}, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := IsResponsesTransportEventType(tc.eventType); got != tc.want { + t.Fatalf("IsResponsesTransportEventType(%q) = %v, want %v", tc.eventType, got, tc.want) + } + }) + } +} diff --git a/dto/responses_token_count_test.go b/dto/responses_token_count_test.go new file mode 100644 index 000000000000..897462fd5444 --- /dev/null +++ b/dto/responses_token_count_test.go @@ -0,0 +1,59 @@ +package dto + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestResponsesTokenCountMetaIncludesFunctionCallOutput guards pre-consume +// sizing for tool-result turns. A Responses turn can consist of nothing but a +// function_call_output item; counting only `content` estimated such a turn at +// zero tokens, which under-reserved quota and skipped the prompt sensitive-word +// check. This is most visible over the WebSocket transport, where upstream +// keeps conversation state and each turn carries only the new items. +func TestResponsesTokenCountMetaIncludesFunctionCallOutput(t *testing.T) { + toolResult := "the weather in Shanghai is 31C and humid" + + var request OpenAIResponsesRequest + require.NoError(t, common.Unmarshal([]byte(`{ + "model": "gpt-4o", + "input": [ + {"type": "function_call_output", "call_id": "call_1", "output": "`+toolResult+`"} + ] + }`), &request)) + + meta := request.GetTokenCountMeta() + require.NotNil(t, meta) + assert.Contains(t, meta.CombineText, toolResult, "function_call_output must reach token counting") +} + +func TestResponsesTokenCountMetaIncludesStructuredFunctionCallOutput(t *testing.T) { + var request OpenAIResponsesRequest + require.NoError(t, common.Unmarshal([]byte(`{ + "model": "gpt-4o", + "input": [ + {"type": "function_call_output", "call_id": "call_1", "output": {"temperature": "31C", "city": "Shanghai"}} + ] + }`), &request)) + + meta := request.GetTokenCountMeta() + require.NotNil(t, meta) + assert.Contains(t, meta.CombineText, "Shanghai", "structured tool output must reach token counting") +} + +func TestResponsesTokenCountMetaStillCountsPlainContent(t *testing.T) { + var request OpenAIResponsesRequest + require.NoError(t, common.Unmarshal([]byte(`{ + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", "content": "hello there"} + ] + }`), &request)) + + meta := request.GetTokenCountMeta() + require.NotNil(t, meta) + assert.Contains(t, meta.CombineText, "hello there") +} diff --git a/main.go b/main.go index dbbf44a1826b..a3e5159dbe6d 100644 --- a/main.go +++ b/main.go @@ -2,6 +2,7 @@ package main import ( "bytes" + "context" "embed" "fmt" "log" @@ -19,6 +20,7 @@ import ( "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/oauth" + "github.com/QuantumNous/new-api/pkg/wsmanager" "github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/router" "github.com/QuantumNous/new-api/service" @@ -89,6 +91,7 @@ func main() { go model.SyncChannelCache(common.SyncFrequency) } + wsmanager.StartSubscriber(context.Background()) // 热更新配置 go model.SyncOptions(common.SyncFrequency) diff --git a/middleware/auth.go b/middleware/auth.go index 23d933fbe0c1..ef394788221b 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -276,20 +276,7 @@ func TokenAuthReadOnly() func(c *gin.Context) { func TokenAuth() func(c *gin.Context) { return func(c *gin.Context) { // 先检测是否为ws - if c.Request.Header.Get("Sec-WebSocket-Protocol") != "" { - // Sec-WebSocket-Protocol: realtime, openai-insecure-api-key.sk-xxx, openai-beta.realtime-v1 - // read sk from Sec-WebSocket-Protocol - key := c.Request.Header.Get("Sec-WebSocket-Protocol") - parts := strings.Split(key, ",") - for _, part := range parts { - part = strings.TrimSpace(part) - if strings.HasPrefix(part, "openai-insecure-api-key") { - key = strings.TrimPrefix(part, "openai-insecure-api-key.") - break - } - } - c.Request.Header.Set("Authorization", "Bearer "+key) - } + applyWebSocketSubprotocolAuthorization(c.Request.Header) // 检查path包含/v1/messages 或 /v1/models if strings.Contains(c.Request.URL.Path, "/v1/messages") || strings.Contains(c.Request.URL.Path, "/v1/models") { anthropicKey := c.Request.Header.Get("x-api-key") @@ -406,6 +393,31 @@ func TokenAuth() func(c *gin.Context) { } } +func applyWebSocketSubprotocolAuthorization(header http.Header) bool { + key, ok := apiKeyFromWebSocketSubprotocol(header.Get("Sec-WebSocket-Protocol")) + if !ok { + return false + } + header.Set("Authorization", "Bearer "+key) + return true +} + +func apiKeyFromWebSocketSubprotocol(protocols string) (string, bool) { + if protocols == "" { + return "", false + } + const insecureAPIKeyPrefix = "openai-insecure-api-key." + parts := strings.Split(protocols, ",") + for _, part := range parts { + part = strings.TrimSpace(part) + if strings.HasPrefix(part, insecureAPIKeyPrefix) { + key := strings.TrimPrefix(part, insecureAPIKeyPrefix) + return key, key != "" + } + } + return "", false +} + func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) error { if token == nil { return fmt.Errorf("token is nil") diff --git a/middleware/auth_test.go b/middleware/auth_test.go new file mode 100644 index 000000000000..edef280b5996 --- /dev/null +++ b/middleware/auth_test.go @@ -0,0 +1,86 @@ +package middleware + +import ( + "net/http" + "testing" +) + +func TestAPIKeyFromWebSocketSubprotocol(t *testing.T) { + tests := []struct { + name string + protocols string + wantKey string + wantOK bool + }{ + { + name: "responses protocol only", + protocols: "responses", + wantOK: false, + }, + { + name: "realtime protocol only", + protocols: "realtime", + wantOK: false, + }, + { + name: "responses with insecure key", + protocols: "responses, openai-insecure-api-key.sk-test", + wantKey: "sk-test", + wantOK: true, + }, + { + name: "realtime with beta and insecure key", + protocols: "realtime, openai-insecure-api-key.sk-realtime, openai-beta.realtime-v1", + wantKey: "sk-realtime", + wantOK: true, + }, + { + name: "empty insecure key", + protocols: "responses, openai-insecure-api-key.", + wantOK: false, + }, + { + name: "bare insecure marker is not a key", + protocols: "openai-insecure-api-key", + wantOK: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotKey, gotOK := apiKeyFromWebSocketSubprotocol(tt.protocols) + if gotOK != tt.wantOK { + t.Fatalf("ok = %v, want %v", gotOK, tt.wantOK) + } + if gotKey != tt.wantKey { + t.Fatalf("key = %q, want %q", gotKey, tt.wantKey) + } + }) + } +} + +func TestApplyWebSocketSubprotocolAuthorizationDoesNotOverrideProtocolOnly(t *testing.T) { + header := http.Header{} + header.Set("Authorization", "Bearer sk-original") + header.Set("Sec-WebSocket-Protocol", "responses") + + if applyWebSocketSubprotocolAuthorization(header) { + t.Fatal("authorization was unexpectedly applied") + } + if got := header.Get("Authorization"); got != "Bearer sk-original" { + t.Fatalf("Authorization = %q, want original bearer", got) + } +} + +func TestApplyWebSocketSubprotocolAuthorizationOverridesWithInsecureKey(t *testing.T) { + header := http.Header{} + header.Set("Authorization", "Bearer sk-original") + header.Set("Sec-WebSocket-Protocol", "responses, openai-insecure-api-key.sk-from-protocol") + + if !applyWebSocketSubprotocolAuthorization(header) { + t.Fatal("authorization was not applied") + } + if got := header.Get("Authorization"); got != "Bearer sk-from-protocol" { + t.Fatalf("Authorization = %q, want protocol bearer", got) + } +} diff --git a/middleware/model-rate-limit.go b/middleware/model-rate-limit.go index 80a3995df097..c84a2e6ad4a1 100644 --- a/middleware/model-rate-limit.go +++ b/middleware/model-rate-limit.go @@ -5,12 +5,14 @@ import ( "fmt" "net/http" "strconv" + "strings" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common/limiter" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/go-redis/redis/v8" @@ -21,6 +23,8 @@ const ( ModelRequestRateLimitSuccessCountMark = "MRRLS" ) +type ModelRequestRateLimitCommit func(success bool) + // 检查Redis中的请求限制 func checkRedisRateLimit(ctx context.Context, rdb *redis.Client, key string, maxCount int, duration int64) (bool, error) { // 如果maxCount为0,表示不限制 @@ -74,30 +78,56 @@ func recordRedisRequest(ctx context.Context, rdb *redis.Client, key string, maxC rdb.Expire(ctx, key, time.Duration(setting.ModelRequestRateLimitDurationMinutes)*time.Minute) } -// Redis限流处理器 -func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { - return func(c *gin.Context) { - userId := strconv.Itoa(c.GetInt("id")) +func modelRequestRateLimitConfig(c *gin.Context) (duration int64, totalMaxCount int, successMaxCount int) { + duration = int64(setting.ModelRequestRateLimitDurationMinutes * 60) + totalMaxCount = setting.ModelRequestRateLimitCount + successMaxCount = setting.ModelRequestRateLimitSuccessCount + + group := common.GetContextKeyString(c, constant.ContextKeyTokenGroup) + if group == "" { + group = common.GetContextKeyString(c, constant.ContextKeyUserGroup) + } + groupTotalCount, groupSuccessCount, found := setting.GetGroupRateLimit(group) + if found { + totalMaxCount = groupTotalCount + successMaxCount = groupSuccessCount + } + return duration, totalMaxCount, successMaxCount +} + +func newModelRateLimitError(message string, statusCode int) *types.NewAPIError { + return types.NewErrorWithStatusCode( + fmt.Errorf("%s", message), + types.ErrorCodeInvalidRequest, + statusCode, + types.ErrOptionWithSkipRetry(), + types.ErrOptionWithNoRecordErrorLog(), + ) +} + +func CheckModelRequestRateLimit(c *gin.Context) (ModelRequestRateLimitCommit, *types.NewAPIError) { + if !setting.ModelRequestRateLimitEnabled { + return func(bool) {}, nil + } + + duration, totalMaxCount, successMaxCount := modelRequestRateLimitConfig(c) + userId := strconv.Itoa(c.GetInt("id")) + + if common.RedisEnabled { ctx := context.Background() rdb := common.RDB - - // 1. 检查成功请求数限制 successKey := fmt.Sprintf("rateLimit:%s:%s", ModelRequestRateLimitSuccessCountMark, userId) allowed, err := checkRedisRateLimit(ctx, rdb, successKey, successMaxCount, duration) if err != nil { fmt.Println("检查成功请求数限制失败:", err.Error()) - abortWithOpenAiMessage(c, http.StatusInternalServerError, "rate_limit_check_failed") - return + return nil, newModelRateLimitError("rate_limit_check_failed", http.StatusInternalServerError) } if !allowed { - abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到请求数限制:%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount)) - return + return nil, newModelRateLimitError(fmt.Sprintf("您已达到请求数限制:%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount), http.StatusTooManyRequests) } - //2.检查总请求数限制并记录总请求(当totalMaxCount为0时会自动跳过,使用令牌桶限流器 if totalMaxCount > 0 { totalKey := fmt.Sprintf("rateLimit:%s", userId) - // 初始化 tb := limiter.New(ctx, rdb) allowed, err = tb.Allow( ctx, @@ -106,95 +136,66 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g limiter.WithRate(int64(totalMaxCount)), limiter.WithRequested(duration), ) - if err != nil { fmt.Println("检查总请求数限制失败:", err.Error()) - abortWithOpenAiMessage(c, http.StatusInternalServerError, "rate_limit_check_failed") - return + return nil, newModelRateLimitError("rate_limit_check_failed", http.StatusInternalServerError) } - if !allowed { - abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到总请求数限制:%d分钟内最多请求%d次,包括失败次数,请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount)) + return nil, newModelRateLimitError(fmt.Sprintf("您已达到总请求数限制:%d分钟内最多请求%d次,包括失败次数,请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount), http.StatusTooManyRequests) } } - // 4. 处理请求 - c.Next() - - // 5. 如果请求成功,记录成功请求 - if c.Writer.Status() < 400 { - recordRedisRequest(ctx, rdb, successKey, successMaxCount) - } + return func(success bool) { + if success { + recordRedisRequest(ctx, rdb, successKey, successMaxCount) + } + }, nil } -} -// 内存限流处理器 -func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { inMemoryRateLimiter.Init(time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute) + totalKey := ModelRequestRateLimitCountMark + userId + successKey := ModelRequestRateLimitSuccessCountMark + userId - return func(c *gin.Context) { - userId := strconv.Itoa(c.GetInt("id")) - totalKey := ModelRequestRateLimitCountMark + userId - successKey := ModelRequestRateLimitSuccessCountMark + userId - - // 1. 检查总请求数限制(当totalMaxCount为0时跳过) - if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) { - c.Status(http.StatusTooManyRequests) - c.Abort() - return - } - - // 2. 检查成功请求数限制 - // 使用一个临时key来检查限制,这样可以避免实际记录 - checkKey := successKey + "_check" - if !inMemoryRateLimiter.Request(checkKey, successMaxCount, duration) { - c.Status(http.StatusTooManyRequests) - c.Abort() - return - } - - // 3. 处理请求 - c.Next() + if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) { + return nil, newModelRateLimitError(fmt.Sprintf("您已达到总请求数限制:%d分钟内最多请求%d次,包括失败次数,请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount), http.StatusTooManyRequests) + } + if successMaxCount > 0 && !inMemoryRateLimiter.Check(successKey, successMaxCount, duration) { + return nil, newModelRateLimitError(fmt.Sprintf("您已达到请求数限制:%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount), http.StatusTooManyRequests) + } - // 4. 如果请求成功,记录到实际的成功请求计数中 - if c.Writer.Status() < 400 { + return func(success bool) { + if success && successMaxCount > 0 { inMemoryRateLimiter.Request(successKey, successMaxCount, duration) } - } + }, nil +} + +func isResponsesWebSocketHandshake(c *gin.Context) bool { + return c != nil && + c.Request != nil && + c.Request.Method == http.MethodGet && + c.Request.URL != nil && + c.Request.URL.Path == "/v1/responses" && + strings.EqualFold(c.Request.Header.Get("Upgrade"), "websocket") } // ModelRequestRateLimit 模型请求限流中间件 func ModelRequestRateLimit() func(c *gin.Context) { return func(c *gin.Context) { - // 在每个请求时检查是否启用限流 if !setting.ModelRequestRateLimitEnabled { c.Next() return } - - // 计算限流参数 - duration := int64(setting.ModelRequestRateLimitDurationMinutes * 60) - totalMaxCount := setting.ModelRequestRateLimitCount - successMaxCount := setting.ModelRequestRateLimitSuccessCount - - // 获取分组 - group := common.GetContextKeyString(c, constant.ContextKeyTokenGroup) - if group == "" { - group = common.GetContextKeyString(c, constant.ContextKeyUserGroup) - } - - //获取分组的限流配置 - groupTotalCount, groupSuccessCount, found := setting.GetGroupRateLimit(group) - if found { - totalMaxCount = groupTotalCount - successMaxCount = groupSuccessCount + if isResponsesWebSocketHandshake(c) { + c.Next() + return } - - // 根据存储类型选择并执行限流处理器 - if common.RedisEnabled { - redisRateLimitHandler(duration, totalMaxCount, successMaxCount)(c) - } else { - memoryRateLimitHandler(duration, totalMaxCount, successMaxCount)(c) + commit, apiErr := CheckModelRequestRateLimit(c) + if apiErr != nil { + abortWithOpenAiMessage(c, apiErr.StatusCode, apiErr.Error(), apiErr.GetErrorCode()) + return } + c.Next() + commit(c.Writer.Status() < 400) } } diff --git a/pkg/wsmanager/wsmanager.go b/pkg/wsmanager/wsmanager.go new file mode 100644 index 000000000000..5d05981acdda --- /dev/null +++ b/pkg/wsmanager/wsmanager.go @@ -0,0 +1,241 @@ +package wsmanager + +import ( + "context" + "fmt" + "os" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" +) + +const ( + KindRealtime = "realtime" + KindResponses = "responses" + + DefaultCloseReason = "channel disabled or deleted" + redisChannel = "new-api:wsmanager:channel-close" +) + +type entry struct { + id uint64 + channelID int + kind string + close func(reason string) +} + +type closeEvent struct { + ChannelIDs []int `json:"channel_ids"` + Reason string `json:"reason"` + Origin string `json:"origin"` +} + +var ( + mu sync.Mutex + nextID uint64 + registry = map[int]map[uint64]*entry{} + + originOnce sync.Once + originID string + + subscriberOnce sync.Once +) + +func Register(channelID int, kind string, close func(reason string)) func() { + if channelID <= 0 || close == nil { + return func() {} + } + var closeOnce sync.Once + safeClose := func(reason string) { + closeOnce.Do(func() { + close(reason) + }) + } + id := atomic.AddUint64(&nextID, 1) + e := &entry{ + id: id, + channelID: channelID, + kind: kind, + close: safeClose, + } + + mu.Lock() + if registry[channelID] == nil { + registry[channelID] = map[uint64]*entry{} + } + registry[channelID][id] = e + mu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + mu.Lock() + defer mu.Unlock() + entries := registry[channelID] + if entries == nil { + return + } + delete(entries, id) + if len(entries) == 0 { + delete(registry, channelID) + } + }) + } +} + +func CloseChannel(channelID int, reason string) int { + return CloseChannels([]int{channelID}, reason) +} + +func CloseChannels(channelIDs []int, reason string) int { + reason = normalizeReason(reason) + entries := takeEntries(channelIDs) + for _, e := range entries { + e.close(reason) + } + if len(entries) > 0 { + common.SysLog(fmt.Sprintf("closed %d active websocket connection(s), channels=%v, kinds=%v, reason=%s", len(entries), entryChannelIDs(entries), entryKindCounts(entries), reason)) + } + return len(entries) +} + +func CloseChannelsAndBroadcast(channelIDs []int, reason string) int { + count := CloseChannels(channelIDs, reason) + if err := PublishCloseChannels(context.Background(), channelIDs, reason); err != nil { + common.SysLog(fmt.Sprintf("failed to publish websocket close event: %v", err)) + } + return count +} + +func StartSubscriber(ctx context.Context) { + if !common.RedisEnabled || common.RDB == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + subscriberOnce.Do(func() { + go subscribe(ctx) + }) +} + +func PublishCloseChannels(ctx context.Context, channelIDs []int, reason string) error { + if !common.RedisEnabled || common.RDB == nil { + return nil + } + ids := uniqueChannelIDs(channelIDs) + if len(ids) == 0 { + return nil + } + payload, err := common.Marshal(closeEvent{ + ChannelIDs: ids, + Reason: normalizeReason(reason), + Origin: getOriginID(), + }) + if err != nil { + return err + } + return common.RDB.Publish(ctx, redisChannel, payload).Err() +} + +func subscribe(ctx context.Context) { + pubsub := common.RDB.Subscribe(ctx, redisChannel) + defer pubsub.Close() + + ch := pubsub.Channel() + for { + select { + case <-ctx.Done(): + return + case msg, ok := <-ch: + if !ok { + return + } + var event closeEvent + if err := common.Unmarshal([]byte(msg.Payload), &event); err != nil { + common.SysLog(fmt.Sprintf("failed to unmarshal websocket close event: %v", err)) + continue + } + if event.Origin == getOriginID() { + continue + } + CloseChannels(event.ChannelIDs, event.Reason) + } + } +} + +func takeEntries(channelIDs []int) []*entry { + ids := uniqueChannelIDs(channelIDs) + if len(ids) == 0 { + return nil + } + + mu.Lock() + defer mu.Unlock() + + var entries []*entry + for _, channelID := range ids { + for _, e := range registry[channelID] { + entries = append(entries, e) + } + delete(registry, channelID) + } + return entries +} + +func entryChannelIDs(entries []*entry) []int { + ids := make([]int, 0, len(entries)) + for _, e := range entries { + if e != nil { + ids = append(ids, e.channelID) + } + } + return uniqueChannelIDs(ids) +} + +func entryKindCounts(entries []*entry) map[string]int { + counts := make(map[string]int) + for _, e := range entries { + if e == nil { + continue + } + counts[e.kind]++ + } + return counts +} + +func uniqueChannelIDs(channelIDs []int) []int { + seen := make(map[int]struct{}, len(channelIDs)) + ids := make([]int, 0, len(channelIDs)) + for _, id := range channelIDs { + if id <= 0 { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + ids = append(ids, id) + } + return ids +} + +func normalizeReason(reason string) string { + if reason == "" { + return DefaultCloseReason + } + return reason +} + +func getOriginID() string { + originOnce.Do(func() { + name := common.NodeName + if name == "" { + name = "node" + } + originID = fmt.Sprintf("%s-%d-%d", name, os.Getpid(), time.Now().UnixNano()) + }) + return originID +} diff --git a/pkg/wsmanager/wsmanager_test.go b/pkg/wsmanager/wsmanager_test.go new file mode 100644 index 000000000000..162e2bea1f82 --- /dev/null +++ b/pkg/wsmanager/wsmanager_test.go @@ -0,0 +1,108 @@ +package wsmanager + +import ( + "context" + "sync" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func resetRegistryForTest() { + mu.Lock() + defer mu.Unlock() + registry = map[int]map[uint64]*entry{} + nextID = 0 +} + +func TestCloseChannelClosesRegisteredConnectionsOnce(t *testing.T) { + resetRegistryForTest() + + var mu sync.Mutex + calls := 0 + Register(10, KindRealtime, func(reason string) { + mu.Lock() + defer mu.Unlock() + require.Equal(t, "test reason", reason, "close callback should receive the provided reason") + calls++ + }) + Register(10, KindResponses, func(reason string) { + mu.Lock() + defer mu.Unlock() + calls++ + }) + + require.Equal(t, 2, CloseChannel(10, "test reason"), "CloseChannel should close every registered connection for the channel") + require.Equal(t, 0, CloseChannel(10, "test reason"), "CloseChannel should not close already removed connections") + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, 2, calls, "all registered close callbacks should be called once") +} + +func TestCloseChannelDoesNotCloseOtherChannels(t *testing.T) { + resetRegistryForTest() + + calls := map[int]int{} + Register(10, KindRealtime, func(reason string) { + calls[10]++ + }) + Register(20, KindRealtime, func(reason string) { + calls[20]++ + }) + + require.Equal(t, 1, CloseChannel(10, "test"), "CloseChannel should only close the requested channel") + assert.Equal(t, 1, calls[10], "requested channel callback should run") + assert.Equal(t, 0, calls[20], "other channel callback should not run") +} + +func TestUnregisterPreventsClose(t *testing.T) { + resetRegistryForTest() + + calls := 0 + unregister := Register(10, KindRealtime, func(reason string) { + calls++ + }) + unregister() + + require.Equal(t, 0, CloseChannel(10, "test"), "unregistered connections should not be closed") + assert.Equal(t, 0, calls, "unregistered close callback should not run") +} + +func TestRegisteredCloseIsIdempotent(t *testing.T) { + resetRegistryForTest() + + calls := 0 + Register(10, KindRealtime, func(reason string) { + calls++ + }) + + mu.Lock() + var registered *entry + for _, e := range registry[10] { + registered = e + } + mu.Unlock() + require.NotNil(t, registered, "registered entry should exist") + + registered.close("test") + registered.close("test") + assert.Equal(t, 1, calls, "registered close callback should be idempotent") +} + +func TestPublishCloseChannelsNoopsWhenRedisDisabled(t *testing.T) { + resetRegistryForTest() + + oldEnabled := common.RedisEnabled + oldRDB := common.RDB + common.RedisEnabled = false + common.RDB = nil + defer func() { + common.RedisEnabled = oldEnabled + common.RDB = oldRDB + }() + + require.NoError(t, PublishCloseChannels(context.Background(), []int{10}, "test"), "publishing should no-op without Redis") +} diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index fa8234523c77..e177e56dab14 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -567,12 +567,14 @@ func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextRe } choice.SetStringContent(responseText) if len(responseThinking) > 0 { - choice.ReasoningContent = responseThinking + choice.ReasoningContent = &responseThinking } if len(tools) > 0 { choice.Message.SetToolCalls(tools) } - choice.Message.ReasoningContent = thinkingContent + if thinkingContent != "" { + choice.Message.ReasoningContent = &thinkingContent + } fullTextResponse.Model = claudeResponse.Model choices = append(choices, choice) fullTextResponse.Choices = choices diff --git a/relay/channel/gemini/adaptor.go b/relay/channel/gemini/adaptor.go index 680c4ee484ec..2bd5c8797673 100644 --- a/relay/channel/gemini/adaptor.go +++ b/relay/channel/gemini/adaptor.go @@ -7,6 +7,7 @@ import ( "net/http" "strings" + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel/openai" @@ -23,6 +24,41 @@ import ( type Adaptor struct { } +// resolveClaudeThinkingConfig translates only Claude-native thinking controls. +// It deliberately lives in the Gemini adaptor so the generic Claude-to-OpenAI +// conversion remains provider-neutral. +func resolveClaudeThinkingConfig(req *dto.ClaudeRequest) (*dto.GeminiThinkingConfig, string) { + if req.Thinking != nil && req.Thinking.Type == "disabled" { + return &dto.GeminiThinkingConfig{IncludeThoughts: common.GetPointer(false)}, "" + } + + if len(req.OutputConfig) > 0 { + if effort := req.GetEfforts(); effort != "" { + return &dto.GeminiThinkingConfig{ + IncludeThoughts: common.GetPointer(true), + ThinkingLevel: common.GetPointer(effort), + }, effort + } + } + + if req.Thinking == nil { + return nil, "" + } + + switch req.Thinking.Type { + case "enabled": + config := &dto.GeminiThinkingConfig{IncludeThoughts: common.GetPointer(true)} + if req.Thinking.BudgetTokens != nil { + config.ThinkingBudget = req.Thinking.BudgetTokens + } + return config, "" + case "adaptive": + return &dto.GeminiThinkingConfig{IncludeThoughts: common.GetPointer(true)}, "" + default: + return nil, "" + } +} + func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) { if len(request.Contents) > 0 { for i, content := range request.Contents { @@ -49,7 +85,19 @@ func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayIn if err != nil { return nil, err } - return a.ConvertOpenAIRequest(c, info, oaiReq.(*dto.GeneralOpenAIRequest)) + converted, err := a.ConvertOpenAIRequest(c, info, oaiReq.(*dto.GeneralOpenAIRequest)) + if err != nil { + return nil, err + } + geminiRequest, ok := converted.(*dto.GeminiChatRequest) + if !ok { + return converted, nil + } + if config, effort := resolveClaudeThinkingConfig(req); config != nil { + geminiRequest.GenerationConfig.ThinkingConfig = config + info.ReasoningEffort = effort + } + return geminiRequest, nil } func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) { @@ -129,8 +177,7 @@ func (a *Adaptor) Init(info *relaycommon.RelayInfo) { func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { - if model_setting.GetGeminiSettings().ThinkingAdapterEnabled && - !model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) { + if shouldApplyGeminiThinkingAdapter(info) { // 新增逻辑:处理 -thinking- 格式 if strings.Contains(info.UpstreamModelName, "-thinking-") { parts := strings.Split(info.UpstreamModelName, "-thinking-") diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 21641e483861..bea46769b5f9 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -131,8 +131,14 @@ func clampThinkingBudgetByEffort(modelName string, effort string) int { return clampThinkingBudget(modelName, maxBudget) } +func shouldApplyGeminiThinkingAdapter(info *relaycommon.RelayInfo) bool { + return model_setting.GetGeminiSettings().ThinkingAdapterEnabled && + !model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) && + !model_setting.ShouldPreserveThinkingSuffix(info.UpstreamModelName) +} + func ThinkingAdaptor(geminiRequest *dto.GeminiChatRequest, info *relaycommon.RelayInfo, oaiRequest ...dto.GeneralOpenAIRequest) { - if model_setting.GetGeminiSettings().ThinkingAdapterEnabled { + if shouldApplyGeminiThinkingAdapter(info) { modelName := info.UpstreamModelName isNew25Pro := strings.HasPrefix(modelName, "gemini-2.5-pro") && !strings.HasPrefix(modelName, "gemini-2.5-pro-preview-05-06") && @@ -145,7 +151,7 @@ func ThinkingAdaptor(geminiRequest *dto.GeminiChatRequest, info *relaycommon.Rel clampedBudget := clampThinkingBudget(modelName, budgetTokens) geminiRequest.GenerationConfig.ThinkingConfig = &dto.GeminiThinkingConfig{ ThinkingBudget: common.GetPointer(clampedBudget), - IncludeThoughts: true, + IncludeThoughts: common.GetPointer(true), } } } @@ -164,11 +170,11 @@ func ThinkingAdaptor(geminiRequest *dto.GeminiChatRequest, info *relaycommon.Rel if isUnsupported { geminiRequest.GenerationConfig.ThinkingConfig = &dto.GeminiThinkingConfig{ - IncludeThoughts: true, + IncludeThoughts: common.GetPointer(true), } } else { geminiRequest.GenerationConfig.ThinkingConfig = &dto.GeminiThinkingConfig{ - IncludeThoughts: true, + IncludeThoughts: common.GetPointer(true), } if geminiRequest.GenerationConfig.MaxOutputTokens != nil && *geminiRequest.GenerationConfig.MaxOutputTokens > 0 { budgetTokens := model_setting.GetGeminiSettings().ThinkingAdapterBudgetTokensPercentage * float64(*geminiRequest.GenerationConfig.MaxOutputTokens) @@ -189,8 +195,8 @@ func ThinkingAdaptor(geminiRequest *dto.GeminiChatRequest, info *relaycommon.Rel } } else if _, level, ok := reasoning.TrimEffortSuffix(info.UpstreamModelName); ok && level != "" { geminiRequest.GenerationConfig.ThinkingConfig = &dto.GeminiThinkingConfig{ - IncludeThoughts: true, - ThinkingLevel: level, + IncludeThoughts: common.GetPointer(true), + ThinkingLevel: common.GetPointer(level), } info.ReasoningEffort = level } @@ -238,7 +244,7 @@ func CovertOpenAI2Gemini(c *gin.Context, textRequest dto.GeneralOpenAIRequest, i geminiRequest.GenerationConfig.StopSequences = stopSequences } - adaptorWithExtraBody := false + explicitThinkingConfig := false // patch extra_body if len(textRequest.ExtraBody) > 0 { @@ -249,69 +255,70 @@ func CovertOpenAI2Gemini(c *gin.Context, textRequest dto.GeneralOpenAIRequest, i // eg. {"google":{"thinking_config":{"thinking_budget":5324,"include_thoughts":true}}} if googleBody, ok := extraBody["google"].(map[string]interface{}); ok { - if !strings.HasSuffix(info.UpstreamModelName, "-nothinking") { - adaptorWithExtraBody = true - // check error param name like thinkingConfig, should be thinking_config - if _, hasErrorParam := googleBody["thinkingConfig"]; hasErrorParam { - return nil, errors.New("extra_body.google.thinkingConfig is not supported, use extra_body.google.thinking_config instead") + // check error param name like thinkingConfig, should be thinking_config + if _, hasErrorParam := googleBody["thinkingConfig"]; hasErrorParam { + return nil, errors.New("extra_body.google.thinkingConfig is not supported, use extra_body.google.thinking_config instead") + } + + if thinkingConfig, ok := googleBody["thinking_config"].(map[string]interface{}); ok { + // check error param name like thinkingBudget, should be thinking_budget + if _, hasErrorParam := thinkingConfig["thinkingBudget"]; hasErrorParam { + return nil, errors.New("extra_body.google.thinking_config.thinkingBudget is not supported, use extra_body.google.thinking_config.thinking_budget instead") } + var hasThinkingConfig bool + var tempThinkingConfig dto.GeminiThinkingConfig - if thinkingConfig, ok := googleBody["thinking_config"].(map[string]interface{}); ok { - // check error param name like thinkingBudget, should be thinking_budget - if _, hasErrorParam := thinkingConfig["thinkingBudget"]; hasErrorParam { - return nil, errors.New("extra_body.google.thinking_config.thinkingBudget is not supported, use extra_body.google.thinking_config.thinking_budget instead") - } - var hasThinkingConfig bool - var tempThinkingConfig dto.GeminiThinkingConfig - - if thinkingBudget, exists := thinkingConfig["thinking_budget"]; exists { - switch v := thinkingBudget.(type) { - case float64: - budgetInt := int(v) - tempThinkingConfig.ThinkingBudget = common.GetPointer(budgetInt) - if budgetInt > 0 { - // 有正数预算 - tempThinkingConfig.IncludeThoughts = true - } else { - // 存在但为0或负数,禁用思考 - tempThinkingConfig.IncludeThoughts = false - } - hasThinkingConfig = true - default: + if thinkingBudget, exists := thinkingConfig["thinking_budget"]; exists { + switch v := thinkingBudget.(type) { + case float64: + if v != float64(int(v)) { return nil, errors.New("extra_body.google.thinking_config.thinking_budget must be an integer") } + budgetInt := int(v) + tempThinkingConfig.ThinkingBudget = common.GetPointer(budgetInt) + // 有正数预算则启用思考;为0或负数则禁用 + tempThinkingConfig.IncludeThoughts = common.GetPointer(budgetInt > 0) + hasThinkingConfig = true + default: + return nil, errors.New("extra_body.google.thinking_config.thinking_budget must be an integer") } + } - if includeThoughts, exists := thinkingConfig["include_thoughts"]; exists { - if v, ok := includeThoughts.(bool); ok { - tempThinkingConfig.IncludeThoughts = v - hasThinkingConfig = true - } else { - return nil, errors.New("extra_body.google.thinking_config.include_thoughts must be a boolean") - } + if includeThoughts, exists := thinkingConfig["include_thoughts"]; exists { + if v, ok := includeThoughts.(bool); ok { + tempThinkingConfig.IncludeThoughts = common.GetPointer(v) + hasThinkingConfig = true + } else { + return nil, errors.New("extra_body.google.thinking_config.include_thoughts must be a boolean") } - if thinkingLevel, exists := thinkingConfig["thinking_level"]; exists { - if v, ok := thinkingLevel.(string); ok { - tempThinkingConfig.ThinkingLevel = v - hasThinkingConfig = true - } else { - return nil, errors.New("extra_body.google.thinking_config.thinking_level must be a string") - } + } + if thinkingLevel, exists := thinkingConfig["thinking_level"]; exists { + if v, ok := thinkingLevel.(string); ok { + tempThinkingConfig.ThinkingLevel = common.GetPointer(v) + hasThinkingConfig = true + } else { + return nil, errors.New("extra_body.google.thinking_config.thinking_level must be a string") } + } - if hasThinkingConfig { - // 避免 panic: 仅在获得配置时分配,防止后续赋值时空指针 - if geminiRequest.GenerationConfig.ThinkingConfig == nil { - geminiRequest.GenerationConfig.ThinkingConfig = &tempThinkingConfig - } else { - // 如果已分配,则合并内容 - if tempThinkingConfig.ThinkingBudget != nil { - geminiRequest.GenerationConfig.ThinkingConfig.ThinkingBudget = tempThinkingConfig.ThinkingBudget - } + if hasThinkingConfig { + if tempThinkingConfig.ThinkingBudget != nil && tempThinkingConfig.ThinkingLevel != nil { + return nil, errors.New("extra_body.google.thinking_config cannot contain both thinking_budget and thinking_level") + } + explicitThinkingConfig = true + // 避免 panic: 仅在获得配置时分配,防止后续赋值时空指针 + if geminiRequest.GenerationConfig.ThinkingConfig == nil { + geminiRequest.GenerationConfig.ThinkingConfig = &tempThinkingConfig + } else { + // 如果已分配,则合并内容 + if tempThinkingConfig.ThinkingBudget != nil { + geminiRequest.GenerationConfig.ThinkingConfig.ThinkingBudget = tempThinkingConfig.ThinkingBudget + } + if tempThinkingConfig.IncludeThoughts != nil { geminiRequest.GenerationConfig.ThinkingConfig.IncludeThoughts = tempThinkingConfig.IncludeThoughts - if tempThinkingConfig.ThinkingLevel != "" { - geminiRequest.GenerationConfig.ThinkingConfig.ThinkingLevel = tempThinkingConfig.ThinkingLevel - } + } + if tempThinkingConfig.ThinkingLevel != nil { + geminiRequest.GenerationConfig.ThinkingConfig.ThinkingLevel = tempThinkingConfig.ThinkingLevel } } } @@ -352,8 +359,20 @@ func CovertOpenAI2Gemini(c *gin.Context, textRequest dto.GeneralOpenAIRequest, i } } - if !adaptorWithExtraBody { - ThinkingAdaptor(&geminiRequest, info, textRequest) + if !explicitThinkingConfig { + // OpenAI reasoning_effort is a native Gemini thinking level. Preserve the + // exact caller value (including xhigh/max) and do not apply the legacy + // budget catalog or clamp. + if textRequest.ReasoningEffort != "" { + level := textRequest.ReasoningEffort + geminiRequest.GenerationConfig.ThinkingConfig = &dto.GeminiThinkingConfig{ + IncludeThoughts: common.GetPointer(true), + ThinkingLevel: common.GetPointer(level), + } + info.ReasoningEffort = level + } else { + ThinkingAdaptor(&geminiRequest, info, textRequest) + } } safetySettings := make([]dto.GeminiChatSafetySettings, 0, len(SafetySettingList)) @@ -991,6 +1010,18 @@ func unescapeMapOrSlice(data interface{}) interface{} { return data } +// collectThoughts extracts non-empty thought part texts and joins them into a +// single reasoning string, keeping thought content separate from text and tool calls. +func collectThoughts(parts []dto.GeminiPart) string { + var thoughtTexts []string + for _, part := range parts { + if part.Thought && part.Text != "" { + thoughtTexts = append(thoughtTexts, part.Text) + } + } + return strings.Join(thoughtTexts, "\n") +} + func getResponseToolCall(item *dto.GeminiPart) *dto.ToolCallResponse { var argsBytes []byte var err error @@ -1097,7 +1128,7 @@ func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) toolCalls = append(toolCalls, *call) } } else if part.Thought { - choice.Message.ReasoningContent = part.Text + // 思考内容由 collectThoughts 单独收集,不与文本/工具调用混合 } else { if part.ExecutableCode != nil { texts = append(texts, "```"+part.ExecutableCode.Language+"\n"+part.ExecutableCode.Code+"\n```") @@ -1115,6 +1146,9 @@ func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) choice.Message.SetToolCalls(toolCalls) isToolCall = true } + if reasoning := collectThoughts(candidate.Content.Parts); reasoning != "" { + choice.Message.ReasoningContent = &reasoning + } choice.Message.SetStringContent(strings.Join(texts, "\n")) } @@ -1171,7 +1205,6 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d } var texts []string isTools := false - isThought := false if candidate.FinishReason != nil { // Map Gemini FinishReason to OpenAI finish_reason switch *candidate.FinishReason { @@ -1218,8 +1251,7 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d } } else if part.Thought { - isThought = true - texts = append(texts, part.Text) + // 思考内容由 collectThoughts 单独收集,不与文本/工具调用混合 } else { if part.ExecutableCode != nil { texts = append(texts, "```"+part.ExecutableCode.Language+"\n"+part.ExecutableCode.Code+"\n```\n") @@ -1232,9 +1264,10 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d } } } - if isThought { - choice.Delta.SetReasoningContent(strings.Join(texts, "\n")) - } else { + if reasoning := collectThoughts(candidate.Content.Parts); reasoning != "" { + choice.Delta.SetReasoningContent(reasoning) + } + if len(texts) > 0 { choice.Delta.SetContentString(strings.Join(texts, "\n")) } if isTools { @@ -1339,6 +1372,14 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * response.Id = id response.Created = createAt response.Model = info.UpstreamModelName + if response.IsToolCall() { + finishReason = constant.FinishReasonToolCalls + if info.RelayFormat == types.RelayFormatClaude { + for choiceIdx := range response.Choices { + response.Choices[choiceIdx].FinishReason = nil + } + } + } for choiceIdx := range response.Choices { choiceKey := response.Choices[choiceIdx].Index for toolIdx := range response.Choices[choiceIdx].Delta.ToolCalls { @@ -1376,7 +1417,6 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * } emptyResponse.Choices[0].Delta.ToolCalls = copiedToolCalls } - finishReason = constant.FinishReasonToolCalls err := handleStream(c, info, emptyResponse) if err != nil { logger.LogError(c, err.Error()) @@ -1399,7 +1439,9 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * logger.LogError(c, err.Error()) } if isStop { - _ = handleStream(c, info, helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, finishReason)) + if info.RelayFormat != types.RelayFormatClaude { + _ = handleStream(c, info, helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, finishReason)) + } } return true }) @@ -1409,6 +1451,10 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * } response := helper.GenerateFinalUsageResponse(id, createAt, info.UpstreamModelName, *usage) + if info.RelayFormat == types.RelayFormatClaude && info.ClaudeConvertInfo != nil && !info.ClaudeConvertInfo.Done { + response = helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, finishReason) + response.Usage = usage + } handleErr := handleFinalStream(c, info, response) if handleErr != nil { common.SysLog("send final response failed: " + handleErr.Error()) diff --git a/relay/channel/gemini/relay_gemini_thinking_test.go b/relay/channel/gemini/relay_gemini_thinking_test.go new file mode 100644 index 000000000000..12a6ad3a74c6 --- /dev/null +++ b/relay/channel/gemini/relay_gemini_thinking_test.go @@ -0,0 +1,323 @@ +package gemini + +import ( + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/setting/model_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newGeminiThinkingTestInfo(model string) *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + OriginModelName: model, + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelBaseUrl: "https://example.test", + UpstreamModelName: model, + }, + } +} + +func newGeminiThinkingTestContext() *gin.Context { + return gin.CreateTestContextOnly(httptest.NewRecorder(), gin.New()) +} + +func TestResolveClaudeThinkingConfig(t *testing.T) { + budget := 1024 + cases := []struct { + name string + request dto.ClaudeRequest + wantConfig bool + wantThoughts *bool + wantBudget *int + wantLevel *string + wantTraceEffort string + }{ + { + name: "manual budget", + request: dto.ClaudeRequest{Thinking: &dto.Thinking{ + Type: "enabled", BudgetTokens: &budget, + }}, + wantConfig: true, + wantThoughts: common.GetPointer(true), + wantBudget: &budget, + }, + { + name: "effort wins over budget", + request: dto.ClaudeRequest{ + Thinking: &dto.Thinking{Type: "enabled", BudgetTokens: &budget}, + OutputConfig: []byte(`{"effort":"xhigh"}`), + }, + wantConfig: true, + wantThoughts: common.GetPointer(true), + wantLevel: common.GetPointer("xhigh"), + wantTraceEffort: "xhigh", + }, + { + name: "disabled wins over effort", + request: dto.ClaudeRequest{ + Thinking: &dto.Thinking{Type: "disabled"}, + OutputConfig: []byte(`{"effort":"high"}`), + }, + wantConfig: true, + wantThoughts: common.GetPointer(false), + }, + { + name: "adaptive enables visible thoughts without a level", + request: dto.ClaudeRequest{Thinking: &dto.Thinking{Type: "adaptive"}}, + wantConfig: true, + wantThoughts: common.GetPointer(true), + }, + { + name: "no client intent", + request: dto.ClaudeRequest{}, + wantConfig: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + config, traceEffort := resolveClaudeThinkingConfig(&tc.request) + if !tc.wantConfig { + require.Nil(t, config) + require.Empty(t, traceEffort) + return + } + + require.NotNil(t, config) + require.NotNil(t, config.IncludeThoughts) + require.Equal(t, *tc.wantThoughts, *config.IncludeThoughts) + if tc.wantBudget == nil { + require.Nil(t, config.ThinkingBudget) + } else { + require.Equal(t, *tc.wantBudget, *config.ThinkingBudget) + } + if tc.wantLevel == nil { + require.Nil(t, config.ThinkingLevel) + } else { + require.Equal(t, *tc.wantLevel, *config.ThinkingLevel) + } + require.Equal(t, tc.wantTraceEffort, traceEffort) + }) + } +} + +func TestConvertClaudeRequestAppliesThinkingIntentWithTools(t *testing.T) { + budget := 1024 + request := &dto.ClaudeRequest{ + Model: "gemini-3.7-flash-high", + Messages: []dto.ClaudeMessage{{ + Role: "user", + Content: "Look up the weather in Tokyo.", + }}, + Tools: []dto.Tool{{ + Name: "lookup_weather", + InputSchema: map[string]any{"type": "object"}, + }}, + Thinking: &dto.Thinking{Type: "enabled", BudgetTokens: &budget}, + OutputConfig: []byte(`{"effort":"high"}`), + } + info := newGeminiThinkingTestInfo("gemini-3.7-flash-high") + + converted, err := (&Adaptor{}).ConvertClaudeRequest(newGeminiThinkingTestContext(), info, request) + require.NoError(t, err) + geminiRequest, ok := converted.(*dto.GeminiChatRequest) + require.True(t, ok) + require.NotNil(t, geminiRequest.GenerationConfig.ThinkingConfig) + require.NotNil(t, geminiRequest.GenerationConfig.ThinkingConfig.IncludeThoughts) + assert.True(t, *geminiRequest.GenerationConfig.ThinkingConfig.IncludeThoughts) + require.NotNil(t, geminiRequest.GenerationConfig.ThinkingConfig.ThinkingLevel) + assert.Equal(t, "high", *geminiRequest.GenerationConfig.ThinkingConfig.ThinkingLevel) + assert.Nil(t, geminiRequest.GenerationConfig.ThinkingConfig.ThinkingBudget) + assert.Equal(t, "high", info.ReasoningEffort) + require.NotEmpty(t, geminiRequest.Tools) +} + +func TestCovertOpenAI2GeminiResolvesThinkingIntent(t *testing.T) { + for _, level := range []string{"low", "xhigh", "max"} { + t.Run("reasoning effort "+level, func(t *testing.T) { + info := newGeminiThinkingTestInfo("gemini-3.7-flash-high") + request := dto.GeneralOpenAIRequest{ReasoningEffort: level} + + converted, err := CovertOpenAI2Gemini(newGeminiThinkingTestContext(), request, info) + require.NoError(t, err) + require.NotNil(t, converted.GenerationConfig.ThinkingConfig) + require.NotNil(t, converted.GenerationConfig.ThinkingConfig.IncludeThoughts) + require.True(t, *converted.GenerationConfig.ThinkingConfig.IncludeThoughts) + require.NotNil(t, converted.GenerationConfig.ThinkingConfig.ThinkingLevel) + require.Equal(t, level, *converted.GenerationConfig.ThinkingConfig.ThinkingLevel) + require.Nil(t, converted.GenerationConfig.ThinkingConfig.ThinkingBudget) + require.Equal(t, level, info.ReasoningEffort) + }) + } + + t.Run("explicit google config wins and preserves false", func(t *testing.T) { + extraBody, err := common.Marshal(map[string]any{ + "google": map[string]any{ + "thinking_config": map[string]any{ + "thinking_level": "low", + "include_thoughts": false, + }, + }, + }) + require.NoError(t, err) + + converted, err := CovertOpenAI2Gemini( + newGeminiThinkingTestContext(), + dto.GeneralOpenAIRequest{ReasoningEffort: "max", ExtraBody: extraBody}, + newGeminiThinkingTestInfo("gemini-3.7-flash-high"), + ) + require.NoError(t, err) + config := converted.GenerationConfig.ThinkingConfig + require.NotNil(t, config) + require.NotNil(t, config.IncludeThoughts) + assert.False(t, *config.IncludeThoughts) + require.NotNil(t, config.ThinkingLevel) + assert.Equal(t, "low", *config.ThinkingLevel) + assert.Nil(t, config.ThinkingBudget) + }) + + t.Run("unrelated google body does not suppress reasoning effort", func(t *testing.T) { + extraBody, err := common.Marshal(map[string]any{ + "google": map[string]any{ + "image_config": map[string]any{"aspect_ratio": "1:1"}, + }, + }) + require.NoError(t, err) + + converted, err := CovertOpenAI2Gemini( + newGeminiThinkingTestContext(), + dto.GeneralOpenAIRequest{ReasoningEffort: "high", ExtraBody: extraBody}, + newGeminiThinkingTestInfo("gemini-3.7-flash-high"), + ) + require.NoError(t, err) + require.NotNil(t, converted.GenerationConfig.ThinkingConfig) + require.NotNil(t, converted.GenerationConfig.ThinkingConfig.ThinkingLevel) + assert.Equal(t, "high", *converted.GenerationConfig.ThinkingConfig.ThinkingLevel) + }) + + for _, tc := range []struct { + name string + config map[string]any + message string + }{ + { + name: "level and budget conflict", + config: map[string]any{"thinking_level": "high", "thinking_budget": 1024}, + message: "cannot contain both", + }, + { + name: "decimal budget", + config: map[string]any{"thinking_budget": 12.5}, + message: "must be an integer", + }, + } { + t.Run(tc.name, func(t *testing.T) { + extraBody, err := common.Marshal(map[string]any{ + "google": map[string]any{"thinking_config": tc.config}, + }) + require.NoError(t, err) + + _, err = CovertOpenAI2Gemini( + newGeminiThinkingTestContext(), + dto.GeneralOpenAIRequest{ExtraBody: extraBody}, + newGeminiThinkingTestInfo("gemini-3.7-flash-high"), + ) + require.Error(t, err) + assert.Contains(t, err.Error(), tc.message) + }) + } +} + +func TestNativeGeminiModelDoesNotUseLegacySuffixAdapter(t *testing.T) { + settings := model_setting.GetGeminiSettings() + previousAdapterState := settings.ThinkingAdapterEnabled + settings.ThinkingAdapterEnabled = true + t.Cleanup(func() { + settings.ThinkingAdapterEnabled = previousAdapterState + }) + + info := newGeminiThinkingTestInfo("gemini-3.7-flash-high") + require.False(t, shouldApplyGeminiThinkingAdapter(info)) + + converted, err := CovertOpenAI2Gemini(newGeminiThinkingTestContext(), dto.GeneralOpenAIRequest{}, info) + require.NoError(t, err) + assert.Nil(t, converted.GenerationConfig.ThinkingConfig) + + url, err := (&Adaptor{}).GetRequestURL(info) + require.NoError(t, err) + assert.Equal(t, "https://example.test/v1beta/models/gemini-3.7-flash-high:generateContent", url) + assert.Equal(t, "gemini-3.7-flash-high", info.UpstreamModelName) +} + +func TestCovertOpenAI2GeminiDoesNotInjectThinkingWithoutIntent(t *testing.T) { + info := newGeminiThinkingTestInfo("gemini-3.7-flash-high") + request := dto.GeneralOpenAIRequest{ + Tools: []dto.ToolCallRequest{{ + Type: "function", + Function: dto.FunctionRequest{ + Name: "lookup_weather", + Parameters: map[string]any{"type": "object"}, + }, + }}, + } + + converted, err := CovertOpenAI2Gemini(newGeminiThinkingTestContext(), request, info) + require.NoError(t, err) + assert.Nil(t, converted.GenerationConfig.ThinkingConfig) +} + +func TestStreamResponseGeminiChat2OpenAISeparatesThoughtTextAndTool(t *testing.T) { + response, _ := streamResponseGeminiChat2OpenAI(&dto.GeminiChatResponse{ + Candidates: []dto.GeminiChatCandidate{{ + Content: dto.GeminiChatContent{Parts: []dto.GeminiPart{ + {Thought: true, Text: "consider the forecast"}, + {Text: "I will check it."}, + {FunctionCall: &dto.FunctionCall{ + FunctionName: "lookup_weather", + Arguments: map[string]any{"city": "Tokyo"}, + }}, + }}, + }}, + }) + + require.Len(t, response.Choices, 1) + choice := response.Choices[0] + assert.Equal(t, "consider the forecast", choice.Delta.GetReasoningContent()) + assert.Equal(t, "I will check it.", choice.Delta.GetContentString()) + require.Len(t, choice.Delta.ToolCalls, 1) + assert.Equal(t, "lookup_weather", choice.Delta.ToolCalls[0].Function.Name) + require.NotNil(t, choice.FinishReason) + assert.Equal(t, "tool_calls", *choice.FinishReason) + assert.False(t, strings.Contains(choice.Delta.GetReasoningContent(), "I will check it.")) +} + +func TestResponseGeminiChat2OpenAISeparatesThoughtTextAndTool(t *testing.T) { + response := responseGeminiChat2OpenAI(newGeminiThinkingTestContext(), &dto.GeminiChatResponse{ + Candidates: []dto.GeminiChatCandidate{{ + Content: dto.GeminiChatContent{Parts: []dto.GeminiPart{ + {Thought: true, Text: "consider the forecast"}, + {Text: "I will check it."}, + {FunctionCall: &dto.FunctionCall{ + FunctionName: "lookup_weather", + Arguments: map[string]any{"city": "Tokyo"}, + }}, + }}, + }}, + }) + + require.Len(t, response.Choices, 1) + choice := response.Choices[0] + assert.Equal(t, "consider the forecast", choice.Message.GetReasoningContent()) + assert.Equal(t, "I will check it.", choice.Message.StringContent()) + toolCalls := choice.Message.ParseToolCalls() + require.Len(t, toolCalls, 1) + assert.Equal(t, "lookup_weather", toolCalls[0].Function.Name) + assert.Equal(t, "tool_calls", choice.FinishReason) +} diff --git a/relay/channel/ollama/stream.go b/relay/channel/ollama/stream.go index 2a264b27e467..43e024deafd7 100644 --- a/relay/channel/ollama/stream.go +++ b/relay/channel/ollama/stream.go @@ -273,7 +273,7 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R msg := dto.Message{Role: "assistant", Content: contentPtr(content)} if rc := reasoningBuilder.String(); rc != "" { - msg.ReasoningContent = rc + msg.ReasoningContent = &rc } full := dto.OpenAITextResponse{ Id: common.GetUUID(), diff --git a/relay/channel/openai/chat_via_responses.go b/relay/channel/openai/chat_via_responses.go index 2c0752275daa..e1547f00a318 100644 --- a/relay/channel/openai/chat_via_responses.go +++ b/relay/channel/openai/chat_via_responses.go @@ -308,6 +308,11 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo sr.Error(err) return } + // Ignore transport keepalives from upstream proxies; they are not + // mappable to chat.completion.chunk events. + if dto.IsResponsesTransportEventType(streamResp.Type) { + return + } switch streamResp.Type { case "response.created": diff --git a/relay/channel/openai/chat_via_responses_test.go b/relay/channel/openai/chat_via_responses_test.go new file mode 100644 index 000000000000..093c17f7719c --- /dev/null +++ b/relay/channel/openai/chat_via_responses_test.go @@ -0,0 +1,70 @@ +package openai + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" +) + +func TestOaiResponsesToChatStreamHandlerDropsKeepaliveAndContinues(t *testing.T) { + oldStreamingTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { + constant.StreamingTimeout = oldStreamingTimeout + }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + info := &relaycommon.RelayInfo{ + DisablePing: true, + RelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "test-model"}, + ShouldIncludeUsage: false, + } + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp-1","model":"test-model"}}`, + `data: {"type":"keepalive","sequence_number":1}`, + `data: {"type":"response.output_text.delta","delta":"hello"}`, + `data: {"type":"response.completed","response":{"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}}`, + "data: [DONE]", + }, "\n\n"))), + } + + usage, apiErr := OaiResponsesToChatStreamHandler(c, info, resp) + if apiErr != nil { + t.Fatalf("OaiResponsesToChatStreamHandler() error = %v", apiErr) + } + if usage == nil { + t.Fatal("OaiResponsesToChatStreamHandler() returned nil usage") + } + if usage.InputTokens != 2 || usage.OutputTokens != 1 || usage.TotalTokens != 3 { + t.Fatalf("usage = %+v, want input=2 output=1 total=3", usage) + } + + body := recorder.Body.String() + t.Logf("forwarded Chat Completions SSE: %q", body) + if strings.Contains(body, "keepalive") || strings.Contains(body, "sequence_number") { + t.Fatalf("transport keepalive leaked into Chat Completions SSE: %q", body) + } + if !strings.Contains(body, "hello") { + t.Fatalf("converted stream is missing output text: %q", body) + } + if !strings.Contains(body, `"finish_reason":"stop"`) { + t.Fatalf("converted stream is missing terminal stop chunk: %q", body) + } + if !strings.Contains(body, "data: [DONE]") { + t.Fatalf("converted stream is missing SSE termination: %q", body) + } +} diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index 08811a77205a..369d2c266405 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -257,5 +257,8 @@ func sendResponsesStreamData(c *gin.Context, streamResponse dto.ResponsesStreamR if data == "" { return } + if dto.IsResponsesTransportEventType(streamResponse.Type) { + return + } helper.ResponseChunkData(c, streamResponse, data) } diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index d33c5555f267..a85751844c0b 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -245,7 +245,7 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo completionTokens := simpleResponse.Usage.CompletionTokens if completionTokens == 0 { for _, choice := range simpleResponse.Choices { - ctkm := service.CountTextToken(choice.Message.StringContent()+choice.Message.ReasoningContent+choice.Message.Reasoning, info.UpstreamModelName) + ctkm := service.CountTextToken(choice.Message.StringContent()+choice.Message.GetReasoningContent(), info.UpstreamModelName) completionTokens += ctkm } } diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 2665b8d027e9..046540b7e02a 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -46,12 +46,7 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http // compute usage usage := dto.Usage{} if responsesResponse.Usage != nil { - usage.PromptTokens = responsesResponse.Usage.InputTokens - usage.CompletionTokens = responsesResponse.Usage.OutputTokens - usage.TotalTokens = responsesResponse.Usage.TotalTokens - if responsesResponse.Usage.InputTokensDetails != nil { - usage.PromptTokensDetails.CachedTokens = responsesResponse.Usage.InputTokensDetails.CachedTokens - } + service.ApplyResponsesUsage(&usage, responsesResponse.Usage) } if info == nil || info.ResponsesUsageInfo == nil || info.ResponsesUsageInfo.BuiltInTools == nil { return &usage, nil @@ -88,23 +83,17 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp sr.Error(err) return } + // Drop transport-only keepalives from upstream proxies. Strict clients + // (serde enums, etc.) fail the whole stream on unknown event types. + if dto.IsResponsesTransportEventType(streamResponse.Type) { + return + } sendResponsesStreamData(c, streamResponse, data) switch streamResponse.Type { case "response.completed": if streamResponse.Response != nil { if streamResponse.Response.Usage != nil { - if streamResponse.Response.Usage.InputTokens != 0 { - usage.PromptTokens = streamResponse.Response.Usage.InputTokens - } - if streamResponse.Response.Usage.OutputTokens != 0 { - usage.CompletionTokens = streamResponse.Response.Usage.OutputTokens - } - if streamResponse.Response.Usage.TotalTokens != 0 { - usage.TotalTokens = streamResponse.Response.Usage.TotalTokens - } - if streamResponse.Response.Usage.InputTokensDetails != nil { - usage.PromptTokensDetails.CachedTokens = streamResponse.Response.Usage.InputTokensDetails.CachedTokens - } + service.ApplyResponsesUsage(usage, streamResponse.Response.Usage) } if streamResponse.Response.HasImageGenerationCall() { c.Set("image_generation_call", true) diff --git a/relay/channel/openai/relay_responses_test.go b/relay/channel/openai/relay_responses_test.go new file mode 100644 index 000000000000..3b7d52b9f04e --- /dev/null +++ b/relay/channel/openai/relay_responses_test.go @@ -0,0 +1,74 @@ +package openai + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + relaycommon "github.com/QuantumNous/new-api/relay/common" + + "github.com/gin-gonic/gin" +) + +func TestOaiResponsesStreamHandlerDropsKeepaliveAndContinues(t *testing.T) { + oldStreamingTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { + constant.StreamingTimeout = oldStreamingTimeout + }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + + info := &relaycommon.RelayInfo{ + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "test-model"}, + } + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp-1","model":"test-model"}}`, + `data: {"type":"keepalive","sequence_number":1}`, + `data: {"type":"response.output_text.delta","delta":"hello"}`, + `data: {"type":"response.completed","response":{"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}}`, + "data: [DONE]", + }, "\n\n"))), + } + + usage, apiErr := OaiResponsesStreamHandler(c, info, resp) + if apiErr != nil { + t.Fatalf("OaiResponsesStreamHandler() error = %v", apiErr) + } + if usage == nil { + t.Fatal("OaiResponsesStreamHandler() returned nil usage") + } + if usage.InputTokens != 2 || usage.OutputTokens != 1 || usage.TotalTokens != 3 { + t.Fatalf("usage = %+v, want input=2 output=1 total=3", usage) + } + + body := recorder.Body.String() + t.Logf("forwarded SSE: %q", body) + if strings.Contains(body, "keepalive") || strings.Contains(body, "sequence_number") { + t.Fatalf("transport keepalive leaked into downstream SSE: %q", body) + } + for _, event := range []string{ + "event: response.created", + "event: response.output_text.delta", + "event: response.completed", + } { + if !strings.Contains(body, event) { + t.Fatalf("downstream SSE is missing %q: %q", event, body) + } + } + if strings.Index(body, "event: response.created") >= strings.Index(body, "event: response.output_text.delta") || + strings.Index(body, "event: response.output_text.delta") >= strings.Index(body, "event: response.completed") { + t.Fatalf("semantic event order was not preserved: %q", body) + } + if info.StreamStatus == nil || info.StreamStatus.EndReason != relaycommon.StreamEndReasonDone { + t.Fatalf("stream status = %+v, want done", info.StreamStatus) + } +} diff --git a/relay/common/relay_utils.go b/relay/common/relay_utils.go index 18df77a645d6..ec408e45375e 100644 --- a/relay/common/relay_utils.go +++ b/relay/common/relay_utils.go @@ -3,6 +3,7 @@ package common import ( "fmt" "net/http" + "net/url" "strconv" "strings" @@ -36,6 +37,53 @@ func GetFullRequestURL(baseURL string, requestURL string, channelType int) strin return fullRequestURL } +// SanitizeURLForLog masks credentials carried in query parameters while +// preserving enough of the upstream URL to diagnose routing failures. +func SanitizeURLForLog(rawURL string) string { + if rawURL == "" { + return rawURL + } + + parsedURL, err := url.Parse(rawURL) + if err != nil { + return rawURL + } + + query := parsedURL.Query() + if len(query) == 0 { + return rawURL + } + + changed := false + for key := range query { + if isSensitiveURLQueryKey(key) { + query.Set(key, "***masked***") + changed = true + } + } + if !changed { + return rawURL + } + + parsedURL.RawQuery = query.Encode() + return parsedURL.String() +} + +func isSensitiveURLQueryKey(key string) bool { + normalized := strings.ToLower(strings.TrimSpace(key)) + switch normalized { + case "key", "api_key", "api-key", "apikey", "x-api-key", + "access_token", "refresh_token", "id_token", "token", + "authorization", "auth", "client_secret", "secret", + "password", "passwd", "signature", "sig", "awsaccesskeyid", + "x-amz-credential", "x-amz-security-token", "x-amz-signature": + return true + } + return strings.Contains(normalized, "token") || + strings.Contains(normalized, "secret") || + strings.Contains(normalized, "signature") +} + func GetAPIVersion(c *gin.Context) string { query := c.Request.URL.Query() apiVersion := query.Get("api-version") diff --git a/relay/common/websocket_idle.go b/relay/common/websocket_idle.go new file mode 100644 index 000000000000..68438e66dfe1 --- /dev/null +++ b/relay/common/websocket_idle.go @@ -0,0 +1,43 @@ +package common + +import ( + "errors" + "net" + "time" + + appcommon "github.com/QuantumNous/new-api/common" + + "github.com/gorilla/websocket" +) + +const WebSocketIdleCloseReason = "websocket idle timeout" + +// WebSocketIdleTimeoutMinutes is the client WebSocket idle timeout in minutes. +// A non-positive value disables the timeout. +var WebSocketIdleTimeoutMinutes = appcommon.GetEnvOrDefault("WEBSOCKET_IDLE_TIMEOUT_MINUTES", 10) + +func GetWebSocketIdleTimeout() time.Duration { + if WebSocketIdleTimeoutMinutes <= 0 { + return 0 + } + return time.Duration(WebSocketIdleTimeoutMinutes) * time.Minute +} + +// RefreshClientWebSocketReadDeadline counts data messages as activity. Gorilla +// handles Ping/Pong control frames inside ReadMessage, so heartbeats alone do +// not extend this deadline. +func RefreshClientWebSocketReadDeadline(conn *websocket.Conn) error { + if conn == nil { + return errors.New("websocket connection is nil") + } + timeout := GetWebSocketIdleTimeout() + if timeout <= 0 { + return conn.SetReadDeadline(time.Time{}) + } + return conn.SetReadDeadline(time.Now().Add(timeout)) +} + +func IsWebSocketIdleTimeout(err error) bool { + var netErr net.Error + return errors.As(err, &netErr) && netErr.Timeout() +} diff --git a/relay/helper/max_tokens_bounds_test.go b/relay/helper/max_tokens_bounds_test.go new file mode 100644 index 000000000000..68e2e3b45974 --- /dev/null +++ b/relay/helper/max_tokens_bounds_test.go @@ -0,0 +1,34 @@ +package helper + +import ( + "testing" + + "github.com/QuantumNous/new-api/dto" + "github.com/stretchr/testify/require" +) + +// TestMaxTokensBounds guards the billing invariant that user-supplied max +// token fields are bounded on every relay format. These values feed +// pre-consume quota math (preConsumedTokens * ratio); a huge or +// wrapped-negative value (e.g. 18446744073686646784 parsed into *uint) must +// be rejected at validation instead of corrupting the pre-charge. +func TestResponsesMaxTokensBounds(t *testing.T) { + // The Responses WebSocket relay parses response.create events itself and + // never goes through GetAndValidateResponsesRequest, so it must reach the + // same bound through the shared validator. + t.Run("responses websocket transport shares the bound", func(t *testing.T) { + huge := uint(18446744073686646784) + err := ValidateResponsesRequest(&dto.OpenAIResponsesRequest{Model: "gpt-4o", MaxOutputTokens: &huge}) + require.Error(t, err) + require.Contains(t, err.Error(), "max_output_tokens is invalid") + + normal := uint(8192) + require.NoError(t, ValidateResponsesRequest(&dto.OpenAIResponsesRequest{Model: "gpt-4o", MaxOutputTokens: &normal})) + }) + + t.Run("responses websocket transport requires model", func(t *testing.T) { + err := ValidateResponsesRequest(&dto.OpenAIResponsesRequest{}) + require.Error(t, err) + require.Contains(t, err.Error(), "model is required") + }) +} diff --git a/relay/helper/stream_scanner.go b/relay/helper/stream_scanner.go index a9bc5e16a720..91c896a9160e 100644 --- a/relay/helper/stream_scanner.go +++ b/relay/helper/stream_scanner.go @@ -104,7 +104,8 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon logger.LogError(c, "timeout waiting for goroutines to exit") } - close(stopChan) + // stopChan is intentionally left open after all workers are drained. + // Worker cleanup sends are allowed to complete after wg.Done(). }() scanner.Buffer(make([]byte, InitialScannerBufferSize), getScannerBufferSize()) diff --git a/relay/helper/valid_request.go b/relay/helper/valid_request.go index c5477ccead65..f983a6bd0d53 100644 --- a/relay/helper/valid_request.go +++ b/relay/helper/valid_request.go @@ -113,15 +113,47 @@ func GetAndValidateEmbeddingRequest(c *gin.Context, relayMode int) (*dto.Embeddi return embeddingRequest, nil } +// maxTokensLimit bounds user-supplied max token fields. These values feed +// pre-consume quota math (preConsumedTokens * ratio); an unbounded value can +// overflow the conversion and corrupt billing. +const maxTokensLimit = math.MaxInt32 / 2 + +func exceedsMaxTokensLimit(values ...*uint) bool { + for _, v := range values { + if lo.FromPtrOr(v, uint(0)) > maxTokensLimit { + return true + } + } + return false +} + +// ValidateResponsesRequest enforces the bounds that protect pre-consume quota +// math for a Responses call. Every transport that builds one — HTTP and the +// WebSocket relay alike — must run it before pricing, so an unvalidated +// max_output_tokens can never reach the billing multiplication. +func ValidateResponsesRequest(request *dto.OpenAIResponsesRequest) error { + if request == nil { + return errors.New("request is required") + } + if request.Model == "" { + return errors.New("model is required") + } + if exceedsMaxTokensLimit(request.MaxOutputTokens) { + return errors.New("max_output_tokens is invalid") + } + return nil +} func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest, error) { request := &dto.OpenAIResponsesRequest{} err := common.UnmarshalBodyReusable(c, request) if err != nil { return nil, err } - if request.Model == "" { - return nil, errors.New("model is required") + if err := ValidateResponsesRequest(request); err != nil { + return nil, err } + // Only the HTTP transport requires input on every call: a WebSocket session + // keeps conversation state upstream, so an incremental turn may omit it. if request.Input == nil { return nil, errors.New("input is required") } @@ -261,7 +293,7 @@ func GetAndValidateTextRequest(c *gin.Context, relayMode int) (*dto.GeneralOpenA textRequest.Model = c.Param("model") } - if lo.FromPtrOr(textRequest.MaxTokens, uint(0)) > math.MaxInt32/2 { + if exceedsMaxTokensLimit(textRequest.MaxTokens) { return nil, errors.New("max_tokens is invalid") } if textRequest.Model == "" { diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index ee48ca64b10b..b1c6f24dc7f4 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -586,7 +586,9 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt common.SysLog("get_channel_null: " + err.Error()) } if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled { - model.UpdateChannelStatus(midjourneyTask.ChannelId, "", 2, "No available account instance") + if model.UpdateChannelStatus(midjourneyTask.ChannelId, "", common.ChannelStatusManuallyDisabled, "No available account instance") { + service.CloseActiveWebSocketsForChannel(midjourneyTask.ChannelId, service.ChannelDisabledCloseReason) + } } } if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 { diff --git a/relay/responses_websocket.go b/relay/responses_websocket.go new file mode 100644 index 000000000000..189417f867b5 --- /dev/null +++ b/relay/responses_websocket.go @@ -0,0 +1,1122 @@ +package relay + +import ( + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + appconstant "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/middleware" + appmodel "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/pkg/wsmanager" + relaychannel "github.com/QuantumNous/new-api/relay/channel" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +const responsesWSEventTypeResponseCreate = "response.create" + +// responsesWSWriteTimeout bounds a single blocked write so a peer that stops +// reading cannot pin a connection forever. Without it the write never returns, +// and idle timeout, channel disable and shutdown all block behind it. +const responsesWSWriteTimeout = 30 * time.Second + +// responsesWSMaxMessageBytes bounds one inbound WebSocket message. The HTTP +// body limit does not cover WebSocket frames, so without it a valid key can +// stream unbounded data into memory. It follows MAX_REQUEST_BODY_MB so the same +// payload is accepted over both transports of /v1/responses, and only diverges +// when WEBSOCKET_MAX_MESSAGE_MB is set explicitly. +// +// Read lazily: constant.MaxRequestBodyMB is populated by InitEnv from main, so +// a package-level var here would capture zero. +// +// NOTE: gorilla enforces this against the compressed wire length (conn.go:924, +// before the decompression reader is attached at conn.go:1019). It is a real +// memory bound only while permessage-deflate stays disabled — see the upgrader +// in controller/relay.go. +func responsesWSMaxMessageBytes() int64 { + maxMB := common.GetEnvOrDefault("WEBSOCKET_MAX_MESSAGE_MB", 0) + if maxMB <= 0 { + maxMB = appconstant.MaxRequestBodyMB + } + if maxMB <= 0 { + maxMB = 32 + } + return int64(maxMB) << 20 +} + +// responsesWSMaxPerUser caps concurrent Responses WebSocket sessions per user; +// 0 disables the cap. Idle sessions hold a goroutine, a socket and an upstream +// connection for up to WEBSOCKET_IDLE_TIMEOUT_MINUTES. +var responsesWSMaxPerUser = common.GetEnvOrDefault("RESPONSES_WEBSOCKET_MAX_PER_USER", 8) + +var ( + responsesWSCountMu sync.Mutex + responsesWSCounts = map[int]int{} +) + +// responsesWSCallOutcome decides how a finished call is billed. +type responsesWSCallOutcome int + +const ( + // responsesWSCallAborted means upstream never accepted the request payload, + // so nothing was generated and the pre-consumed quota is returned in full. + responsesWSCallAborted responsesWSCallOutcome = iota + // responsesWSCallSettled means upstream accepted the request: bill what was + // observed, whether the stream ended normally, failed, or was cut short by a + // disconnect. This matches the HTTP and realtime relays, which also settle + // on mid-stream client disconnect rather than refunding generated output. + responsesWSCallSettled +) + +type responsesWSCreateEvent struct { + Type string `json:"type"` + EventID string `json:"event_id,omitempty"` + Request common.RawMessage `json:"response,omitempty"` +} + +type responsesWSCreateRequest struct { + Request dto.OpenAIResponsesRequest + Generate common.RawMessage +} + +type responsesWSErrorEvent struct { + Type string `json:"type"` + Status int `json:"status"` + EventID string `json:"event_id,omitempty"` + Error *types.OpenAIError `json:"error"` +} + +type responsesWSCallState struct { + info *relaycommon.RelayInfo + commitRate middleware.ModelRequestRateLimitCommit + + // mu guards usage and outputText. The upstream reader goroutine appends to + // them while the client goroutine may win the race to finish the call on a + // disconnect and read them for settlement. + mu sync.Mutex + usage *dto.Usage + outputText strings.Builder +} + +type responsesWSSession struct { + c *gin.Context + client *websocket.Conn + target *websocket.Conn + unregister func() + lockedModel string + lockedChannel *appmodel.Channel + nextEventIndex int + closeOnce sync.Once + + clientWriteMu sync.Mutex + // targetMu guards target and unregister. It is never held across network + // I/O, so closing the session cannot block behind an in-flight write. + targetMu sync.Mutex + // targetWriteMu only serializes writes, as gorilla allows a single writer. + targetWriteMu sync.Mutex + stateMu sync.Mutex + current *responsesWSCallState +} + +func ResponsesWebSocketHelper(c *gin.Context, client *websocket.Conn) *types.NewAPIError { + return responsesWebSocketHelper(c, client, relaycommon.RefreshClientWebSocketReadDeadline) +} + +func responsesWebSocketHelper(c *gin.Context, client *websocket.Conn, refreshReadDeadline func(*websocket.Conn) error) *types.NewAPIError { + userId := common.GetContextKeyInt(c, appconstant.ContextKeyUserId) + if !acquireResponsesWSSlot(userId) { + return types.NewErrorWithStatusCode( + fmt.Errorf("too many concurrent responses websocket connections (limit %d)", responsesWSMaxPerUser), + types.ErrorCodeInvalidRequest, + http.StatusTooManyRequests, + types.ErrOptionWithSkipRetry(), + ) + } + defer releaseResponsesWSSlot(userId) + session := &responsesWSSession{ + c: c, + client: client, + } + defer session.closeTarget() + defer session.settleCurrent() + client.SetReadLimit(responsesWSMaxMessageBytes()) + if err := refreshReadDeadline(client); err != nil { + return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry()) + } + + for { + messageType, message, err := client.ReadMessage() + if err != nil { + if relaycommon.IsWebSocketIdleTimeout(err) { + logger.LogInfo(c, "responses websocket closed after idle timeout") + session.closeForIdleTimeout() + return nil + } + if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + return nil + } + return types.NewError(err, types.ErrorCodeBadRequestBody, types.ErrOptionWithSkipRetry()) + } + if err := refreshReadDeadline(client); err != nil { + return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry()) + } + + eventType, eventErr := responsesWSEventType(message) + if eventErr != nil { + session.sendError("", newResponsesWSInvalidRequestError(eventErr)) + continue + } + + if eventType != responsesWSEventTypeResponseCreate { + if !session.hasTarget() { + session.sendError("", newResponsesWSInvalidRequestError(errors.New("first responses websocket event must be response.create"))) + continue + } + if err := session.writeTarget(messageType, message); err != nil { + return session.handleControlEventWriteFailure(err) + } + continue + } + + create, eventID, err := normalizeResponsesWSCreateEvent(message) + if err != nil { + session.sendError("", newResponsesWSInvalidRequestError(err)) + continue + } + if err := helper.ValidateResponsesRequest(&create.Request); err != nil { + session.sendError(eventID, newResponsesWSInvalidRequestError(err)) + continue + } + if err := session.handleResponseCreate(create, eventID); err != nil { + session.sendError(eventID, err) + } + } +} + +func responsesWSEventType(message []byte) (string, error) { + var event struct { + Type string `json:"type"` + } + if err := common.Unmarshal(message, &event); err != nil { + return "", fmt.Errorf("invalid websocket event json: %w", err) + } + if strings.TrimSpace(event.Type) == "" { + return "", errors.New("websocket event type is required") + } + return event.Type, nil +} + +func newResponsesWSInvalidRequestError(err error) *types.NewAPIError { + return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) +} + +func normalizeResponsesWSCreateEvent(message []byte) (responsesWSCreateRequest, string, error) { + var event responsesWSCreateEvent + if err := common.Unmarshal(message, &event); err != nil { + return responsesWSCreateRequest{}, "", err + } + if event.Type != responsesWSEventTypeResponseCreate { + return responsesWSCreateRequest{}, event.EventID, fmt.Errorf("unsupported event type %q", event.Type) + } + + var generate common.RawMessage + var raw map[string]common.RawMessage + rawErr := common.Unmarshal(message, &raw) + if rawErr == nil { + if generateRaw, ok := raw["generate"]; ok { + generate = generateRaw + } + } + + payload := event.Request + if len(payload) == 0 { + if rawErr != nil { + return responsesWSCreateRequest{}, event.EventID, rawErr + } + delete(raw, "type") + delete(raw, "event_id") + delete(raw, "generate") + stripResponsesWSTransportFields(raw) + var err error + payload, err = common.Marshal(raw) + if err != nil { + return responsesWSCreateRequest{}, event.EventID, err + } + } else { + var responseMap map[string]common.RawMessage + if err := common.Unmarshal(payload, &responseMap); err == nil { + if len(generate) == 0 { + if generateRaw, ok := responseMap["generate"]; ok { + generate = generateRaw + } + } + if _, exists := responseMap["generate"]; exists { + delete(responseMap, "generate") + if merged, err := common.Marshal(responseMap); err == nil { + payload = merged + } + } + } + } + + var req dto.OpenAIResponsesRequest + if err := common.Unmarshal(payload, &req); err != nil { + return responsesWSCreateRequest{}, event.EventID, err + } + req.Stream = nil + req.StreamOptions = nil + return responsesWSCreateRequest{ + Request: req, + Generate: generate, + }, event.EventID, nil +} + +func (s *responsesWSSession) handleResponseCreate(create responsesWSCreateRequest, eventID string) *types.NewAPIError { + req := create.Request + if s.lockedModel != "" && req.Model != s.lockedModel { + return types.NewErrorWithStatusCode( + fmt.Errorf("responses websocket connection is locked to model %q; got %q", s.lockedModel, req.Model), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + + if s.hasCurrent() { + return types.NewErrorWithStatusCode( + errors.New("another response.create is already in progress on this websocket connection"), + types.ErrorCodeInvalidRequest, + http.StatusConflict, + types.ErrOptionWithSkipRetry(), + ) + } + + commitRate, apiErr := middleware.CheckModelRequestRateLimit(s.c) + if apiErr != nil { + return apiErr + } + + if !s.hasTarget() { + return s.connectAndSendFirst(create, commitRate) + } + + state, payload, apiErr := s.prepareCall(create, commitRate) + if apiErr != nil { + commitRate(false) + return apiErr + } + if !s.tryReserveCurrent(state) { + state.refund(s.c) + commitRate(false) + return types.NewErrorWithStatusCode( + errors.New("another response.create is already in progress on this websocket connection"), + types.ErrorCodeInvalidRequest, + http.StatusConflict, + types.ErrOptionWithSkipRetry(), + ) + } + if err := s.writeTarget(websocket.TextMessage, payload); err != nil { + return s.handleTargetWriteFailureWithState(state, err) + } + return nil +} + +func (s *responsesWSSession) handleControlEventWriteFailure(err error) *types.NewAPIError { + apiErr := s.handleTargetWriteFailure(err) + s.sendError("", apiErr) + return nil +} + +func (s *responsesWSSession) handleTargetWriteFailure(err error) *types.NewAPIError { + s.closeTarget() + apiErr := types.NewError(err, types.ErrorCodeBadResponse) + apiErr, _ = s.processChannelError(s.lockedChannel, apiErr, nil) + return apiErr +} + +func (s *responsesWSSession) handleTargetWriteFailureWithState(state *responsesWSCallState, err error) *types.NewAPIError { + s.finishCall(state, responsesWSCallAborted) + return s.handleTargetWriteFailure(err) +} + +func (s *responsesWSSession) connectAndSendFirst(create responsesWSCreateRequest, commitRate middleware.ModelRequestRateLimitCommit) *types.NewAPIError { + req := create.Request + if err := checkResponsesWSModelAccess(s.c, req.Model); err != nil { + commitRate(false) + return err + } + + retryParam := &service.RetryParam{ + Ctx: s.c, + TokenGroup: common.GetContextKeyString(s.c, appconstant.ContextKeyUsingGroup), + ModelName: req.Model, + Retry: common.GetPointer(0), + } + if retryParam.TokenGroup == "" { + retryParam.TokenGroup = common.GetContextKeyString(s.c, appconstant.ContextKeyTokenGroup) + } + + var lastErr *types.NewAPIError + for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { + channel, apiErr := selectResponsesWSChannel(s.c, req.Model, retryParam) + if apiErr != nil { + lastErr = apiErr + break + } + addResponsesWSUsedChannel(s.c, channel.Id) + + if channel.Type != appconstant.ChannelTypeOpenAI && channel.Type != appconstant.ChannelTypeCodex { + lastErr = types.NewErrorWithStatusCode( + fmt.Errorf("responses websocket only supports OpenAI and Codex channels, got channel type %d", channel.Type), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + continue + } + + state, payload, apiErr := s.prepareCall(create, commitRate) + if apiErr != nil { + commitRate(false) + return apiErr + } + + adaptor := GetAdaptor(state.info.ApiType) + if adaptor == nil { + state.refund(s.c) + apiErr = types.NewError(fmt.Errorf("invalid api type: %d", state.info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry()) + var shouldRetry bool + lastErr, shouldRetry = s.processChannelError(channel, apiErr, retryParam) + if !shouldRetry { + break + } + continue + } + adaptor.Init(state.info) + target, apiErr := dialResponsesWebSocketUpstream(s.c, adaptor, state.info) + if apiErr != nil { + state.refund(s.c) + var shouldRetry bool + lastErr, shouldRetry = s.processChannelError(channel, apiErr, retryParam) + if !shouldRetry { + break + } + continue + } + + s.setTarget(target) + if !s.tryReserveCurrent(state) { + s.closeTarget() + state.refund(s.c) + commitRate(false) + return types.NewErrorWithStatusCode(errors.New("another response.create is already in progress on this websocket connection"), types.ErrorCodeInvalidRequest, http.StatusConflict, types.ErrOptionWithSkipRetry()) + } + if err := s.writeTarget(websocket.TextMessage, payload); err != nil { + s.finishCall(state, responsesWSCallAborted) + s.closeTarget() + apiErr = types.NewError(err, types.ErrorCodeBadResponse) + var shouldRetry bool + lastErr, shouldRetry = s.processChannelError(channel, apiErr, retryParam) + if !shouldRetry { + break + } + continue + } + + s.lockedModel = req.Model + s.lockedChannel = channel + s.registerChannelClose(channel.Id) + service.RecordChannelAffinity(s.c, channel.Id) + s.startTargetReader() + return nil + } + + if lastErr == nil { + lastErr = types.NewError(errors.New("failed to connect responses websocket upstream"), types.ErrorCodeDoRequestFailed, types.ErrOptionWithSkipRetry()) + } + commitRate(false) + return lastErr +} + +func (s *responsesWSSession) processChannelError(channel *appmodel.Channel, apiErr *types.NewAPIError, retryParam *service.RetryParam) (*types.NewAPIError, bool) { + if apiErr == nil { + return nil, false + } + apiErr = service.NormalizeViolationFeeError(apiErr) + statusCodeMapping := "" + if s.c != nil { + statusCodeMapping = s.c.GetString("status_code_mapping") + } + service.ResetStatusCode(apiErr, statusCodeMapping) + if channel != nil && s.c != nil { + service.ProcessChannelError(s.c, *types.NewChannelError( + channel.Id, + channel.Type, + channel.Name, + channel.ChannelInfo.IsMultiKey, + common.GetContextKeyString(s.c, appconstant.ContextKeyChannelKey), + channel.GetAutoBan(), + ), apiErr) + } + if retryParam == nil { + return apiErr, false + } + return apiErr, service.ShouldRetryRelayError(s.c, apiErr, common.RetryTimes-retryParam.GetRetry()) +} + +func (s *responsesWSSession) prepareCall(create responsesWSCreateRequest, commitRate middleware.ModelRequestRateLimitCommit) (*responsesWSCallState, []byte, *types.NewAPIError) { + req := create.Request + common.SetContextKey(s.c, appconstant.ContextKeyRequestStartTime, time.Now()) + relayInfo := relaycommon.GenRelayInfoResponses(s.c, &req) + // The stream field is stripped from the frame before parsing, so + // GenRelayInfoResponses sees stream=nil and would record the call as + // non-stream, which also hides first-response time in the usage-log UI. + // WebSocket delivery is inherently incremental; mark it streaming like the + // realtime relay does. + relayInfo.IsStream = true + relayInfo.ClientWs = s.client + s.c.Set(string(appconstant.ContextKeyIsStream), true) + relayInfo.RequestId = fmt.Sprintf("%s-ws-%d", relayInfo.RequestId, s.nextEventIndex) + s.nextEventIndex++ + + meta := req.GetTokenCountMeta() + if setting.ShouldCheckPromptSensitive() && meta != nil { + contains, words := service.CheckSensitiveText(meta.CombineText) + if contains { + return nil, nil, types.NewError(fmt.Errorf("user sensitive words detected: %s", strings.Join(words, ", ")), types.ErrorCodeSensitiveWordsDetected, types.ErrOptionWithSkipRetry()) + } + } + + tokens, err := service.EstimateRequestToken(s.c, meta, relayInfo) + if err != nil { + return nil, nil, types.NewError(err, types.ErrorCodeCountTokenFailed) + } + relayInfo.SetEstimatePromptTokens(tokens) + + priceData, err := helper.ModelPriceHelper(s.c, relayInfo, tokens, meta) + if err != nil { + return nil, nil, types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest)) + } + if !priceData.FreeModel { + if apiErr := service.PreConsumeBilling(s.c, priceData.QuotaToPreConsume, relayInfo); apiErr != nil { + return nil, nil, apiErr + } + } + + payload, apiErr := buildResponsesWSCreatePayload(s.c, relayInfo, req, create.Generate) + if apiErr != nil { + if relayInfo.Billing != nil { + relayInfo.Billing.Refund(s.c) + } + return nil, nil, apiErr + } + + return &responsesWSCallState{ + info: relayInfo, + usage: &dto.Usage{}, + commitRate: commitRate, + }, payload, nil +} + +func buildResponsesWSCreatePayload(c *gin.Context, relayInfo *relaycommon.RelayInfo, req dto.OpenAIResponsesRequest, generate common.RawMessage) ([]byte, *types.NewAPIError) { + relayInfo.InitChannelMeta(c) + request, err := common.DeepCopy(&req) + if err != nil { + return nil, types.NewError(fmt.Errorf("failed to copy responses request: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) + } + if err := helper.ModelMappedHelper(c, relayInfo, request); err != nil { + return nil, types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry()) + } + + adaptor := GetAdaptor(relayInfo.ApiType) + if adaptor == nil { + return nil, types.NewError(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry()) + } + adaptor.Init(relayInfo) + convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, relayInfo, *request) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + relaycommon.AppendRequestConversionFromRequest(relayInfo, convertedRequest) + jsonData, err := common.Marshal(convertedRequest) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + jsonData, err = relaycommon.RemoveDisabledFields(jsonData, relayInfo.ChannelOtherSettings, relayInfo.ChannelSetting.PassThroughBodyEnabled) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + jsonData, err = removeResponsesWSTransportFields(jsonData) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + if len(relayInfo.ParamOverride) > 0 { + jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, relayInfo) + if err != nil { + return nil, newAPIErrorFromParamOverride(err) + } + } + + event, err := buildResponsesWSCreateEvent(jsonData, generate) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + return event, nil +} + +func buildResponsesWSCreateEvent(jsonData []byte, generate common.RawMessage) ([]byte, error) { + var event map[string]common.RawMessage + if err := common.Unmarshal(jsonData, &event); err != nil { + return nil, err + } + typeData, err := common.Marshal(responsesWSEventTypeResponseCreate) + if err != nil { + return nil, err + } + event["type"] = typeData + delete(event, "event_id") + stripResponsesWSTransportFields(event) + if len(generate) > 0 { + event["generate"] = generate + } + return common.Marshal(event) +} + +func removeResponsesWSTransportFields(jsonData []byte) ([]byte, error) { + var data map[string]common.RawMessage + if err := common.Unmarshal(jsonData, &data); err != nil { + return jsonData, err + } + stripResponsesWSTransportFields(data) + return common.Marshal(data) +} + +func stripResponsesWSTransportFields(data map[string]common.RawMessage) { + delete(data, "stream") + delete(data, "stream_options") + delete(data, "background") +} + +func dialResponsesWebSocketUpstream(c *gin.Context, adaptor relaychannel.Adaptor, info *relaycommon.RelayInfo) (*websocket.Conn, *types.NewAPIError) { + fullRequestURL, err := adaptor.GetRequestURL(info) + if err != nil { + return nil, types.NewError(fmt.Errorf("get request url failed: %w", err), types.ErrorCodeDoRequestFailed) + } + fullRequestURL = toWebSocketURL(fullRequestURL) + + targetHeader := http.Header{} + if err := adaptor.SetupRequestHeader(c, &targetHeader, info); err != nil { + return nil, types.NewError(fmt.Errorf("setup request header failed: %w", err), types.ErrorCodeDoRequestFailed) + } + headerOverride, err := relaychannel.ResolveHeaderOverride(info, c) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid) + } + for key, value := range headerOverride { + targetHeader.Set(key, value) + } + prepareResponsesWebSocketHeaders(c, &targetHeader) + + targetConn, resp, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader) + if err != nil { + statusCode := http.StatusInternalServerError + if resp != nil { + statusCode = resp.StatusCode + } + return nil, types.NewErrorWithStatusCode(fmt.Errorf("dial failed to %s: %w", relaycommon.SanitizeURLForLog(fullRequestURL), err), types.ErrorCodeDoRequestFailed, statusCode) + } + targetConn.SetReadLimit(responsesWSMaxMessageBytes()) + return targetConn, nil +} + +// prepareResponsesWebSocketHeaders removes HTTP/SSE negotiation that may have +// been inferred from RelayInfo.IsStream and preserves the client's Responses +// WebSocket beta contract for the upstream handshake. +func prepareResponsesWebSocketHeaders(c *gin.Context, header *http.Header) { + if header == nil { + return + } + header.Del("Accept") + if c == nil || c.Request == nil { + return + } + if beta := strings.TrimSpace(c.Request.Header.Get("OpenAI-Beta")); beta != "" { + header.Set("OpenAI-Beta", beta) + } +} + +func toWebSocketURL(raw string) string { + switch { + case strings.HasPrefix(raw, "https://"): + return "wss://" + strings.TrimPrefix(raw, "https://") + case strings.HasPrefix(raw, "http://"): + return "ws://" + strings.TrimPrefix(raw, "http://") + default: + return raw + } +} + +func (s *responsesWSSession) startTargetReader() { + target := s.getTarget() + if target == nil { + return + } + go func() { + for { + messageType, message, err := target.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + logger.LogError(s.c, "responses websocket upstream read failed: "+err.Error()) + } + s.settleCurrent() + _ = s.client.Close() + return + } + // Upstream traffic also counts as activity: a long generation can + // stream for minutes while the client only listens, and that must + // not trip the client idle timeout. + // + // Calling this from the reader goroutine while the client goroutine + // blocks in ReadMessage is safe despite gorilla listing + // SetReadDeadline as a read method: it is a bare passthrough to + // net.Conn.SetReadDeadline (conn.go:1105), which is documented to be + // callable concurrently with a blocked Read. + _ = relaycommon.RefreshClientWebSocketReadDeadline(s.client) + + // Drop transport-only keepalives (e.g. from CLI API proxies). Strict + // clients reject unknown event.type variants and abort the turn. + if isResponsesWSTransportKeepalive(message) { + continue + } + + s.observeUpstreamMessage(message) + if err := s.writeClient(messageType, message); err != nil { + logger.LogError(s.c, "responses websocket client write failed: "+err.Error()) + s.settleCurrent() + s.closeTarget() + return + } + } + }() +} + +// isResponsesWSTransportKeepalive peeks at event.type without full parsing so +// the reader can drop non-semantic frames before billing observation/forwarding. +func isResponsesWSTransportKeepalive(message []byte) bool { + var envelope struct { + Type string `json:"type"` + } + if err := common.Unmarshal(message, &envelope); err != nil { + return false + } + return dto.IsResponsesTransportEventType(envelope.Type) +} + +func (s *responsesWSSession) observeUpstreamMessage(message []byte) { + state := s.getCurrent() + if state == nil { + return + } + state.info.SetFirstResponseTime() + + var streamResponse dto.ResponsesStreamResponse + if err := common.Unmarshal(message, &streamResponse); err != nil { + return + } + + switch streamResponse.Type { + case "response.completed", "response.done", "response.incomplete", + "response.failed", "response.cancelled", "response.canceled": + // A terminal event carries the authoritative usage even when the response + // failed, so settle on it instead of discarding what upstream generated + // and already billed us for. + s.applyTerminalResponseUsage(state, streamResponse.Response) + s.finishCall(state, responsesWSCallSettled) + case "response.output_text.delta": + state.mu.Lock() + state.outputText.WriteString(streamResponse.Delta) + state.mu.Unlock() + case dto.ResponsesOutputTypeItemDone: + if streamResponse.Item != nil && streamResponse.Item.Type == dto.BuildInCallWebSearchCall { + if state.info != nil && state.info.ResponsesUsageInfo != nil && state.info.ResponsesUsageInfo.BuiltInTools != nil { + if webSearchTool, exists := state.info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; exists && webSearchTool != nil { + webSearchTool.CallCount++ + } + } + } + case "error": + s.finishCall(state, responsesWSCallSettled) + } +} + +func (s *responsesWSSession) applyTerminalResponseUsage(state *responsesWSCallState, response *dto.OpenAIResponsesResponse) { + if state == nil || response == nil { + return + } + if response.Usage != nil { + state.mu.Lock() + service.ApplyResponsesUsage(state.usage, response.Usage) + state.mu.Unlock() + } + if response.HasImageGenerationCall() { + s.c.Set("image_generation_call", true) + s.c.Set("image_generation_call_quality", response.GetQuality()) + s.c.Set("image_generation_call_size", response.GetSize()) + } +} + +func (s *responsesWSSession) finishCall(state *responsesWSCallState, outcome responsesWSCallOutcome) { + if state == nil { + return + } + if !s.clearCurrent(state) { + return + } + // Refund only when upstream produced nothing: either it never accepted the + // request, or it accepted but no usage and no output text were observed. + // Anything actually generated gets billed, otherwise disconnecting just + // before the terminal event would yield free output. + if outcome == responsesWSCallAborted || !finalizeResponsesWSUsage(state) { + state.refund(s.c) + if state.commitRate != nil { + state.commitRate(false) + } + return + } + + // Bill a snapshot: the goroutine that lost the clearCurrent race may still + // be applying a late terminal event to state.usage under state.mu. + state.mu.Lock() + usage := *state.usage + state.mu.Unlock() + service.PostTextConsumeQuota(s.c, state.info, &usage, nil) + if state.commitRate != nil { + state.commitRate(true) + } +} + +// finalizeResponsesWSUsage fills in what upstream did not report — the usual +// case for a stream cut short — and reports whether anything is billable. +func finalizeResponsesWSUsage(state *responsesWSCallState) bool { + if state == nil || state.usage == nil || state.info == nil { + return false + } + state.mu.Lock() + defer state.mu.Unlock() + if state.usage.CompletionTokens == 0 { + if output := state.outputText.String(); output != "" { + state.usage.CompletionTokens = service.CountTextToken(output, state.info.UpstreamModelName) + } + } + if state.usage.PromptTokens == 0 && state.usage.CompletionTokens != 0 { + state.usage.PromptTokens = state.info.GetEstimatePromptTokens() + } + if state.usage.TotalTokens == 0 { + state.usage.TotalTokens = state.usage.PromptTokens + state.usage.CompletionTokens + } + return state.usage.TotalTokens > 0 +} + +func (state *responsesWSCallState) refund(c *gin.Context) { + if state != nil && state.info != nil && state.info.Billing != nil { + state.info.Billing.Refund(c) + } +} + +func (s *responsesWSSession) tryReserveCurrent(state *responsesWSCallState) bool { + s.stateMu.Lock() + defer s.stateMu.Unlock() + if s.current != nil { + return false + } + s.current = state + return true +} + +func (s *responsesWSSession) hasCurrent() bool { + s.stateMu.Lock() + defer s.stateMu.Unlock() + return s.current != nil +} + +func (s *responsesWSSession) clearCurrent(state *responsesWSCallState) bool { + s.stateMu.Lock() + defer s.stateMu.Unlock() + if state != nil && s.current != state { + return false + } + s.current = nil + return true +} + +func (s *responsesWSSession) getCurrent() *responsesWSCallState { + s.stateMu.Lock() + defer s.stateMu.Unlock() + return s.current +} + +// settleCurrent ends an in-flight call that was interrupted rather than +// completed — client disconnect, upstream read failure, idle timeout or channel +// shutdown. The payload already reached upstream, so it settles on observed +// usage; finishCall still refunds if nothing was generated. +func (s *responsesWSSession) settleCurrent() { + state := s.getCurrent() + if state != nil { + s.finishCall(state, responsesWSCallSettled) + } +} + +func acquireResponsesWSSlot(userId int) bool { + if responsesWSMaxPerUser <= 0 || userId == 0 { + return true + } + responsesWSCountMu.Lock() + defer responsesWSCountMu.Unlock() + if responsesWSCounts[userId] >= responsesWSMaxPerUser { + return false + } + responsesWSCounts[userId]++ + return true +} + +func releaseResponsesWSSlot(userId int) { + if responsesWSMaxPerUser <= 0 || userId == 0 { + return + } + responsesWSCountMu.Lock() + defer responsesWSCountMu.Unlock() + if responsesWSCounts[userId] <= 1 { + delete(responsesWSCounts, userId) + return + } + responsesWSCounts[userId]-- +} + +func (s *responsesWSSession) writeClient(messageType int, message []byte) error { + s.clientWriteMu.Lock() + defer s.clientWriteMu.Unlock() + if err := s.client.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil { + return err + } + return s.client.WriteMessage(messageType, message) +} + +func (s *responsesWSSession) hasTarget() bool { + s.targetMu.Lock() + defer s.targetMu.Unlock() + return s.target != nil +} + +func (s *responsesWSSession) getTarget() *websocket.Conn { + s.targetMu.Lock() + defer s.targetMu.Unlock() + return s.target +} + +func (s *responsesWSSession) setTarget(target *websocket.Conn) { + s.targetMu.Lock() + defer s.targetMu.Unlock() + s.target = target +} + +func (s *responsesWSSession) writeTarget(messageType int, message []byte) error { + // Resolve the target under targetMu, then release it before writing: a slow + // upstream must not be able to block closeTarget or the idle/policy paths. + target := s.getTarget() + if target == nil { + return errors.New("responses websocket upstream is not connected") + } + s.targetWriteMu.Lock() + defer s.targetWriteMu.Unlock() + if err := target.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil { + return err + } + return target.WriteMessage(messageType, message) +} + +func (s *responsesWSSession) sendError(eventID string, apiErr *types.NewAPIError) { + if apiErr == nil { + return + } + payload, err := buildResponsesWSErrorPayload(eventID, apiErr) + if err != nil { + return + } + _ = s.writeClient(websocket.TextMessage, payload) +} + +func buildResponsesWSErrorPayload(eventID string, apiErr *types.NewAPIError) ([]byte, error) { + if apiErr == nil { + return nil, errors.New("api error is nil") + } + status := apiErr.StatusCode + if status == 0 { + status = http.StatusInternalServerError + } + openaiErr := apiErr.ToOpenAIError() + return common.Marshal(&responsesWSErrorEvent{ + Type: "error", + Status: status, + EventID: eventID, + Error: &openaiErr, + }) +} + +func (s *responsesWSSession) closeTarget() { + var target *websocket.Conn + var unregister func() + s.targetMu.Lock() + target = s.target + s.target = nil + unregister = s.unregister + s.unregister = nil + s.targetMu.Unlock() + if unregister != nil { + unregister() + } + if target != nil { + _ = target.Close() + } +} + +func (s *responsesWSSession) registerChannelClose(channelID int) { + unregister := wsmanager.Register(channelID, wsmanager.KindResponses, func(reason string) { + s.closeForPolicy(reason) + }) + s.targetMu.Lock() + if s.unregister != nil { + s.unregister() + } + s.unregister = unregister + s.targetMu.Unlock() +} + +func (s *responsesWSSession) closeForPolicy(reason string) { + s.closeWithCode(websocket.ClosePolicyViolation, reason) +} + +func (s *responsesWSSession) closeForIdleTimeout() { + s.closeWithCode(websocket.CloseGoingAway, relaycommon.WebSocketIdleCloseReason) +} + +func (s *responsesWSSession) closeWithCode(code int, reason string) { + s.closeOnce.Do(func() { + s.settleCurrent() + deadline := time.Now().Add(time.Second) + closeMessage := websocket.FormatCloseMessage(code, reason) + _ = s.client.WriteControl(websocket.CloseMessage, closeMessage, deadline) + if target := s.getTarget(); target != nil { + _ = target.WriteControl(websocket.CloseMessage, closeMessage, deadline) + } + s.closeTarget() + _ = s.client.Close() + }) +} + +func checkResponsesWSModelAccess(c *gin.Context, modelName string) *types.NewAPIError { + if !common.GetContextKeyBool(c, appconstant.ContextKeyTokenModelLimitEnabled) { + return nil + } + raw, ok := common.GetContextKey(c, appconstant.ContextKeyTokenModelLimit) + if !ok { + return types.NewErrorWithStatusCode(errors.New("token has no model access"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + tokenModelLimit, ok := raw.(map[string]bool) + if !ok { + tokenModelLimit = map[string]bool{} + } + matchName := ratio_setting.FormatMatchingModelName(modelName) + if _, ok := tokenModelLimit[matchName]; !ok { + return types.NewErrorWithStatusCode(fmt.Errorf("token is not allowed to use model %s", modelName), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + return nil +} + +func selectResponsesWSChannel(c *gin.Context, modelName string, retryParam *service.RetryParam) (*appmodel.Channel, *types.NewAPIError) { + if channelIdRaw, ok := common.GetContextKey(c, appconstant.ContextKeyTokenSpecificChannelId); ok { + channelID, ok := channelIdRaw.(string) + if !ok { + return nil, types.NewErrorWithStatusCode(errors.New("invalid specified channel id"), types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + id, err := strconv.Atoi(channelID) + if err != nil { + return nil, types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + channel, err := appmodel.GetChannelById(id, true) + if err != nil { + return nil, types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + if channel.Status != common.ChannelStatusEnabled { + return nil, types.NewErrorWithStatusCode(errors.New("specified channel is disabled"), types.ErrorCodeGetChannelFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + if err := middleware.SetupContextForSelectedChannel(c, channel, modelName); err != nil { + return nil, err + } + return channel, nil + } + + usingGroup := common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup) + if usingGroup == "" { + usingGroup = retryParam.TokenGroup + } + + if retryParam.GetRetry() == 0 { + if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelName, usingGroup); found { + preferred, err := appmodel.CacheGetChannel(preferredChannelID) + if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { + if usingGroup == "auto" { + userGroup := common.GetContextKeyString(c, appconstant.ContextKeyUserGroup) + for _, g := range service.GetUserAutoGroup(userGroup) { + if appmodel.IsChannelEnabledForGroupModel(g, modelName, preferred.Id) { + common.SetContextKey(c, appconstant.ContextKeyAutoGroup, g) + service.MarkChannelAffinityUsed(c, g, preferred.Id) + if err := middleware.SetupContextForSelectedChannel(c, preferred, modelName); err != nil { + return nil, err + } + return preferred, nil + } + } + } else if appmodel.IsChannelEnabledForGroupModel(usingGroup, modelName, preferred.Id) { + service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + if err := middleware.SetupContextForSelectedChannel(c, preferred, modelName); err != nil { + return nil, err + } + return preferred, nil + } + } + } + } + + channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam) + if err != nil { + return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败(retry): %s", selectGroup, modelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + } + if channel == nil { + return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在(retry)", selectGroup, modelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + } + if err := middleware.SetupContextForSelectedChannel(c, channel, modelName); err != nil { + return nil, err + } + return channel, nil +} + +func addResponsesWSUsedChannel(c *gin.Context, channelId int) { + useChannel := c.GetStringSlice("use_channel") + useChannel = append(useChannel, fmt.Sprintf("%d", channelId)) + c.Set("use_channel", useChannel) +} diff --git a/relay/responses_websocket_test.go b/relay/responses_websocket_test.go new file mode 100644 index 000000000000..4eefc192e6f4 --- /dev/null +++ b/relay/responses_websocket_test.go @@ -0,0 +1,533 @@ +package relay + +import ( + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNormalizeResponsesWSCreateEventWrapper(t *testing.T) { + message := []byte(`{ + "type": "response.create", + "event_id": "evt_1", + "generate": false, + "response": { + "model": "gpt-5.3-codex-spark", + "input": "hi", + "store": false, + "stream": true, + "stream_options": {"include_usage": true} + } + }`) + + create, eventID, err := normalizeResponsesWSCreateEvent(message) + if err != nil { + t.Fatalf("normalizeResponsesWSCreateEvent() error = %v", err) + } + req := create.Request + if eventID != "evt_1" { + t.Fatalf("eventID = %q, want evt_1", eventID) + } + if req.Model != "gpt-5.3-codex-spark" { + t.Fatalf("model = %q", req.Model) + } + if strings.TrimSpace(string(create.Generate)) != "false" { + t.Fatalf("generate = %s, want false", create.Generate) + } + if req.Stream != nil { + t.Fatalf("stream = %v, want nil", req.Stream) + } + if req.StreamOptions != nil { + t.Fatalf("stream_options = %#v, want nil", req.StreamOptions) + } + if strings.TrimSpace(string(req.Store)) != "false" { + t.Fatalf("store = %s, want false", req.Store) + } +} + +func TestNormalizeResponsesWSCreateEventFlat(t *testing.T) { + message := []byte(`{ + "type": "response.create", + "event_id": "evt_2", + "model": "gpt-5.3-codex-spark", + "input": "hi", + "generate": false, + "stream": true, + "background": true, + "stream_options": {"include_usage": true} + }`) + + create, eventID, err := normalizeResponsesWSCreateEvent(message) + if err != nil { + t.Fatalf("normalizeResponsesWSCreateEvent() error = %v", err) + } + req := create.Request + if eventID != "evt_2" { + t.Fatalf("eventID = %q, want evt_2", eventID) + } + if req.Model != "gpt-5.3-codex-spark" { + t.Fatalf("model = %q", req.Model) + } + if strings.TrimSpace(string(create.Generate)) != "false" { + t.Fatalf("generate = %s, want false", create.Generate) + } + if req.Stream != nil { + t.Fatalf("stream = %v, want nil", req.Stream) + } + if req.StreamOptions != nil { + t.Fatalf("stream_options = %#v, want nil", req.StreamOptions) + } +} + +func TestBuildResponsesWSCreateEventIsFlat(t *testing.T) { + payload := []byte(`{ + "model": "gpt-5.3-codex-spark", + "input": "hi", + "store": false, + "event_id": "evt_upstream", + "stream": true, + "background": true, + "stream_options": {"include_usage": true} + }`) + + got, err := buildResponsesWSCreateEvent(payload, common.RawMessage(`false`)) + if err != nil { + t.Fatalf("buildResponsesWSCreateEvent() error = %v", err) + } + var data map[string]any + if err := common.Unmarshal(got, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data["type"] != responsesWSEventTypeResponseCreate { + t.Fatalf("type = %#v", data["type"]) + } + if data["model"] != "gpt-5.3-codex-spark" || data["input"] != "hi" || data["store"] != false { + t.Fatalf("unexpected flat event fields: %s", got) + } + if data["generate"] != false { + t.Fatalf("generate = %#v, want false", data["generate"]) + } + for _, key := range []string{"response", "event_id", "stream", "background", "stream_options"} { + if _, ok := data[key]; ok { + t.Fatalf("field %q should not be present in upstream event: %s", key, got) + } + } +} + +func TestHTTPResponsesRequestDoesNotMarshalGenerate(t *testing.T) { + var req dto.OpenAIResponsesRequest + if err := common.Unmarshal([]byte(`{"model":"gpt-5.3-codex-spark","input":"hi","generate":false}`), &req); err != nil { + t.Fatalf("unmarshal request: %v", err) + } + got, err := common.Marshal(req) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + var data map[string]any + if err := common.Unmarshal(got, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if _, ok := data["generate"]; ok { + t.Fatalf("generate leaked into HTTP request JSON: %s", got) + } +} + +func TestBuildResponsesWSErrorPayloadIncludesStatus(t *testing.T) { + payload, err := buildResponsesWSErrorPayload("evt_err", types.NewErrorWithStatusCode( + errors.New("model is required"), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + )) + if err != nil { + t.Fatalf("buildResponsesWSErrorPayload() error = %v", err) + } + var data struct { + Type string `json:"type"` + Status int `json:"status"` + EventID string `json:"event_id"` + Error *types.OpenAIError `json:"error"` + } + if err := common.Unmarshal(payload, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data.Type != "error" || data.Status != http.StatusBadRequest || data.EventID != "evt_err" { + t.Fatalf("unexpected error event: %s", payload) + } + if data.Error == nil || data.Error.Code != string(types.ErrorCodeInvalidRequest) { + t.Fatalf("unexpected error body: %#v", data.Error) + } +} + +func TestResponsesWSInvalidRequestErrorUsesBadRequestStatus(t *testing.T) { + payload, err := buildResponsesWSErrorPayload("", newResponsesWSInvalidRequestError(errors.New("bad event"))) + if err != nil { + t.Fatalf("buildResponsesWSErrorPayload() error = %v", err) + } + var data struct { + Status int `json:"status"` + } + if err := common.Unmarshal(payload, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data.Status != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", data.Status, http.StatusBadRequest) + } +} + +func TestRemoveResponsesWSTransportFields(t *testing.T) { + payload := []byte(`{ + "model": "gpt-5.3-codex-spark", + "stream": true, + "background": true, + "stream_options": {"include_usage": true}, + "store": false + }`) + + got, err := removeResponsesWSTransportFields(payload) + if err != nil { + t.Fatalf("removeResponsesWSTransportFields() error = %v", err) + } + var data map[string]any + if err := common.Unmarshal(got, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + for _, key := range []string{"stream", "background", "stream_options"} { + if _, ok := data[key]; ok { + t.Fatalf("transport field %q still present in %s", key, got) + } + } + if data["store"] != false { + t.Fatalf("store = %#v, want false", data["store"]) + } +} + +func TestToWebSocketURL(t *testing.T) { + tests := map[string]string{ + "https://api.openai.com/v1/responses": "wss://api.openai.com/v1/responses", + "http://127.0.0.1:3000/v1/responses": "ws://127.0.0.1:3000/v1/responses", + "wss://chatgpt.com/backend-api/codex/responses": "wss://chatgpt.com/backend-api/codex/responses", + "ws://127.0.0.1:3000/backend-api/codex/responses": "ws://127.0.0.1:3000/backend-api/codex/responses", + } + + for input, want := range tests { + if got := toWebSocketURL(input); got != want { + t.Fatalf("toWebSocketURL(%q) = %q, want %q", input, got, want) + } + } +} + +func TestPrepareResponsesWebSocketHeaders(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + c.Request.Header.Set("OpenAI-Beta", "responses_websockets=2026-02-06") + header := http.Header{ + "Accept": []string{"text/event-stream"}, + "Content-Type": []string{"application/json"}, + } + + prepareResponsesWebSocketHeaders(c, &header) + + assert.Empty(t, header.Get("Accept"), "SSE Accept must not leak into a WebSocket handshake") + assert.Equal(t, "application/json", header.Get("Content-Type")) + assert.Equal(t, "responses_websockets=2026-02-06", header.Get("OpenAI-Beta")) +} + +func TestPrepareResponsesWebSocketHeadersWithoutClientBetaPreservesAdaptorBeta(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + header := http.Header{} + header.Set("Accept", "text/event-stream") + header.Set("OpenAI-Beta", "responses=experimental") + + prepareResponsesWebSocketHeaders(c, &header) + + assert.Empty(t, header.Get("Accept")) + assert.Equal(t, "responses=experimental", header.Get("OpenAI-Beta")) +} + +func TestHandleTargetWriteFailureWithStateReleasesCurrentAndClearsTarget(t *testing.T) { + target, cleanup := newTestResponsesWSTarget(t) + defer cleanup() + + var committed *bool + session := &responsesWSSession{target: target} + state := &responsesWSCallState{ + info: &relaycommon.RelayInfo{}, + commitRate: func(success bool) { + committed = &success + }, + } + session.current = state + + apiErr := session.handleTargetWriteFailureWithState(state, errors.New("write failed")) + + if apiErr == nil { + t.Fatal("apiErr is nil") + } + if session.target != nil { + t.Fatal("target was not cleared") + } + if session.getCurrent() != nil { + t.Fatal("current response was not released") + } + if committed == nil || *committed { + t.Fatalf("commit success = %v, want false", committed) + } +} + +func TestHandleControlEventWriteFailureSendsResponsesError(t *testing.T) { + clientConn, serverConn, cleanupClient := newTestWebSocketPair(t) + defer cleanupClient() + target, cleanupTarget := newTestResponsesWSTarget(t) + defer cleanupTarget() + + session := &responsesWSSession{ + client: serverConn, + target: target, + } + apiErr := session.handleControlEventWriteFailure(errors.New("write failed")) + if apiErr != nil { + t.Fatalf("handleControlEventWriteFailure() error = %v", apiErr) + } + if session.target != nil { + t.Fatal("target was not cleared") + } + + if err := clientConn.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, payload, err := clientConn.ReadMessage() + if err != nil { + t.Fatalf("read responses error event: %v", err) + } + var data struct { + Type string `json:"type"` + Status int `json:"status"` + } + if err := common.Unmarshal(payload, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data.Type != "error" || data.Status == 0 { + t.Fatalf("unexpected error event: %s", payload) + } +} + +// TestFinalizeResponsesWSUsageBillsInterruptedStream pins the billing policy +// for a stream that never reached its terminal event — the client disconnected, +// upstream died, or the idle timeout fired. Upstream already generated (and +// charged us for) that output, so it must be billable from the observed delta +// text, not refunded in full. +func TestFinalizeResponsesWSUsageBillsInterruptedStream(t *testing.T) { + info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-sonnet-4"}} + info.SetEstimatePromptTokens(123) + state := &responsesWSCallState{info: info, usage: &dto.Usage{}} + state.outputText.WriteString("partial answer streamed before the client vanished") + + require.True(t, finalizeResponsesWSUsage(state), "generated output must be billable") + assert.Positive(t, state.usage.CompletionTokens, "completion tokens should be counted from observed output") + assert.Equal(t, 123, state.usage.PromptTokens, "prompt tokens should fall back to the pre-consume estimate") + assert.Equal(t, state.usage.PromptTokens+state.usage.CompletionTokens, state.usage.TotalTokens) +} + +func TestFinalizeResponsesWSUsageReportsNothingBillableWithoutOutput(t *testing.T) { + state := &responsesWSCallState{ + info: &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-sonnet-4"}}, + usage: &dto.Usage{}, + } + + assert.False(t, finalizeResponsesWSUsage(state), "a call that produced nothing must stay refundable") +} + +// TestFinishCallAbortedRefundsDespiteObservedOutput guards the other side of the +// policy: when the request never reached upstream there is nothing to pay for, +// even if stale state carries text. +func TestFinishCallAbortedRefundsDespiteObservedOutput(t *testing.T) { + var committed *bool + session := &responsesWSSession{} + state := &responsesWSCallState{ + info: &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-sonnet-4"}}, + usage: &dto.Usage{}, + commitRate: func(success bool) { + committed = &success + }, + } + state.outputText.WriteString("never sent upstream") + session.current = state + + session.finishCall(state, responsesWSCallAborted) + + assert.Nil(t, session.getCurrent(), "current response was not released") + require.NotNil(t, committed, "commit was not invoked") + assert.False(t, *committed, "an aborted call must not be committed as a successful request") +} + +// TestApplyTerminalResponseUsageRecordsFailedResponseUsage covers the fix for +// terminal failure events: upstream reports real usage on response.failed, and +// discarding it meant billing nothing for output the provider already charged. +func TestApplyTerminalResponseUsageRecordsFailedResponseUsage(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + session := &responsesWSSession{c: c} + state := &responsesWSCallState{ + info: &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-sonnet-4"}}, + usage: &dto.Usage{}, + } + + session.applyTerminalResponseUsage(state, &dto.OpenAIResponsesResponse{ + Usage: &dto.Usage{InputTokens: 40, OutputTokens: 9, TotalTokens: 49}, + }) + + assert.Equal(t, 40, state.usage.PromptTokens) + assert.Equal(t, 9, state.usage.CompletionTokens) +} + +func TestResponsesWSSlotCapIsPerUserAndReleased(t *testing.T) { + original := responsesWSMaxPerUser + responsesWSMaxPerUser = 2 + defer func() { responsesWSMaxPerUser = original }() + + const userId = 4242 + require.True(t, acquireResponsesWSSlot(userId)) + require.True(t, acquireResponsesWSSlot(userId)) + assert.False(t, acquireResponsesWSSlot(userId), "third concurrent session must be rejected") + assert.True(t, acquireResponsesWSSlot(userId+1), "cap must be scoped per user") + + releaseResponsesWSSlot(userId) + assert.True(t, acquireResponsesWSSlot(userId), "releasing must free a slot") + + releaseResponsesWSSlot(userId) + releaseResponsesWSSlot(userId) + releaseResponsesWSSlot(userId + 1) + + responsesWSCountMu.Lock() + defer responsesWSCountMu.Unlock() + assert.Empty(t, responsesWSCounts, "fully released users must not leak counter entries") +} + +func TestObserveUpstreamFailedReleasesCurrent(t *testing.T) { + var committed *bool + session := &responsesWSSession{} + state := &responsesWSCallState{ + info: &relaycommon.RelayInfo{}, + commitRate: func(success bool) { + committed = &success + }, + } + session.current = state + + session.observeUpstreamMessage([]byte(`{"type":"response.failed"}`)) + + if session.getCurrent() != nil { + t.Fatal("current response was not released") + } + if committed == nil || *committed { + t.Fatalf("commit success = %v, want false", committed) + } +} + +func TestIsResponsesWSTransportKeepalive(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + message string + want bool + }{ + {name: "keepalive", message: `{"type":"keepalive"}`, want: true}, + {name: "keep_alive with fields", message: `{"type":"keep_alive","ts":1}`, want: true}, + {name: "heartbeat", message: `{"type":"heartbeat"}`, want: true}, + {name: "response delta", message: `{"type":"response.output_text.delta","delta":"hi"}`, want: false}, + {name: "invalid json", message: `not-json`, want: false}, + {name: "missing type", message: `{"delta":"x"}`, want: false}, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := isResponsesWSTransportKeepalive([]byte(tc.message)); got != tc.want { + t.Fatalf("isResponsesWSTransportKeepalive(%s) = %v, want %v", tc.message, got, tc.want) + } + }) + } +} + +func TestStartTargetReaderDropsKeepalive(t *testing.T) { + clientSide, clientServer, cleanupClient := newTestWebSocketPair(t) + defer cleanupClient() + targetSide, targetServer, cleanupTarget := newTestWebSocketPair(t) + defer cleanupTarget() + + // Session reads from targetSide (as if connected to upstream) and writes to clientServer. + session := &responsesWSSession{ + c: &gin.Context{}, + client: clientServer, + target: targetSide, + } + session.startTargetReader() + + // Upstream injects a transport keepalive then a real event. + require.NoError(t, targetServer.WriteMessage(websocket.TextMessage, []byte(`{"type":"keepalive"}`))) + require.NoError(t, targetServer.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created"}`))) + + // Client must only receive the semantic event. + _ = clientSide.SetReadDeadline(time.Now().Add(2 * time.Second)) + _, message, err := clientSide.ReadMessage() + require.NoError(t, err) + assert.JSONEq(t, `{"type":"response.created"}`, string(message)) + + // Ensure no extra keepalive frame is waiting. + _ = clientSide.SetReadDeadline(time.Now().Add(150 * time.Millisecond)) + _, _, err = clientSide.ReadMessage() + require.Error(t, err, "client must not receive a second message after keepalive was dropped") + + session.closeTarget() + _ = clientServer.Close() +} + +func newTestResponsesWSTarget(t *testing.T) (*websocket.Conn, func()) { + t.Helper() + target, _, cleanup := newTestWebSocketPair(t) + return target, cleanup +} + +func newTestWebSocketPair(t *testing.T) (*websocket.Conn, *websocket.Conn, func()) { + t.Helper() + upgrader := websocket.Upgrader{} + serverConnCh := make(chan *websocket.Conn, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + serverConnCh <- conn + })) + + targetURL := "ws" + strings.TrimPrefix(server.URL, "http") + target, _, err := websocket.DefaultDialer.Dial(targetURL, nil) + if err != nil { + server.Close() + t.Fatalf("dial websocket: %v", err) + } + serverConn := <-serverConnCh + cleanup := func() { + _ = target.Close() + _ = serverConn.Close() + server.Close() + } + return target, serverConn, cleanup +} diff --git a/relay/websocket.go b/relay/websocket.go index 57a51895b006..ac3d8f52ee50 100644 --- a/relay/websocket.go +++ b/relay/websocket.go @@ -2,8 +2,10 @@ package relay import ( "fmt" + "time" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/pkg/wsmanager" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/types" @@ -33,6 +35,15 @@ func WssHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types. if resp != nil { info.TargetWs = resp.(*websocket.Conn) defer info.TargetWs.Close() + unregister := wsmanager.Register(info.ChannelId, wsmanager.KindRealtime, func(reason string) { + deadline := time.Now().Add(time.Second) + closeMessage := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, reason) + _ = info.ClientWs.WriteControl(websocket.CloseMessage, closeMessage, deadline) + _ = info.TargetWs.WriteControl(websocket.CloseMessage, closeMessage, deadline) + _ = info.ClientWs.Close() + _ = info.TargetWs.Close() + }) + defer unregister() } usage, newAPIError := adaptor.DoResponse(c, nil, info) diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd6..59487c7d6c64 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -71,6 +71,11 @@ func SetRelayRouter(router *gin.Engine) { relayV1Router.Use(middleware.SystemPerformanceCheck()) relayV1Router.Use(middleware.TokenAuth()) relayV1Router.Use(middleware.ModelRequestRateLimit()) + { + // Responses WebSocket route. Channel selection happens after the first + // response.create event because the model is in the WebSocket payload. + relayV1Router.GET("/responses", controller.ResponsesWebSocket) + } { // WebSocket 路由(统一到 Relay) wsRouter := relayV1Router.Group("") diff --git a/scripts/validate-fork-release-tag.sh b/scripts/validate-fork-release-tag.sh new file mode 100755 index 000000000000..7ecca8c1fc98 --- /dev/null +++ b/scripts/validate-fork-release-tag.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash + +set -euo pipefail + +tag=${1:-} +if [[ ! $tag =~ ^v[0-9]+\.[0-9]+\.[0-9]+(-[0-9A-Za-z.]+)?-[0-9]+$ ]]; then + echo "invalid fork release tag: '$tag'" >&2 + echo "expected an upstream version plus numeric fork suffix, for example v1.0.0-rc.23-0 or v0.13.2-0" >&2 + exit 64 +fi + +case "$tag" in + v1.*|v0.13.*) + ;; + *) + echo "unsupported fork release track: '$tag'" >&2 + echo "supported tracks are v1.x and v0.13.x" >&2 + exit 64 + ;; +esac + +printf '%s\n' "$tag" diff --git a/scripts/validate-fork-release-tag_test.sh b/scripts/validate-fork-release-tag_test.sh new file mode 100755 index 000000000000..25b237c78f45 --- /dev/null +++ b/scripts/validate-fork-release-tag_test.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +validator="$script_dir/validate-fork-release-tag.sh" + +for tag in v1.0.0-rc.23-0 v1.0.0-rc.23-12 v0.13.2-0 v0.13.2-9; do + "$validator" "$tag" >/dev/null +done + +for tag in v1.0.0-rc.23 v0.13.2 latest nightly v2.0.0-0; do + if "$validator" "$tag" >/dev/null 2>&1; then + echo "validator unexpectedly accepted '$tag'" >&2 + exit 1 + fi +done + +echo "fork release tag validation passed" diff --git a/service/channel.go b/service/channel.go index 3fde6e207b68..a8ed9b7723cd 100644 --- a/service/channel.go +++ b/service/channel.go @@ -15,6 +15,15 @@ func formatNotifyType(channelId int, status int) string { return fmt.Sprintf("%s_%d_%d", dto.NotifyTypeChannelUpdate, channelId, status) } +func shouldCloseActiveWebSocketsAfterDisable(channelId int) bool { + channel, err := model.GetChannelById(channelId, true) + if err != nil { + common.SysLog(fmt.Sprintf("failed to check channel status before closing active websockets: channel_id=%d, error=%v", channelId, err)) + return true + } + return channel.Status != common.ChannelStatusEnabled +} + // disable & notify func DisableChannel(channelError types.ChannelError, reason string) { common.SysLog(fmt.Sprintf("通道「%s」(#%d)发生错误,准备禁用,原因:%s", channelError.ChannelName, channelError.ChannelId, reason)) @@ -27,6 +36,9 @@ func DisableChannel(channelError types.ChannelError, reason string) { success := model.UpdateChannelStatus(channelError.ChannelId, channelError.UsingKey, common.ChannelStatusAutoDisabled, reason) if success { + if shouldCloseActiveWebSocketsAfterDisable(channelError.ChannelId) { + CloseActiveWebSocketsForChannel(channelError.ChannelId, ChannelDisabledCloseReason) + } subject := fmt.Sprintf("通道「%s」(#%d)已被禁用", channelError.ChannelName, channelError.ChannelId) content := fmt.Sprintf("通道「%s」(#%d)已被禁用,原因:%s", channelError.ChannelName, channelError.ChannelId, reason) NotifyRootUser(formatNotifyType(channelError.ChannelId, common.ChannelStatusAutoDisabled), subject, content) diff --git a/service/convert.go b/service/convert.go index 95acf835ee46..30ebddeb9ca7 100644 --- a/service/convert.go +++ b/service/convert.go @@ -258,13 +258,6 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon } var claudeResponses []*dto.ClaudeResponse - // stopOpenBlocks emits the required content_block_stop event(s) for the currently open block(s) - // according to Anthropic's SSE streaming state machine: - // content_block_start -> content_block_delta* -> content_block_stop (per index). - // - // For text/thinking, there is at most one open block at info.ClaudeConvertInfo.Index. - // For tools, OpenAI tool_calls can stream multiple parallel tool_use blocks (indexed from 0), - // so we may have multiple open blocks and must stop each one explicitly. stopOpenBlocks := func() { switch info.ClaudeConvertInfo.LastMessagesType { case relaycommon.LastMessageTypeText, relaycommon.LastMessageTypeThinking: @@ -276,13 +269,9 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon } } } - // stopOpenBlocksAndAdvance closes the currently open block(s) and advances the content block index - // to the next available slot for subsequent content_block_start events. - // - // This prevents invalid streams where a content_block_delta (e.g. thinking_delta) is emitted for an - // index whose active content_block type is different (the typical cause of "Mismatched content block type"). stopOpenBlocksAndAdvance := func() { - if info.ClaudeConvertInfo.LastMessagesType == relaycommon.LastMessageTypeNone { + if info.ClaudeConvertInfo.LastMessagesType == "" || + info.ClaudeConvertInfo.LastMessagesType == relaycommon.LastMessageTypeNone { return } stopOpenBlocks() @@ -296,309 +285,167 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon } info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeNone } - if info.SendResponseCount == 1 { - msg := &dto.ClaudeMediaMessage{ - Id: openAIResponse.Id, - Model: openAIResponse.Model, - Type: "message", - Role: "assistant", - Usage: &dto.ClaudeUsage{ - InputTokens: info.GetEstimatePromptTokens(), - OutputTokens: 0, + finish := func(usage *dto.Usage, defaultStopReason string) { + stopOpenBlocks() + stopReason := stopReasonOpenAI2Claude(info.FinishReason) + if stopReason == "" { + stopReason = defaultStopReason + } + claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ + Type: "message_delta", + Usage: buildClaudeUsageFromOpenAIUsage(usage), + Delta: &dto.ClaudeMediaMessage{ + StopReason: common.GetPointer(stopReason), }, + }) + claudeResponses = append(claudeResponses, &dto.ClaudeResponse{Type: "message_stop"}) + info.ClaudeConvertInfo.Done = true + } + emitThinking := func(reasoning string) { + if reasoning == "" { + return } - msg.SetContent(make([]any, 0)) + if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeThinking { + stopOpenBlocksAndAdvance() + idx := info.ClaudeConvertInfo.Index + claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ + Index: &idx, + Type: "content_block_start", + ContentBlock: &dto.ClaudeMediaMessage{ + Type: "thinking", + Thinking: common.GetPointer(""), + }, + }) + info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeThinking + } + idx := info.ClaudeConvertInfo.Index claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Type: "message_start", - Message: msg, + Index: &idx, + Type: "content_block_delta", + Delta: &dto.ClaudeMediaMessage{ + Type: "thinking_delta", + Thinking: common.GetPointer(reasoning), + }, }) - //claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - // Type: "ping", - //}) - if openAIResponse.IsToolCall() { - info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeTools - info.ClaudeConvertInfo.ToolCallBaseIndex = 0 - info.ClaudeConvertInfo.ToolCallMaxIndexOffset = 0 - var toolCall dto.ToolCallResponse - if len(openAIResponse.Choices) > 0 && len(openAIResponse.Choices[0].Delta.ToolCalls) > 0 { - toolCall = openAIResponse.Choices[0].Delta.ToolCalls[0] - } else { - first := openAIResponse.GetFirstToolCall() - if first != nil { - toolCall = *first - } else { - toolCall = dto.ToolCallResponse{} - } - } - resp := &dto.ClaudeResponse{ - Type: "content_block_start", + } + emitText := func(content string) { + if content == "" { + return + } + if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeText { + stopOpenBlocksAndAdvance() + idx := info.ClaudeConvertInfo.Index + claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ + Index: &idx, + Type: "content_block_start", ContentBlock: &dto.ClaudeMediaMessage{ - Id: toolCall.ID, - Type: "tool_use", - Name: toolCall.Function.Name, - Input: map[string]interface{}{}, + Type: "text", + Text: common.GetPointer(""), }, - } - resp.SetIndex(0) - claudeResponses = append(claudeResponses, resp) - // 首块包含工具 delta,则追加 input_json_delta - if toolCall.Function.Arguments != "" { - idx := 0 - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Index: &idx, - Type: "content_block_delta", - Delta: &dto.ClaudeMediaMessage{ - Type: "input_json_delta", - PartialJson: &toolCall.Function.Arguments, - }, - }) - } - } else { - + }) + info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeText + } + idx := info.ClaudeConvertInfo.Index + claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ + Index: &idx, + Type: "content_block_delta", + Delta: &dto.ClaudeMediaMessage{ + Type: "text_delta", + Text: common.GetPointer(content), + }, + }) + } + emitTools := func(toolCalls []dto.ToolCallResponse) { + if len(toolCalls) == 0 { + return + } + if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeTools { + stopOpenBlocksAndAdvance() + info.ClaudeConvertInfo.ToolCallBaseIndex = info.ClaudeConvertInfo.Index + info.ClaudeConvertInfo.ToolCallMaxIndexOffset = 0 + info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeTools } - // 判断首个响应是否存在内容(非标准的 OpenAI 响应) - if len(openAIResponse.Choices) > 0 { - reasoning := openAIResponse.Choices[0].Delta.GetReasoningContent() - content := openAIResponse.Choices[0].Delta.GetContentString() - if reasoning != "" { - if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeThinking { - stopOpenBlocksAndAdvance() - } - idx := info.ClaudeConvertInfo.Index + base := info.ClaudeConvertInfo.ToolCallBaseIndex + maxOffset := info.ClaudeConvertInfo.ToolCallMaxIndexOffset + for i, toolCall := range toolCalls { + offset := i + if toolCall.Index != nil { + offset = *toolCall.Index + } + if offset > maxOffset { + maxOffset = offset + } + idx := base + offset + if toolCall.Function.Name != "" { claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ Index: &idx, Type: "content_block_start", ContentBlock: &dto.ClaudeMediaMessage{ - Type: "thinking", - Thinking: common.GetPointer[string](""), - }, - }) - idx2 := idx - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Index: &idx2, - Type: "content_block_delta", - Delta: &dto.ClaudeMediaMessage{ - Type: "thinking_delta", - Thinking: &reasoning, + Id: toolCall.ID, + Type: "tool_use", + Name: toolCall.Function.Name, + Input: map[string]interface{}{}, }, }) - info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeThinking - } else if content != "" { - if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeText { - stopOpenBlocksAndAdvance() - } - idx := info.ClaudeConvertInfo.Index + } + if toolCall.Function.Arguments != "" { + arguments := toolCall.Function.Arguments claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ Index: &idx, - Type: "content_block_start", - ContentBlock: &dto.ClaudeMediaMessage{ - Type: "text", - Text: common.GetPointer[string](""), - }, - }) - idx2 := idx - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Index: &idx2, Type: "content_block_delta", Delta: &dto.ClaudeMediaMessage{ - Type: "text_delta", - Text: common.GetPointer[string](content), + Type: "input_json_delta", + PartialJson: &arguments, }, }) - info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeText } } + info.ClaudeConvertInfo.ToolCallMaxIndexOffset = maxOffset + info.ClaudeConvertInfo.Index = base + maxOffset + } - // 如果首块就带 finish_reason,需要立即发送停止块 - if len(openAIResponse.Choices) > 0 && openAIResponse.Choices[0].FinishReason != nil && *openAIResponse.Choices[0].FinishReason != "" { - info.FinishReason = *openAIResponse.Choices[0].FinishReason - stopOpenBlocks() - oaiUsage := openAIResponse.Usage - if oaiUsage == nil { - oaiUsage = info.ClaudeConvertInfo.Usage - } - if oaiUsage != nil { - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Type: "message_delta", - Usage: buildClaudeUsageFromOpenAIUsage(oaiUsage), - Delta: &dto.ClaudeMediaMessage{ - StopReason: common.GetPointer[string](stopReasonOpenAI2Claude(info.FinishReason)), - }, - }) - } - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Type: "message_stop", - }) - info.ClaudeConvertInfo.Done = true + if info.SendResponseCount == 1 { + msg := &dto.ClaudeMediaMessage{ + Id: openAIResponse.Id, + Model: openAIResponse.Model, + Type: "message", + Role: "assistant", + Usage: &dto.ClaudeUsage{ + InputTokens: info.GetEstimatePromptTokens(), + OutputTokens: 0, + }, } - return claudeResponses + msg.SetContent(make([]any, 0)) + claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ + Type: "message_start", + Message: msg, + }) } if len(openAIResponse.Choices) == 0 { - // Some OpenAI-compatible upstreams end with a usage-only SSE chunk. - oaiUsage := openAIResponse.Usage - if oaiUsage == nil { - oaiUsage = info.ClaudeConvertInfo.Usage - } - if oaiUsage != nil { - stopOpenBlocks() - stopReason := stopReasonOpenAI2Claude(info.FinishReason) - if stopReason == "" { - stopReason = "end_turn" + if openAIResponse.Usage != nil || info.FinishReason != "" { + usage := openAIResponse.Usage + if usage == nil { + usage = info.ClaudeConvertInfo.Usage } - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Type: "message_delta", - Usage: buildClaudeUsageFromOpenAIUsage(oaiUsage), - Delta: &dto.ClaudeMediaMessage{ - StopReason: common.GetPointer[string](stopReason), - }, - }) - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Type: "message_stop", - }) - info.ClaudeConvertInfo.Done = true + finish(usage, "end_turn") } return claudeResponses - } else { - chosenChoice := openAIResponse.Choices[0] - doneChunk := chosenChoice.FinishReason != nil && *chosenChoice.FinishReason != "" - if doneChunk { - info.FinishReason = *chosenChoice.FinishReason - oaiUsage := openAIResponse.Usage - if oaiUsage == nil { - oaiUsage = info.ClaudeConvertInfo.Usage - // Some upstreams emit finish_reason first, then send a final usage-only chunk. - // Defer closing until usage is available so the final message_delta carries it. - return claudeResponses - } - } - - var claudeResponse dto.ClaudeResponse - var isEmpty bool - claudeResponse.Type = "content_block_delta" - if len(chosenChoice.Delta.ToolCalls) > 0 { - toolCalls := chosenChoice.Delta.ToolCalls - if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeTools { - stopOpenBlocksAndAdvance() - info.ClaudeConvertInfo.ToolCallBaseIndex = info.ClaudeConvertInfo.Index - info.ClaudeConvertInfo.ToolCallMaxIndexOffset = 0 - } - info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeTools - base := info.ClaudeConvertInfo.ToolCallBaseIndex - maxOffset := info.ClaudeConvertInfo.ToolCallMaxIndexOffset - - for i, toolCall := range toolCalls { - offset := 0 - if toolCall.Index != nil { - offset = *toolCall.Index - } else { - offset = i - } - if offset > maxOffset { - maxOffset = offset - } - blockIndex := base + offset - - idx := blockIndex - if toolCall.Function.Name != "" { - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Index: &idx, - Type: "content_block_start", - ContentBlock: &dto.ClaudeMediaMessage{ - Id: toolCall.ID, - Type: "tool_use", - Name: toolCall.Function.Name, - Input: map[string]interface{}{}, - }, - }) - } - - if len(toolCall.Function.Arguments) > 0 { - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Index: &idx, - Type: "content_block_delta", - Delta: &dto.ClaudeMediaMessage{ - Type: "input_json_delta", - PartialJson: &toolCall.Function.Arguments, - }, - }) - } - } - info.ClaudeConvertInfo.ToolCallMaxIndexOffset = maxOffset - info.ClaudeConvertInfo.Index = base + maxOffset - } else { - reasoning := chosenChoice.Delta.GetReasoningContent() - textContent := chosenChoice.Delta.GetContentString() - if reasoning != "" || textContent != "" { - if reasoning != "" { - if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeThinking { - stopOpenBlocksAndAdvance() - idx := info.ClaudeConvertInfo.Index - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Index: &idx, - Type: "content_block_start", - ContentBlock: &dto.ClaudeMediaMessage{ - Type: "thinking", - Thinking: common.GetPointer[string](""), - }, - }) - } - info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeThinking - claudeResponse.Delta = &dto.ClaudeMediaMessage{ - Type: "thinking_delta", - Thinking: &reasoning, - } - } else { - if info.ClaudeConvertInfo.LastMessagesType != relaycommon.LastMessageTypeText { - stopOpenBlocksAndAdvance() - idx := info.ClaudeConvertInfo.Index - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Index: &idx, - Type: "content_block_start", - ContentBlock: &dto.ClaudeMediaMessage{ - Type: "text", - Text: common.GetPointer[string](""), - }, - }) - } - info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeText - claudeResponse.Delta = &dto.ClaudeMediaMessage{ - Type: "text_delta", - Text: common.GetPointer[string](textContent), - } - } - } else { - isEmpty = true - } - } + } - claudeResponse.Index = common.GetPointer[int](info.ClaudeConvertInfo.Index) - if !isEmpty && claudeResponse.Delta != nil { - claudeResponses = append(claudeResponses, &claudeResponse) - } + chosenChoice := openAIResponse.Choices[0] + emitThinking(chosenChoice.Delta.GetReasoningContent()) + emitText(chosenChoice.Delta.GetContentString()) + emitTools(chosenChoice.Delta.ToolCalls) - if doneChunk || info.ClaudeConvertInfo.Done { - stopOpenBlocks() - oaiUsage := openAIResponse.Usage - if oaiUsage == nil { - oaiUsage = info.ClaudeConvertInfo.Usage - } - if oaiUsage != nil { - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Type: "message_delta", - Usage: buildClaudeUsageFromOpenAIUsage(oaiUsage), - Delta: &dto.ClaudeMediaMessage{ - StopReason: common.GetPointer[string](stopReasonOpenAI2Claude(info.FinishReason)), - }, - }) - } - claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ - Type: "message_stop", - }) - info.ClaudeConvertInfo.Done = true - return claudeResponses + if chosenChoice.FinishReason != nil && *chosenChoice.FinishReason != "" { + info.FinishReason = *chosenChoice.FinishReason + usage := openAIResponse.Usage + if usage == nil { + usage = info.ClaudeConvertInfo.Usage } + finish(usage, "end_turn") } return claudeResponses @@ -615,6 +462,16 @@ func ResponseOpenAI2Claude(openAIResponse *dto.OpenAITextResponse, info *relayco } for _, choice := range openAIResponse.Choices { stopReason = stopReasonOpenAI2Claude(choice.FinishReason) + if reasoning := choice.Message.GetReasoningContent(); reasoning != "" { + claudeContent := dto.ClaudeMediaMessage{Type: "thinking"} + claudeContent.Thinking = &reasoning + contents = append(contents, claudeContent) + } + if text := choice.Message.StringContent(); text != "" { + claudeContent := dto.ClaudeMediaMessage{Type: "text"} + claudeContent.SetText(text) + contents = append(contents, claudeContent) + } if choice.FinishReason == "tool_calls" { for _, toolUse := range choice.Message.ParseToolCalls() { claudeContent := dto.ClaudeMediaMessage{} @@ -629,11 +486,6 @@ func ResponseOpenAI2Claude(openAIResponse *dto.OpenAITextResponse, info *relayco } contents = append(contents, claudeContent) } - } else { - claudeContent := dto.ClaudeMediaMessage{} - claudeContent.Type = "text" - claudeContent.SetText(choice.Message.StringContent()) - contents = append(contents, claudeContent) } } claudeResponse.Content = contents diff --git a/service/convert_thinking_test.go b/service/convert_thinking_test.go new file mode 100644 index 000000000000..73331045cdea --- /dev/null +++ b/service/convert_thinking_test.go @@ -0,0 +1,178 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newClaudeConversionInfo(sendResponseCount int) *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + SendResponseCount: sendResponseCount, + ClaudeConvertInfo: &relaycommon.ClaudeConvertInfo{ + LastMessagesType: relaycommon.LastMessageTypeNone, + }, + } +} + +func claudeResponseTypes(responses []*dto.ClaudeResponse) []string { + types := make([]string, 0, len(responses)) + for _, response := range responses { + types = append(types, response.Type) + } + return types +} + +func TestStreamResponseOpenAI2ClaudeOrdersThinkingTextAndToolUse(t *testing.T) { + finishReason := "tool_calls" + reasoning := "I need current weather." + content := "Checking Tokyo now." + info := newClaudeConversionInfo(1) + + responses := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl-test", + Model: "gemini-3.7-flash-high", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + FinishReason: &finishReason, + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ + ReasoningContent: &reasoning, + Content: &content, + ToolCalls: []dto.ToolCallResponse{{ + ID: "tool-1", + Type: "function", + Function: dto.FunctionResponse{ + Name: "lookup_weather", + Arguments: `{"city":"Tokyo"}`, + }, + }}, + }, + }}, + }, info) + + require.Equal(t, []string{ + "message_start", + "content_block_start", "content_block_delta", "content_block_stop", + "content_block_start", "content_block_delta", "content_block_stop", + "content_block_start", "content_block_delta", "content_block_stop", + "message_delta", "message_stop", + }, claudeResponseTypes(responses)) + + assert.Equal(t, "thinking", responses[1].ContentBlock.Type) + assert.Equal(t, "thinking_delta", responses[2].Delta.Type) + assert.Equal(t, reasoning, *responses[2].Delta.Thinking) + assert.Equal(t, 0, *responses[1].Index) + assert.Equal(t, "text", responses[4].ContentBlock.Type) + assert.Equal(t, content, *responses[5].Delta.Text) + assert.Equal(t, 1, *responses[4].Index) + assert.Equal(t, "tool_use", responses[7].ContentBlock.Type) + assert.Equal(t, "lookup_weather", responses[7].ContentBlock.Name) + assert.Equal(t, 2, *responses[7].Index) + assert.Equal(t, "tool_use", *responses[10].Delta.StopReason) + assert.Nil(t, responses[10].Usage) + assert.True(t, info.ClaudeConvertInfo.Done) +} + +func TestStreamResponseOpenAI2ClaudeClosesWithoutUsageAfterFirstChunk(t *testing.T) { + finishReason := "stop" + reasoning := "brief thought" + info := newClaudeConversionInfo(1) + + responses := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + FinishReason: &finishReason, + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ReasoningContent: &reasoning}, + }}, + }, info) + + require.Equal(t, []string{ + "message_start", "content_block_start", "content_block_delta", "content_block_stop", "message_delta", "message_stop", + }, claudeResponseTypes(responses)) + assert.Equal(t, "end_turn", *responses[4].Delta.StopReason) + assert.Nil(t, responses[4].Usage) + assert.True(t, info.ClaudeConvertInfo.Done) +} + +func TestStreamResponseOpenAI2ClaudeClosesUsageOnlyTerminalChunk(t *testing.T) { + info := newClaudeConversionInfo(2) + info.ClaudeConvertInfo.LastMessagesType = relaycommon.LastMessageTypeText + info.ClaudeConvertInfo.Index = 3 + info.FinishReason = "stop" + + responses := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Usage: &dto.Usage{PromptTokens: 5, CompletionTokens: 7}, + }, info) + + require.Equal(t, []string{"content_block_stop", "message_delta", "message_stop"}, claudeResponseTypes(responses)) + assert.Equal(t, 3, *responses[0].Index) + require.NotNil(t, responses[1].Usage) + assert.Equal(t, 5, responses[1].Usage.InputTokens) + assert.Equal(t, 7, responses[1].Usage.OutputTokens) + assert.Equal(t, "end_turn", *responses[1].Delta.StopReason) +} + +func TestResponseOpenAI2ClaudeIncludesThinkingBeforeText(t *testing.T) { + reasoning := "reason through it" + response := ResponseOpenAI2Claude(&dto.OpenAITextResponse{ + Id: "chatcmpl-test", + Model: "gemini-3.7-flash-high", + Choices: []dto.OpenAITextResponseChoice{{ + FinishReason: "stop", + Message: dto.Message{ + Role: "assistant", + Content: "final answer", + ReasoningContent: &reasoning, + }, + }}, + }, &relaycommon.RelayInfo{}) + + require.Len(t, response.Content, 2) + assert.Equal(t, "thinking", response.Content[0].Type) + assert.Equal(t, reasoning, *response.Content[0].Thinking) + assert.Equal(t, "text", response.Content[1].Type) + assert.Equal(t, "final answer", *response.Content[1].Text) +} + +func TestResponseOpenAI2ClaudeKeepsThinkingTextAndToolUse(t *testing.T) { + reasoning := "I need to use the weather tool." + message := dto.Message{Role: "assistant", ReasoningContent: &reasoning} + message.SetStringContent("Checking Tokyo now.") + message.SetToolCalls([]dto.ToolCallResponse{{ + ID: "tool-1", + Type: "function", + Function: dto.FunctionResponse{ + Name: "lookup_weather", + Arguments: `{"city":"Tokyo"}`, + }, + }}) + + response := ResponseOpenAI2Claude(&dto.OpenAITextResponse{ + Choices: []dto.OpenAITextResponseChoice{{ + FinishReason: "tool_calls", + Message: message, + }}, + }, &relaycommon.RelayInfo{}) + + require.Len(t, response.Content, 3) + assert.Equal(t, "thinking", response.Content[0].Type) + assert.Equal(t, reasoning, *response.Content[0].Thinking) + assert.Equal(t, "text", response.Content[1].Type) + assert.Equal(t, "Checking Tokyo now.", *response.Content[1].Text) + assert.Equal(t, "tool_use", response.Content[2].Type) + assert.Equal(t, "lookup_weather", response.Content[2].Name) +} + +func TestStreamResponseOpenAI2ClaudeDoesNotDuplicateTerminalEvents(t *testing.T) { + finishReason := "stop" + info := newClaudeConversionInfo(2) + + responses := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{FinishReason: &finishReason}}, + }, info) + require.Equal(t, []string{"message_delta", "message_stop"}, claudeResponseTypes(responses)) + + assert.Empty(t, StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{}, info)) + assert.True(t, info.ClaudeConvertInfo.Done) +} diff --git a/service/log_info_generate.go b/service/log_info_generate.go index 54448d59d673..56c0196cd0fb 100644 --- a/service/log_info_generate.go +++ b/service/log_info_generate.go @@ -44,6 +44,9 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m other["model_price"] = modelPrice other["user_group_ratio"] = userGroupRatio other["frt"] = float64(relayInfo.FirstResponseTime.UnixMilli() - relayInfo.StartTime.UnixMilli()) + if relayInfo.ClientWs != nil { + other["ws"] = true + } if relayInfo.ReasoningEffort != "" { other["reasoning_effort"] = relayInfo.ReasoningEffort } diff --git a/service/log_info_generate_test.go b/service/log_info_generate_test.go new file mode 100644 index 000000000000..11d06f18b254 --- /dev/null +++ b/service/log_info_generate_test.go @@ -0,0 +1,45 @@ +package service + +import ( + "net/http/httptest" + "testing" + "time" + + relaycommon "github.com/QuantumNous/new-api/relay/common" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +func TestGenerateTextOtherInfoMarksWebSocketTransport(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + now := time.Now() + relayInfo := &relaycommon.RelayInfo{ + ClientWs: &websocket.Conn{}, + StartTime: now, + FirstResponseTime: now, + ChannelMeta: &relaycommon.ChannelMeta{}, + } + + other := GenerateTextOtherInfo(ctx, relayInfo, 1, 1, 1, 0, 0, 0, 1) + + require.Equal(t, true, other["ws"]) +} + +func TestGenerateTextOtherInfoOmitsWebSocketTransportForHTTP(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + now := time.Now() + relayInfo := &relaycommon.RelayInfo{ + StartTime: now, + FirstResponseTime: now, + ChannelMeta: &relaycommon.ChannelMeta{}, + } + + other := GenerateTextOtherInfo(ctx, relayInfo, 1, 1, 1, 0, 0, 0, 1) + + _, ok := other["ws"] + require.False(t, ok) +} diff --git a/service/relay_error.go b/service/relay_error.go new file mode 100644 index 000000000000..74b7a3956547 --- /dev/null +++ b/service/relay_error.go @@ -0,0 +1,96 @@ +package service + +import ( + "fmt" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/bytedance/gopkg/util/gopool" + "github.com/gin-gonic/gin" +) + +func ShouldRetryRelayError(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { + if openaiErr == nil { + return false + } + if ShouldSkipRetryAfterChannelAffinityFailure(c) { + return false + } + if c != nil { + if _, ok := c.Get("specific_channel_id"); ok { + return false + } + } + if types.IsChannelError(openaiErr) { + return true + } + if types.IsSkipRetryError(openaiErr) { + return false + } + if retryTimes <= 0 { + return false + } + code := openaiErr.StatusCode + if code >= 200 && code < 300 { + return false + } + if code < 100 || code > 599 { + return true + } + if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) { + return false + } + return operation_setting.ShouldRetryByStatusCode(code) +} + +func ProcessChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) { + if err == nil { + return + } + logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.MaskSensitiveErrorWithStatusCode()))) + if ShouldDisableChannel(err) && channelError.AutoBan { + gopool.Go(func() { + DisableChannel(channelError, err.ErrorWithStatusCode()) + }) + } + + if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) { + userId := c.GetInt("id") + tokenName := c.GetString("token_name") + modelName := c.GetString("original_model") + tokenId := c.GetInt("token_id") + userGroup := c.GetString("group") + channelId := c.GetInt("channel_id") + other := make(map[string]interface{}) + if c.Request != nil && c.Request.URL != nil { + other["request_path"] = c.Request.URL.Path + } + other["error_type"] = err.GetErrorType() + other["error_code"] = err.GetErrorCode() + other["status_code"] = err.StatusCode + other["channel_id"] = channelId + other["channel_name"] = c.GetString("channel_name") + other["channel_type"] = c.GetInt("channel_type") + adminInfo := make(map[string]interface{}) + adminInfo["use_channel"] = c.GetStringSlice("use_channel") + isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey) + if isMultiKey { + adminInfo["is_multi_key"] = true + adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex) + } + AppendChannelAffinityAdminInfo(c, adminInfo) + other["admin_info"] = adminInfo + startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime) + if startTime.IsZero() { + startTime = time.Now() + } + useTimeSeconds := int(time.Since(startTime).Seconds()) + model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other) + } +} diff --git a/service/relay_error_test.go b/service/relay_error_test.go new file mode 100644 index 000000000000..4ed854fa2927 --- /dev/null +++ b/service/relay_error_test.go @@ -0,0 +1,21 @@ +package service + +import ( + "errors" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" +) + +func TestShouldRetryRelayErrorSpecificChannelSkipsChannelError(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set("specific_channel_id", "1") + err := types.NewError(errors.New("channel failed"), types.ErrorCodeChannelNoAvailableKey) + + if ShouldRetryRelayError(c, err, 1) { + t.Fatal("specific channel channel error should not retry") + } +} diff --git a/service/responses_usage.go b/service/responses_usage.go new file mode 100644 index 000000000000..c241af466813 --- /dev/null +++ b/service/responses_usage.go @@ -0,0 +1,43 @@ +package service + +import "github.com/QuantumNous/new-api/dto" + +func ApplyResponsesUsage(dst *dto.Usage, src *dto.Usage) { + if dst == nil || src == nil { + return + } + if src.InputTokens != 0 { + dst.PromptTokens = src.InputTokens + dst.InputTokens = src.InputTokens + } + if src.OutputTokens != 0 { + dst.CompletionTokens = src.OutputTokens + dst.OutputTokens = src.OutputTokens + } + if src.TotalTokens != 0 { + dst.TotalTokens = src.TotalTokens + } + if src.InputTokensDetails != nil { + inputDetails := *src.InputTokensDetails + dst.InputTokensDetails = &inputDetails + dst.PromptTokensDetails = inputDetails + } + outputDetails := src.CompletionTokenDetails + if src.OutputTokensDetails != nil { + outputDetails = *src.OutputTokensDetails + } + if !isZeroOutputTokenDetails(outputDetails) { + dst.CompletionTokenDetails = outputDetails + dst.OutputTokensDetails = &outputDetails + } + dst.PromptCacheHitTokens = src.PromptCacheHitTokens + dst.UsageSemantic = src.UsageSemantic + dst.UsageSource = src.UsageSource +} + +func isZeroOutputTokenDetails(details dto.OutputTokenDetails) bool { + return details.TextTokens == 0 && + details.AudioTokens == 0 && + details.ImageTokens == 0 && + details.ReasoningTokens == 0 +} diff --git a/service/responses_usage_test.go b/service/responses_usage_test.go new file mode 100644 index 000000000000..9f7f7a5e961f --- /dev/null +++ b/service/responses_usage_test.go @@ -0,0 +1,87 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/dto" +) + +func TestApplyResponsesUsageCopiesTokenDetails(t *testing.T) { + dst := &dto.Usage{} + src := &dto.Usage{ + InputTokens: 11, + OutputTokens: 7, + TotalTokens: 18, + InputTokensDetails: &dto.InputTokenDetails{ + CachedTokens: 3, + CachedCreationTokens: 2, + TextTokens: 6, + AudioTokens: 4, + ImageTokens: 5, + }, + OutputTokensDetails: &dto.OutputTokenDetails{ + TextTokens: 1, + AudioTokens: 2, + ImageTokens: 3, + ReasoningTokens: 4, + }, + PromptCacheHitTokens: 3, + UsageSemantic: "openai", + UsageSource: "upstream", + } + + ApplyResponsesUsage(dst, src) + + if dst.PromptTokens != 11 || dst.CompletionTokens != 7 || dst.TotalTokens != 18 { + t.Fatalf("usage tokens = %#v", dst) + } + if dst.InputTokensDetails == nil { + t.Fatal("InputTokensDetails is nil") + } + if dst.PromptTokensDetails.CachedTokens != 3 || + dst.PromptTokensDetails.CachedCreationTokens != 2 || + dst.PromptTokensDetails.TextTokens != 6 || + dst.PromptTokensDetails.AudioTokens != 4 || + dst.PromptTokensDetails.ImageTokens != 5 { + t.Fatalf("prompt details = %#v", dst.PromptTokensDetails) + } + if dst.CompletionTokenDetails.TextTokens != 1 || + dst.CompletionTokenDetails.AudioTokens != 2 || + dst.CompletionTokenDetails.ImageTokens != 3 || + dst.CompletionTokenDetails.ReasoningTokens != 4 { + t.Fatalf("completion details = %#v", dst.CompletionTokenDetails) + } + if dst.OutputTokensDetails == nil { + t.Fatal("OutputTokensDetails is nil") + } + if dst.OutputTokensDetails.TextTokens != 1 || + dst.OutputTokensDetails.AudioTokens != 2 || + dst.OutputTokensDetails.ImageTokens != 3 || + dst.OutputTokensDetails.ReasoningTokens != 4 { + t.Fatalf("output details = %#v", dst.OutputTokensDetails) + } + if dst.UsageSemantic != "openai" || dst.UsageSource != "upstream" { + t.Fatalf("usage metadata = %#v", dst) + } +} + +func TestApplyResponsesUsageFallsBackToCompletionTokenDetails(t *testing.T) { + dst := &dto.Usage{} + src := &dto.Usage{ + CompletionTokenDetails: dto.OutputTokenDetails{ + ReasoningTokens: 9, + }, + } + + ApplyResponsesUsage(dst, src) + + if dst.CompletionTokenDetails.ReasoningTokens != 9 { + t.Fatalf("reasoning tokens = %d, want 9", dst.CompletionTokenDetails.ReasoningTokens) + } + if dst.OutputTokensDetails == nil { + t.Fatal("OutputTokensDetails is nil") + } + if dst.OutputTokensDetails.ReasoningTokens != 9 { + t.Fatalf("output reasoning tokens = %d, want 9", dst.OutputTokensDetails.ReasoningTokens) + } +} diff --git a/service/ws_close.go b/service/ws_close.go new file mode 100644 index 000000000000..5298e3224053 --- /dev/null +++ b/service/ws_close.go @@ -0,0 +1,13 @@ +package service + +import "github.com/QuantumNous/new-api/pkg/wsmanager" + +const ChannelDisabledCloseReason = wsmanager.DefaultCloseReason + +func CloseActiveWebSocketsForChannel(channelID int, reason string) int { + return wsmanager.CloseChannelsAndBroadcast([]int{channelID}, reason) +} + +func CloseActiveWebSocketsForChannels(channelIDs []int, reason string) int { + return wsmanager.CloseChannelsAndBroadcast(channelIDs, reason) +} diff --git a/setting/model_setting/global.go b/setting/model_setting/global.go index d0c4d312893c..1f8cbfd49b61 100644 --- a/setting/model_setting/global.go +++ b/setting/model_setting/global.go @@ -69,6 +69,12 @@ func ShouldPreserveThinkingSuffix(modelName string) bool { if target == "" { return false } + // gemini-3.7-flash-high is a canonical native Gemini model name; its + // trailing -high is not the legacy effort suffix and must reach upstream + // unchanged even when the global adapter is enabled. + if target == "gemini-3.7-flash-high" { + return true + } for _, entry := range globalSettings.ThinkingModelBlacklist { if strings.TrimSpace(entry) == target { diff --git a/web/src/components/table/usage-logs/UsageLogsColumnDefs.jsx b/web/src/components/table/usage-logs/UsageLogsColumnDefs.jsx index 07e6fbb9896d..5a09c01a1c18 100644 --- a/web/src/components/table/usage-logs/UsageLogsColumnDefs.jsx +++ b/web/src/components/table/usage-logs/UsageLogsColumnDefs.jsx @@ -144,10 +144,7 @@ function renderType(type, t) { function buildStreamStatusTooltip(ss, t) { if (!ss) return null; - const lines = [ - t('流状态') + ':' + t('异常'), - (ss.end_reason || 'unknown'), - ]; + const lines = [t('流状态') + ':' + t('异常'), ss.end_reason || 'unknown']; if (ss.error_count > 0) { lines.push(`${t('软错误')}: ${ss.error_count}`); } @@ -163,14 +160,14 @@ function buildStreamStatusTooltip(ss, t) { ); } -function renderIsStream(bool, t, streamStatus) { +function renderIsStream(bool, t, streamStatus, isWebSocket = false) { const isError = streamStatus && streamStatus.status !== 'ok'; if (bool) { return ( - {t('流')} + {isWebSocket ? t('WebSocket') : t('流')} {isError && ( @@ -185,11 +182,7 @@ function renderIsStream(bool, t, streamStatus) { userSelect: 'none', }} > - + )} @@ -461,7 +454,11 @@ function getUsageLogDetailSummary(record, text, billingDisplayMode, t) { }; } - const summaryOpts = { ...other, displayMode: billingDisplayMode, outputMode: 'segments' }; + const summaryOpts = { + ...other, + displayMode: billingDisplayMode, + outputMode: 'segments', + }; if (other?.billing_mode === 'tiered_expr') { return { segments: renderTieredModelPriceSimple(summaryOpts) }; @@ -709,7 +706,12 @@ export const getLogsColumns = ({ {renderUseTime(text, t)} {renderFirstUseTime(other?.frt, t)} - {renderIsStream(record.is_stream, t, other?.stream_status)} + {renderIsStream( + record.is_stream, + t, + other?.stream_status, + other?.ws === true, + )} ); diff --git a/web/src/i18n/locales/en.json b/web/src/i18n/locales/en.json index dc8ad6cb9464..5a8e01694640 100644 --- a/web/src/i18n/locales/en.json +++ b/web/src/i18n/locales/en.json @@ -3787,6 +3787,7 @@ "见上方动态计费详情": "See dynamic pricing details above", "含时间条件": "Time rules", "含请求条件": "Request rules", - "(当前仅支持易支付接口,默认使用上方服务器地址作为回调地址!)": "(Currently only supports Epay interface, the default callback address is the server address above!)" + "(当前仅支持易支付接口,默认使用上方服务器地址作为回调地址!)": "(Currently only supports Epay interface, the default callback address is the server address above!)", + "WebSocket": "WebSocket" } } diff --git a/web/src/i18n/locales/fr.json b/web/src/i18n/locales/fr.json index 8e7d143d0954..3856c7c2872b 100644 --- a/web/src/i18n/locales/fr.json +++ b/web/src/i18n/locales/fr.json @@ -3642,6 +3642,7 @@ "默认折叠侧边栏": "Réduire la barre latérale par défaut", "默认测试模型": "Modèle de test par défaut", "默认用户消息": "Bonjour", - "默认补全倍率": "Taux de complétion par défaut" + "默认补全倍率": "Taux de complétion par défaut", + "WebSocket": "WebSocket" } } diff --git a/web/src/i18n/locales/ja.json b/web/src/i18n/locales/ja.json index 88d2899e17c6..bc4aae296cbd 100644 --- a/web/src/i18n/locales/ja.json +++ b/web/src/i18n/locales/ja.json @@ -3611,6 +3611,7 @@ "默认折叠侧边栏": "サイドバーをデフォルトで折りたたむ", "默认测试模型": "デフォルトテストモデル", "默认用户消息": "こんにちは", - "默认补全倍率": "デフォルト補完倍率" + "默认补全倍率": "デフォルト補完倍率", + "WebSocket": "WebSocket" } } diff --git a/web/src/i18n/locales/ru.json b/web/src/i18n/locales/ru.json index 2980af179f6f..3a04d2f88a8b 100644 --- a/web/src/i18n/locales/ru.json +++ b/web/src/i18n/locales/ru.json @@ -3662,6 +3662,7 @@ "默认折叠侧边栏": "Сворачивать боковую панель по умолчанию", "默认测试模型": "Модель для тестирования по умолчанию", "默认用户消息": "Здравствуйте", - "默认补全倍率": "Коэффициент завершения по умолчанию" + "默认补全倍率": "Коэффициент завершения по умолчанию", + "WebSocket": "WebSocket" } } diff --git a/web/src/i18n/locales/vi.json b/web/src/i18n/locales/vi.json index 4ca1a77f3122..b8f1460c26cd 100644 --- a/web/src/i18n/locales/vi.json +++ b/web/src/i18n/locales/vi.json @@ -4176,6 +4176,7 @@ "默认折叠侧边栏": "Mặc định thu gọn thanh bên", "默认测试模型": "Mô hình kiểm tra mặc định", "默认用户消息": "Xin chào", - "默认补全倍率": "Tỷ lệ hoàn thành mặc định" + "默认补全倍率": "Tỷ lệ hoàn thành mặc định", + "WebSocket": "WebSocket" } } diff --git a/web/src/i18n/locales/zh-CN.json b/web/src/i18n/locales/zh-CN.json index e54a1c0f9114..b048f85a3751 100644 --- a/web/src/i18n/locales/zh-CN.json +++ b/web/src/i18n/locales/zh-CN.json @@ -3771,6 +3771,7 @@ "缓存创建-1h": "缓存创建-1h", "见上方动态计费详情": "见上方动态计费详情", "含时间条件": "含时间条件", - "含请求条件": "含请求条件" + "含请求条件": "含请求条件", + "WebSocket": "WebSocket" } } diff --git a/web/src/i18n/locales/zh-TW.json b/web/src/i18n/locales/zh-TW.json index b31c9e1e0eac..bfc340a101ed 100644 --- a/web/src/i18n/locales/zh-TW.json +++ b/web/src/i18n/locales/zh-TW.json @@ -3635,6 +3635,7 @@ "默认折叠侧边栏": "預設摺疊側邊欄", "默认测试模型": "預設測試模型", "默认用户消息": "你好", - "默认补全倍率": "預設補全倍率" + "默认补全倍率": "預設補全倍率", + "WebSocket": "WebSocket" } }