new-api/service/channel_monitor_checker.go

package service

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

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
)

const channelMonitorResponseSnippetLimit = 500

type channelMonitorCheckOptions struct {
	client      *http.Client
	validateURL func(string) error
	now         func() time.Time
}

func RunMonitorCheck(monitor *model.ChannelMonitor) error {
	return runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{})
}

func runMonitorCheckWithOptions(monitor *model.ChannelMonitor, options channelMonitorCheckOptions) error {
	if monitor == nil {
		return errors.New("channel monitor is nil")
	}

	checkNow := options.now
	if checkNow == nil {
		checkNow = time.Now
	}

	template := SelectChannelMonitorRequestTemplate(monitor.RequestTemplates, nil)
	models, modelParseErr := channelMonitorModels(monitor)

	var checkErrs []error
	if modelParseErr != nil {
		checkErrs = append(checkErrs, modelParseErr)
	}
	for _, modelName := range models {
		if err := checkChannelMonitorModel(monitor, modelName, template, options, checkNow); err != nil {
			checkErrs = append(checkErrs, err)
		}
	}

	now := checkNow()
	if err := monitor.UpdateLastCheckedAt(now); err != nil {
		checkErrs = append(checkErrs, fmt.Errorf("update monitor last_checked_at: %w", err))
	}
	return errors.Join(checkErrs...)
}

func channelMonitorModels(monitor *model.ChannelMonitor) ([]string, error) {
	models := make([]string, 0, 1)
	primaryModel := strings.TrimSpace(monitor.PrimaryModel)
	if primaryModel != "" {
		models = append(models, primaryModel)
	}

	rawExtraModels := strings.TrimSpace(monitor.ExtraModels)
	if rawExtraModels == "" {
		return models, nil
	}

	var extraModels []string
	if err := common.Unmarshal([]byte(rawExtraModels), &extraModels); err != nil {
		return models, fmt.Errorf("invalid extra_models json array ignored: %w", err)
	}
	for _, modelName := range extraModels {
		modelName = strings.TrimSpace(modelName)
		if modelName != "" {
			models = append(models, modelName)
		}
	}
	return models, nil
}

func checkChannelMonitorModel(monitor *model.ChannelMonitor, modelName string, template string, options channelMonitorCheckOptions, now func() time.Time) error {
	start := now()
	requestBody, err := buildChannelMonitorRequestBody(monitor, modelName, template)
	record := &model.ChannelMonitorHistory{
		MonitorId:   monitor.Id,
		ChannelId:   monitor.ChannelId,
		ModelName:   modelName,
		RequestBody: requestBody,
		CheckedAt:   start,
	}
	if err != nil {
		return insertChannelMonitorErrorHistory(record, fmt.Errorf("build request body: %w", err))
	}

	validateURL := options.validateURL
	if validateURL == nil {
		validateURL = ValidateChannelMonitorTargetURL
	}
	if err := validateURL(monitor.TargetUrl); err != nil {
		return insertChannelMonitorErrorHistory(record, fmt.Errorf("SSRF validation failed: %w", err))
	}

	headers, err := parseChannelMonitorCustomHeaders(monitor.CustomHeaders)
	if err != nil {
		return insertChannelMonitorErrorHistory(record, err)
	}

	timeout := channelMonitorTimeout(monitor.Timeout)
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, monitor.TargetUrl, bytes.NewReader([]byte(requestBody)))
	if err != nil {
		return insertChannelMonitorErrorHistory(record, fmt.Errorf("create request: %w", err))
	}
	req.Header.Set("Content-Type", "application/json")
	for key, value := range headers {
		req.Header.Set(key, value)
	}

	client := options.client
	if client == nil {
		client = newChannelMonitorHTTPClient(timeout, validateURL)
	}

	resp, err := client.Do(req)
	record.LatencyMs = int(time.Since(start).Milliseconds())
	if err != nil {
		return insertChannelMonitorErrorHistory(record, fmt.Errorf("request failed: %w", err))
	}
	defer resp.Body.Close()

	record.StatusCode = resp.StatusCode
	snippet, readErr := readChannelMonitorResponseSnippet(resp.Body)
	record.ResponseSnippet = snippet
	if readErr != nil {
		return insertChannelMonitorErrorHistory(record, fmt.Errorf("read response: %w", readErr))
	}

	record.Status = channelMonitorStatus(resp.StatusCode, snippet, record.LatencyMs, timeout)
	if record.Status == "failed" && resp.StatusCode >= 200 && resp.StatusCode < 300 && strings.TrimSpace(snippet) == "" {
		record.ErrorMessage = "empty response body"
	}
	if err := record.Insert(); err != nil {
		return fmt.Errorf("insert monitor history: %w", err)
	}
	return nil
}

