Codex
aa6c3f50f4
Fixed-group tokens with a backup group now take precedence over the session affinity skip-retry-on-failure rule, so retryable upstream errors (503) fall through to the backup group instead of stopping on the primary channel. Non-retryable errors, explicit channel pins, and auto groups are unchanged. Adds controller regression tests for the affinity override and a build-time Go test gate.
59 lines
2.3 KiB
Go
59 lines
2.3 KiB
Go
package controller
|
|
|
|
import (
|
|
"errors"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestFixedBackupRetryOverridesAffinityOnly(t *testing.T) {
|
|
configureTokenAutoGroupsTest(t, "2", "[]")
|
|
for _, tc := range []struct {
|
|
name string
|
|
group string
|
|
backups []string
|
|
cross bool
|
|
status int
|
|
skip bool
|
|
pinned bool
|
|
want bool
|
|
}{
|
|
{"503 with backup", "default", []string{"vip"}, true, 503, false, false, true},
|
|
{"no backup", "default", nil, true, 503, false, false, false},
|
|
{"same group", "default", []string{"default"}, true, 503, false, false, false},
|
|
{"revoked backup", "default", []string{"unavailable"}, true, 503, false, false, false},
|
|
{"retry disabled", "default", []string{"vip"}, false, 503, false, false, false},
|
|
{"auto unchanged", "auto", []string{"vip"}, true, 503, false, false, false},
|
|
{"invalid request", "default", []string{"vip"}, true, 400, true, false, false},
|
|
{"skip retry", "default", []string{"vip"}, true, 503, true, false, false},
|
|
{"explicit channel pin", "default", []string{"vip"}, true, 503, false, true, false},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
c := newPinRetryContext()
|
|
c.Set("channel_affinity_skip_retry_on_failure", true)
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
|
|
common.SetContextKey(c, constant.ContextKeyTokenGroup, tc.group)
|
|
common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, tc.backups)
|
|
common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, tc.cross)
|
|
if tc.pinned {
|
|
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
|
|
ChannelId: 13, Source: dto.PinSourceToken,
|
|
Rank: dto.PinRankToken, RetryMode: dto.PinRetrySingleAttempt,
|
|
})
|
|
}
|
|
err := types.NewOpenAIError(errors.New("upstream unavailable"), types.ErrorCodeBadResponseStatusCode, tc.status)
|
|
if tc.skip {
|
|
err = types.NewOpenAIError(errors.New("do not replay"), types.ErrorCodeBadResponseStatusCode, tc.status, types.ErrOptionWithSkipRetry())
|
|
}
|
|
require.Equal(t, tc.want, shouldRetry(c, err, 1))
|
|
require.False(t, shouldRetry(c, err, 0), "exhausted retry budget stays bounded")
|
|
})
|
|
}
|
|
}
|