iop/packages/go/agentpolicy/quota.go

package agentpolicy

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"iop/packages/go/agentprovider/cli/status"
	"iop/packages/go/agentruntime"
)

// QuotaState is the only quota result admitted to a durable work-attempt
// observation. It intentionally has no provider output or diagnostic field.
type QuotaState string

const (
	QuotaStateAvailable     QuotaState = "available"
	QuotaStateExhausted     QuotaState = "exhausted"
	QuotaStateUnknown       QuotaState = "unknown"
	QuotaStateNotApplicable QuotaState = "not_applicable"
)

// ObservationValidity distinguishes a valid unknown quota from evidence that
// is too old or malformed to make a continuation decision.
type ObservationValidity string

const (
	ObservationValid   ObservationValidity = "valid"
	ObservationStale   ObservationValidity = "stale"
	ObservationCorrupt ObservationValidity = "corrupt"
)

// QuotaObservation is the safe, immutable projection of one provider quota
// snapshot. It excludes raw status output, checker errors, and cap details.
type QuotaObservation struct {
	SnapshotID          string
	Adapter             string
	Target              string
	State               QuotaState
	CheckedAt           time.Time
	Validity            ObservationValidity
	Reasons             []string
	projectionIntegrity string
}

// quotaObservationJSON is the strict durable representation of a quota
// projection. The integrity value is serialized without exposing a caller-settable
// Go field, and is verified before the value is returned from JSON decoding.
type quotaObservationJSON struct {
	SnapshotID          string              `json:"snapshot_id"`
	Adapter             string              `json:"adapter"`
	Target              string              `json:"target"`
	State               QuotaState          `json:"state"`
	CheckedAt           time.Time           `json:"checked_at"`
	Validity            ObservationValidity `json:"validity"`
	Reasons             []string            `json:"reasons"`
	ProjectionIntegrity string              `json:"projection_integrity"`
}

// FailureObservation is the safe, policy-relevant projection of a runtime
// failure. Message and Metadata are deliberately excluded from durable state.
type FailureObservation struct {
	Code      agentruntime.FailureCode
	Retryable bool
}

// AttemptObservation binds safe quota and failure evidence to one work
// attempt. Callers retain the returned value; this package never shares input
// slices or maps with it.
type AttemptObservation struct {
	ObservedAt time.Time
	Quota      QuotaObservation
	Failure    FailureObservation
}

// NormalizeAttemptObservation converts quota and failure input into the only
// observation shape a continuation policy can consume. It never copies a
// failure message, failure metadata, provider output, or checker error.
func NormalizeAttemptObservation(
	snapshot status.QuotaSnapshot,
	failure *agentruntime.Failure,
	observedAt time.Time,
	maxAge time.Duration,
) AttemptObservation {
	return AttemptObservation{
		ObservedAt: observedAt.UTC(),
		Quota:      NormalizeQuotaObservation(snapshot, observedAt, maxAge),
		Failure:    NormalizeFailureObservation(failure),
	}
}

// NormalizeQuotaObservation creates a safe quota projection. Invalid source
// identity and timestamps are marked corrupt; an otherwise valid but old
// snapshot is marked stale. Both states must block continuation.
func NormalizeQuotaObservation(
	snapshot status.QuotaSnapshot,
	observedAt time.Time,
	maxAge time.Duration,
) QuotaObservation {
	if observedAt.IsZero() || status.ValidateQuotaSnapshot(snapshot) != nil {
		return CorruptQuotaObservation()
	}
	checkedAt, err := time.Parse(time.RFC3339Nano, snapshot.CheckedAt)
	if err != nil || checkedAt.IsZero() || checkedAt.After(observedAt.UTC()) {
		return CorruptQuotaObservation()
	}
	target := snapshot.Targets[0]
	state, ok := parseQuotaState(target.Status)
	if !ok {
		return CorruptQuotaObservation()
	}
	observation := QuotaObservation{
		SnapshotID: snapshot.SnapshotID,
		Adapter:    target.Adapter,
		Target:     target.Target,
		State:      state,
		CheckedAt:  checkedAt.UTC(),
		Validity:   ObservationValid,
		Reasons:    cloneStrings(snapshot.ReasonCodes),
	}
	if maxAge > 0 && observation.CheckedAt.Add(maxAge).Before(observedAt.UTC()) {
		observation.Validity = ObservationStale
	}
	observation.projectionIntegrity = quotaObservationIntegrity(observation)
	return observation
}