func buildChannelMonitorRequestBody(monitor *model.ChannelMonitor, modelName string, template string) (string, error) {
	if strings.TrimSpace(monitor.CustomBody) != "" {
		return replaceChannelMonitorTemplateVariables(monitor.CustomBody, modelName, template), nil
	}

	body := map[string]any{
		"model": modelName,
		"messages": []map[string]string{
			{"role": "user", "content": template},
		},
	}
	encoded, err := common.Marshal(body)
	if err != nil {
		return "", err
	}
	return string(encoded), nil
}

func replaceChannelMonitorTemplateVariables(raw string, modelName string, template string) string {
	replacer := strings.NewReplacer(
		"{{model}}", modelName,
		"{{model_name}}", modelName,
		"${model}", modelName,
		"${model_name}", modelName,
		"{{message}}", template,
		"{{template}}", template,
		"${message}", template,
		"${template}", template,
	)
	return replacer.Replace(raw)
}

func parseChannelMonitorCustomHeaders(raw string) (map[string]string, error) {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return nil, nil
	}

	var headers map[string]string
	if err := common.Unmarshal([]byte(raw), &headers); err != nil {
		return nil, fmt.Errorf("invalid custom_headers json object: %w", err)
	}
	return headers, nil
}

func newChannelMonitorHTTPClient(timeout time.Duration, validateURL func(string) error) *http.Client {
	dialer := &net.Dialer{
		Timeout:   timeout,
		KeepAlive: 30 * time.Second,
	}
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			DialContext:         newChannelMonitorProtectedDialContext(dialer),
			ForceAttemptHTTP2:   true,
			MaxIdleConns:        100,
			IdleConnTimeout:     90 * time.Second,
			TLSHandshakeTimeout: 10 * time.Second,
		},
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) >= 10 {
				return errors.New("stopped after 10 redirects")
			}
			if validateURL == nil {
				return nil
			}
			if err := validateURL(req.URL.String()); err != nil {
				return fmt.Errorf("redirect target rejected: %w", err)
			}
			return nil
		},
	}
}

func channelMonitorTimeout(timeoutSeconds int) time.Duration {
	if timeoutSeconds <= 0 {
		timeoutSeconds = 30
	}
	return time.Duration(timeoutSeconds) * time.Second
}

func channelMonitorStatus(statusCode int, snippet string, latencyMs int, timeout time.Duration) string {
	if statusCode < 200 || statusCode >= 300 {
		return "failed"
	}
	if strings.TrimSpace(snippet) == "" {
		return "failed"
	}
	if latencyMs >= int(timeout.Milliseconds()) {
		return "degraded"
	}
	return "pass"
}

func readChannelMonitorResponseSnippet(reader io.Reader) (string, error) {
	body, err := io.ReadAll(io.LimitReader(reader, channelMonitorResponseSnippetLimit))
	if err != nil {
		return "", err
	}
	return string(body), nil
}

func insertChannelMonitorErrorHistory(record *model.ChannelMonitorHistory, err error) error {
	record.Status = "error"
	record.ErrorMessage = err.Error()
	if insertErr := record.Insert(); insertErr != nil {
		return fmt.Errorf("%w; insert monitor history: %v", err, insertErr)
	}
	return err
}