new-api-custom/relay/common/response_model_guard.go

package common

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	appcommon "github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/relaykit/types"
)

const ResponseGuardMaxAttempts = 3
const ResponseGuardErrorCode types.ErrorCode = "response_model_rejected"

// The guard is request-scoped, unlike channel metadata which is reset on retry.
type ResponseModelGuard struct {
	Attempts []*ResponseGuardAttempt `json:"attempts"`
	Deadline time.Time               `json:"-"`
}

type ResponseGuardUsage struct {
	Input  int `json:"input_tokens"`
	Output int `json:"output_tokens"`
	Total  int `json:"total_tokens"`
}

type ResponseGuardAttempt struct {
	Timing     *UpstreamTiming    `json:"-"`
	OnEvent    func([]byte, bool) `json:"-"`
	OnClose    func()             `json:"-"`
	closeOnce  sync.Once
	Drain      func()              `json:"-"`
	Number     int                 `json:"number"`
	ChannelID  int                 `json:"channel_id"`
	Expected   string              `json:"expected_model"`
	Returned   string              `json:"returned_model,omitempty"`
	Reason     string              `json:"reason"`
	ElapsedMS  int64               `json:"elapsed_ms"`
	Usage      *ResponseGuardUsage `json:"usage,omitempty"`
	Delivered  bool                `json:"delivered"`
	Status     int                 `json:"status,omitempty"`
	ResponseID string              `json:"-"`
	Sequence   int                 `json:"-"`
	started    time.Time
	timeout    time.Duration
	confirmed  bool
	err        error
	cancel     context.CancelFunc
	timer      *time.Timer
	stopClient func() bool
}

func (g *ResponseModelGuard) Last() *ResponseGuardAttempt {
	if g == nil || len(g.Attempts) == 0 {
		return nil
	}
	return g.Attempts[len(g.Attempts)-1]
}

func (g *ResponseModelGuard) Begin(ctx context.Context, number, channel int) (context.Context, *ResponseGuardAttempt) {
	upstream, cancel := context.WithCancel(context.WithoutCancel(ctx))
	a := &ResponseGuardAttempt{Number: number, ChannelID: channel, Expected: "gpt-6-astra", Reason: "pending", started: time.Now(), cancel: cancel}
	// Model confirmation must allow the same initial wait as an ordinary stream.
	// A separate 10-second cutoff rejects valid slow upstream responses.
	a.timeout = time.Duration(constant.StreamingTimeout) * time.Second
	if a.timeout <= 0 {
		a.timeout = 300 * time.Second
	}
	if !g.Deadline.IsZero() {
		a.timeout = max(time.Until(g.Deadline), time.Nanosecond)
	}
	a.timer = time.AfterFunc(a.timeout, cancel)
	a.stopClient = context.AfterFunc(ctx, cancel)
	g.Attempts = append(g.Attempts, a)
	return upstream, a
}

func (a *ResponseGuardAttempt) Finish(delivered bool) {
	a.ElapsedMS = time.Since(a.started).Milliseconds()
	a.Delivered = delivered
}

func (a *ResponseGuardAttempt) Close() {
	a.closeOnce.Do(func() {
		if a.OnClose != nil {
			a.OnClose()
		}
		a.timer.Stop()
		a.stopClient()
		a.cancel()
	})
}

func (a *ResponseGuardAttempt) Reject(reason string) error {
	if a.err == nil {
		a.Reason = reason
		a.err = errors.New("response model guard: " + reason)
	}
	return a.err
}

func (a *ResponseGuardAttempt) APIError() *types.NewAPIError {
	if a == nil || a.err == nil {
		return nil
	}
	return types.NewErrorWithStatusCode(errors.New(a.PublicErrorMessage()), ResponseGuardErrorCode, http.StatusConflict, types.ErrOptionWithSkipRetry())
}

