diff --git a/model/channel_monitor.go b/model/channel_monitor.go index d0b20fd7..650558b5 100644 --- a/model/channel_monitor.go +++ b/model/channel_monitor.go @@ -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 diff --git a/service/channel_monitor_checker.go b/service/channel_monitor_checker.go new file mode 100644 index 00000000..baad47fb --- /dev/null +++ b/service/channel_monitor_checker.go @@ -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 +} diff --git a/service/channel_monitor_checker_test.go b/service/channel_monitor_checker_test.go new file mode 100644 index 00000000..4f7fd209 --- /dev/null +++ b/service/channel_monitor_checker_test.go @@ -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") +} diff --git a/service/channel_monitor_runner.go b/service/channel_monitor_runner.go new file mode 100644 index 00000000..d22e8805 --- /dev/null +++ b/service/channel_monitor_runner.go @@ -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) +}