Agent-System-Backend/internal/store/workflow_step_source.go

package store

import (
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"agent-system-backend/pkg/extension"
)

type WorkflowStepSource struct {
	RunID, StepID, Owner, ProjectID, StepAttemptID                      string
	OutputSHA256                                                        [32]byte
	AttemptID, AuthorizationEpoch, PublicationEpoch, Watermark, Version string
}

// UnbindWorkflowStepSource removes an earlier attempt before the engine begins
// another execution of the same step. A failed/unknown retry cannot inherit a
// previous output's proof, even if the output bytes happen to be identical.
func (s *Store) UnbindWorkflowStepSource(owner, runID, stepID string) error {
	if owner == "" || runID == "" || stepID == "" {
		return errors.New("workflow step source scope is incomplete")
	}
	_, err := s.db.Exec(s.q(`DELETE FROM workflow_step_sources WHERE owner=? AND run_id=? AND step_id=?`), owner, runID, stepID)
	return err
}

// BindWorkflowStepSource associates a materialized prompt step with only the
// responded model receipts created during that step. The caller must have
// excluded parallel prompt execution, whose receipt attribution is ambiguous.
func (s *Store) BindWorkflowStepSource(owner, runID, stepID, stepAttemptID, kind string, startedAt time.Time, output []byte) (WorkflowStepSource, error) {
	if kind != "prompt" || owner == "" || runID == "" || stepID == "" || stepAttemptID == "" || startedAt.IsZero() {
		return WorkflowStepSource{}, errors.New("workflow step source requires a started prompt step")
	}
	run, err := s.GetRunByID(runID)
	if err != nil || run.Owner != owner || run.ProjectID == "" {
		return WorkflowStepSource{}, ErrNotFound
	}
	rows, err := s.db.Query(s.q(`SELECT created_at,manifest_json FROM context_manifests WHERE owner=? AND target_id=? AND delivery='responded'`), owner, runID)
	if err != nil {
		return WorkflowStepSource{}, err
	}
	defer rows.Close()
	var latest extension.ProjectContextManifest
	var latestAt time.Time
	for rows.Next() {
		var atRaw, raw string
		if err := rows.Scan(&atRaw, &raw); err != nil {
			return WorkflowStepSource{}, err
		}
		at, err := time.Parse(time.RFC3339Nano, atRaw)
		if err != nil {
			return WorkflowStepSource{}, err
		}
		if at.Before(startedAt) {
			continue
		}
		var m extension.ProjectContextManifest
		if err := json.Unmarshal([]byte(raw), &m); err != nil {
			return WorkflowStepSource{}, err
		}
		if m.WorkflowStepAttemptID != stepAttemptID {
			continue
		}
		if m.ProjectID != run.ProjectID || m.AuthorizationEpoch == "" || m.PublicationEpoch == "" || m.AttemptID == "" {
			return WorkflowStepSource{}, errors.New("workflow step model receipt lacks trusted project epochs")
		}
		if !latestAt.IsZero() && (m.ProjectID != latest.ProjectID || m.AuthorizationEpoch != latest.AuthorizationEpoch || m.PublicationEpoch != latest.PublicationEpoch || m.Watermark != latest.Watermark) {
			return WorkflowStepSource{}, errors.New("workflow step model receipts span project authorization revisions")
		}
		if latestAt.IsZero() || at.After(latestAt) {
			latest, latestAt = m, at
		}
	}
	if err := rows.Err(); err != nil {
		return WorkflowStepSource{}, err
	}
	if latestAt.IsZero() {
		return WorkflowStepSource{}, errors.New("workflow prompt output has no responded model receipt")
	}
	b := WorkflowStepSource{RunID: runID, StepID: stepID, Owner: owner, ProjectID: string(run.ProjectID), StepAttemptID: stepAttemptID, OutputSHA256: sha256.Sum256(output),
		AttemptID: latest.AttemptID, AuthorizationEpoch: latest.AuthorizationEpoch, PublicationEpoch: latest.PublicationEpoch, Watermark: latest.Watermark, Version: latest.Version}
	_, err = s.db.Exec(s.q(`INSERT INTO workflow_step_sources(run_id,step_id,owner,project_id,step_attempt_id,output_sha256,attempt_id,authorization_epoch,publication_epoch,watermark,version,created_at) VALUES(?,?,?,?,?,?,?,?,?,?,?,?) ON CONFLICT(run_id,step_id) DO UPDATE SET owner=excluded.owner,project_id=excluded.project_id,step_attempt_id=excluded.step_attempt_id,output_sha256=excluded.output_sha256,attempt_id=excluded.attempt_id,authorization_epoch=excluded.authorization_epoch,publication_epoch=excluded.publication_epoch,watermark=excluded.watermark,version=excluded.version,created_at=excluded.created_at`),
		runID, stepID, owner, b.ProjectID, b.StepAttemptID, hex.EncodeToString(b.OutputSHA256[:]), b.AttemptID, b.AuthorizationEpoch, b.PublicationEpoch, b.Watermark, b.Version, nowISO())
	if err != nil {
		return WorkflowStepSource{}, fmt.Errorf("bind workflow step output: %w", err)
	}
	return b, nil
}

func (s *Store) WorkflowStepSource(owner, runID, stepID string) (WorkflowStepSource, error) {
	var b WorkflowStepSource
	var hash string
	err := s.db.QueryRow(s.q(`SELECT run_id,step_id,owner,project_id,step_attempt_id,output_sha256,attempt_id,authorization_epoch,publication_epoch,watermark,version FROM workflow_step_sources WHERE owner=? AND run_id=? AND step_id=?`),
		owner, runID, stepID).Scan(&b.RunID, &b.StepID, &b.Owner, &b.ProjectID, &b.StepAttemptID, &hash, &b.AttemptID, &b.AuthorizationEpoch, &b.PublicationEpoch, &b.Watermark, &b.Version)
	if errors.Is(err, sql.ErrNoRows) {
		return WorkflowStepSource{}, ErrNotFound
	}
	if err != nil {
		return WorkflowStepSource{}, err
	}
	decoded, err := hex.DecodeString(hash)
	if err != nil || len(decoded) != len(b.OutputSHA256) {
		return WorkflowStepSource{}, errors.New("workflow step source has invalid output hash")
	}
	copy(b.OutputSHA256[:], decoded)
	return b, nil
}

func (s *Store) RespondedContextManifest(owner, runID, attemptID string) (extension.ProjectContextManifest, error) {
	var raw string
	err := s.db.QueryRow(s.q(`SELECT manifest_json FROM context_manifests WHERE owner=? AND target_id=? AND attempt_id=? AND delivery='responded'`), owner, runID, attemptID).Scan(&raw)
	if errors.Is(err, sql.ErrNoRows) {
		return extension.ProjectContextManifest{}, ErrNotFound
	}
	if err != nil {
		return extension.ProjectContextManifest{}, err
	}
	var m extension.ProjectContextManifest
	if err := json.Unmarshal([]byte(raw), &m); err != nil {
		return extension.ProjectContextManifest{}, err
	}
	return m, nil
}