Agent-System/internal/app/plan_owner.go

package app

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"sync"

	"chat-agent-system/internal/app/intentflow"
	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/taskrun"
)

// TaskRunAdapter is the U8 generate adapter. It records a stable idempotency
// key and result reference without changing single-intent taskrun HTTP.
type TaskRunAdapter struct {
	mu       sync.Mutex
	receipts map[string]intentflow.Receipt
	starts   int
	Denied   bool
	Starter  func(ctx context.Context, exec intentflow.StepExec) (resultRef string, err error)
}

func NewTaskRunAdapter() *TaskRunAdapter {
	return &TaskRunAdapter{receipts: map[string]intentflow.Receipt{}}
}

func (a *TaskRunAdapter) Authorize(_ context.Context, exec intentflow.StepExec) error {
	if a != nil && a.Denied {
		return intentflow.ErrPlanActionForbidden
	}
	if exec.ActionKey != "task.run" && exec.ActionKey != "task.run.adhoc" {
		return intentflow.ErrPlanActionForbidden
	}
	return nil
}

func (a *TaskRunAdapter) Execute(ctx context.Context, exec intentflow.StepExec) (intentflow.Receipt, error) {
	if err := a.Authorize(ctx, exec); err != nil {
		return intentflow.Receipt{}, err
	}
	a.mu.Lock()
	if rec, ok := a.receipts[exec.IdempotencyKey]; ok {
		a.mu.Unlock()
		if rec.ActionDigest != exec.ActionDigest {
			return intentflow.Receipt{}, intentflow.ErrPlanDecisionConflict
		}
		return rec, nil
	}
	a.mu.Unlock()
	ref := "taskrun:" + exec.IdempotencyKey
	if a.Starter != nil {
		got, err := a.Starter(ctx, exec)
		if err != nil {
			return intentflow.Receipt{}, err
		}
		if got != "" {
			ref = got
		}
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	if rec, ok := a.receipts[exec.IdempotencyKey]; ok {
		if rec.ActionDigest != exec.ActionDigest {
			return intentflow.Receipt{}, intentflow.ErrPlanDecisionConflict
		}
		return rec, nil
	}
	a.starts++
	rec := intentflow.Receipt{
		IdempotencyKey: exec.IdempotencyKey,
		ActionDigest:   exec.ActionDigest,
		ResultRef:      ref,
		ResultDigest:   "task:" + exec.ActionDigest,
		SafeSummary:    "task.run accepted",
	}
	if a.receipts == nil {
		a.receipts = map[string]intentflow.Receipt{}
	}
	a.receipts[exec.IdempotencyKey] = rec
	return rec, nil
}

// LiveTaskStarter claims a taskrun with the plan step idempotency key.
// When live is false it only occupies the run row and does not start PI.
func LiveTaskStarter(mgr *taskrun.Manager, live bool) func(context.Context, intentflow.StepExec) (string, error) {
	return func(ctx context.Context, exec intentflow.StepExec) (string, error) {
		if mgr == nil {
			return "", fmt.Errorf("task generate adapter: taskrun manager missing")
		}
		if existing, err := mgr.FindByIdempotencyKey(ctx, exec.IdempotencyKey); err != nil {
			return "", err
		} else if existing != nil && existing.RunID != "" {
			return existing.RunID, nil
		}
		taskKey := "intake"
		if exec.ActionKey == "task.run.adhoc" {
			taskKey = "adhoc"
		}
		rec := &taskrun.RunRecord{
			RunID:          planRunID(exec.IdempotencyKey),
			SessionID:      exec.Principal.SessionID,
			ProjectID:      exec.Principal.ProjectID,
			UserID:         exec.Principal.UserID,
			IdempotencyKey: exec.IdempotencyKey,
			TaskKey:        taskKey,
			Status:         "pending",
		}
		if !live {
			claimed, _, err := mgr.ClaimRun(ctx, rec)
			if err != nil {
				return "", err
			}
			return claimed.RunID, nil
		}
		created, err := mgr.CreateAndStart(ctx, rec.RunID, rec, bizdev.TaskRunRequest{
			TaskKey:   taskKey,
			UserID:    exec.Principal.UserID,
			ProjectID: exec.Principal.ProjectID,
			SessionID: exec.Principal.SessionID,
		}, bizdev.StageRunRequest{}, "")
		if err != nil {
			return "", err
		}
		return created.RunID, nil
	}
}

func planRunID(idempotencyKey string) string {
	sum := sha256.Sum256([]byte(idempotencyKey))
	return "run_plan_" + hex.EncodeToString(sum[:8])
}

func (a *TaskRunAdapter) LookupReceipt(_ context.Context, idempotencyKey string) (intentflow.Receipt, bool, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	rec, ok := a.receipts[idempotencyKey]
	return rec, ok, nil
}

func (a *TaskRunAdapter) StartCount() int {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.starts
}