295 lines
16 KiB
Go
295 lines
16 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"git.cin.su/m/Assistant/contracts/work"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
type WorkRepository struct{ pool *pgxpool.Pool }
|
|
|
|
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) 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) 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
|
|
}
|