// CorruptQuotaObservation is the canonical secret-free replacement for quota
// evidence that cannot be validated. It deliberately retains no caller data.
func CorruptQuotaObservation() QuotaObservation {
	return QuotaObservation{Validity: ObservationCorrupt}
}

// MarshalJSON writes only a valid, sealed quota projection. This prevents an
// invalid in-memory value from becoming durable evidence through a generic
// ManagerState JSON encode.
func (o QuotaObservation) MarshalJSON() ([]byte, error) {
	if !validateQuotaObservation(o) {
		return nil, fmt.Errorf("agentpolicy: cannot encode corrupt quota observation")
	}
	return json.Marshal(quotaObservationJSON{
		SnapshotID:          o.SnapshotID,
		Adapter:             o.Adapter,
		Target:              o.Target,
		State:               o.State,
		CheckedAt:           o.CheckedAt.UTC(),
		Validity:            o.Validity,
		Reasons:             cloneStrings(o.Reasons),
		ProjectionIntegrity: o.projectionIntegrity,
	})
}

// UnmarshalJSON accepts one exact, sealed durable quota projection. It
// rejects unknown fields, malformed data, and any projection-integrity drift
// before the containing manager state can be used.
func (o *QuotaObservation) UnmarshalJSON(data []byte) error {
	var encoded quotaObservationJSON
	if err := decodeStrictJSON(data, &encoded); err != nil {
		return fmt.Errorf("agentpolicy: decode quota observation: %w", err)
	}
	next := QuotaObservation{
		SnapshotID:          encoded.SnapshotID,
		Adapter:             encoded.Adapter,
		Target:              encoded.Target,
		State:               encoded.State,
		CheckedAt:           encoded.CheckedAt.UTC(),
		Validity:            encoded.Validity,
		Reasons:             cloneStrings(encoded.Reasons),
		projectionIntegrity: encoded.ProjectionIntegrity,
	}
	if !validateQuotaObservation(next) {
		return fmt.Errorf("agentpolicy: invalid quota observation")
	}
	*o = next
	return nil
}

// SanitizeQuotaObservation returns a defensive copy of valid durable evidence
// or the canonical corrupt observation. It is safe to use at untrusted port
// boundaries before persistence.
func SanitizeQuotaObservation(observation QuotaObservation) QuotaObservation {
	observation.Reasons = cloneStrings(observation.Reasons)
	if !validateQuotaObservation(observation) {
		return CorruptQuotaObservation()
	}
	return observation
}

// SanitizeAttemptObservation canonicalizes an untrusted invocation
// observation before the manager stores or evaluates it. A missing timestamp
// uses the manager-supplied fallback; malformed quota data retains no source
// identity, reason, cap, or diagnostic.
func SanitizeAttemptObservation(
	observation AttemptObservation,
	fallbackObservedAt time.Time,
) AttemptObservation {
	observedAt := observation.ObservedAt.UTC()
	if observedAt.IsZero() {
		observedAt = fallbackObservedAt.UTC()
	}
	if observedAt.IsZero() {
		observedAt = time.Unix(0, 0).UTC()
	}
	failure := observation.Failure
	if !knownFailureCode(failure.Code) {
		failure = FailureObservation{Code: agentruntime.FailureCodeUnknown}
	}
	quota := SanitizeQuotaObservation(observation.Quota)
	if quota.Validity != ObservationCorrupt && quota.CheckedAt.After(observedAt) {
		quota = CorruptQuotaObservation()
	}
	return AttemptObservation{
		ObservedAt: observedAt,
		Quota:      quota,
		Failure:    failure,
	}
}

// NormalizeFailureObservation converts unknown future failure codes to the
// explicit unknown category without retaining provider diagnostics.
func NormalizeFailureObservation(failure *agentruntime.Failure) FailureObservation {
	if failure == nil || !knownFailureCode(failure.Code) {
		return FailureObservation{Code: agentruntime.FailureCodeUnknown}
	}
	return FailureObservation{Code: failure.Code, Retryable: failure.Retryable}
}

// Clone returns an independent attempt value suitable for state snapshots.
func (o AttemptObservation) Clone() AttemptObservation {
	out := o
	out.Quota.Reasons = cloneStrings(o.Quota.Reasons)
	return out
}

