superidou
9501751792
三轮审查修复(60+ 项),相对远端 main(b536672)的关键变更:
- 安全: 数据面 SSRF 拨号防护(防 DNS rebinding)/上游凭据剥离/登录防枚举
与锁定态统一/可信代理(X-Forwarded-For)限流加固/会话版本失效机制/
撤销即时传播/弱密钥拒绝启动/脱敏字节级重写(保签名契约)
- 业务逻辑: 裸 body 上传 panic/bootstrap 审计管线卡死/定价通配符优先级/
全局工具可见性/调度器停机补跑/TOTP 挑战令牌消费顺序/熔断探针语义/
>4MB 响应 token 计量/管理员重置密码作废会话 等
- 前端: 新 logo(语枢 AI 网关主题)/Provider 凭据异常警示/删除入口/
后端错误消息透传/localStorage 敏感数据收敛
- 部署: CREDENTIAL_MASTER_KEY 持久化与弱值拒绝/Provider DELETE 接口/
nginx 安全头/worker 内存限制
- 新增迁移 000029(key_hash 索引)/000030(usage_daily 币种维度)
404 lines
13 KiB
Go
404 lines
13 KiB
Go
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)
|
||
}
|