iop/apps/edge/internal/openai/single_request_quality_gate.go

package openai

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/binary"
	"encoding/json"
	"errors"
	"sync"

	edgeservice "iop/apps/edge/internal/service"
)

var errSingleRequestClosedTerminal = errors.New("single-request stage reached a closed terminal")

// singleRequestTerminalFailure transports only one validated closed
// disposition. The cause is always a package sentinel used for errors.Is; raw
// provider/tool errors are deliberately not retained.
type singleRequestTerminalFailure struct {
	disposition edgeservice.SingleRequestTerminalDisposition
	cause       error
}

func (e *singleRequestTerminalFailure) Error() string { return errSingleRequestClosedTerminal.Error() }
func (e *singleRequestTerminalFailure) Unwrap() error { return e.cause }

func singleRequestTerminalDisposition(err error) (edgeservice.SingleRequestTerminalDisposition, bool) {
	var terminal *singleRequestTerminalFailure
	if !errors.As(err, &terminal) || terminal == nil || terminal.disposition.Validate() != nil {
		return edgeservice.SingleRequestTerminalDisposition{}, false
	}
	return terminal.disposition, true
}

type singleRequestQualityGate struct {
	mu         sync.Mutex
	toolCycles map[string]map[[sha256.Size]byte]struct{}
}

func newSingleRequestQualityGate() *singleRequestQualityGate {
	return &singleRequestQualityGate{toolCycles: make(map[string]map[[sha256.Size]byte]struct{})}
}

func singleRequestQualityGateOrNew(gate *singleRequestQualityGate) *singleRequestQualityGate {
	if gate != nil {
		return gate
	}
	return newSingleRequestQualityGate()
}

func (g *singleRequestQualityGate) failure(kind edgeservice.SingleRequestTerminalKind, class edgeservice.SingleRequestTerminalErrorClass, cause error) error {
	disposition := edgeservice.SingleRequestTerminalDisposition{Kind: kind, ErrorClass: class}
	if disposition.Validate() != nil {
		disposition = edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}
	}
	return &singleRequestTerminalFailure{disposition: disposition, cause: cause}
}

func (g *singleRequestQualityGate) reclassify(err, cause error) error {
	if disposition, ok := singleRequestTerminalDisposition(err); ok {
		return g.failure(disposition.Kind, disposition.ErrorClass, cause)
	}
	return g.providerFailure(context.Background(), err, cause)
}

func (g *singleRequestQualityGate) providerFailure(ctx context.Context, err, cause error) error {
	if disposition, ok := singleRequestTerminalDisposition(err); ok {
		return g.failure(disposition.Kind, disposition.ErrorClass, cause)
	}
	switch {
	case ctx != nil && errors.Is(ctx.Err(), context.Canceled):
		return g.failure(edgeservice.SingleRequestTerminalCancelled, "", cause)
	case ctx != nil && errors.Is(ctx.Err(), context.DeadlineExceeded), errors.Is(err, context.DeadlineExceeded):
		return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorTimeout, cause)
	case errors.Is(err, errProviderStageOutputLimit):
		return g.failure(edgeservice.SingleRequestTerminalLength, "", cause)
	case errors.Is(err, errProviderStageContextLimit):
		return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorContext, cause)
	case errors.Is(err, errProviderStageMalformed):
		return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorMalformed, cause)
	default:
		return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorProvider, cause)
	}
}

func (g *singleRequestQualityGate) validation(cause error) error {
	return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorValidation, cause)
}

func (g *singleRequestQualityGate) malformed(cause error) error {
	return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorMalformed, cause)
}

func (g *singleRequestQualityGate) budget(cause error) error {
	return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorBudget, cause)
}

func (g *singleRequestQualityGate) length(cause error) error {
	return g.failure(edgeservice.SingleRequestTerminalLength, "", cause)
}

func (g *singleRequestQualityGate) contextLimit(cause error) error {
	return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorContext, cause)
}

func (g *singleRequestQualityGate) internalTool(cause error) error {
	return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorInternalTool, cause)
}

