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

package intentflow

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"
	"unicode/utf8"
)

const SegmentProposalSchemaV1 = "intentflow.segment_proposal.v1"

var (
	ErrSegmentProposalDisabled    = errors.New("SEGMENT_PROPOSAL_DISABLED")
	ErrSegmentProposalMalformed   = errors.New("SEGMENT_PROPOSAL_MALFORMED")
	ErrSegmentProposalTimeout     = errors.New("SEGMENT_PROPOSAL_TIMEOUT")
	ErrSegmentProposalOutOfBounds = errors.New("SEGMENT_PROPOSAL_OUT_OF_BOUNDS")
)

// SegmentSpan is a half-open byte range into the original utterance.
type SegmentSpan struct {
	Start int    `json:"start"`
	End   int    `json:"end"`
	Text  string `json:"text,omitempty"`
}

// SegmentProposalRequest is the minimal U9 input.
type SegmentProposalRequest struct {
	Text          string
	RequestDigest string
	MaxSegments   int
	SchemaVersion string
}

// SegmentProposal is a versioned, digest-bound suggestion.
type SegmentProposal struct {
	SchemaVersion string        `json:"schema_version"`
	Digest        string        `json:"digest"`
	Spans         []SegmentSpan `json:"spans"`
}

// SegmentProposer proposes spans only when rules cannot decide (U9).
type SegmentProposer interface {
	Propose(ctx context.Context, req SegmentProposalRequest) (*SegmentProposal, error)
}

type disabledSegmentProposer struct{}

func (disabledSegmentProposer) Propose(context.Context, SegmentProposalRequest) (*SegmentProposal, error) {
	return nil, ErrSegmentProposalDisabled
}

// ActiveSegmentProposer returns disabled when MULTI_INTENT_U9_SEGMENT is off.
func ActiveSegmentProposer() SegmentProposer {
	if !CurrentFeatureFlags().U9Segment {
		return disabledSegmentProposer{}
	}
	return NewStubSegmentProposer(0)
}

// StubSegmentProposer is scaffolding (no LLM): schema/span validation + timeout fail-closed.
type StubSegmentProposer struct {
	ArtificialDelay time.Duration
}

func NewStubSegmentProposer(delay time.Duration) *StubSegmentProposer {
	return &StubSegmentProposer{ArtificialDelay: delay}
}

func (s *StubSegmentProposer) Propose(ctx context.Context, req SegmentProposalRequest) (*SegmentProposal, error) {
	if err := ctx.Err(); err != nil {
		return nil, ErrSegmentProposalTimeout
	}
	if s != nil && s.ArtificialDelay > 0 {
		timer := time.NewTimer(s.ArtificialDelay)
		defer timer.Stop()
		select {
		case <-ctx.Done():
			return nil, ErrSegmentProposalTimeout
		case <-timer.C:
		}
	}
	if err := ctx.Err(); err != nil {
		return nil, ErrSegmentProposalTimeout
	}

	schema := req.SchemaVersion
	if schema == "" {
		schema = SegmentProposalSchemaV1
	}
	if schema != SegmentProposalSchemaV1 {
		return nil, ErrSegmentProposalMalformed
	}
	text := req.Text
	if text == "" || !utf8.ValidString(text) {
		return nil, ErrSegmentProposalMalformed
	}
	maxSeg := req.MaxSegments
	if maxSeg <= 0 {
		maxSeg = 4
	}
	if maxSeg > 4 || req.RequestDigest == "" {
		return nil, ErrSegmentProposalMalformed
	}

	span := SegmentSpan{Start: 0, End: len(text), Text: text}
	if err := ValidateProposedSpans(text, []SegmentSpan{span}); err != nil {
		return nil, err
	}
	return &SegmentProposal{
		SchemaVersion: schema,
		Digest:        fmt.Sprintf("stub:%s", req.RequestDigest),
		Spans:         []SegmentSpan{span},
	}, nil
}