func (a *ResponseGuardAttempt) PublicErrorMessage() string {
	if a != nil && a.Reason == "model_mismatch" && a.Returned != "" {
		return "Requested model " + a.Expected + ", but the upstream response declared " + a.Returned + ". No verified result could be delivered. Usage from each upstream attempt is charged independently, including rejected results."
	}
	if a != nil && a.Expected != "" {
		return "Could not verify that the upstream response used " + a.Expected + ". Usage from each upstream attempt is charged independently, including rejected results."
	}
	return "Could not verify the upstream response model."
}

func (g *ResponseModelGuard) APIError() *types.NewAPIError {
	if g == nil {
		return nil
	}
	for i := len(g.Attempts) - 1; i >= 0; i-- {
		attempt := g.Attempts[i]
		if attempt != nil && attempt.Reason == "model_mismatch" && attempt.Returned != "" {
			return attempt.APIError()
		}
	}
	last := g.Last()
	if last == nil {
		return nil
	}
	return last.APIError()
}

// Only read the provider's declaration and numeric usage, never generated text.
func (a *ResponseGuardAttempt) observe(payload []byte, stream bool) error {
	if stream {
		a.Timing.Mark("first_sse_data")
	}
	if a.OnEvent != nil {
		a.OnEvent(payload, stream)
	}
	var event struct {
		Model    string              `json:"model"`
		ID       string              `json:"id"`
		Sequence int                 `json:"sequence_number"`
		Usage    *ResponseGuardUsage `json:"usage"`
		Response *struct {
			Model string              `json:"model"`
			ID    string              `json:"id"`
			Usage *ResponseGuardUsage `json:"usage"`
		} `json:"response"`
	}
	if err := appcommon.Unmarshal(payload, &event); err != nil {
		return a.Reject("invalid_response")
	}
	model, id, usage := event.Model, event.ID, event.Usage
	if stream {
		model, id, usage = "", "", nil
		if event.Response != nil {
			model, id, usage = event.Response.Model, event.Response.ID, event.Response.Usage
		}
	}
	if usage != nil {
		a.Usage = usage
	}
	a.Sequence = max(a.Sequence, event.Sequence)
	model = strings.TrimSpace(model)
	if model == "" {
		return nil
	}
	// Bound diagnostics from an untrusted upstream.
	a.Returned = string([]rune(model)[:min(len([]rune(model)), 128)])
	if model != a.Expected {
		return a.Reject("model_mismatch")
	}
	if id != "" {
		a.ResponseID = string([]rune(id)[:min(len([]rune(id)), 256)])
	}
	if !a.confirmed {
		if !a.timer.Stop() {
			return a.Reject("confirmation_timeout")
		}
		if !a.stopClient() {
			return a.Reject("client_cancelled")
		}
		a.confirmed = true
		a.Reason = "accepted"
	}
	return nil
}

// Confirm buffers only the prefix until a provider model is known. No caller
// may commit headers or start a downstream pinger before it returns success.
func (a *ResponseGuardAttempt) Confirm(resp *http.Response, ctx context.Context) error {
	a.Status = resp.StatusCode
	if resp.StatusCode != http.StatusOK {
		a.timer.Stop()
		a.stopClient()
		a.Reason = "http_error"
		resp.Body = &guardCloseBody{Reader: resp.Body, body: resp.Body, close: a.Close}
		return nil
	}
	if !strings.HasPrefix(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
		body, err := io.ReadAll(io.LimitReader(resp.Body, (1<<20)+1))
		if err != nil {
			return a.readError(ctx, err)
		}
		if len(body) > 1<<20 {
			return a.Reject("confirmation_buffer_limit")
		}
		if err := a.observe(body, false); err != nil {
			return err
		}
		if !a.confirmed {
			return a.Reject("missing_model")
		}
		_ = resp.Body.Close()
		resp.Body = &guardCloseBody{Reader: bytes.NewReader(body), body: io.NopCloser(bytes.NewReader(nil)), close: a.Close}
		return nil
	}
	b := &guardStreamBody{reader: bufio.NewReader(resp.Body), body: resp.Body, attempt: a}
	a.Drain = func() {
		b.draining = true
		for {
			if _, err := b.readFrame(8 << 20); err != nil {
				return
			}
		}
	}
	var prefix bytes.Buffer
	for !a.confirmed {
		frame, err := b.readFrame((1 << 20) - prefix.Len())
		if err != nil {
			return a.readError(ctx, err)
		}
		prefix.Write(frame)
	}
	b.pending = prefix.Bytes()
	resp.Body = b
	return nil
}