// ValidateAttemptObservation checks the safe snapshot shape without treating
// unknown, stale, or corrupt evidence as a successful observation.
func ValidateAttemptObservation(observation AttemptObservation) bool {
	if observation.ObservedAt.IsZero() || !knownFailureCode(observation.Failure.Code) {
		return false
	}
	if !validateQuotaObservation(observation.Quota) {
		return false
	}
	return observation.Quota.Validity == ObservationCorrupt ||
		!observation.Quota.CheckedAt.After(observation.ObservedAt.UTC())
}

func parseQuotaState(value string) (QuotaState, bool) {
	switch QuotaState(value) {
	case QuotaStateAvailable, QuotaStateExhausted, QuotaStateUnknown,
		QuotaStateNotApplicable:
		return QuotaState(value), true
	default:
		return "", false
	}
}

func validateQuotaObservation(observation QuotaObservation) bool {
	if observation.Validity == ObservationCorrupt {
		return observation.SnapshotID == "" &&
			observation.Adapter == "" &&
			observation.Target == "" &&
			observation.State == "" &&
			observation.CheckedAt.IsZero() &&
			len(observation.Reasons) == 0 &&
			observation.projectionIntegrity == ""
	}
	if observation.Validity != ObservationValid && observation.Validity != ObservationStale {
		return false
	}
	if !validQuotaSnapshotIdentity(observation.SnapshotID) ||
		!safeObservationIdentity(observation.Adapter) ||
		!safeObservationIdentity(observation.Target) ||
		observation.CheckedAt.IsZero() ||
		status.ValidateQuotaReasonCodes(observation.Reasons) != nil {
		return false
	}
	state, ok := parseQuotaState(string(observation.State))
	if !ok {
		return false
	}
	var reasonsValid bool
	switch state {
	case QuotaStateNotApplicable:
		reasonsValid = len(observation.Reasons) == 1 &&
			observation.Reasons[0] == "quota_not_applicable"
	case QuotaStateUnknown:
		reasonsValid = len(observation.Reasons) == 1 &&
			(observation.Reasons[0] == "checker_error" ||
				observation.Reasons[0] == "cap_evidence_unknown")
	case QuotaStateAvailable, QuotaStateExhausted:
		reasonsValid = len(observation.Reasons) == 0
	default:
		return false
	}
	if !reasonsValid {
		return false
	}
	return constantBytesEqual(
		observation.projectionIntegrity,
		quotaObservationIntegrity(observation),
	)
}

// quotaObservationIntegrity seals every policy-visible projection field after
// snapshot validation and final valid/stale classification. Ordered reasons
// are deliberate: their sequence is part of the original evidence.
func quotaObservationIntegrity(observation QuotaObservation) string {
	parts := []string{
		"quota-observation-v1",
		observation.SnapshotID,
		observation.Adapter,
		observation.Target,
		string(observation.State),
		observation.CheckedAt.UTC().Format(time.RFC3339Nano),
		string(observation.Validity),
	}
	parts = append(parts, observation.Reasons...)
	return digestParts(parts...)
}

func knownFailureCode(code agentruntime.FailureCode) bool {
	switch code {
	case agentruntime.FailureCodeUnknown,
		agentruntime.FailureCodeCancelled,
		agentruntime.FailureCodeDeadlineExceeded,
		agentruntime.FailureCodeInvalidRequest,
		agentruntime.FailureCodeSessionNotFound,
		agentruntime.FailureCodeUnavailable,
		agentruntime.FailureCodeQuotaExhausted,
		agentruntime.FailureCodeProcessExit,
		agentruntime.FailureCodeProvider,
		agentruntime.FailureCodeInternal:
		return true
	default:
		return false
	}
}

func safeObservationIdentity(value string) bool {
	return value != "" && strings.TrimSpace(value) == value &&
		!strings.ContainsAny(value, "\x00\r\n")
}

func validQuotaSnapshotIdentity(value string) bool {
	const prefix = "quota-"
	if len(value) != len(prefix)+64 || !strings.HasPrefix(value, prefix) {
		return false
	}
	for _, character := range value[len(prefix):] {
		if (character < '0' || character > '9') &&
			(character < 'a' || character > 'f') {
			return false
		}
	}
	return true
}

func cloneStrings(input []string) []string {
	if len(input) == 0 {
		return nil
	}
	return append([]string(nil), input...)
}