weichenandCursor
aac42936bc
Normalize empty TEXT[] binds, isolate step-changed aggregates from plan revision uniqueness, propagate AdvanceEligible errors, and emit taskrun.completed to the registered notify consumer so finalize no longer rolls back live receipts. Co-authored-by: Cursor <cursoragent@cursor.com>
495 lines
14 KiB
Go
495 lines
14 KiB
Go
package intentflow
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/app/uow"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
)
|
|
|
|
// Repository is the PostgreSQL owner of intent_plans*.
|
|
type Repository struct{}
|
|
|
|
func NewRepository() Repository { return Repository{} }
|
|
|
|
func txFrom(scope uow.TxScope) (pgx.Tx, error) {
|
|
return uow.AsPG(scope)
|
|
}
|
|
|
|
func isUniqueViolation(err error) bool {
|
|
var pgErr *pgconn.PgError
|
|
return errors.As(err, &pgErr) && pgErr.Code == "23505"
|
|
}
|
|
|
|
func uniqueConstraint(err error) string {
|
|
var pgErr *pgconn.PgError
|
|
if errors.As(err, &pgErr) {
|
|
return pgErr.ConstraintName
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (Repository) InsertPlan(ctx context.Context, scope uow.TxScope, plan *Plan) error {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
snap := plan.DisplaySnapshot
|
|
if len(snap) == 0 {
|
|
snap = nil
|
|
}
|
|
_, err = tx.Exec(ctx, `
|
|
INSERT INTO intent_plans (
|
|
plan_id, environment, project_id, session_id, user_id, client_turn_id, request_digest,
|
|
segmenter_version, policy_version, status, revision, expires_at, last_fencing_token,
|
|
write_plan, execution_mode, safe_summary, display_snapshot, created_at, updated_at
|
|
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19)`,
|
|
plan.PlanID, plan.Principal.Environment, plan.Principal.ProjectID, plan.Principal.SessionID,
|
|
plan.Principal.UserID, plan.ClientTurnID, plan.RequestDigest, plan.SegmenterVersion,
|
|
plan.PolicyVersion, string(plan.Status), plan.Revision, plan.ExpiresAt, plan.LastFencingToken,
|
|
plan.WritePlan, string(plan.ExecutionMode), plan.SafeSummary, jsonbOrNil(snap),
|
|
plan.CreatedAt, plan.UpdatedAt,
|
|
)
|
|
if err != nil {
|
|
if isUniqueViolation(err) {
|
|
if uniqueConstraint(err) == "uq_intent_plans_one_active_write" {
|
|
return uow.ErrBusy
|
|
}
|
|
return errTurnExists
|
|
}
|
|
return err
|
|
}
|
|
for _, s := range plan.Steps {
|
|
if err := insertStep(ctx, tx, s); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func insertStep(ctx context.Context, tx pgx.Tx, s Step) error {
|
|
slots := s.Slots
|
|
if len(slots) == 0 {
|
|
slots = []byte("{}")
|
|
}
|
|
dependsOn := textArrayOrEmpty(s.DependsOn)
|
|
_, err := tx.Exec(ctx, `
|
|
INSERT INTO intent_plan_steps (
|
|
step_id, plan_id, ordinal, action_key, slots, depends_on, status, expected_resource_version,
|
|
action_digest, idempotency_key, attempt, result_ref, result_digest, safe_summary,
|
|
lease_owner, lease_until, fencing_token, version, execution_mode, 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)`,
|
|
s.StepID, s.PlanID, s.Ordinal, s.ActionKey, slots, dependsOn, string(s.Status),
|
|
s.ExpectedResourceVersion, s.ActionDigest, s.IdempotencyKey, s.Attempt, s.ResultRef,
|
|
s.ResultDigest, s.SafeSummary, s.LeaseOwner, s.LeaseUntil, s.FencingToken, s.Version,
|
|
string(s.ExecutionMode), s.CreatedAt, s.UpdatedAt,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (r Repository) GetPlan(ctx context.Context, scope uow.TxScope, planID string) (*Plan, error) {
|
|
return r.loadPlan(ctx, scope, planID, false)
|
|
}
|
|
|
|
func (r Repository) GetPlanForUpdate(ctx context.Context, scope uow.TxScope, planID string) (*Plan, error) {
|
|
return r.loadPlan(ctx, scope, planID, true)
|
|
}
|
|
|
|
func (Repository) loadPlan(ctx context.Context, scope uow.TxScope, planID string, forUpdate bool) (*Plan, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
q := planSelect + " WHERE plan_id=$1"
|
|
if forUpdate {
|
|
q += " FOR UPDATE"
|
|
}
|
|
plan, err := scanPlan(tx.QueryRow(ctx, q, planID))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
steps, err := loadSteps(ctx, tx, plan.PlanID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
plan.Steps = steps
|
|
return plan, nil
|
|
}
|
|
|
|
func (Repository) GetByTurn(ctx context.Context, scope uow.TxScope, env, sessionID, clientTurnID string) (*Plan, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
plan, err := scanPlan(tx.QueryRow(ctx, planSelect+`
|
|
WHERE environment=$1 AND session_id=$2 AND client_turn_id=$3`, env, sessionID, clientTurnID))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
steps, err := loadSteps(ctx, tx, plan.PlanID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
plan.Steps = steps
|
|
return plan, nil
|
|
}
|
|
|
|
func (Repository) GetActiveWrite(ctx context.Context, scope uow.TxScope, p Principal) (*Plan, error) {
|
|
return loadActive(ctx, scope, p, true)
|
|
}
|
|
|
|
func (Repository) GetActive(ctx context.Context, scope uow.TxScope, p Principal) (*Plan, error) {
|
|
return loadActive(ctx, scope, p, false)
|
|
}
|
|
|
|
func loadActive(ctx context.Context, scope uow.TxScope, p Principal, writeOnly bool) (*Plan, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
q := planSelect + `
|
|
WHERE environment=$1 AND project_id=$2 AND session_id=$3 AND user_id=$4
|
|
AND status = ANY($5)`
|
|
args := []any{p.Environment, p.ProjectID, p.SessionID, p.UserID, activeStatusStrings()}
|
|
if writeOnly {
|
|
q += ` AND write_plan`
|
|
}
|
|
q += ` ORDER BY updated_at DESC LIMIT 1`
|
|
plan, err := scanPlan(tx.QueryRow(ctx, q, args...))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
steps, err := loadSteps(ctx, tx, plan.PlanID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
plan.Steps = steps
|
|
return plan, nil
|
|
}
|
|
|
|
func (Repository) CASPlan(ctx context.Context, scope uow.TxScope, planID string, expectedRev int64, from []PlanStatus, to PlanStatus, now time.Time) (int64, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
var newRev int64
|
|
err = tx.QueryRow(ctx, `
|
|
UPDATE intent_plans
|
|
SET status=$1, revision=revision+1, updated_at=$2
|
|
WHERE plan_id=$3 AND revision=$4 AND status = ANY($5)
|
|
RETURNING revision`, string(to), now, planID, expectedRev, statusStrings(from)).Scan(&newRev)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return 0, nil
|
|
}
|
|
return newRev, err
|
|
}
|
|
|
|
func (Repository) InsertDecision(ctx context.Context, scope uow.TxScope, d *Decision) error {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = tx.Exec(ctx, `
|
|
INSERT INTO intent_plan_decisions (
|
|
decision_id, plan_id, nonce, plan_revision, actor_user_id, kind, step_ids,
|
|
exclude_unselected, decision_digest, outcome, created_at
|
|
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`,
|
|
d.DecisionID, d.PlanID, d.Nonce, d.PlanRevision, d.ActorUserID, string(d.Kind),
|
|
textArrayOrEmpty(d.StepIDs), d.ExcludeUnselected, d.DecisionDigest, d.Outcome, d.CreatedAt,
|
|
)
|
|
if isUniqueViolation(err) {
|
|
return uow.ErrDigestConflict
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (Repository) GetDecisionByNonce(ctx context.Context, scope uow.TxScope, planID, nonce string) (*Decision, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var d Decision
|
|
var kind string
|
|
err = tx.QueryRow(ctx, `
|
|
SELECT decision_id, plan_id, nonce, plan_revision, actor_user_id, kind, step_ids,
|
|
exclude_unselected, decision_digest, outcome, created_at
|
|
FROM intent_plan_decisions WHERE plan_id=$1 AND nonce=$2`, planID, nonce).Scan(
|
|
&d.DecisionID, &d.PlanID, &d.Nonce, &d.PlanRevision, &d.ActorUserID, &kind,
|
|
&d.StepIDs, &d.ExcludeUnselected, &d.DecisionDigest, &d.Outcome, &d.CreatedAt,
|
|
)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
d.Kind = DecisionKind(kind)
|
|
return &d, nil
|
|
}
|
|
|
|
func (Repository) ReplaceSteps(ctx context.Context, scope uow.TxScope, planID string, steps []Step) error {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, s := range steps {
|
|
if s.PlanID == "" {
|
|
s.PlanID = planID
|
|
}
|
|
_, err := tx.Exec(ctx, `
|
|
UPDATE intent_plan_steps
|
|
SET status=$1, depends_on=$2, expected_resource_version=$3, action_digest=$4,
|
|
idempotency_key=$5, attempt=$6, result_ref=$7, result_digest=$8, safe_summary=$9,
|
|
lease_owner=$10, lease_until=$11, fencing_token=$12, version=$13, execution_mode=$14,
|
|
updated_at=$15, slots=$16
|
|
WHERE step_id=$17 AND plan_id=$18`,
|
|
string(s.Status), textArrayOrEmpty(s.DependsOn), s.ExpectedResourceVersion, s.ActionDigest,
|
|
s.IdempotencyKey, s.Attempt, s.ResultRef, s.ResultDigest, s.SafeSummary,
|
|
s.LeaseOwner, s.LeaseUntil, s.FencingToken, s.Version, string(s.ExecutionMode),
|
|
s.UpdatedAt, jsonbOrNil(s.Slots), s.StepID, planID,
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (Repository) ClaimStep(ctx context.Context, scope uow.TxScope, planID, stepID, worker string, until, now time.Time, idempotencyKey string) (*Step, int64, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
var token int64
|
|
err = tx.QueryRow(ctx, `
|
|
UPDATE intent_plans
|
|
SET last_fencing_token = last_fencing_token + 1, updated_at = $2
|
|
WHERE plan_id = $1
|
|
AND EXISTS (
|
|
SELECT 1 FROM intent_plan_steps s
|
|
WHERE s.plan_id = $1 AND s.step_id = $3
|
|
AND (
|
|
s.status = 'eligible'
|
|
OR (s.status = 'running' AND (s.lease_until IS NULL OR s.lease_until < $2))
|
|
)
|
|
)
|
|
RETURNING last_fencing_token`, planID, now, stepID).Scan(&token)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, 0, nil
|
|
}
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
row := tx.QueryRow(ctx, `
|
|
UPDATE intent_plan_steps
|
|
SET status='running',
|
|
lease_owner=$1,
|
|
lease_until=$2,
|
|
attempt=attempt+1,
|
|
version=version+1,
|
|
fencing_token=$3,
|
|
idempotency_key=CASE WHEN idempotency_key='' THEN $4 ELSE idempotency_key END,
|
|
updated_at=$5
|
|
WHERE step_id=$6 AND plan_id=$7
|
|
AND (
|
|
status='eligible'
|
|
OR (status='running' AND (lease_until IS NULL OR lease_until < $5))
|
|
)
|
|
RETURNING `+stepColumns, worker, until, token, idempotencyKey, now, stepID, planID)
|
|
step, err := scanStep(row)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, 0, nil
|
|
}
|
|
return step, token, err
|
|
}
|
|
|
|
func (Repository) WriteStepReceipt(ctx context.Context, scope uow.TxScope, worker string, token int64, step Step, now time.Time) (bool, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
tag, err := tx.Exec(ctx, `
|
|
UPDATE intent_plan_steps
|
|
SET status='succeeded', result_ref=$1, result_digest=$2, safe_summary=$3,
|
|
lease_owner='', lease_until=NULL, version=version+1, updated_at=$4
|
|
WHERE step_id=$5 AND plan_id=$6 AND fencing_token=$7 AND lease_owner=$8 AND status='running'`,
|
|
step.ResultRef, step.ResultDigest, step.SafeSummary, now, step.StepID, step.PlanID, token, worker)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if tag.RowsAffected() == 1 {
|
|
return true, nil
|
|
}
|
|
var status, digest string
|
|
err = tx.QueryRow(ctx, `SELECT status, result_digest FROM intent_plan_steps WHERE step_id=$1 AND plan_id=$2`,
|
|
step.StepID, step.PlanID).Scan(&status, &digest)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return false, nil
|
|
}
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
return status == string(StepSucceeded) && digest == step.ResultDigest, nil
|
|
}
|
|
|
|
func (Repository) ExpireWaiting(ctx context.Context, scope uow.TxScope, now time.Time) (int, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
rows, err := tx.Query(ctx, `
|
|
UPDATE intent_plans
|
|
SET status='expired', revision=revision+1, updated_at=$1
|
|
WHERE status = ANY($2) AND expires_at <= $1
|
|
RETURNING plan_id`, now, waitingStatusStrings())
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer rows.Close()
|
|
var ids []string
|
|
for rows.Next() {
|
|
var id string
|
|
if err := rows.Scan(&id); err != nil {
|
|
return 0, err
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
if len(ids) == 0 {
|
|
return 0, nil
|
|
}
|
|
_, err = tx.Exec(ctx, `
|
|
UPDATE intent_plan_steps
|
|
SET status='cancelled', updated_at=$1
|
|
WHERE plan_id = ANY($2) AND status IN ('pending','eligible')`, now, ids)
|
|
return len(ids), err
|
|
}
|
|
|
|
func (Repository) PurgeClosed(ctx context.Context, scope uow.TxScope, before time.Time) (int, error) {
|
|
tx, err := txFrom(scope)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
tag, err := tx.Exec(ctx, `
|
|
DELETE FROM intent_plans
|
|
WHERE updated_at < $1
|
|
AND status IN ('completed','failed','cancelled','expired','superseded')`, before)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return int(tag.RowsAffected()), nil
|
|
}
|
|
|
|
const planSelect = `
|
|
SELECT plan_id, environment, project_id, session_id, user_id, client_turn_id, request_digest,
|
|
segmenter_version, policy_version, status, revision, expires_at, last_fencing_token,
|
|
write_plan, execution_mode, safe_summary, display_snapshot, created_at, updated_at
|
|
FROM intent_plans`
|
|
|
|
const stepColumns = `
|
|
step_id, plan_id, ordinal, action_key, slots, depends_on, status, expected_resource_version,
|
|
action_digest, idempotency_key, attempt, result_ref, result_digest, safe_summary,
|
|
lease_owner, lease_until, fencing_token, version, execution_mode, created_at, updated_at`
|
|
|
|
func loadSteps(ctx context.Context, tx pgx.Tx, planID string) ([]Step, error) {
|
|
rows, err := tx.Query(ctx, `SELECT `+stepColumns+` FROM intent_plan_steps WHERE plan_id=$1 ORDER BY ordinal`, planID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []Step
|
|
for rows.Next() {
|
|
s, err := scanStep(rows)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, *s)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
type scannable interface {
|
|
Scan(dest ...any) error
|
|
}
|
|
|
|
func scanPlan(row scannable) (*Plan, error) {
|
|
var p Plan
|
|
var status, mode string
|
|
var snap []byte
|
|
err := row.Scan(
|
|
&p.PlanID, &p.Principal.Environment, &p.Principal.ProjectID, &p.Principal.SessionID,
|
|
&p.Principal.UserID, &p.ClientTurnID, &p.RequestDigest, &p.SegmenterVersion,
|
|
&p.PolicyVersion, &status, &p.Revision, &p.ExpiresAt, &p.LastFencingToken,
|
|
&p.WritePlan, &mode, &p.SafeSummary, &snap, &p.CreatedAt, &p.UpdatedAt,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
p.Status = PlanStatus(status)
|
|
p.ExecutionMode = ExecutionMode(mode)
|
|
p.DisplaySnapshot = snap
|
|
return &p, nil
|
|
}
|
|
|
|
func scanStep(row scannable) (*Step, error) {
|
|
var s Step
|
|
var status, mode string
|
|
err := row.Scan(
|
|
&s.StepID, &s.PlanID, &s.Ordinal, &s.ActionKey, &s.Slots, &s.DependsOn, &status,
|
|
&s.ExpectedResourceVersion, &s.ActionDigest, &s.IdempotencyKey, &s.Attempt,
|
|
&s.ResultRef, &s.ResultDigest, &s.SafeSummary, &s.LeaseOwner, &s.LeaseUntil,
|
|
&s.FencingToken, &s.Version, &mode, &s.CreatedAt, &s.UpdatedAt,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
s.Status = StepStatus(status)
|
|
s.ExecutionMode = ExecutionMode(mode)
|
|
return &s, nil
|
|
}
|
|
|
|
func jsonbOrNil(raw []byte) []byte {
|
|
if len(raw) == 0 {
|
|
return []byte("{}")
|
|
}
|
|
if !json.Valid(raw) {
|
|
return []byte("{}")
|
|
}
|
|
return raw
|
|
}
|
|
|
|
func textArrayOrEmpty(vals []string) []string {
|
|
if vals == nil {
|
|
return []string{}
|
|
}
|
|
return vals
|
|
}
|
|
|
|
func statusStrings(in []PlanStatus) []string {
|
|
out := make([]string, len(in))
|
|
for i, s := range in {
|
|
out[i] = string(s)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func activeStatusStrings() []string { return statusStrings(activeStatuses()) }
|
|
|
|
func waitingStatusStrings() []string { return statusStrings(waitingStatuses()) }
|