package scheduler import ( "bytes" "context" "encoding/json" "errors" "fmt" "io" "log/slog" "net/http" "strings" "time" "unicode/utf8" platformid "aigateway.local/core/internal/platform/id" ) const maxExecutionResponseBytes = 2 << 20 type Engine struct { service *Service baseURL string client *http.Client workerID string batchSize int maxAttempts int executionTTL time.Duration logger *slog.Logger } func NewEngine(service *Service, baseURL, workerID string, batchSize, maxAttempts int, executionTTL time.Duration, logger *slog.Logger) *Engine { if batchSize < 1 { batchSize = 10 } if maxAttempts < 1 { maxAttempts = 3 } if executionTTL < time.Minute { executionTTL = 5 * time.Minute } if logger == nil { logger = slog.Default() } return &Engine{service: service, baseURL: strings.TrimRight(baseURL, "/"), client: &http.Client{Timeout: executionTTL}, workerID: workerID, batchSize: batchSize, maxAttempts: maxAttempts, executionTTL: executionTTL, logger: logger} } // Tick coalesces overdue schedules into durable pending runs, reclaims stale // executions, and processes a bounded batch. SKIP LOCKED makes it safe for // multiple scheduler replicas to call Tick concurrently. func (e *Engine) Tick(ctx context.Context, now time.Time) (int, error) { if err := e.scheduleDue(ctx, now.UTC()); err != nil { return 0, err } if err := e.reclaim(ctx, now.UTC()); err != nil { return 0, err } runs, err := e.claim(ctx) if err != nil { return 0, err } for _, run := range runs { task, getErr := e.service.Get(ctx, run.TaskID) if getErr != nil { _ = e.complete(ctx, run, Task{ID: run.TaskID, Code: run.TaskCode}, nil, getErr) continue } response, executeErr := e.execute(ctx, task, run) if executeErr != nil && run.Attempts < e.maxAttempts { if retryErr := e.retry(ctx, run, executeErr); retryErr != nil { return len(runs), retryErr } continue } if completeErr := e.complete(ctx, run, task, response, executeErr); completeErr != nil { return len(runs), completeErr } } return len(runs), nil } func (e *Engine) scheduleDue(ctx context.Context, now time.Time) error { tx, err := e.service.pool.Begin(ctx) if err != nil { return err } defer func() { _ = tx.Rollback(ctx) }() rows, err := tx.Query(ctx, `SELECT id::text,cron_expression,timezone,next_run_at FROM gateway.scheduled_tasks WHERE enabled AND next_run_at <= $1 ORDER BY next_run_at,id FOR UPDATE SKIP LOCKED LIMIT $2`, now, e.batchSize) if err != nil { return err } type due struct { id, expression, timezone string scheduledFor time.Time } items := []due{} for rows.Next() { var item due if err = rows.Scan(&item.id, &item.expression, &item.timezone, &item.scheduledFor); err != nil { rows.Close() return err } items = append(items, item) } rows.Close() if err = rows.Err(); err != nil { return err } for _, item := range items { schedule, parseErr := ParseCron(item.expression) location, locationErr := time.LoadLocation(item.timezone) if parseErr != nil || locationErr != nil { _, err = tx.Exec(ctx, `UPDATE gateway.scheduled_tasks SET enabled=false,next_run_at=NULL,last_status='failed',last_error='cron 或时区配置无效',updated_at=clock_timestamp() WHERE id=$1`, item.id) if err != nil { return err } continue } // 补跑停机期间漏掉的执行:从旧的 next_run_at 起逐个 occurrence // 落一条 pending run,直到越过 now,而不是只补最新一次。否则调度器 // 宕机超过一个周期后,中间所有计划执行被静默丢弃。 runID, idErr := platformid.NewUUID() if idErr != nil { return idErr } if _, err = tx.Exec(ctx, `INSERT INTO gateway.scheduled_task_runs(id,task_id,trigger_type,scheduled_for) VALUES($1,$2,'scheduled',$3) ON CONFLICT(task_id,trigger_type,scheduled_for) DO NOTHING`, runID, item.id, item.scheduledFor); err != nil { return err } nextRun := item.scheduledFor var nextErr error inserted := 1 for nextRun.Before(now) || nextRun.Equal(now) { if inserted >= catchUpLimit { break } nextRun, nextErr = schedule.Next(nextRun, location) if nextErr != nil { return nextErr } if nextRun.After(now) { break } runID, idErr = platformid.NewUUID() if idErr != nil { return idErr } if _, err = tx.Exec(ctx, `INSERT INTO gateway.scheduled_task_runs(id,task_id,trigger_type,scheduled_for) VALUES($1,$2,'scheduled',$3) ON CONFLICT(task_id,trigger_type,scheduled_for) DO NOTHING`, runID, item.id, nextRun); err != nil { return err } inserted++ } next, nextErr := schedule.Next(now, location) if nextErr != nil { return nextErr } if _, err = tx.Exec(ctx, `UPDATE gateway.scheduled_tasks SET next_run_at=$2,updated_at=clock_timestamp() WHERE id=$1`, item.id, next); err != nil { return err } } return tx.Commit(ctx) } // catchUpLimit 是单任务单次补跑的最大执行数;超过部分丢弃,防止调度器长期 // 停机后瞬间插入海量补跑记录。 const catchUpLimit = 100 func (e *Engine) reclaim(ctx context.Context, now time.Time) error { cutoff := now.Add(-e.executionTTL) tx, err := e.service.pool.Begin(ctx) if err != nil { return err } defer func() { _ = tx.Rollback(ctx) }() if _, err = tx.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status='pending',worker_id='',started_at=NULL,error='上次执行超时,已回收重试' WHERE status='running' AND started_at < $1 AND attempts < $2`, cutoff, e.maxAttempts); err != nil { return err } rows, err := tx.Query(ctx, `WITH picked AS (SELECT id FROM gateway.scheduled_task_runs WHERE status='running' AND started_at < $1 AND attempts >= $2 ORDER BY started_at,id FOR UPDATE SKIP LOCKED LIMIT $3) UPDATE gateway.scheduled_task_runs r SET worker_id=$4,started_at=$5 FROM picked WHERE r.id=picked.id RETURNING r.id::text`, cutoff, e.maxAttempts, e.batchSize, e.workerID, now) if err != nil { return err } ids := []string{} for rows.Next() { var id string if err = rows.Scan(&id); err != nil { rows.Close() return err } ids = append(ids, id) } rows.Close() if err = rows.Err(); err != nil { return err } if err = tx.Commit(ctx); err != nil { return err } for _, id := range ids { run, getErr := e.service.getRun(ctx, id) if getErr != nil { return getErr } task, taskErr := e.service.Get(ctx, run.TaskID) if taskErr != nil { task = Task{ID: run.TaskID, Code: run.TaskCode} } if err = e.complete(ctx, run, task, nil, errors.New("执行超时且达到最大重试次数")); err != nil { return err } } return nil } func (e *Engine) claim(ctx context.Context) ([]Run, error) { tx, err := e.service.pool.Begin(ctx) if err != nil { return nil, err } defer func() { _ = tx.Rollback(ctx) }() rows, err := tx.Query(ctx, `WITH picked AS (SELECT id FROM gateway.scheduled_task_runs WHERE status='pending' ORDER BY created_at,id FOR UPDATE SKIP LOCKED LIMIT $1) UPDATE gateway.scheduled_task_runs r SET status='running',attempts=attempts+1,worker_id=$2,started_at=clock_timestamp(),error='' FROM picked WHERE r.id=picked.id RETURNING r.id::text`, e.batchSize, e.workerID) if err != nil { return nil, err } ids := []string{} for rows.Next() { var id string if err = rows.Scan(&id); err != nil { rows.Close() return nil, err } ids = append(ids, id) } rows.Close() if err = rows.Err(); err != nil { return nil, err } if err = tx.Commit(ctx); err != nil { return nil, err } runs := make([]Run, 0, len(ids)) for _, id := range ids { run, getErr := e.service.getRun(ctx, id) if getErr != nil { return nil, getErr } runs = append(runs, run) } return runs, nil } func (e *Engine) conversationMessages(ctx context.Context, task Task, currentRunID string) ([]map[string]any, error) { messages := []map[string]any{} if task.ConversationID != "" { rows, err := e.service.pool.Query(ctx, `SELECT response FROM gateway.scheduled_task_runs WHERE task_id=$1 AND id<>$2 AND status='success' AND response IS NOT NULL ORDER BY created_at DESC LIMIT 5`, task.ID, currentRunID) if err != nil { return nil, err } defer rows.Close() answers := []string{} for rows.Next() { var response json.RawMessage if err = rows.Scan(&response); err != nil { return nil, err } if answer := responseAnswer(response); answer != "" { answers = append(answers, answer) } } for i := len(answers) - 1; i >= 0; i-- { messages = append(messages, map[string]any{"role": "user", "content": task.Prompt}, map[string]any{"role": "assistant", "content": answers[i]}) } } messages = append(messages, map[string]any{"role": "user", "content": task.Prompt}) return messages, nil } func responseAnswer(raw json.RawMessage) string { var response map[string]any if json.Unmarshal(raw, &response) != nil { return "" } choices, _ := response["choices"].([]any) if len(choices) == 0 { return "" } choice, _ := choices[0].(map[string]any) message, _ := choice["message"].(map[string]any) answer, _ := message["content"].(string) return answer } func (e *Engine) execute(ctx context.Context, task Task, run Run) (json.RawMessage, error) { secret, err := e.service.decryptAPIKey(task) if err != nil { return nil, fmt.Errorf("解密执行 API Key: %w", err) } messages, err := e.conversationMessages(ctx, task, run.ID) if err != nil { return nil, err } var variables map[string]any if json.Unmarshal(task.Variables, &variables) != nil { variables = map[string]any{} } payload := map[string]any{"messages": messages, "variables": variables} path := "/v1/applications/" + task.TargetCode + "/chat/completions" if task.TargetType == "digital_employee" { path = "/v1/digital-employees/" + task.TargetCode + "/chat/completions" payload["skill_ids"] = task.SkillIDs payload["mcp_server_ids"] = task.MCPServerIDs } body, _ := json.Marshal(payload) request, err := http.NewRequestWithContext(ctx, http.MethodPost, e.baseURL+path, bytes.NewReader(body)) if err != nil { return nil, err } request.Header.Set("Content-Type", "application/json") request.Header.Set("X-Gateway-API-Key", secret) request.Header.Set("X-Request-ID", "scheduled-"+run.ID) response, err := e.client.Do(request) if err != nil { return nil, err } defer response.Body.Close() data, err := io.ReadAll(io.LimitReader(response.Body, maxExecutionResponseBytes+1)) if err != nil { return nil, err } if len(data) > maxExecutionResponseBytes { return nil, errors.New("模型响应超过 2MB 上限") } if !json.Valid(data) { return nil, errors.New("模型响应不是合法 JSON") } if response.StatusCode < 200 || response.StatusCode >= 300 { return nil, fmt.Errorf("模型调用失败(HTTP %d): %s", response.StatusCode, truncate(string(data), 1000)) } return json.RawMessage(data), nil } func truncate(value string, maximum int) string { value = strings.TrimSpace(value) if len(value) <= maximum { return value } cut := value[:maximum] // 按字节截断可能把多字节 rune 切半,产生无效 UTF-8;PostgreSQL text 列 // 会拒绝写入,导致任务永远停在重试循环。回退到最近的 rune 边界。 for len(cut) > 0 && !utf8.RuneStart(cut[len(cut)-1]) { cut = cut[:len(cut)-1] } return cut } func (e *Engine) retry(ctx context.Context, run Run, executeErr error) error { errorText := truncate(executeErr.Error(), 4000) tag, err := e.service.pool.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status='pending',worker_id='',started_at=NULL,response=NULL,error=$2 WHERE id=$1 AND status='running' AND worker_id=$3`, run.ID, errorText, e.workerID) if err != nil { return err } if tag.RowsAffected() == 0 { return errors.New("定时任务执行租约已失效") } e.logger.Warn("scheduled task execution will retry", "task", run.TaskCode, "run_id", run.ID, "attempt", run.Attempts, "error", executeErr) return nil } func (e *Engine) complete(ctx context.Context, run Run, task Task, response json.RawMessage, executeErr error) error { status := "success" errorText := "" eventType := "scheduled_task.completed" if executeErr != nil { status = "failed" eventType = "scheduled_task.failed" errorText = truncate(executeErr.Error(), 4000) e.logger.Warn("scheduled task execution failed", "task", task.Code, "run_id", run.ID, "error", executeErr) } tx, err := e.service.pool.Begin(ctx) if err != nil { return err } defer func() { _ = tx.Rollback(ctx) }() tag, err := tx.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status=$2,response=$3,error=$4,finished_at=clock_timestamp() WHERE id=$1 AND status='running' AND worker_id=$5`, run.ID, status, response, errorText, e.workerID) if err != nil { return err } if tag.RowsAffected() == 0 { return errors.New("定时任务执行租约已失效") } _, err = tx.Exec(ctx, `UPDATE gateway.scheduled_tasks SET last_run_at=clock_timestamp(),last_status=$2,last_error=$3,updated_at=clock_timestamp() WHERE id=$1`, task.ID, status, errorText) if err != nil { return err } eventID, err := platformid.NewUUID() if err != nil { return err } payload, _ := json.Marshal(map[string]any{"scheduled_task_id": task.ID, "task_code": task.Code, "run_id": run.ID, "status": status, "error": errorText, "actor_id": task.CreatedBy, "notification_channel_id": task.NotificationChannelID}) _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,'scheduled_task',$3,$4)`, eventID, eventType, run.ID, payload) if err != nil { return err } return tx.Commit(ctx) }