func (a *ResponseGuardAttempt) readError(ctx context.Context, err error) error {
	if a.err != nil {
		return a.err
	}
	if ctx.Err() != nil {
		return a.Reject("client_cancelled")
	}
	if time.Since(a.started) >= a.timeout {
		return a.Reject("confirmation_timeout")
	}
	if errors.Is(err, io.EOF) {
		return a.Reject("missing_model")
	}
	return a.Reject("response_read_error")
}

type guardCloseBody struct {
	io.Reader
	body  io.Closer
	close func()
	once  sync.Once
}

func (b *guardCloseBody) Close() error { b.once.Do(b.close); return b.body.Close() }

// Validate complete SSE records before exposing their bytes to the existing
// scanner, so a late wrong-model event cannot race past the check.
type guardStreamBody struct {
	draining bool
	reader   *bufio.Reader
	body     io.ReadCloser
	attempt  *ResponseGuardAttempt
	pending  []byte
	once     sync.Once
}

func (b *guardStreamBody) Close() error { b.once.Do(b.attempt.Close); return b.body.Close() }
func (b *guardStreamBody) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	if len(b.pending) == 0 {
		frame, err := b.readFrame(128 << 20)
		if err != nil {
			return 0, err
		}
		b.pending = frame
	}
	n := copy(p, b.pending)
	b.pending = b.pending[n:]
	return n, nil
}

func (b *guardStreamBody) readFrame(limit int) ([]byte, error) {
	var frame, line, data bytes.Buffer
	for {
		part, err := b.reader.ReadSlice('\n')
		if frame.Len()+len(part) > limit {
			if b.draining {
				return nil, io.ErrShortBuffer
			}
			return nil, b.attempt.Reject("confirmation_buffer_limit")
		}
		frame.Write(part)
		line.Write(part)
		if errors.Is(err, bufio.ErrBufferFull) {
			continue
		}
		if err != nil && !errors.Is(err, io.EOF) {
			return nil, err
		}
		text := strings.TrimSuffix(strings.TrimSuffix(line.String(), "\n"), "\r")
		line.Reset()
		if value, ok := strings.CutPrefix(text, "data:"); ok {
			if data.Len() > 0 {
				data.WriteByte('\n')
			}
			data.WriteString(strings.TrimPrefix(value, " "))
		}
		if text == "" || err != nil {
			if frame.Len() == 0 {
				return nil, io.EOF
			}
			payload := bytes.TrimSpace(data.Bytes())
			if len(payload) > 0 && !bytes.Equal(payload, []byte("[DONE]")) {
				if b.draining {
					if b.attempt.OnEvent != nil {
						b.attempt.OnEvent(payload, true)
					}
				} else if e := b.attempt.observe(payload, true); e != nil {
					return nil, e
				}
			}
			if b.draining && bytes.Equal(payload, []byte("[DONE]")) {
				return nil, io.EOF
			}
			if !b.attempt.confirmed && bytes.Equal(payload, []byte("[DONE]")) {
				return nil, b.attempt.Reject("missing_model")
			}
			return frame.Bytes(), nil
		}
	}
}

// Detach keeps only this upstream attempt alive while a different channel is
// tried; its numeric usage is collected for at most one minute.
func (a *ResponseGuardAttempt) Detach() {
	a.timer.Stop()
	a.stopClient()
	a.timer = time.AfterFunc(60*time.Second, a.cancel)
}