Agent-System/internal/app/intentflow/memory.go

package intentflow

import (
	"context"
	"fmt"
	"slices"
	"sync"
	"time"

	"chat-agent-system/internal/app/uow"
)

type serialScope struct{ ctx context.Context }

func (s serialScope) Context() context.Context { return s.ctx }

// SerialUoW serializes commands. Memory tests use it so CAS races have one winner.
type SerialUoW struct{ mu sync.Mutex }

func (s *SerialUoW) Within(ctx context.Context, fn func(uow.TxScope) error) error {
	if s == nil {
		return fmt.Errorf("intentflow: serial uow required")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return fn(serialScope{ctx: ctx})
}

type MemoryStore struct {
	mu        sync.Mutex
	plans     map[string]*Plan
	turn      map[string]string
	decisions map[string]*Decision
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		plans:     map[string]*Plan{},
		turn:      map[string]string{},
		decisions: map[string]*Decision{},
	}
}

func turnKey(env, session, turn string) string {
	return env + "\x1f" + session + "\x1f" + turn
}

func nonceKey(planID, nonce string) string {
	return planID + "\x1f" + nonce
}

func (m *MemoryStore) InsertPlan(ctx context.Context, scope uow.TxScope, plan *Plan) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if plan == nil || plan.PlanID == "" {
		return fmt.Errorf("intentflow: plan required")
	}
	tk := turnKey(plan.Principal.Environment, plan.Principal.SessionID, plan.ClientTurnID)
	if existingID, ok := m.turn[tk]; ok {
		_ = existingID
		return errTurnExists
	}
	if plan.WritePlan && isActiveStatus(plan.Status) {
		for _, p := range m.plans {
			if sameScope(p.Principal, plan.Principal) && p.WritePlan && isActiveStatus(p.Status) {
				return uow.ErrBusy
			}
		}
	}
	cloned := clonePlan(plan)
	m.plans[plan.PlanID] = cloned
	m.turn[tk] = plan.PlanID
	return nil
}

func (m *MemoryStore) GetPlan(ctx context.Context, scope uow.TxScope, planID string) (*Plan, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	p := m.plans[planID]
	if p == nil {
		return nil, nil
	}
	return clonePlan(p), nil
}

func (m *MemoryStore) GetPlanForUpdate(ctx context.Context, scope uow.TxScope, planID string) (*Plan, error) {
	return m.GetPlan(ctx, scope, planID)
}

func (m *MemoryStore) GetByTurn(ctx context.Context, scope uow.TxScope, env, sessionID, clientTurnID string) (*Plan, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	id := m.turn[turnKey(env, sessionID, clientTurnID)]
	if id == "" {
		return nil, nil
	}
	return clonePlan(m.plans[id]), nil
}

func (m *MemoryStore) GetActiveWrite(ctx context.Context, scope uow.TxScope, p Principal) (*Plan, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, plan := range m.plans {
		if sameScope(plan.Principal, p) && plan.WritePlan && isActiveStatus(plan.Status) {
			return clonePlan(plan), nil
		}
	}
	return nil, nil
}

func (m *MemoryStore) GetActive(ctx context.Context, scope uow.TxScope, p Principal) (*Plan, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var best *Plan
	for _, plan := range m.plans {
		if !sameScope(plan.Principal, p) || !isActiveStatus(plan.Status) {
			continue
		}
		if best == nil || plan.UpdatedAt.After(best.UpdatedAt) {
			cp := clonePlan(plan)
			best = cp
		}
	}
	return best, nil
}

func (m *MemoryStore) CASPlan(ctx context.Context, scope uow.TxScope, planID string, expectedRev int64, from []PlanStatus, to PlanStatus, now time.Time) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	p := m.plans[planID]
	if p == nil || p.Revision != expectedRev || !slices.Contains(from, p.Status) {
		return 0, nil
	}
	p.Status = to
	p.Revision++
	p.UpdatedAt = now
	return p.Revision, nil
}

func (m *MemoryStore) InsertDecision(ctx context.Context, scope uow.TxScope, d *Decision) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	k := nonceKey(d.PlanID, d.Nonce)
	if _, ok := m.decisions[k]; ok {
		return uow.ErrDigestConflict
	}
	cp := *d
	cp.StepIDs = slices.Clone(d.StepIDs)
	m.decisions[k] = &cp
	return nil
}

func (m *MemoryStore) GetDecisionByNonce(ctx context.Context, scope uow.TxScope, planID, nonce string) (*Decision, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	d := m.decisions[nonceKey(planID, nonce)]
	if d == nil {
		return nil, nil
	}
	cp := *d
	cp.StepIDs = slices.Clone(d.StepIDs)
	return &cp, nil
}