// ValidateProposedSpans rejects out-of-bounds or overlapping spans.
func ValidateProposedSpans(text string, spans []SegmentSpan) error {
	if len(spans) == 0 {
		return ErrSegmentProposalMalformed
	}
	n := len(text)
	prevEnd := 0
	for _, sp := range spans {
		if sp.Start < 0 || sp.End < sp.Start || sp.End > n {
			return ErrSegmentProposalOutOfBounds
		}
		if sp.Start < prevEnd {
			return ErrSegmentProposalOutOfBounds
		}
		if sp.Text != "" && sp.Text != text[sp.Start:sp.End] {
			return ErrSegmentProposalMalformed
		}
		prevEnd = sp.End
	}
	return nil
}

var mutateActionKeys = map[string]struct{}{
	"project.fact.supplement": {},
}

var docGovActionKeys = map[string]struct{}{
	"document.claim":           {},
	"document.review.submit":   {},
	"document.review.approve":  {},
	"document.review.reject":   {},
	"document.final.approve":   {},
	"document.final.reject":    {},
	"document.proposal.decide": {},
}

var runControlActionKeys = map[string]struct{}{
	"workflow.switch": {},
	"workflow.run":    {},
	"run.resume":      {},
	"run.restart":     {},
}

// DryRunAdapter is a fail-closed / dry-run ActionOwner for U10–U12 scaffolding.
type DryRunAdapter struct {
	mu         sync.Mutex
	allowed    map[string]struct{}
	receipts   map[string]Receipt
	Denied     bool
	Label      string
	RequireRev bool
}

func NewMutateDryRunAdapter() *DryRunAdapter {
	return newDryRunAdapter("mutate", mutateActionKeys, true)
}

func NewDocGovDryRunAdapter() *DryRunAdapter {
	return newDryRunAdapter("docgov", docGovActionKeys, true)
}

func NewRunControlDryRunAdapter() *DryRunAdapter {
	return newDryRunAdapter("run_control", runControlActionKeys, false)
}

func newDryRunAdapter(label string, keys map[string]struct{}, requireRev bool) *DryRunAdapter {
	cp := make(map[string]struct{}, len(keys))
	for k := range keys {
		cp[k] = struct{}{}
	}
	return &DryRunAdapter{
		allowed:    cp,
		receipts:   map[string]Receipt{},
		Label:      label,
		RequireRev: requireRev,
	}
}

func (a *DryRunAdapter) allows(key string) bool {
	if a == nil {
		return false
	}
	_, ok := a.allowed[key]
	return ok
}

func (a *DryRunAdapter) Authorize(_ context.Context, exec StepExec) error {
	if a == nil || a.Denied || !a.allows(exec.ActionKey) {
		return ErrPlanActionForbidden
	}
	if a.RequireRev && exec.ExpectedResourceVersion == "" {
		return ErrPlanActionVersionConflict
	}
	return nil
}

func (a *DryRunAdapter) Execute(ctx context.Context, exec StepExec) (Receipt, error) {
	if err := a.Authorize(ctx, exec); err != nil {
		return Receipt{}, err
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	if existing, ok := a.receipts[exec.IdempotencyKey]; ok {
		if existing.ActionDigest != exec.ActionDigest {
			return Receipt{}, ErrPlanDecisionConflict
		}
		return existing, nil
	}
	rec := Receipt{
		IdempotencyKey: exec.IdempotencyKey,
		ActionDigest:   exec.ActionDigest,
		ResultRef:      "dryrun:" + a.Label + ":" + exec.StepID,
		ResultDigest:   "validated_only:" + exec.ActionDigest,
		SafeSummary:    a.Label + " dry-run accepted",
	}
	a.receipts[exec.IdempotencyKey] = rec
	return rec, nil
}

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

// BuildOwnerMux registers U10–U12 dry-run owners only when their flags are on.
func BuildOwnerMux(query, task ActionOwner) OwnerMux {
	m := OwnerMux{Query: query, Task: task}
	f := CurrentFeatureFlags()
	if f.U10Mutate {
		m.Mutate = NewMutateDryRunAdapter()
	}
	if f.U11DocGov {
		m.DocGov = NewDocGovDryRunAdapter()
	}
	if f.U12RunControl {
		m.RunControl = NewRunControlDryRunAdapter()
	}
	return m
}