feat: sync deployed v1.0.0-rc.38-custom.42

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.

controller/relay_attempt_test.go

@@ -0,0 +1,66 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	hostdto "github.com/QuantumNous/new-api/dto"
+	"github.com/QuantumNous/new-api/model"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/relaykit/types"
+	"github.com/QuantumNous/new-api/service"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestRelayAttemptSwitchSafety(t *testing.T) {
+	for _, tc := range []struct {
+		name, reason, previous, conversation string
+		count                                int
+		pin, cancel, written, disabled, want bool
+	}{
+		{name: "explicit mismatch", reason: "model_mismatch", count: 1, want: true},
+		{name: "token disabled", reason: "model_mismatch", count: 1, disabled: true},
+		{name: "third channel allowed", reason: "model_mismatch", count: 2, want: true},
+		{name: "third channel final", reason: "model_mismatch", count: 3},
+		{name: "missing model", reason: "missing_model", count: 1},
+		{name: "confirmation timeout", reason: "confirmation_timeout", count: 1},
+		{name: "pinned", reason: "model_mismatch", count: 1, pin: true},
+		{name: "cancelled", reason: "model_mismatch", count: 1, cancel: true},
+		{name: "already delivered", reason: "model_mismatch", count: 1, written: true},
+		{name: "upstream response id", reason: "model_mismatch", count: 1, previous: "resp_private"},
+		{name: "upstream conversation", reason: "model_mismatch", count: 1, conversation: `"conv_private"`},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			c, _ := gin.CreateTestContext(httptest.NewRecorder())
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			c.Request = httptest.NewRequest("POST", "/v1/responses", nil).WithContext(ctx)
+			c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "a"}}})
+			if tc.pin {
+				service.GetChannelConstraints(c).AddPin(hostdto.ChannelPin{ChannelId: 1, Source: hostdto.PinSourceToken, Rank: hostdto.PinRankToken, RetryMode: hostdto.PinRetrySingleAttempt})
+			}
+			if tc.cancel {
+				cancel()
+			}
+			if tc.written {
+				_, _ = c.Writer.WriteString("already sent")
+			}
+			info := &relaycommon.RelayInfo{DurableAttemptBilling: true, UpstreamRequests: tc.count, Request: &dto.OpenAIResponsesRequest{PreviousResponseID: tc.previous, Conversation: []byte(tc.conversation)}, ResponseGuard: &relaycommon.ResponseModelGuard{Attempts: []*relaycommon.ResponseGuardAttempt{{Reason: tc.reason}}}}
+			err := types.NewError(errors.New("guard rejected"), relaycommon.ResponseGuardErrorCode)
+			info.AstraAutoSwitch = !tc.disabled
+			require.Equal(t, tc.want, canSwitchRejectedAstra(c, info, err))
+		})
+	}
+}
+
+func TestRelayAttemptGuardResult(t *testing.T) {
+	require.Equal(t, "mismatch_rejected", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "model_mismatch", Number: 1}))
+	require.Equal(t, "switched_success", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 2}))
+	require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 1}))
+	require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra", ReturnedModel: "gpt-6-astra"}))
+	require.Equal(t, "unchecked", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra"}))
+	require.Equal(t, "transport_error", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "transport_error", Number: 1}))
+}