Agent-System/internal/app/query_owner.go

package app

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

	"chat-agent-system/internal/app/intentflow"
	"chat-agent-system/internal/kb"
)

// ProjectAccess checks membership for the plan principal before query execute.
type ProjectAccess func(ctx context.Context, projectID, userID int64) error

// QueryOwner is the production auto_read query ActionOwner. It re-auths the
// principal, optionally checks expected resource versions, runs read-only KB
// search when online, and stamps receipts with an auth-scope digest. It never
// performs business writes.
type QueryOwner struct {
	Access          ProjectAccess
	KB              *kb.Client
	CurrentVersion  func(ctx context.Context, exec intentflow.StepExec) (string, error)
	Denied          bool
	RequiredVersion map[string]string

	mu       sync.Mutex
	receipts map[string]intentflow.Receipt
	execs    map[string]int
}

func NewQueryOwner(access ProjectAccess, client *kb.Client) *QueryOwner {
	return &QueryOwner{
		Access:   access,
		KB:       client,
		receipts: map[string]intentflow.Receipt{},
		execs:    map[string]int{},
	}
}

func (q *QueryOwner) Authorize(ctx context.Context, exec intentflow.StepExec) error {
	if q == nil || q.Denied {
		return intentflow.ErrPlanActionForbidden
	}
	if !isAutoReadQueryKey(exec.ActionKey) {
		return intentflow.ErrPlanActionForbidden
	}
	if exec.Principal.UserID <= 0 || exec.Principal.ProjectID <= 0 || exec.Principal.SessionID == "" {
		return intentflow.ErrPlanActionForbidden
	}
	if q.Access != nil {
		if err := q.Access(ctx, exec.Principal.ProjectID, exec.Principal.UserID); err != nil {
			return intentflow.ErrPlanActionForbidden
		}
	}
	if want, ok := q.RequiredVersion[exec.ActionKey]; ok {
		if exec.ExpectedResourceVersion == "" || exec.ExpectedResourceVersion != want {
			return intentflow.ErrPlanActionVersionConflict
		}
	}
	if exec.ExpectedResourceVersion != "" && q.CurrentVersion != nil {
		got, err := q.CurrentVersion(ctx, exec)
		if err != nil {
			return err
		}
		if got != exec.ExpectedResourceVersion {
			return intentflow.ErrPlanActionVersionConflict
		}
	}
	return nil
}

func (q *QueryOwner) Execute(ctx context.Context, exec intentflow.StepExec) (intentflow.Receipt, error) {
	if err := q.Authorize(ctx, exec); err != nil {
		return intentflow.Receipt{}, err
	}
	q.mu.Lock()
	if rec, ok := q.receipts[exec.IdempotencyKey]; ok {
		q.mu.Unlock()
		if rec.ActionDigest != exec.ActionDigest {
			return intentflow.Receipt{}, intentflow.ErrPlanDecisionConflict
		}
		return rec, nil
	}
	q.mu.Unlock()

	scope := AuthScopeDigest(exec)
	resultVer := exec.ExpectedResourceVersion
	if resultVer == "" {
		resultVer = "none"
	}
	hits := 0
	queryText := queryTextFromExec(exec)
	if q.KB != nil && q.KB.Online() && strings.TrimSpace(queryText) != "" {
		domains := domainsForAction(exec.ActionKey)
		kbHits, err := q.KB.Search(ctx, queryText, kb.SearchOptions{Domains: domains, Limit: 5})
		if err == nil {
			hits = len(kbHits)
		}
	}
	rec := intentflow.Receipt{
		IdempotencyKey: exec.IdempotencyKey,
		ActionDigest:   exec.ActionDigest,
		ResultRef:      fmt.Sprintf("query:%s:%s", exec.ActionKey, exec.IdempotencyKey),
		ResultDigest: fmt.Sprintf(
			"query:%s:scope=%s:ver=%s:recover=%s:hits=%d",
			exec.ActionKey, scope, resultVer, shortKey(exec.IdempotencyKey), hits,
		),
		SafeSummary: fmt.Sprintf("%s authorized (scope=%s)", exec.ActionKey, shortKey(scope)),
	}
	q.mu.Lock()
	defer q.mu.Unlock()
	if existing, ok := q.receipts[exec.IdempotencyKey]; ok {
		if existing.ActionDigest != exec.ActionDigest {
			return intentflow.Receipt{}, intentflow.ErrPlanDecisionConflict
		}
		return existing, nil
	}
	q.execs[exec.IdempotencyKey]++
	q.receipts[exec.IdempotencyKey] = rec
	return rec, nil
}

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

func (q *QueryOwner) ExecCount(idempotencyKey string) int {
	q.mu.Lock()
	defer q.mu.Unlock()
	return q.execs[idempotencyKey]
}

// AuthScopeDigest binds tenant/project/session/environment/action for receipts.
func AuthScopeDigest(exec intentflow.StepExec) string {
	h := sha256.New()
	_, _ = fmt.Fprintf(h, "u=%d|p=%d|s=%s|e=%s|a=%s",
		exec.Principal.UserID, exec.Principal.ProjectID,
		exec.Principal.SessionID, exec.Principal.Environment, exec.ActionKey)
	return hex.EncodeToString(h.Sum(nil)[:16])
}

func isAutoReadQueryKey(key string) bool {
	row, err := intentflow.DefaultBatchActionPolicy().Lookup(key)
	if err != nil || row.RejectPlan || row.RequiresBatchConfirm {
		return false
	}
	return intentflow.TierAutoRead.Allows(row.OpenFromTier)
}

func queryTextFromExec(exec intentflow.StepExec) string {
	if s := strings.TrimSpace(exec.SafeSummary); s != "" {
		return s
	}
	if len(exec.Slots) > 0 {
		return strings.TrimSpace(string(exec.Slots))
	}
	return exec.ActionKey
}

func domainsForAction(actionKey string) []kb.KnowledgeDomain {
	switch actionKey {
	case "case.query":
		return []kb.KnowledgeDomain{kb.DomainProjectAsset}
	default:
		return []kb.KnowledgeDomain{kb.DomainRegulatory}
	}
}

func shortKey(s string) string {
	if len(s) <= 12 {
		return s
	}
	sum := sha256.Sum256([]byte(s))
	return hex.EncodeToString(sum[:6])
}