517 lines
29 KiB
Go
517 lines
29 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"git.cin.su/m/Assistant/contracts/work"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
type WorkRepository struct{ pool *pgxpool.Pool }
|
|
|
|
var (
|
|
ErrDispatchIntentConflict = errors.New("dispatch intent idempotency conflict")
|
|
ErrAgentLeaseUnavailable = errors.New("agent lease unavailable")
|
|
ErrAgentCallbackRejected = errors.New("agent callback rejected")
|
|
)
|
|
|
|
func OpenWork(ctx context.Context, config ContinuityConfig) (*WorkRepository, error) {
|
|
pool, err := openPool(ctx, config)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &WorkRepository{pool}, nil
|
|
}
|
|
func (r *WorkRepository) Close() { r.pool.Close() }
|
|
func (r *WorkRepository) Ping(ctx context.Context) error { return r.pool.Ping(ctx) }
|
|
func (r *WorkRepository) Create(ctx context.Context, target work.TaskTarget) error {
|
|
_, err := r.pool.Exec(ctx, `INSERT INTO work.tasks (task_id,goal_id,scope_id,revision) VALUES ($1,$2,$3,$4) ON CONFLICT (task_id) DO NOTHING`, target.TaskID, target.GoalID, target.ScopeID, target.Revision)
|
|
return err
|
|
}
|
|
func (r *WorkRepository) Commit(ctx context.Context, target work.TaskTarget) (work.TaskRevision, error) {
|
|
var revision work.TaskRevision
|
|
err := r.pool.QueryRow(ctx, `UPDATE work.tasks SET revision=revision+1 WHERE task_id=$1 AND goal_id=$2 AND scope_id=$3 AND revision=$4 RETURNING revision`, target.TaskID, target.GoalID, target.ScopeID, target.Revision).Scan(&revision)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("commit task: %w", err)
|
|
}
|
|
return revision, nil
|
|
}
|
|
func (r *WorkRepository) Load(ctx context.Context, id work.TaskID) (work.TaskTarget, error) {
|
|
var target work.TaskTarget
|
|
err := r.pool.QueryRow(ctx, `SELECT task_id, goal_id, scope_id, revision FROM work.tasks WHERE task_id=$1`, id).Scan(&target.TaskID, &target.GoalID, &target.ScopeID, &target.Revision)
|
|
if err != nil {
|
|
return work.TaskTarget{}, fmt.Errorf("load task: %w", err)
|
|
}
|
|
return target, nil
|
|
}
|
|
|
|
func (r *WorkRepository) SaveDispatchIntent(ctx context.Context, intent work.AgentDispatchIntent) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
_, err = tx.Exec(ctx, `INSERT INTO work.agent_dispatch_intents(dispatch_id,idempotency_key,delegation_id,run_id,implementation_id,implementation_revision,source_revision,worktree_id,verification_plan_ref,created_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) ON CONFLICT(idempotency_key) DO NOTHING`, intent.DispatchID, intent.IdempotencyKey, intent.DelegationID, intent.RunID, intent.ImplementationID, intent.ImplementationRevision, intent.SourceRevision, intent.WorktreeID, intent.VerificationPlanRef, intent.CreatedAtUnixMillis)
|
|
if err != nil {
|
|
return fmt.Errorf("save dispatch intent: %w", err)
|
|
}
|
|
var existing work.AgentDispatchIntent
|
|
err = tx.QueryRow(ctx, `SELECT dispatch_id,idempotency_key,delegation_id,run_id,implementation_id,implementation_revision,source_revision,worktree_id,verification_plan_ref,created_at_unix_millis FROM work.agent_dispatch_intents WHERE idempotency_key=$1`, intent.IdempotencyKey).Scan(&existing.DispatchID, &existing.IdempotencyKey, &existing.DelegationID, &existing.RunID, &existing.ImplementationID, &existing.ImplementationRevision, &existing.SourceRevision, &existing.WorktreeID, &existing.VerificationPlanRef, &existing.CreatedAtUnixMillis)
|
|
if err != nil {
|
|
return fmt.Errorf("read dispatch intent: %w", err)
|
|
}
|
|
if existing != intent {
|
|
return ErrDispatchIntentConflict
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) LoadDispatchIntent(ctx context.Context, dispatchID string) (work.AgentDispatchIntent, error) {
|
|
var intent work.AgentDispatchIntent
|
|
err := r.pool.QueryRow(ctx, `SELECT dispatch_id,idempotency_key,delegation_id,run_id,implementation_id,implementation_revision,source_revision,worktree_id,verification_plan_ref,created_at_unix_millis FROM work.agent_dispatch_intents WHERE dispatch_id=$1`, dispatchID).Scan(&intent.DispatchID, &intent.IdempotencyKey, &intent.DelegationID, &intent.RunID, &intent.ImplementationID, &intent.ImplementationRevision, &intent.SourceRevision, &intent.WorktreeID, &intent.VerificationPlanRef, &intent.CreatedAtUnixMillis)
|
|
if err != nil {
|
|
return work.AgentDispatchIntent{}, fmt.Errorf("load dispatch intent: %w", err)
|
|
}
|
|
return intent, nil
|
|
}
|
|
|
|
func (r *WorkRepository) EnqueueDispatch(ctx context.Context, item work.AgentDispatchQueueItem) error {
|
|
_, err := r.pool.Exec(ctx, `INSERT INTO work.agent_dispatch_queue(dispatch_id,state,available_at_unix_millis,attempts,claimed_by,updated_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6) ON CONFLICT(dispatch_id) DO NOTHING`, item.DispatchID, item.State, item.AvailableAtUnixMillis, item.Attempts, item.ClaimedBy, item.UpdatedAtUnixMillis)
|
|
return err
|
|
}
|
|
|
|
func (r *WorkRepository) ClaimNextDispatch(ctx context.Context, workerID string, nowUnixMillis int64) (work.AgentDispatchQueueItem, error) {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return work.AgentDispatchQueueItem{}, err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
var item work.AgentDispatchQueueItem
|
|
err = tx.QueryRow(ctx, `SELECT dispatch_id,state,available_at_unix_millis,attempts,claimed_by,updated_at_unix_millis FROM work.agent_dispatch_queue WHERE state='QUEUED' AND available_at_unix_millis <= $1 ORDER BY available_at_unix_millis,dispatch_id FOR UPDATE SKIP LOCKED LIMIT 1`, nowUnixMillis).Scan(&item.DispatchID, &item.State, &item.AvailableAtUnixMillis, &item.Attempts, &item.ClaimedBy, &item.UpdatedAtUnixMillis)
|
|
if err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return work.AgentDispatchQueueItem{}, nil
|
|
}
|
|
return work.AgentDispatchQueueItem{}, err
|
|
}
|
|
item.State = work.AgentDispatchClaimed
|
|
item.Attempts++
|
|
item.ClaimedBy = workerID
|
|
item.UpdatedAtUnixMillis = nowUnixMillis
|
|
_, err = tx.Exec(ctx, `UPDATE work.agent_dispatch_queue SET state=$1,attempts=$2,claimed_by=$3,updated_at_unix_millis=$4 WHERE dispatch_id=$5 AND state='QUEUED'`, item.State, item.Attempts, item.ClaimedBy, item.UpdatedAtUnixMillis, item.DispatchID)
|
|
if err != nil {
|
|
return work.AgentDispatchQueueItem{}, err
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return work.AgentDispatchQueueItem{}, err
|
|
}
|
|
return item, nil
|
|
}
|
|
|
|
func (r *WorkRepository) WaitForDispatch(ctx context.Context, interval time.Duration) error {
|
|
if interval <= 0 {
|
|
return fmt.Errorf("invalid worker poll interval")
|
|
}
|
|
_, err := r.pool.Exec(ctx, `SELECT pg_sleep($1)`, interval.Seconds())
|
|
return err
|
|
}
|
|
|
|
func (r *WorkRepository) CompleteDispatch(ctx context.Context, dispatchID, workerID string, state work.AgentDispatchQueueState, nowUnixMillis int64) error {
|
|
if state != work.AgentDispatchCompleted && state != work.AgentDispatchFailed {
|
|
return fmt.Errorf("invalid terminal dispatch state: %s", state)
|
|
}
|
|
command, err := r.pool.Exec(ctx, `UPDATE work.agent_dispatch_queue SET state=$1,updated_at_unix_millis=$2 WHERE dispatch_id=$3 AND state='CLAIMED' AND claimed_by=$4`, state, nowUnixMillis, dispatchID, workerID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if command.RowsAffected() != 1 {
|
|
return ErrAgentCallbackRejected
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *WorkRepository) SaveWorkGraph(ctx context.Context, graph work.WorkGraph) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
if _, err := tx.Exec(ctx, `INSERT INTO work.graphs(graph_id,revision,state,max_parallel,max_attempts) VALUES($1,$2,$3,$4,$5) ON CONFLICT(graph_id,revision) DO NOTHING`, graph.ID, graph.Revision, graph.State, graph.Budget.MaxParallel, graph.Budget.MaxAttempts); err != nil {
|
|
return err
|
|
}
|
|
for _, node := range graph.Nodes {
|
|
if _, err := tx.Exec(ctx, `INSERT INTO work.graph_nodes(graph_id,revision,node_id,dispatch_id,idempotency_key,state,attempts) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(graph_id,revision,node_id) DO NOTHING`, graph.ID, graph.Revision, node.ID, node.DispatchID, node.IdempotencyKey, node.State, node.Attempts); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
for _, edge := range graph.Edges {
|
|
if _, err := tx.Exec(ctx, `INSERT INTO work.graph_edges(graph_id,revision,from_node_id,to_node_id) VALUES($1,$2,$3,$4) ON CONFLICT(graph_id,revision,from_node_id,to_node_id) DO NOTHING`, graph.ID, graph.Revision, edge.From, edge.To); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) LoadWorkGraph(ctx context.Context, graphID string, revision int64) (work.WorkGraph, error) {
|
|
var graph work.WorkGraph
|
|
err := r.pool.QueryRow(ctx, `SELECT graph_id,revision,state,max_parallel,max_attempts FROM work.graphs WHERE graph_id=$1 AND revision=$2`, graphID, revision).Scan(&graph.ID, &graph.Revision, &graph.State, &graph.Budget.MaxParallel, &graph.Budget.MaxAttempts)
|
|
if err != nil {
|
|
return work.WorkGraph{}, fmt.Errorf("load graph: %w", err)
|
|
}
|
|
rows, err := r.pool.Query(ctx, `SELECT node_id,dispatch_id,idempotency_key,state,attempts FROM work.graph_nodes WHERE graph_id=$1 AND revision=$2 ORDER BY node_id`, graphID, revision)
|
|
if err != nil {
|
|
return work.WorkGraph{}, err
|
|
}
|
|
for rows.Next() {
|
|
var node work.GraphNode
|
|
if err := rows.Scan(&node.ID, &node.DispatchID, &node.IdempotencyKey, &node.State, &node.Attempts); err != nil {
|
|
rows.Close()
|
|
return work.WorkGraph{}, err
|
|
}
|
|
graph.Nodes = append(graph.Nodes, node)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
rows.Close()
|
|
return work.WorkGraph{}, err
|
|
}
|
|
rows.Close()
|
|
edges, err := r.pool.Query(ctx, `SELECT from_node_id,to_node_id FROM work.graph_edges WHERE graph_id=$1 AND revision=$2 ORDER BY from_node_id,to_node_id`, graphID, revision)
|
|
if err != nil {
|
|
return work.WorkGraph{}, err
|
|
}
|
|
for edges.Next() {
|
|
var edge work.GraphEdge
|
|
if err := edges.Scan(&edge.From, &edge.To); err != nil {
|
|
edges.Close()
|
|
return work.WorkGraph{}, err
|
|
}
|
|
graph.Edges = append(graph.Edges, edge)
|
|
}
|
|
if err := edges.Err(); err != nil {
|
|
edges.Close()
|
|
return work.WorkGraph{}, err
|
|
}
|
|
edges.Close()
|
|
return graph, nil
|
|
}
|
|
|
|
func (r *WorkRepository) SaveGraphSynthesis(ctx context.Context, conflict work.ConflictSet, candidates []work.ResultCandidate, synthesis work.GraphSynthesisCandidate) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
for _, candidate := range candidates {
|
|
if _, err := tx.Exec(ctx, `INSERT INTO work.graph_result_candidates(result_candidate_id,graph_id,graph_revision,node_id,author_run_id,artifact_refs,evidence_summary,summary) VALUES($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT(result_candidate_id) DO NOTHING`, candidate.ID, candidate.GraphID, candidate.GraphRevision, candidate.NodeID, candidate.AuthorRunID, candidate.ArtifactRefs, candidate.EvidenceSummary, candidate.Summary); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if _, err := tx.Exec(ctx, `INSERT INTO work.graph_conflicts(conflict_id,graph_id,graph_revision,node_id,candidate_ids,resolved) VALUES($1,$2,$3,$4,$5,$6) ON CONFLICT(conflict_id) DO NOTHING`, conflict.ID, conflict.GraphID, conflict.GraphRevision, conflict.NodeID, conflict.CandidateIDs, conflict.Resolved); err != nil {
|
|
return err
|
|
}
|
|
if _, err := tx.Exec(ctx, `INSERT INTO work.graph_synthesis_candidates(synthesis_id,graph_id,graph_revision,node_id,candidate_ids,selected_candidate_id,dissent_candidate_ids,verifier_run_id,verifier_identity,evidence_refs,summary,state) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) ON CONFLICT(synthesis_id) DO NOTHING`, synthesis.ID, synthesis.GraphID, synthesis.GraphRevision, synthesis.NodeID, synthesis.CandidateIDs, synthesis.SelectedCandidateID, synthesis.DissentCandidateIDs, synthesis.VerifierRunID, synthesis.VerifierIdentity, synthesis.EvidenceRefs, synthesis.Summary, synthesis.State); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) RecoverStaleClaims(ctx context.Context, nowUnixMillis, staleAfterMillis int64) (int64, error) {
|
|
command, err := r.pool.Exec(ctx, `UPDATE work.agent_dispatch_queue SET state='QUEUED',claimed_by='',updated_at_unix_millis=$1 WHERE state='CLAIMED' AND updated_at_unix_millis <= $2`, nowUnixMillis, nowUnixMillis-staleAfterMillis)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return command.RowsAffected(), nil
|
|
}
|
|
|
|
func (r *WorkRepository) SaveAgentLease(ctx context.Context, lease work.AgentLease) error {
|
|
err := r.pool.QueryRow(ctx, `INSERT INTO work.agent_leases(lease_id,run_id,worker_id,fencing_token,state,acquired_at_unix_millis,expires_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(run_id) DO UPDATE SET lease_id=EXCLUDED.lease_id,worker_id=EXCLUDED.worker_id,fencing_token=EXCLUDED.fencing_token,state=EXCLUDED.state,acquired_at_unix_millis=EXCLUDED.acquired_at_unix_millis,expires_at_unix_millis=EXCLUDED.expires_at_unix_millis WHERE work.agent_leases.state='EXPIRED' OR work.agent_leases.expires_at_unix_millis <= EXCLUDED.acquired_at_unix_millis RETURNING run_id`, lease.ID, lease.RunID, lease.WorkerID, lease.FencingToken, lease.State, lease.AcquiredAt.UnixMilli(), lease.ExpiresAt.UnixMilli()).Scan(new(string))
|
|
if err != nil {
|
|
return ErrAgentLeaseUnavailable
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *WorkRepository) RenewAgentLease(ctx context.Context, lease work.AgentLease, expiresAt time.Time) error {
|
|
command, err := r.pool.Exec(ctx, `UPDATE work.agent_leases SET expires_at_unix_millis=$1 WHERE run_id=$2 AND worker_id=$3 AND fencing_token=$4 AND state='ACTIVE' AND expires_at_unix_millis > $5`, expiresAt.UnixMilli(), lease.RunID, lease.WorkerID, lease.FencingToken, time.Now().UnixMilli())
|
|
if err != nil || command.RowsAffected() != 1 {
|
|
return ErrAgentLeaseUnavailable
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *WorkRepository) AcceptAgentCallback(ctx context.Context, callback work.AgentCallback) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
var runID string
|
|
err = tx.QueryRow(ctx, `SELECT run_id FROM work.agent_leases WHERE run_id=$1 AND worker_id=$2 AND fencing_token=$3 AND state='ACTIVE' AND expires_at_unix_millis > $4 FOR UPDATE`, callback.RunID, callback.WorkerID, callback.FencingToken, callback.ObservedAt.UnixMilli()).Scan(&runID)
|
|
if err != nil {
|
|
return ErrAgentCallbackRejected
|
|
}
|
|
command, err := tx.Exec(ctx, `INSERT INTO work.agent_callback_heads(run_id,last_sequence,updated_at_unix_millis) VALUES($1,$2,$3) ON CONFLICT(run_id) DO UPDATE SET last_sequence=EXCLUDED.last_sequence,updated_at_unix_millis=EXCLUDED.updated_at_unix_millis WHERE work.agent_callback_heads.last_sequence < EXCLUDED.last_sequence`, runID, callback.Sequence, callback.ObservedAt.UnixMilli())
|
|
if err != nil || command.RowsAffected() != 1 {
|
|
return ErrAgentCallbackRejected
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) SavePlanningResult(ctx context.Context, result work.PlanningResult) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
goal := result.Goal
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.goals(goal_id,source_intent_id,source_utterance_id,actor_id,principal_id,scope_id,project_id,project_display_name,desired_outcome,constraints_text,state,provenance,revision,created_at_unix_millis,updated_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15) ON CONFLICT(goal_id) DO NOTHING`, goal.ID, goal.SourceIntentID, goal.SourceUtteranceID, goal.ActorID, goal.PrincipalID, goal.ScopeID, goal.ProjectID, goal.ProjectDisplayName, goal.DesiredOutcome, goal.Constraints, goal.State, goal.Provenance, goal.Revision, goal.CreatedAtUnixMillis, goal.UpdatedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
candidate := result.Candidate
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.plan_candidates(plan_candidate_id,goal_id,source,state,summary,validation_summary,requires_clarification,accepted_at_unix_millis,recorded_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9) ON CONFLICT(plan_candidate_id) DO NOTHING`, candidate.ID, candidate.GoalID, candidate.Source, candidate.State, candidate.Summary, candidate.ValidationSummary, candidate.RequiresClarification, candidate.AcceptedAtUnixMillis, candidate.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
plan := result.Plan
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.work_plans(work_plan_id,goal_id,accepted_candidate_id,revision,state,summary,recorded_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(work_plan_id) DO NOTHING`, plan.ID, plan.GoalID, plan.AcceptedCandidateID, plan.Revision, plan.State, plan.Summary, plan.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
for _, requirement := range candidate.Requirements {
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.capability_requirements(capability_requirement_id,plan_candidate_id,capability_kind,capability_id,purpose,state,epistemic_precondition,recorded_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT(capability_requirement_id) DO NOTHING`, requirement.ID, requirement.PlanCandidateID, requirement.CapabilityKind, requirement.CapabilityID, requirement.Purpose, requirement.State, requirement.EpistemicPrecondition, requirement.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
for _, step := range candidate.Steps {
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.plan_steps(plan_step_id,plan_candidate_id,ordinal,kind,description,requires_capability_id,recorded_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(plan_step_id) DO NOTHING`, step.ID, step.PlanCandidateID, step.Ordinal, step.Kind, step.Description, step.RequiresCapabilityID, step.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
for _, delegation := range candidate.Delegations {
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.delegation_candidates(delegation_candidate_id,plan_candidate_id,capability_requirement_id,kind,executor_hint,state,reason,recorded_at_unix_millis) VALUES($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT(delegation_candidate_id) DO NOTHING`, delegation.ID, delegation.PlanCandidateID, delegation.CapabilityRequirementID, delegation.Kind, delegation.ExecutorHint, delegation.State, delegation.Reason, delegation.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) SaveDelegationBundle(ctx context.Context, delegation work.Delegation, run work.AgentRun) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.delegations(
|
|
delegation_id,goal_id,work_plan_id,plan_candidate_id,capability_requirement_id,capability_kind,
|
|
selected_implementation_id,selected_implementation_revision,selected_implementation_kind,protocol,
|
|
source_intent_id,source_utterance_id,conversation_id,project_id,project_display_name,desired_outcome,
|
|
constraint_summary,verification_summary,context_summary,context_refs,state,latest_run_id,latest_result_id,
|
|
created_at_unix_millis,updated_at_unix_millis)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23,$24,$25)
|
|
ON CONFLICT(delegation_id) DO NOTHING`,
|
|
delegation.ID, delegation.GoalID, delegation.PlanID, delegation.PlanCandidateID, delegation.CapabilityRequirementID,
|
|
delegation.CapabilityKind, delegation.SelectedImplementationID, delegation.SelectedImplementationRev,
|
|
delegation.SelectedImplementationKind, delegation.Protocol, delegation.SourceIntentID, delegation.SourceUtteranceID,
|
|
delegation.ConversationID, delegation.ProjectID, delegation.ProjectDisplayName, delegation.DesiredOutcome,
|
|
delegation.ConstraintSummary, delegation.VerificationSummary, delegation.ContextSummary, delegation.ContextRefs,
|
|
delegation.State, delegation.LatestRunID, delegation.LatestResultID, delegation.CreatedAtUnixMillis, delegation.UpdatedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.agent_runs(
|
|
run_id,delegation_id,implementation_id,implementation_revision,protocol,state,derived_instruction_summary,
|
|
started_at_unix_millis,updated_at_unix_millis,completed_at_unix_millis)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) ON CONFLICT(run_id) DO NOTHING`,
|
|
run.ID, run.DelegationID, run.ImplementationID, run.ImplementationRevision, run.Protocol, run.State,
|
|
run.DerivedInstructionSummary, run.StartedAtUnixMillis, run.UpdatedAtUnixMillis, run.CompletedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) SaveAgentCheckpoint(ctx context.Context, run work.AgentRun, checkpoint work.AgentCheckpoint) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
if _, err = tx.Exec(ctx, `UPDATE work.agent_runs
|
|
SET state=$2,updated_at_unix_millis=$3,completed_at_unix_millis=$4
|
|
WHERE run_id=$1`, run.ID, run.State, run.UpdatedAtUnixMillis, run.CompletedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `UPDATE work.delegations
|
|
SET state=$2,updated_at_unix_millis=$3,latest_run_id=$4
|
|
WHERE delegation_id=$1`, checkpoint.DelegationID, delegationStateForRun(run.State), checkpoint.RecordedAtUnixMillis, run.ID); err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.agent_checkpoints(
|
|
checkpoint_id,delegation_id,run_id,state,blocker_class,completed_summary,remaining_summary,test_summary,
|
|
branch_name,commit_hash,next_step,artifact_refs,recorded_at_unix_millis)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13) ON CONFLICT(checkpoint_id) DO NOTHING`,
|
|
checkpoint.ID, checkpoint.DelegationID, checkpoint.RunID, checkpoint.State, checkpoint.BlockerClass,
|
|
checkpoint.CompletedSummary, checkpoint.RemainingSummary, checkpoint.TestSummary, checkpoint.BranchName,
|
|
checkpoint.CommitHash, checkpoint.NextStep, checkpoint.ArtifactRefs, checkpoint.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) SaveAgentResult(ctx context.Context, run work.AgentRun, result work.AgentResult) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
if _, err = tx.Exec(ctx, `UPDATE work.agent_runs
|
|
SET state=$2,updated_at_unix_millis=$3,completed_at_unix_millis=$4
|
|
WHERE run_id=$1`, run.ID, run.State, run.UpdatedAtUnixMillis, run.CompletedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `UPDATE work.delegations
|
|
SET state=$2,updated_at_unix_millis=$3,latest_run_id=$4,latest_result_id=$5
|
|
WHERE delegation_id=$1`, result.DelegationID, work.DelegationVerifying, result.RecordedAtUnixMillis, run.ID, result.ID); err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.agent_results(
|
|
result_id,delegation_id,run_id,state,summary,outcome_ref,verification_required,verification_summary,artifact_refs,recorded_at_unix_millis)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) ON CONFLICT(result_id) DO NOTHING`,
|
|
result.ID, result.DelegationID, result.RunID, result.State, result.Summary, result.OutcomeRef, result.VerificationRequired,
|
|
result.VerificationSummary, result.ArtifactRefs, result.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) SaveResumedRun(ctx context.Context, delegation work.Delegation, run work.AgentRun) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
if _, err = tx.Exec(ctx, `UPDATE work.delegations
|
|
SET state=$2,updated_at_unix_millis=$3,latest_run_id=$4
|
|
WHERE delegation_id=$1`, delegation.ID, delegation.State, delegation.UpdatedAtUnixMillis, delegation.LatestRunID); err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.agent_runs(
|
|
run_id,delegation_id,implementation_id,implementation_revision,protocol,state,derived_instruction_summary,
|
|
started_at_unix_millis,updated_at_unix_millis,completed_at_unix_millis)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) ON CONFLICT(run_id) DO NOTHING`,
|
|
run.ID, run.DelegationID, run.ImplementationID, run.ImplementationRevision, run.Protocol, run.State,
|
|
run.DerivedInstructionSummary, run.StartedAtUnixMillis, run.UpdatedAtUnixMillis, run.CompletedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) SaveDelegationVerification(ctx context.Context, delegation work.Delegation, verification work.DelegationVerification) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
if _, err = tx.Exec(ctx, `UPDATE work.delegations
|
|
SET state=$2,updated_at_unix_millis=$3
|
|
WHERE delegation_id=$1`, delegation.ID, delegation.State, delegation.UpdatedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `INSERT INTO work.delegation_verifications(
|
|
verification_id,delegation_id,run_id,result_id,verified,summary,evidence_refs,recorded_at_unix_millis)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT(verification_id) DO NOTHING`,
|
|
verification.ID, verification.DelegationID, verification.RunID, verification.ResultID, verification.Verified,
|
|
verification.Summary, verification.EvidenceRefs, verification.RecordedAtUnixMillis); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *WorkRepository) LatestAgentResult(ctx context.Context, delegationID work.DelegationID) (work.AgentResult, error) {
|
|
var value work.AgentResult
|
|
err := r.pool.QueryRow(ctx, `SELECT result_id,delegation_id,run_id,state,summary,outcome_ref,verification_required,verification_summary,artifact_refs,recorded_at_unix_millis
|
|
FROM work.agent_results WHERE delegation_id=$1 ORDER BY recorded_at_unix_millis DESC, result_id DESC LIMIT 1`, delegationID).
|
|
Scan(&value.ID, &value.DelegationID, &value.RunID, &value.State, &value.Summary, &value.OutcomeRef, &value.VerificationRequired, &value.VerificationSummary, &value.ArtifactRefs, &value.RecordedAtUnixMillis)
|
|
if err != nil {
|
|
return work.AgentResult{}, fmt.Errorf("load latest agent result: %w", err)
|
|
}
|
|
return value, nil
|
|
}
|
|
|
|
func (r *WorkRepository) CountAgentRuns(ctx context.Context, delegationID work.DelegationID) (int, error) {
|
|
var count int
|
|
err := r.pool.QueryRow(ctx, `SELECT COUNT(*) FROM work.agent_runs WHERE delegation_id=$1`, delegationID).Scan(&count)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("count agent runs: %w", err)
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
func (r *WorkRepository) LatestAgentCheckpoint(ctx context.Context, delegationID work.DelegationID) (work.AgentCheckpoint, error) {
|
|
var value work.AgentCheckpoint
|
|
err := r.pool.QueryRow(ctx, `SELECT checkpoint_id,delegation_id,run_id,state,blocker_class,completed_summary,remaining_summary,test_summary,branch_name,commit_hash,next_step,artifact_refs,recorded_at_unix_millis
|
|
FROM work.agent_checkpoints WHERE delegation_id=$1 ORDER BY recorded_at_unix_millis DESC, checkpoint_id DESC LIMIT 1`, delegationID).
|
|
Scan(&value.ID, &value.DelegationID, &value.RunID, &value.State, &value.BlockerClass, &value.CompletedSummary, &value.RemainingSummary,
|
|
&value.TestSummary, &value.BranchName, &value.CommitHash, &value.NextStep, &value.ArtifactRefs, &value.RecordedAtUnixMillis)
|
|
if err != nil {
|
|
return work.AgentCheckpoint{}, err
|
|
}
|
|
return value, nil
|
|
}
|
|
|
|
func (r *WorkRepository) LoadDelegation(ctx context.Context, id work.DelegationID) (work.Delegation, error) {
|
|
var value work.Delegation
|
|
err := r.pool.QueryRow(ctx, `SELECT delegation_id,goal_id,work_plan_id,plan_candidate_id,capability_requirement_id,capability_kind,
|
|
selected_implementation_id,selected_implementation_revision,selected_implementation_kind,protocol,
|
|
source_intent_id,source_utterance_id,conversation_id,project_id,project_display_name,desired_outcome,constraint_summary,
|
|
verification_summary,context_summary,context_refs,state,latest_run_id,latest_result_id,created_at_unix_millis,updated_at_unix_millis
|
|
FROM work.delegations WHERE delegation_id=$1`, id).
|
|
Scan(&value.ID, &value.GoalID, &value.PlanID, &value.PlanCandidateID, &value.CapabilityRequirementID, &value.CapabilityKind,
|
|
&value.SelectedImplementationID, &value.SelectedImplementationRev, &value.SelectedImplementationKind, &value.Protocol,
|
|
&value.SourceIntentID, &value.SourceUtteranceID, &value.ConversationID, &value.ProjectID, &value.ProjectDisplayName,
|
|
&value.DesiredOutcome, &value.ConstraintSummary, &value.VerificationSummary, &value.ContextSummary, &value.ContextRefs,
|
|
&value.State, &value.LatestRunID, &value.LatestResultID, &value.CreatedAtUnixMillis, &value.UpdatedAtUnixMillis)
|
|
if err != nil {
|
|
return work.Delegation{}, fmt.Errorf("load delegation: %w", err)
|
|
}
|
|
return value, nil
|
|
}
|
|
|
|
func (r *WorkRepository) LoadAgentRun(ctx context.Context, id work.AgentRunID) (work.AgentRun, error) {
|
|
var value work.AgentRun
|
|
err := r.pool.QueryRow(ctx, `SELECT run_id,delegation_id,implementation_id,implementation_revision,protocol,state,derived_instruction_summary,started_at_unix_millis,updated_at_unix_millis,completed_at_unix_millis
|
|
FROM work.agent_runs WHERE run_id=$1`, id).
|
|
Scan(&value.ID, &value.DelegationID, &value.ImplementationID, &value.ImplementationRevision, &value.Protocol, &value.State, &value.DerivedInstructionSummary,
|
|
&value.StartedAtUnixMillis, &value.UpdatedAtUnixMillis, &value.CompletedAtUnixMillis)
|
|
if err != nil {
|
|
return work.AgentRun{}, fmt.Errorf("load agent run: %w", err)
|
|
}
|
|
return value, nil
|
|
}
|
|
|
|
func delegationStateForRun(state work.AgentRunState) work.DelegationState {
|
|
switch state {
|
|
case work.AgentRunStarted, work.AgentRunRunning:
|
|
return work.DelegationRunning
|
|
case work.AgentRunWaiting, work.AgentRunCheckpointed, work.AgentRunExecutionWindowExhausted:
|
|
return work.DelegationWaiting
|
|
case work.AgentRunBlocked:
|
|
return work.DelegationBlocked
|
|
case work.AgentRunCompleted:
|
|
return work.DelegationVerifying
|
|
case work.AgentRunCancelled:
|
|
return work.DelegationCancelled
|
|
case work.AgentRunFailed:
|
|
return work.DelegationFailed
|
|
default:
|
|
return work.DelegationPending
|
|
}
|
|
}
|
|
|
|
func latestRowErr(err error) error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|