new-api-custom/controller/detection.go

package controller

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/middleware"
	"github.com/QuantumNous/new-api/model"
	benchmark "github.com/QuantumNous/new-api/pkg/model_benchmark"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/relay"
	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/tidwall/gjson"
)

var detectionKeyPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]{16,96}$`)
var detectionCalls sync.Map

type detectionRequest struct {
	Group        string `json:"group"`
	Kind         string `json:"kind"`
	Key          string `json:"key"`
	PriceVersion string `json:"price_version"`
}

func detectionInput(c *gin.Context) (detectionRequest, *model.User, service.DetectionPrice, int, error) {
	var input detectionRequest
	if err := c.ShouldBindJSON(&input); err != nil {
		return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection request")
	}
	if input.Kind != "benchmark" && input.Kind != "fingerprint" {
		return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection type")
	}
	user, err := model.GetUserById(c.GetInt("id"), false)
	if err != nil {
		return input, nil, service.DetectionPrice{}, 0, err
	}
	if !service.IsUserSelectableGroup(user.Group, input.Group) {
		return input, nil, service.DetectionPrice{}, 0, errors.New("group unavailable")
	}
	price, reserve, err := service.DetectionPricing(user.Group, input.Group, input.Kind)
	return input, user, price, reserve, err
}

func EstimateDetection(c *gin.Context) {
	input, _, price, reserve, err := detectionInput(c)
	if err != nil {
		common.ApiError(c, err)
		return
	}
	encoded, _ := common.Marshal(price)
	calls := benchmark.ProbeCount
	if input.Kind == "fingerprint" {
		calls = fingerprint.ProbeCount
	}
	c.JSON(200, gin.H{"success": true, "data": gin.H{"reserved": reserve, "ratio": price.Snapshot.GroupRatio, "calls": calls, "price_version": model.DetectionSecretHash(string(encoded))}})
}

func CreateUserDetection(c *gin.Context) {
	input, user, price, reserve, err := detectionInput(c)
	if err != nil {
		common.ApiError(c, err)
		return
	}
	if !detectionKeyPattern.MatchString(input.Key) {
		common.ApiError(c, errors.New("invalid idempotency key"))
		return
	}
	var previous []model.DetectionRun
	key := fmt.Sprintf("user:%d:%s", user.Id, input.Key)
	if err = model.DB.Where("request_key = ?", key).Find(&previous).Error; err != nil {
		common.ApiError(c, err)
		return
	}
	if len(previous) > 0 {
		if previous[0].GroupName != input.Group || previous[0].Kind != input.Kind {
			common.ApiError(c, errors.New("idempotency key conflict"))
			return
		}
		c.JSON(200, gin.H{"success": true, "data": gin.H{"id": previous[0].ID}})
		return
	}
	encoded, _ := common.Marshal(price)
	if input.PriceVersion != model.DetectionSecretHash(string(encoded)) {
		common.ApiError(c, errors.New("pricing changed; request a new estimate"))
		return
	}
	now := time.Now()
	price.Request.At = &now
	encoded, _ = common.Marshal(price)
	channel, err := model.GetRandomSatisfiedChannel(input.Group, fingerprint.Model, 0, nil)
	if err != nil || channel == nil || !service.ChannelSupportsAstra(channel) {
		common.ApiError(c, errors.New("no Astra channel available"))
		return
	}
	id, err := common.GenerateRandomCharsKey(32)
	if err != nil {
		common.ApiError(c, err)
		return
	}
	run := &model.DetectionRun{ID: id, UserID: user.Id, RequestKey: fmt.Sprintf("user:%d:%s", user.Id, input.Key), ChannelID: channel.Id, GroupName: input.Group, Kind: input.Kind, Status: "pending", CreatedAt: time.Now().Unix(), Reserved: reserve, Price: string(encoded)}
	if err = model.CreateDetection(run); err != nil {
		common.ApiError(c, err)
		return
	}
	c.JSON(200, gin.H{"success": true, "data": gin.H{"id": run.ID}})
}

func GetUserDetection(c *gin.Context) {
	var run model.DetectionRun
	query := model.DB.Where("user_id = ?", c.GetInt("id"))
	if c.Param("id") == "latest" {
		query = query.Where("group_name = ?", c.Query("group")).Order("created_at desc")
	} else {
		query = query.Where("id = ?", c.Param("id"))
	}
	if err := query.First(&run).Error; err != nil {
		if c.Param("id") == "latest" {
			c.JSON(200, gin.H{"success": true, "data": nil})
		} else {
			c.JSON(404, gin.H{"success": false, "message": "Detection not found"})
		}
		return
	}
	var calls []model.DetectionCall
	if err := model.DB.Where("run_id = ?", run.ID).Find(&calls).Error; err != nil {
		common.ApiError(c, err)
		return
	}
	usage := []any{}
	for _, call := range calls {
		var parsed any
		_ = common.UnmarshalJsonStr(call.Usage, &parsed)
		if parsed != nil {
			usage = append(usage, parsed)
		}
		if call.Error != "" && run.Error == "" {
			run.Error = "usage_or_generation_incomplete"
		}
	}
	c.JSON(200, gin.H{"success": true, "data": gin.H{"id": run.ID, "kind": run.Kind, "status": run.Status, "quality": run.Quality, "reserved": run.Reserved, "charged": run.Charged, "settled": run.Settled, "usage": usage, "error": run.Error}})
}

func GetDetectionArtwork(c *gin.Context) {
	getDetectionArtwork(c, false)
}

func GetSiteDetectionArtwork(c *gin.Context) {
	getDetectionArtwork(c, true)
}

func getDetectionArtwork(c *gin.Context, allGroups bool) {
	var run model.DetectionRun
	if err := model.DB.Select("html", "channel_id", "created_at").Where("id = ? AND created_at > ?", c.Param("id"), time.Now().Add(-6*time.Hour).Unix()).First(&run).Error; err != nil || run.HTML == "" {
		c.Status(404)
		return
	}
	settings, err := model.GetStatusPageGroups(c.Request.Context())
	if err != nil {
		c.Status(http.StatusServiceUnavailable)
		return
	}
	channels, err := model.GetStatusChannels(c.Request.Context())
	if err != nil {
		c.Status(http.StatusServiceUnavailable)
		return
	}
	usable := service.GetUserUsableGroups("")
	visible := false
	for _, channel := range channels {
		if channel.Id != run.ChannelID {
			continue
		}
		for group := range strings.SplitSeq(channel.Group, ",") {
			group = strings.TrimSpace(group)
			if _, ok := usable[group]; (allGroups || (ok && !settings[group].Hidden)) && run.CreatedAt > settings[group].ResetBefore {
				visible = true
			}
		}
	}
	if !visible {
		c.Status(http.StatusNotFound)
		return
	}
	c.Header("Content-Security-Policy", "sandbox allow-scripts; default-src 'none'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; img-src data:; form-action 'none'; base-uri 'none'; frame-ancestors 'self'")
	c.Header("X-Content-Type-Options", "nosniff")
	c.Header("Cache-Control", "no-store")
	c.Data(200, "text/html; charset=utf-8", []byte(sanitizeDetectionHTML(run.HTML)))
}

// Decode only generation fields: external credentials cannot forward tools,
// previous responses, billing modifiers or arbitrary provider parameters.
func detectionProxyRequest(data []byte) (*dto.OpenAIResponsesRequest, error) {
	var fields map[string]any
	if len(data) > 64<<10 || common.Unmarshal(data, &fields) != nil {
		return nil, errors.New("invalid request")
	}
	if name, ok := fields["model"].(string); !ok || name != fingerprint.Model {
		return nil, errors.New("model not allowed")
	}
	for key := range fields {
		switch key {
		case "model", "input", "instructions", "stream", "max_output_tokens", "reasoning", "store", "tools", "tool_choice":
		default:
			return nil, errors.New("unsupported request field")
		}
	}
	input, err := common.Marshal(fields["input"])
	if err != nil || len(input) > 16384 {
		return nil, errors.New("input too large")
	}
	var checkText func(any) bool
	checkText = func(v any) bool {
		switch value := v.(type) {
		case string:
			return true
		case []any:
			for _, item := range value {
				if !checkText(item) {
					return false
				}
			}
			return true
		case map[string]any:
			for k, item := range value {
				switch k {
				case "role":
					role, ok := item.(string)
					if !ok || (role != "user" && role != "system" && role != "developer") {
						return false
					}
				case "type":
					kind, ok := item.(string)
					if !ok || (kind != "input_text" && kind != "message") {
						return false
					}
				case "content", "text":
					if !checkText(item) {
						return false
					}
				default:
					return false
				}
			}
			return true
		}
		return false
	}
	if !checkText(fields["input"]) {
		return nil, errors.New("text input required")
	}
	instructions, _ := fields["instructions"].(string)
	if len(input)+len(instructions) > 16384 {
		return nil, errors.New("input too large")
	}
	// Service-selected bounds cannot be overridden by the bearer.
	fields = map[string]any{"model": fingerprint.Model, "input": fields["input"], "instructions": instructions, "stream": true, "max_output_tokens": 16384, "reasoning": map[string]string{"effort": "medium"}, "store": false, "tools": []any{}}
	encoded, _ := common.Marshal(fields)
	var request dto.OpenAIResponsesRequest
	err = common.Unmarshal(encoded, &request)
	return &request, err
}

func DetectionProxy(c *gin.Context) {
	secret := strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ")
	if !strings.HasPrefix(secret, "det_") {
		c.AbortWithStatus(401)
		return
	}
	data, err := io.ReadAll(io.LimitReader(c.Request.Body, (64<<10)+1))
	if err != nil {
		c.AbortWithStatus(400)
		return
	}
	request, err := detectionProxyRequest(data)
	if err != nil {
		c.JSON(400, gin.H{"error": gin.H{"message": err.Error()}})
		return
	}
	call, err := model.ClaimDetectionCall(model.DetectionSecretHash(secret))
	if err != nil {
		c.AbortWithStatus(401)
		return
	}
	var run model.DetectionRun
	if err = model.DB.First(&run, "id = ?", call.RunID).Error; err != nil {
		c.AbortWithStatus(503)
		return
	}
	ctx, cancel := context.WithDeadline(c.Request.Context(), time.Unix(run.Deadline-10, 0))
	detectionCalls.Store(call.ID, cancel)
	defer func() { cancel(); detectionCalls.Delete(call.ID) }()
	channel, err := model.GetChannelById(run.ChannelID, true)
	var raw []byte
	var contentType string
	// Flush SSE comments while generation is buffered for accounting. Without
	// these bytes, the public edge times out a valid long reasoning request.
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("X-Accel-Buffering", "no")
	c.Writer.WriteHeader(http.StatusOK)
	_, _ = c.Writer.WriteString(": detection started\n\n")
	c.Writer.Flush()
	if err == nil && channel.Status == common.ChannelStatusEnabled {
		raw, contentType, err = waitDetectionResponse(c, ctx, channel, request)
	} else {
		err = errors.New("channel unavailable")
	}
	success := err == nil
	call.Available = &success
	if err != nil {
		call.Error = "channel_call_failed"
		if ctx.Err() != nil {
			call.Available = nil
			call.Error = "interrupted_usage_unknown"
		}
	} else {
		usage, complete := detectionResponseUsage(raw, contentType)
		if !complete {
			call.Error = "incomplete_response"
			call.Available = nil
		} else {
			call.Usage = usage.Raw
			quota, quotaErr := service.DetectionUsageQuota(run.Price, usage.Get("input_tokens").Int(), usage.Get("output_tokens").Int(), usage.Get("input_tokens_details.cached_tokens").Int(), max(usage.Get("input_tokens_details.cached_creation_tokens").Int(), usage.Get("input_tokens_details.cache_creation_tokens").Int()))
			if quotaErr != nil || quota > run.Reserved {
				call.Error = "settlement_unavailable"
			} else if run.UserID > 0 {
				call.Quota = quota
			}
		}
	}
	if err = model.FinishDetectionCall(call); err != nil {
		success = false
	}
	if !success {
		_, _ = c.Writer.WriteString("event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"Detection generation failed\"}}\n\n")
		c.Writer.Flush()
		return
	}
	if strings.Contains(contentType, "text/event-stream") {
		_, _ = c.Writer.Write(raw)
	} else {
		event, _ := common.Marshal(map[string]any{"type": "response.completed", "response": gjson.ParseBytes(raw).Value()})
		_, _ = fmt.Fprintf(c.Writer, "event: response.completed\ndata: %s\n\n", event)
	}
	c.Writer.Flush()
}

func waitDetectionResponse(c *gin.Context, ctx context.Context, channel *model.Channel, request *dto.OpenAIResponsesRequest) ([]byte, string, error) {
	type response struct {
		raw         []byte
		contentType string
		err         error
	}
	done := make(chan response, 1)
	go func() {
		raw, contentType, err := runDetectionResponse(ctx, channel, request)
		done <- response{raw, contentType, err}
	}()
	ticker := time.NewTicker(15 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case result := <-done:
			return result.raw, result.contentType, result.err
		case <-ticker.C:
			_, _ = c.Writer.WriteString(": waiting for generation\n\n")
			c.Writer.Flush()
		}
	}
}

func runDetectionResponse(ctx context.Context, channel *model.Channel, request *dto.OpenAIResponsesRequest) ([]byte, string, error) {
	encoded, err := common.Marshal(request)
	if err != nil {
		return nil, "", err
	}
	c, _ := gin.CreateTestContext(httptest.NewRecorder())
	c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/responses", bytes.NewReader(encoded))
	c.Request.Header.Set("Content-Type", "application/json")
	if middleware.SetupContextForSelectedChannel(c, channel, fingerprint.Model) != nil {
		return nil, "", errors.New("channel setup failed")
	}
	info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIResponses, request, nil)
	if err != nil {
		return nil, "", err
	}
	info.IsChannelTest = true
	info.DisablePing = true
	info.AstraDetectionDisabled = true
	info.AstraAutoSwitch = false
	adaptor, body, closer, apiErr := relay.PrepareResponsesRequest(c, info, request)
	if closer != nil {
		defer closer.Close()
	}
	if apiErr != nil {
		return nil, "", errors.New("preparation failed")
	}
	body, boundedCloser, err := boundDetectionBody(body, request)
	if err != nil {
		return nil, "", err
	}
	defer boundedCloser.Close()
	response, err := adaptor.DoRequest(c, info, body)
	if err != nil {
		return nil, "", err
	}
	resp, ok := response.(*http.Response)
	if !ok || resp == nil {
		return nil, "", errors.New("unsupported transport")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, "", errors.New("upstream request failed")
	}
	raw, err := io.ReadAll(io.LimitReader(resp.Body, (8<<20)+1))
	if len(raw) > 8<<20 {
		return nil, "", errors.New("response too large")
	}
	return raw, resp.Header.Get("Content-Type"), err
}

func detectionResponseUsage(raw []byte, contentType string) (gjson.Result, bool) {
	result := gjson.ParseBytes(raw)
	if strings.Contains(contentType, "text/event-stream") {
		result = gjson.Result{}
		for line := range strings.SplitSeq(string(raw), "\n") {
			if strings.HasPrefix(line, "data:") {
				event := gjson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
				switch event.Get("type").String() {
				case "response.completed", "response.incomplete", "response.failed":
					result = event.Get("response")
				}
			}
		}
	}
	usage := result.Get("usage")
	status := result.Get("status").String()
	return usage, (status == "completed" || status == "incomplete" || status == "failed") && usage.Get("input_tokens").Type == gjson.Number && usage.Get("output_tokens").Type == gjson.Number
}

// Apply billing and safety limits AFTER channel parameter overrides. Never let
// an administrative override turn a bounded paid test into unbounded tool work.
func boundDetectionBody(body common.ReplayableBody, request *dto.OpenAIResponsesRequest) (common.ReplayableBody, io.Closer, error) {
	raw, err := io.ReadAll(io.LimitReader(body, 128<<10))
	if err != nil {
		return nil, nil, err
	}
	var prepared map[string]any
	if common.Unmarshal(raw, &prepared) != nil {
		return nil, nil, errors.New("invalid prepared detection request")
	}
	var input any
	if common.Unmarshal(request.Input, &input) != nil {
		return nil, nil, errors.New("invalid detection input")
	}
	prepared["input"] = input
	var instructions any
	_ = common.Unmarshal(request.Instructions, &instructions)
	prepared["instructions"] = instructions
	prepared["max_output_tokens"] = *request.MaxOutputTokens
	if *request.MaxOutputTokens == 16384 {
		prepared["reasoning"] = map[string]string{"effort": "medium"}
	}
	if *request.MaxOutputTokens == fingerprint.OutputTokens || *request.MaxOutputTokens == benchmark.OutputTokens {
		prepared["reasoning"] = map[string]string{"effort": "low"}
		delete(prepared, "instructions")
		delete(prepared, "temperature")
		delete(prepared, "top_p")
	}
	prepared["tools"] = []any{}
	prepared["tool_choice"] = "none"
	prepared["store"] = false
	prepared["stream"] = true
	for _, key := range []string{"previous_response_id", "conversation", "background", "max_tokens", "max_completion_tokens"} {
		delete(prepared, key)
	}
	encoded, err := common.Marshal(prepared)
	if err != nil {
		return nil, nil, err
	}
	return relaycommon.NewOutboundJSONBody(encoded)
}