weichenandCursor
7466a15319
Use deliverablesRef + limited poll for post-task server docs; pass reconnectRunId across auto-retry; GetRun reads idempotency_key; StartClaimedRun fails closed on UpdateRun; MessageBubble/E2E scope this-turn assistant; correct audit retry/schema19 wording. Co-authored-by: Cursor <cursoragent@cursor.com>
504 lines
14 KiB
Go
504 lines
14 KiB
Go
package taskrun
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
|
|
"chat-agent-system/internal/stream"
|
|
)
|
|
|
|
var ErrStoreNotAvailable = errors.New("taskrun: store not available")
|
|
var ErrIdempotencyConflict = errors.New("taskrun: idempotency key already exists")
|
|
|
|
// Store 持久化 task run 与 SSE 事件，支持重启后查询与重放。
|
|
type Store interface {
|
|
Available() bool
|
|
InsertRun(ctx context.Context, rec *RunRecord) error
|
|
UpdateRun(ctx context.Context, rec *RunRecord) error
|
|
GetByIdempotencyKey(ctx context.Context, key string) (*RunRecord, error)
|
|
AppendEvent(ctx context.Context, runID string, seq int64, ev stream.StreamEvent) error
|
|
GetRun(ctx context.Context, runID string) (*RunRecord, error)
|
|
ListEvents(ctx context.Context, runID string, afterSeq int64) ([]StoredEvent, error)
|
|
AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error
|
|
HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error
|
|
ReleaseLease(ctx context.Context, runID, workerID string) error
|
|
ResetForResume(ctx context.Context, runID string, attempt int) error
|
|
MarkOrphanedRunning(ctx context.Context) (int, error)
|
|
}
|
|
|
|
// PGStore PostgreSQL 实现。
|
|
type PGStore struct {
|
|
pool *pgxpool.Pool
|
|
}
|
|
|
|
func NewPGStore(pool *pgxpool.Pool) *PGStore {
|
|
return &PGStore{pool: pool}
|
|
}
|
|
|
|
func (s *PGStore) Available() bool {
|
|
if s == nil || s.pool == nil {
|
|
return false
|
|
}
|
|
var ok bool
|
|
_ = s.pool.QueryRow(context.Background(), `
|
|
SELECT EXISTS (
|
|
SELECT 1 FROM information_schema.tables
|
|
WHERE table_schema = 'public' AND table_name = 'task_runs'
|
|
)
|
|
`).Scan(&ok)
|
|
return ok
|
|
}
|
|
|
|
func (s *PGStore) hasLeaseColumns(ctx context.Context) bool {
|
|
var ok bool
|
|
_ = s.pool.QueryRow(ctx, `
|
|
SELECT EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id'
|
|
)
|
|
`).Scan(&ok)
|
|
return ok
|
|
}
|
|
|
|
func (s *PGStore) InsertRun(ctx context.Context, rec *RunRecord) error {
|
|
if !s.Available() {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
if rec == nil || rec.RunID == "" {
|
|
return fmt.Errorf("taskrun: run_id required")
|
|
}
|
|
attempt := rec.Attempt
|
|
if attempt <= 0 {
|
|
attempt = 1
|
|
}
|
|
idemKey := nullScopeText(rec.IdempotencyKey)
|
|
if s.hasLeaseColumns(ctx) {
|
|
var leaseHeartbeat *time.Time
|
|
if rec.LeaseExpiresAt != nil {
|
|
now := time.Now()
|
|
leaseHeartbeat = &now
|
|
}
|
|
_, err := s.pool.Exec(ctx, `
|
|
INSERT INTO task_runs (
|
|
run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
|
|
task_key, stage, status, workspace_id, assistant_message_id,
|
|
error_message, worker_id, lease_expires_at, lease_heartbeat_at, attempt,
|
|
workflow_instance_id, stage_instance_id, workflow_stage_key,
|
|
created_at, updated_at
|
|
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21)
|
|
ON CONFLICT (run_id) DO NOTHING
|
|
`,
|
|
rec.RunID, rec.SessionID, nullInt64(rec.ProjectID), rec.UserID, rec.ClientTurnID, idemKey,
|
|
rec.TaskKey, rec.Stage, rec.Status, rec.WorkspaceID, nullInt64(rec.AssistantMsgID),
|
|
rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, leaseHeartbeat, attempt,
|
|
nullInt64(rec.WorkflowInstanceID), nullInt64(rec.StageInstanceID), nullScopeText(rec.WorkflowStageKey),
|
|
rec.CreatedAt, rec.UpdatedAt,
|
|
)
|
|
return mapIdempotencyConflict(err)
|
|
}
|
|
_, err := s.pool.Exec(ctx, `
|
|
INSERT INTO task_runs (
|
|
run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
|
|
task_key, stage, status, workspace_id, assistant_message_id,
|
|
error_message, workflow_instance_id, stage_instance_id, workflow_stage_key,
|
|
created_at, updated_at
|
|
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)
|
|
ON CONFLICT (run_id) DO NOTHING
|
|
`,
|
|
rec.RunID, rec.SessionID, nullInt64(rec.ProjectID), rec.UserID, rec.ClientTurnID, idemKey,
|
|
rec.TaskKey, rec.Stage, rec.Status, rec.WorkspaceID, nullInt64(rec.AssistantMsgID),
|
|
rec.ErrorMessage, nullInt64(rec.WorkflowInstanceID), nullInt64(rec.StageInstanceID), nullScopeText(rec.WorkflowStageKey),
|
|
rec.CreatedAt, rec.UpdatedAt,
|
|
)
|
|
return mapIdempotencyConflict(err)
|
|
}
|
|
|
|
func mapIdempotencyConflict(err error) error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
var pgErr *pgconn.PgError
|
|
if errors.As(err, &pgErr) && pgErr.Code == "23505" {
|
|
return fmt.Errorf("%w: %v", ErrIdempotencyConflict, err)
|
|
}
|
|
msg := strings.ToLower(err.Error())
|
|
if strings.Contains(msg, "idempotency_key") && (strings.Contains(msg, "duplicate") || strings.Contains(msg, "unique") || strings.Contains(msg, "23505")) {
|
|
return fmt.Errorf("%w: %v", ErrIdempotencyConflict, err)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (s *PGStore) UpdateRun(ctx context.Context, rec *RunRecord) error {
|
|
if !s.Available() {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
if rec == nil || rec.RunID == "" {
|
|
return fmt.Errorf("taskrun: run_id required")
|
|
}
|
|
if s.hasLeaseColumns(ctx) {
|
|
_, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs SET
|
|
status = $2,
|
|
workspace_id = $3,
|
|
error_message = $4,
|
|
worker_id = $5,
|
|
lease_expires_at = $6,
|
|
assistant_message_id = $7,
|
|
updated_at = $8
|
|
WHERE run_id = $1
|
|
`, rec.RunID, rec.Status, rec.WorkspaceID, rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, nullInt64(rec.AssistantMsgID), rec.UpdatedAt)
|
|
return err
|
|
}
|
|
_, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs SET
|
|
status = $2,
|
|
workspace_id = $3,
|
|
error_message = $4,
|
|
assistant_message_id = $5,
|
|
updated_at = $6
|
|
WHERE run_id = $1
|
|
`, rec.RunID, rec.Status, rec.WorkspaceID, rec.ErrorMessage, nullInt64(rec.AssistantMsgID), rec.UpdatedAt)
|
|
return err
|
|
}
|
|
|
|
func (s *PGStore) AppendEvent(ctx context.Context, runID string, seq int64, ev stream.StreamEvent) error {
|
|
if !s.Available() {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
payload, err := json.Marshal(ev)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = s.pool.Exec(ctx, `
|
|
INSERT INTO task_run_events (run_id, seq, event_type, payload)
|
|
VALUES ($1, $2, $3, $4::jsonb)
|
|
ON CONFLICT (run_id, seq) DO NOTHING
|
|
`, runID, seq, string(ev.Type), payload)
|
|
return err
|
|
}
|
|
|
|
func (s *PGStore) GetRun(ctx context.Context, runID string) (*RunRecord, error) {
|
|
if !s.Available() {
|
|
return nil, ErrStoreNotAvailable
|
|
}
|
|
var rec RunRecord
|
|
var projectID, assistantMsgID *int64
|
|
var idempotencyKey *string
|
|
if s.hasLeaseColumns(ctx) {
|
|
err := s.pool.QueryRow(ctx, `
|
|
SELECT run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
|
|
task_key, stage, status, workspace_id, assistant_message_id,
|
|
error_message, worker_id, lease_expires_at, attempt, created_at, updated_at
|
|
FROM task_runs
|
|
WHERE run_id = $1
|
|
`, runID).Scan(
|
|
&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, &idempotencyKey,
|
|
&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
|
|
&rec.ErrorMessage, &rec.WorkerID, &rec.LeaseExpiresAt, &rec.Attempt,
|
|
&rec.CreatedAt, &rec.UpdatedAt,
|
|
)
|
|
if err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, fmt.Errorf("run not found")
|
|
}
|
|
return nil, err
|
|
}
|
|
} else {
|
|
err := s.pool.QueryRow(ctx, `
|
|
SELECT run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
|
|
task_key, stage, status, workspace_id, assistant_message_id,
|
|
error_message, created_at, updated_at
|
|
FROM task_runs
|
|
WHERE run_id = $1
|
|
`, runID).Scan(
|
|
&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, &idempotencyKey,
|
|
&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
|
|
&rec.ErrorMessage, &rec.CreatedAt, &rec.UpdatedAt,
|
|
)
|
|
if err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, fmt.Errorf("run not found")
|
|
}
|
|
return nil, err
|
|
}
|
|
}
|
|
if projectID != nil {
|
|
rec.ProjectID = *projectID
|
|
}
|
|
if assistantMsgID != nil {
|
|
rec.AssistantMsgID = *assistantMsgID
|
|
}
|
|
if idempotencyKey != nil {
|
|
rec.IdempotencyKey = *idempotencyKey
|
|
}
|
|
return &rec, nil
|
|
}
|
|
|
|
// GetByIdempotencyKey 按稳定幂等键查找既有 run；无键或未命中返回 (nil, nil)。
|
|
func (s *PGStore) GetByIdempotencyKey(ctx context.Context, key string) (*RunRecord, error) {
|
|
if !s.Available() {
|
|
return nil, ErrStoreNotAvailable
|
|
}
|
|
if key == "" {
|
|
return nil, nil
|
|
}
|
|
var runID string
|
|
err := s.pool.QueryRow(ctx, `
|
|
SELECT run_id FROM task_runs
|
|
WHERE idempotency_key = $1
|
|
LIMIT 1
|
|
`, key).Scan(&runID)
|
|
if err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
return s.GetRun(ctx, runID)
|
|
}
|
|
|
|
func (s *PGStore) ListEvents(ctx context.Context, runID string, afterSeq int64) ([]StoredEvent, error) {
|
|
if !s.Available() {
|
|
return nil, ErrStoreNotAvailable
|
|
}
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT seq, payload
|
|
FROM task_run_events
|
|
WHERE run_id = $1 AND seq > $2
|
|
ORDER BY seq ASC
|
|
`, runID, afterSeq)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
out := make([]StoredEvent, 0)
|
|
for rows.Next() {
|
|
var seq int64
|
|
var payload []byte
|
|
if err := rows.Scan(&seq, &payload); err != nil {
|
|
return nil, err
|
|
}
|
|
var ev stream.StreamEvent
|
|
if err := json.Unmarshal(payload, &ev); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, StoredEvent{Seq: seq, Event: ev})
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (s *PGStore) AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error {
|
|
if !s.Available() {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
if !s.hasLeaseColumns(ctx) {
|
|
return nil
|
|
}
|
|
expires := time.Now().Add(ttl)
|
|
tag, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs
|
|
SET worker_id = $2,
|
|
lease_expires_at = $3,
|
|
lease_heartbeat_at = now(),
|
|
updated_at = now()
|
|
WHERE run_id = $1 AND status = 'running'
|
|
`, runID, workerID, expires)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return fmt.Errorf("taskrun: lease acquire failed for %s", runID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *PGStore) HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error {
|
|
if !s.Available() {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
if !s.hasLeaseColumns(ctx) {
|
|
return nil
|
|
}
|
|
expires := time.Now().Add(ttl)
|
|
tag, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs
|
|
SET lease_expires_at = $3,
|
|
lease_heartbeat_at = now(),
|
|
updated_at = now()
|
|
WHERE run_id = $1 AND worker_id = $2 AND status = 'running'
|
|
`, runID, workerID, expires)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return fmt.Errorf("taskrun: heartbeat rejected for %s", runID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *PGStore) ReleaseLease(ctx context.Context, runID, workerID string) error {
|
|
if !s.Available() {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
if !s.hasLeaseColumns(ctx) {
|
|
return nil
|
|
}
|
|
_, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs
|
|
SET worker_id = '',
|
|
lease_expires_at = NULL,
|
|
lease_heartbeat_at = NULL,
|
|
updated_at = now()
|
|
WHERE run_id = $1 AND (worker_id = $2 OR worker_id = '')
|
|
`, runID, workerID)
|
|
return err
|
|
}
|
|
|
|
func (s *PGStore) ResetForResume(ctx context.Context, runID string, attempt int) error {
|
|
if !s.Available() {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
tag, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs
|
|
SET status = 'running',
|
|
error_message = '',
|
|
attempt = $2,
|
|
updated_at = now()
|
|
WHERE run_id = $1 AND status = 'interrupted'
|
|
`, runID, attempt)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return fmt.Errorf("taskrun: run %s not resumable", runID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *PGStore) MarkOrphanedRunning(ctx context.Context) (int, error) {
|
|
if !s.Available() {
|
|
return 0, ErrStoreNotAvailable
|
|
}
|
|
if s.hasLeaseColumns(ctx) {
|
|
tag, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs
|
|
SET status = 'interrupted',
|
|
error_message = 'worker lease 过期，任务已中断（进程重启或 worker 失联）',
|
|
worker_id = '',
|
|
lease_expires_at = NULL,
|
|
lease_heartbeat_at = NULL,
|
|
updated_at = now()
|
|
WHERE status = 'running'
|
|
AND lease_expires_at IS NOT NULL
|
|
AND lease_expires_at < now() - interval '2 seconds'
|
|
`)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return int(tag.RowsAffected()), nil
|
|
}
|
|
tag, err := s.pool.Exec(ctx, `
|
|
UPDATE task_runs
|
|
SET status = 'interrupted',
|
|
error_message = '进程重启导致任务中断，请重新提交',
|
|
updated_at = now()
|
|
WHERE status = 'running'
|
|
`)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return int(tag.RowsAffected()), nil
|
|
}
|
|
|
|
func nullInt64(v int64) *int64 {
|
|
if v <= 0 {
|
|
return nil
|
|
}
|
|
return &v
|
|
}
|
|
|
|
func nullScopeText(v string) *string {
|
|
if v == "" {
|
|
return nil
|
|
}
|
|
return &v
|
|
}
|
|
|
|
// TablesReady 检查 task_runs 表是否已迁移。
|
|
func TablesReady(ctx context.Context, pool *pgxpool.Pool) bool {
|
|
if pool == nil {
|
|
return false
|
|
}
|
|
var ok bool
|
|
_ = pool.QueryRow(ctx, `
|
|
SELECT EXISTS (
|
|
SELECT 1 FROM information_schema.tables
|
|
WHERE table_schema = 'public' AND table_name = 'task_runs'
|
|
)
|
|
`).Scan(&ok)
|
|
return ok
|
|
}
|
|
|
|
// LeaseColumnsReady 检查 worker lease 列是否已迁移。
|
|
func LeaseColumnsReady(ctx context.Context, pool *pgxpool.Pool) bool {
|
|
if pool == nil {
|
|
return false
|
|
}
|
|
var ok bool
|
|
_ = pool.QueryRow(ctx, `
|
|
SELECT EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id'
|
|
)
|
|
`).Scan(&ok)
|
|
return ok
|
|
}
|
|
|
|
// noopStore 无持久化时的空实现。
|
|
type noopStore struct{}
|
|
|
|
func (noopStore) Available() bool { return false }
|
|
func (noopStore) InsertRun(context.Context, *RunRecord) error {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) UpdateRun(context.Context, *RunRecord) error { return ErrStoreNotAvailable }
|
|
func (noopStore) AppendEvent(context.Context, string, int64, stream.StreamEvent) error {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) GetRun(context.Context, string) (*RunRecord, error) {
|
|
return nil, ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) GetByIdempotencyKey(context.Context, string) (*RunRecord, error) {
|
|
return nil, ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) ListEvents(context.Context, string, int64) ([]StoredEvent, error) {
|
|
return nil, ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) AcquireLease(context.Context, string, string, time.Duration) error {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) HeartbeatLease(context.Context, string, string, time.Duration) error {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) ReleaseLease(context.Context, string, string) error { return ErrStoreNotAvailable }
|
|
func (noopStore) ResetForResume(context.Context, string, int) error {
|
|
return ErrStoreNotAvailable
|
|
}
|
|
func (noopStore) MarkOrphanedRunning(context.Context) (int, error) { return 0, ErrStoreNotAvailable }
|
|
|
|
// EnsureUpdatedAt 若 UpdatedAt 为零则设为当前时间。
|
|
func EnsureUpdatedAt(rec *RunRecord) {
|
|
if rec != nil && rec.UpdatedAt.IsZero() {
|
|
rec.UpdatedAt = time.Now()
|
|
}
|
|
}
|