func (g *singleRequestQualityGate) serviceFailure(ctx context.Context, err, cause error) error {
	if disposition, ok := singleRequestTerminalDisposition(err); ok {
		return g.failure(disposition.Kind, disposition.ErrorClass, cause)
	}
	switch {
	case ctx != nil && errors.Is(ctx.Err(), context.Canceled), errors.Is(err, edgeservice.ErrSingleRequestCancelled):
		return g.failure(edgeservice.SingleRequestTerminalCancelled, "", cause)
	case ctx != nil && errors.Is(ctx.Err(), context.DeadlineExceeded), errors.Is(err, context.DeadlineExceeded):
		return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorTimeout, cause)
	case errors.Is(err, edgeservice.ErrSingleRequestInternalToolBudget):
		return g.budget(cause)
	case errors.Is(err, edgeservice.ErrSingleRequestInternalToolInvalidCall), errors.Is(err, edgeservice.ErrSingleRequestInvalidState), errors.Is(err, edgeservice.ErrSingleRequestInvalidSequence):
		return g.malformed(cause)
	default:
		return g.internalTool(cause)
	}
}

// observeToolCycle hashes one bounded canonical action/result pair and rejects
// the first proven duplicate within the same request stage. Only fixed hashes
// are retained; arguments and tool output never enter guard state.
func (g *singleRequestQualityGate) observeToolCycle(stage, name string, arguments json.RawMessage, result edgeservice.InternalWorkspaceToolResult, cause error) error {
	if g == nil || stage == "" || name == "" || len(arguments) == 0 {
		return singleRequestQualityGateOrNew(g).malformed(cause)
	}
	success := result.Status == "success" && result.ErrorCode == ""
	repairableNotFound := result.Status == "error" && result.ErrorCode == "not_found"
	if !success && !repairableNotFound {
		switch result.Status {
		case "timeout":
			return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorTimeout, cause)
		case "cancelled":
			return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorInternalTool, cause)
		case "invalid":
			return g.malformed(cause)
		default:
			if result.ErrorCode == "invalid_request" {
				return g.malformed(cause)
			}
			return g.internalTool(cause)
		}
	}

	canonicalArguments, err := canonicalSingleRequestJSON(arguments)
	if err != nil {
		return g.malformed(cause)
	}
	hash := sha256.New()
	writeSingleRequestFingerprintPart(hash, []byte(name))
	writeSingleRequestFingerprintPart(hash, canonicalArguments)
	writeSingleRequestFingerprintPart(hash, []byte(result.Status))
	writeSingleRequestFingerprintPart(hash, []byte(result.ErrorCode))
	writeSingleRequestFingerprintPart(hash, result.Content)
	for _, entry := range result.Entries {
		writeSingleRequestFingerprintPart(hash, []byte(entry))
	}
	writeSingleRequestFingerprintPart(hash, result.Stdout)
	writeSingleRequestFingerprintPart(hash, result.Stderr)
	var scalar [5]byte
	binary.BigEndian.PutUint32(scalar[:4], uint32(result.ExitCode))
	if result.Truncated {
		scalar[4] = 1
	}
	writeSingleRequestFingerprintPart(hash, scalar[:])
	var fingerprint [sha256.Size]byte
	copy(fingerprint[:], hash.Sum(nil))

	g.mu.Lock()
	defer g.mu.Unlock()
	seen := g.toolCycles[stage]
	if seen == nil {
		seen = make(map[[sha256.Size]byte]struct{})
		g.toolCycles[stage] = seen
	}
	if _, repeated := seen[fingerprint]; repeated {
		return g.failure(edgeservice.SingleRequestTerminalError, edgeservice.SingleRequestTerminalErrorRepetition, cause)
	}
	seen[fingerprint] = struct{}{}
	return nil
}

func canonicalSingleRequestJSON(raw []byte) ([]byte, error) {
	if validateSingleRequestJSON(raw) != nil {
		return nil, errProviderStageMalformed
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.UseNumber()
	var value any
	if err := decoder.Decode(&value); err != nil {
		return nil, errProviderStageMalformed
	}
	canonical, err := json.Marshal(value)
	if err != nil {
		return nil, errProviderStageMalformed
	}
	return canonical, nil
}

type singleRequestFingerprintWriter interface {
	Write([]byte) (int, error)
}

func writeSingleRequestFingerprintPart(hash singleRequestFingerprintWriter, value []byte) {
	var size [8]byte
	binary.BigEndian.PutUint64(size[:], uint64(len(value)))
	_, _ = hash.Write(size[:])
	_, _ = hash.Write(value)
}