fix: preserve channel permission for cyber policy 403 errors (#615)

Author: zijiren233

core/relay/adaptor/interface.go

@@ -121,6 +121,10 @@ type Error interface {
 	StatusCode() int
 }
 
+type ErrorCodeProvider interface {
+	ErrorCode() any
+}
+
 var ErrGetBalanceNotImplemented = errors.New("get balance not implemented")
 
 type Balancer interface {

core/relay/adaptor/utils.go

@@ -22,6 +22,15 @@ func (e BasicError[T]) StatusCode() int {
 	return e.statusCode
 }
 
+func (e BasicError[T]) ErrorCode() any {
+	provider, ok := any(e.error).(ErrorCodeProvider)
+	if !ok {
+		return nil
+	}
+
+	return provider.ErrorCode()
+}
+
 func (e BasicError[T]) Error() string {
 	return fmt.Sprintf("status code: %d, error: %v", e.statusCode, e.error)
 }

core/relay/model/chat.go

@@ -192,6 +192,10 @@ type OpenAIErrorResponse struct {
 	Error OpenAIError `json:"error"`
 }
 
+func (e OpenAIErrorResponse) ErrorCode() any {
+	return e.Error.Code
+}
+
 type OpenAIError struct {
 	Code    any    `json:"code,omitempty"`
 	Message string `json:"message,omitempty"`

core/relay/plugin/monitor/monitor.go

@@ -50,13 +50,37 @@ var channelNoPermissionStatusCodesMap = map[int]struct{}{
 	http.StatusNotFound:        {},
 }
 
+var forbiddenHasPermissionErrorCodesMap = map[string]struct{}{
+	"session_blocked_by_cyber_policy": {},
+}
+
 func ChannelStatusHasPermission(statusCode int) bool {
 	_, ok := channelNoPermissionStatusCodesMap[statusCode]
 	return !ok
 }
 
 func ChannelHasPermission(relayErr adaptor.Error) bool {
-	return ChannelStatusHasPermission(relayErr.StatusCode())
+	if ChannelStatusHasPermission(relayErr.StatusCode()) {
+		return true
+	}
+
+	if relayErr.StatusCode() != http.StatusForbidden {
+		return false
+	}
+
+	provider, ok := relayErr.(adaptor.ErrorCodeProvider)
+	if !ok {
+		return false
+	}
+
+	errorCode, ok := provider.ErrorCode().(string)
+	if !ok {
+		return false
+	}
+
+	_, ok = forbiddenHasPermissionErrorCodesMap[errorCode]
+
+	return ok
 }
 
 func getRequestDuration(meta *meta.Meta) time.Duration {

core/relay/plugin/monitor/monitor_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/labring/aiproxy/core/common/config"
 	relaymeta "github.com/labring/aiproxy/core/relay/meta"
+	relaymodel "github.com/labring/aiproxy/core/relay/model"
 	"github.com/stretchr/testify/require"
 )
 
@@ -73,3 +74,38 @@ func TestChannelStatusHasPermission(t *testing.T) {
 
 	require.True(t, ChannelStatusHasPermission(http.StatusBadRequest))
 }
+
+func TestChannelHasPermissionForForbiddenErrorCode(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		code string
+		want bool
+	}{
+		{
+			name: "request blocked by cyber policy",
+			code: "session_blocked_by_cyber_policy",
+			want: true,
+		},
+		{
+			name: "channel permission failure",
+			code: "insufficient_user_quota",
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			relayErr := relaymodel.NewOpenAIError(http.StatusForbidden, relaymodel.OpenAIError{
+				Code:    tt.code,
+				Message: "forbidden",
+				Type:    relaymodel.ErrorTypeUpstream,
+			})
+
+			require.Equal(t, tt.want, ChannelHasPermission(relayErr))
+		})
+	}
+}