feat: add channel monitor runner
Implement monitor probe execution, history recording, redirect-safe target checks, and a background runner for due monitors.
This commit is contained in:
zizi
parent
a63300fc1f
commit
8f91ab1eb6
@ -32,6 +32,11 @@ func (m *ChannelMonitor) Update() error {
|
||||
return DB.Save(m).Error
|
||||
}
|
||||
|
||||
func (m *ChannelMonitor) UpdateLastCheckedAt(checkedAt time.Time) error {
|
||||
m.LastCheckedAt = &checkedAt
|
||||
return DB.Model(m).Update("last_checked_at", checkedAt).Error
|
||||
}
|
||||
|
||||
func GetChannelMonitorByID(id int) (*ChannelMonitor, error) {
|
||||
var monitor ChannelMonitor
|
||||
err := DB.First(&monitor, "id = ?", id).Error
|
||||
|
||||
251
service/channel_monitor_checker.go
Normal file
251
service/channel_monitor_checker.go
Normal file
@ -0,0 +1,251 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"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 {
|
||||
return &http.Client{
|
||||
Timeout: timeout,
|
||||
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
|
||||
}
|
||||
182
service/channel_monitor_checker_test.go
Normal file
182
service/channel_monitor_checker_test.go
Normal file
@ -0,0 +1,182 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupChannelMonitorTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
|
||||
common.UsingSQLite = true
|
||||
common.UsingMySQL = false
|
||||
common.UsingPostgreSQL = false
|
||||
|
||||
previousDB := model.DB
|
||||
previousLogDB := model.LOG_DB
|
||||
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.ChannelMonitor{}, &model.ChannelMonitorHistory{}))
|
||||
|
||||
model.DB = db
|
||||
model.LOG_DB = db
|
||||
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
model.LOG_DB = previousLogDB
|
||||
sqlDB, err := db.DB()
|
||||
if err == nil {
|
||||
_ = sqlDB.Close()
|
||||
}
|
||||
})
|
||||
return db
|
||||
}
|
||||
|
||||
func allowTestChannelMonitorURL(string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestRunMonitorCheckPassWithCustomTemplateBodyAndHeaders(t *testing.T) {
|
||||
db := setupChannelMonitorTestDB(t)
|
||||
|
||||
var receivedBody string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, "ok", r.Header.Get("X-Test"))
|
||||
body, err := io.ReadAll(r.Body)
|
||||
require.NoError(t, err)
|
||||
receivedBody = string(body)
|
||||
_, _ = w.Write([]byte(`{"id":"chatcmpl_test"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
monitor := &model.ChannelMonitor{
|
||||
Name: "custom",
|
||||
ChannelId: 7,
|
||||
TargetUrl: server.URL,
|
||||
PrimaryModel: "m-custom",
|
||||
RequestTemplates: `["custom prompt"]`,
|
||||
CustomBody: `{"model":"{{model}}","input":"{{template}}"}`,
|
||||
CustomHeaders: `{"X-Test":"ok"}`,
|
||||
Timeout: 5,
|
||||
Status: 1,
|
||||
}
|
||||
require.NoError(t, db.Create(monitor).Error)
|
||||
|
||||
err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
|
||||
client: server.Client(),
|
||||
validateURL: allowTestChannelMonitorURL,
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, receivedBody, `"model":"m-custom"`)
|
||||
require.Contains(t, receivedBody, `"input":"custom prompt"`)
|
||||
require.NotNil(t, monitor.LastCheckedAt)
|
||||
|
||||
var history model.ChannelMonitorHistory
|
||||
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||
require.Equal(t, monitor.Id, history.MonitorId)
|
||||
require.Equal(t, 7, history.ChannelId)
|
||||
require.Equal(t, "m-custom", history.ModelName)
|
||||
require.Equal(t, "pass", history.Status)
|
||||
require.Equal(t, 200, history.StatusCode)
|
||||
require.GreaterOrEqual(t, history.LatencyMs, 0)
|
||||
require.Empty(t, history.ErrorMessage)
|
||||
require.Contains(t, history.RequestBody, `"model":"m-custom"`)
|
||||
require.Contains(t, history.ResponseSnippet, "chatcmpl_test")
|
||||
require.False(t, history.CheckedAt.IsZero())
|
||||
}
|
||||
|
||||
func TestRunMonitorCheckRecordsFailedStatusCode(t *testing.T) {
|
||||
db := setupChannelMonitorTestDB(t)
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"error":"bad"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
monitor := &model.ChannelMonitor{
|
||||
Name: "failed",
|
||||
TargetUrl: server.URL,
|
||||
PrimaryModel: "m-failed",
|
||||
Timeout: 5,
|
||||
Status: 1,
|
||||
}
|
||||
require.NoError(t, db.Create(monitor).Error)
|
||||
|
||||
err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
|
||||
client: server.Client(),
|
||||
validateURL: allowTestChannelMonitorURL,
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
var history model.ChannelMonitorHistory
|
||||
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||
require.Equal(t, "failed", history.Status)
|
||||
require.Equal(t, http.StatusInternalServerError, history.StatusCode)
|
||||
require.Contains(t, history.ResponseSnippet, "bad")
|
||||
}
|
||||
|
||||
func TestRunMonitorCheckRecordsSSRFFailure(t *testing.T) {
|
||||
db := setupChannelMonitorTestDB(t)
|
||||
|
||||
monitor := &model.ChannelMonitor{
|
||||
Name: "ssrf",
|
||||
TargetUrl: "http://127.0.0.1:1/v1/chat/completions",
|
||||
PrimaryModel: "m-ssrf",
|
||||
Timeout: 5,
|
||||
Status: 1,
|
||||
}
|
||||
require.NoError(t, db.Create(monitor).Error)
|
||||
|
||||
err := RunMonitorCheck(monitor)
|
||||
|
||||
require.Error(t, err)
|
||||
var history model.ChannelMonitorHistory
|
||||
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||
require.Equal(t, "error", history.Status)
|
||||
require.Contains(t, history.ErrorMessage, "SSRF validation failed")
|
||||
require.Contains(t, history.RequestBody, "m-ssrf")
|
||||
require.Equal(t, 0, history.StatusCode)
|
||||
}
|
||||
|
||||
func TestRunMonitorCheckRecordsInvalidCustomHeaders(t *testing.T) {
|
||||
db := setupChannelMonitorTestDB(t)
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Fatal("request should not be sent when custom headers are invalid")
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
monitor := &model.ChannelMonitor{
|
||||
Name: "headers",
|
||||
TargetUrl: server.URL,
|
||||
PrimaryModel: "m-headers",
|
||||
CustomHeaders: `{"X-Test":123}`,
|
||||
Timeout: 5,
|
||||
Status: 1,
|
||||
}
|
||||
require.NoError(t, db.Create(monitor).Error)
|
||||
|
||||
err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
|
||||
client: server.Client(),
|
||||
validateURL: allowTestChannelMonitorURL,
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
var history model.ChannelMonitorHistory
|
||||
require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
|
||||
require.Equal(t, "error", history.Status)
|
||||
require.Contains(t, history.ErrorMessage, "invalid custom_headers json object")
|
||||
}
|
||||
109
service/channel_monitor_runner.go
Normal file
109
service/channel_monitor_runner.go
Normal file
@ -0,0 +1,109 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
)
|
||||
|
||||
var (
|
||||
channelMonitorRunnerOnce sync.Once
|
||||
channelMonitorStopMu sync.Mutex
|
||||
channelMonitorStopCh chan struct{}
|
||||
channelMonitorStopped bool
|
||||
|
||||
channelMonitorRunningMu sync.Mutex
|
||||
channelMonitorRunning = map[int]struct{}{}
|
||||
)
|
||||
|
||||
func StartChannelMonitorRunner() {
|
||||
channelMonitorRunnerOnce.Do(func() {
|
||||
channelMonitorStopMu.Lock()
|
||||
defer channelMonitorStopMu.Unlock()
|
||||
|
||||
channelMonitorStopCh = make(chan struct{})
|
||||
channelMonitorStopped = false
|
||||
go runChannelMonitorLoop(channelMonitorStopCh)
|
||||
})
|
||||
}
|
||||
|
||||
func StopChannelMonitorRunner() {
|
||||
channelMonitorStopMu.Lock()
|
||||
defer channelMonitorStopMu.Unlock()
|
||||
|
||||
if channelMonitorStopCh != nil && !channelMonitorStopped {
|
||||
close(channelMonitorStopCh)
|
||||
channelMonitorStopped = true
|
||||
}
|
||||
}
|
||||
|
||||
func runChannelMonitorLoop(stopCh <-chan struct{}) {
|
||||
ticker := time.NewTicker(time.Minute)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
runDueChannelMonitors()
|
||||
case <-stopCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func runDueChannelMonitors() {
|
||||
monitors, err := model.GetEnabledChannelMonitors()
|
||||
if err != nil {
|
||||
log.Printf("channel monitor runner: get monitors error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
for _, monitor := range monitors {
|
||||
if !channelMonitorDue(&monitor, now) {
|
||||
continue
|
||||
}
|
||||
if !markChannelMonitorRunning(monitor.Id) {
|
||||
continue
|
||||
}
|
||||
|
||||
monitorCopy := monitor
|
||||
go func() {
|
||||
defer unmarkChannelMonitorRunning(monitorCopy.Id)
|
||||
if err := RunMonitorCheck(&monitorCopy); err != nil {
|
||||
log.Printf("channel monitor runner: monitor %d check error: %v", monitorCopy.Id, err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
func channelMonitorDue(monitor *model.ChannelMonitor, now time.Time) bool {
|
||||
if monitor.LastCheckedAt == nil {
|
||||
return true
|
||||
}
|
||||
interval := monitor.CheckInterval
|
||||
if interval <= 0 {
|
||||
interval = 300
|
||||
}
|
||||
return !monitor.LastCheckedAt.Add(time.Duration(interval) * time.Second).After(now)
|
||||
}
|
||||
|
||||
func markChannelMonitorRunning(monitorID int) bool {
|
||||
channelMonitorRunningMu.Lock()
|
||||
defer channelMonitorRunningMu.Unlock()
|
||||
|
||||
if _, ok := channelMonitorRunning[monitorID]; ok {
|
||||
return false
|
||||
}
|
||||
channelMonitorRunning[monitorID] = struct{}{}
|
||||
return true
|
||||
}
|
||||
|
||||
func unmarkChannelMonitorRunning(monitorID int) {
|
||||
channelMonitorRunningMu.Lock()
|
||||
defer channelMonitorRunningMu.Unlock()
|
||||
|
||||
delete(channelMonitorRunning, monitorID)
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user