func (m *MemoryStore) ReplaceSteps(ctx context.Context, scope uow.TxScope, planID string, steps []Step) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	p := m.plans[planID]
	if p == nil {
		return fmt.Errorf("intentflow: plan %s not found", planID)
	}
	p.Steps = cloneSteps(steps)
	p.UpdatedAt = time.Now().UTC()
	return nil
}

func (m *MemoryStore) ClaimStep(ctx context.Context, scope uow.TxScope, planID, stepID, worker string, until, now time.Time, idempotencyKey string) (*Step, int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	p := m.plans[planID]
	if p == nil {
		return nil, 0, nil
	}
	for i := range p.Steps {
		s := &p.Steps[i]
		if s.StepID != stepID {
			continue
		}
		claimable := s.Status == StepEligible ||
			(s.Status == StepRunning && (s.LeaseUntil == nil || s.LeaseUntil.Before(now)))
		if !claimable {
			return nil, 0, nil
		}
		p.LastFencingToken++
		s.Status = StepRunning
		s.LeaseOwner = worker
		untilCopy := until
		s.LeaseUntil = &untilCopy
		s.Attempt++
		s.Version++
		s.FencingToken = p.LastFencingToken
		if s.IdempotencyKey == "" {
			s.IdempotencyKey = idempotencyKey
		}
		s.UpdatedAt = now
		p.UpdatedAt = now
		return cloneStep(s), s.FencingToken, nil
	}
	return nil, 0, nil
}

func (m *MemoryStore) WriteStepReceipt(ctx context.Context, scope uow.TxScope, worker string, token int64, step Step, now time.Time) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	p := m.plans[step.PlanID]
	if p == nil {
		return false, nil
	}
	for i := range p.Steps {
		s := &p.Steps[i]
		if s.StepID != step.StepID {
			continue
		}
		if s.Status == StepSucceeded && s.ResultDigest == step.ResultDigest {
			return true, nil
		}
		if s.Status != StepRunning || s.LeaseOwner != worker || s.FencingToken != token {
			return false, nil
		}
		s.Status = StepSucceeded
		s.ResultRef = step.ResultRef
		s.ResultDigest = step.ResultDigest
		s.SafeSummary = step.SafeSummary
		s.LeaseOwner = ""
		s.LeaseUntil = nil
		s.Version++
		s.UpdatedAt = now
		p.UpdatedAt = now
		return true, nil
	}
	return false, nil
}

func (m *MemoryStore) ExpireWaiting(ctx context.Context, scope uow.TxScope, now time.Time) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	n := 0
	for _, p := range m.plans {
		if !isWaitingStatus(p.Status) || !now.After(p.ExpiresAt) {
			continue
		}
		p.Status = PlanExpired
		p.Revision++
		p.UpdatedAt = now
		for i := range p.Steps {
			if p.Steps[i].Status == StepPending || p.Steps[i].Status == StepEligible {
				p.Steps[i].Status = StepCancelled
				p.Steps[i].UpdatedAt = now
			}
		}
		n++
	}
	return n, nil
}

func (m *MemoryStore) PurgeClosed(ctx context.Context, scope uow.TxScope, before time.Time) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	n := 0
	for id, p := range m.plans {
		switch p.Status {
		case PlanCompleted, PlanFailed, PlanCancelled, PlanExpired, PlanSuperseded:
		default:
			continue
		}
		if !p.UpdatedAt.Before(before) {
			continue
		}
		delete(m.turn, turnKey(p.Principal.Environment, p.Principal.SessionID, p.ClientTurnID))
		for k, d := range m.decisions {
			if d.PlanID == id {
				delete(m.decisions, k)
			}
		}
		delete(m.plans, id)
		n++
	}
	return n, nil
}

func sameScope(a, b Principal) bool {
	return a.Environment == b.Environment && a.ProjectID == b.ProjectID && a.SessionID == b.SessionID && a.UserID == b.UserID
}

func clonePlan(p *Plan) *Plan {
	if p == nil {
		return nil
	}
	cp := *p
	cp.DisplaySnapshot = slices.Clone(p.DisplaySnapshot)
	cp.Steps = cloneSteps(p.Steps)
	return &cp
}

func cloneSteps(in []Step) []Step {
	out := make([]Step, len(in))
	for i := range in {
		out[i] = *cloneStep(&in[i])
	}
	return out
}

func cloneStep(s *Step) *Step {
	cp := *s
	cp.Slots = slices.Clone(s.Slots)
	cp.DependsOn = slices.Clone(s.DependsOn)
	if s.LeaseUntil != nil {
		t := *s.LeaseUntil
		cp.LeaseUntil = &t
	}
	return &cp
}