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

package openai

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"math"
	"sort"
	"strings"
	"sync"
	"time"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/streamgate"
)

const (
	maxToolValidationAttempts = 2

	runtimeMetadataToolValidationAttempt = "openai_tool_validation_attempt"
	runtimeMetadataToolValidationRetryOf = "openai_tool_validation_retry_of"
	runtimeMetadataToolValidationFailure = "openai_tool_validation_failure"

	toolCallOriginNative = "native"
	toolCallOriginText   = "text"
)

type toolValidationContract struct {
	enabled bool
	specs   map[string]textToolSpec
}

type toolCallCandidate struct {
	name      string
	arguments any
	origin    string
	raw       any
}

func toolValidationContractFromRequest(req chatCompletionRequest) toolValidationContract {
	if len(req.Tools) == 0 || toolChoiceDisablesToolValidation(req.ToolChoice) {
		return toolValidationContract{}
	}
	specs := requestedToolSpecs(req.Tools)
	if len(specs) == 0 {
		return toolValidationContract{}
	}
	return toolValidationContract{enabled: true, specs: specs}
}

func toolChoiceDisablesToolValidation(toolChoice any) bool {
	choice, ok := toolChoice.(string)
	return ok && strings.EqualFold(strings.TrimSpace(choice), "none")
}

func toolValidationAttemptMetadata(metadata map[string]string, attempt int, retryOf, failureReason string) map[string]string {
	out := make(map[string]string, len(metadata)+3)
	for key, value := range metadata {
		out[key] = value
	}
	out[runtimeMetadataToolValidationAttempt] = fmt.Sprintf("%d", attempt)
	if strings.TrimSpace(retryOf) != "" {
		out[runtimeMetadataToolValidationRetryOf] = retryOf
	}
	if strings.TrimSpace(failureReason) != "" {
		out[runtimeMetadataToolValidationFailure] = failureReason
	}
	return out
}

func validateToolCallResponse(contract toolValidationContract, toolCalls []any, origin string) error {
	if !contract.enabled || len(toolCalls) == 0 {
		return nil
	}
	candidates, err := toolCallCandidatesFromResponse(toolCalls, origin)
	if err != nil {
		return err
	}
	for i, candidate := range candidates {
		if err := validateToolCallCandidate(contract, candidate); err != nil {
			return fmt.Errorf("tool call %d: %w", i, err)
		}
	}
	return nil
}

func toolCallCandidatesFromResponse(toolCalls []any, origin string) ([]toolCallCandidate, error) {
	candidates := make([]toolCallCandidate, 0, len(toolCalls))
	for i, raw := range toolCalls {
		call, ok := raw.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("tool call %d is not an object", i)
		}
		fn, ok := call["function"].(map[string]any)
		if !ok {
			return nil, fmt.Errorf("tool call %d function is not an object", i)
		}
		name, _ := fn["name"].(string)
		name = strings.TrimSpace(name)
		if name == "" {
			return nil, fmt.Errorf("tool call %d function.name is required", i)
		}
		arguments, err := parseToolCallArguments(fn["arguments"])
		if err != nil {
			return nil, fmt.Errorf("tool call %d function.arguments: %w", i, err)
		}
		candidates = append(candidates, toolCallCandidate{
			name:      name,
			arguments: arguments,
			origin:    origin,
			raw:       raw,
		})
	}
	return candidates, nil
}

func parseToolCallArguments(raw any) (any, error) {
	switch v := raw.(type) {
	case string:
		trimmed := strings.TrimSpace(v)
		if trimmed == "" {
			return nil, fmt.Errorf("must be valid JSON")
		}
		var decoded any
		if err := json.Unmarshal([]byte(trimmed), &decoded); err != nil {
			return nil, fmt.Errorf("must be valid JSON: %w", err)
		}
		return decoded, nil
	case map[string]any:
		return v, nil
	default:
		return nil, fmt.Errorf("must be a JSON string or object")
	}
}

func validateToolCallCandidate(contract toolValidationContract, candidate toolCallCandidate) error {
	spec, ok := contract.specs[candidate.name]
	if !ok {
		return fmt.Errorf("tool %q is not in request tools", candidate.name)
	}
	if err := validateJSONSchemaSubset(candidate.arguments, spec.parameters, "$.arguments"); err != nil {
		return fmt.Errorf("tool %q %s", candidate.name, err.Error())
	}
	return nil
}

func validateJSONSchemaSubset(value any, schema any, path string) error {
	m, ok := schema.(map[string]any)
	if !ok || len(m) == 0 {
		return nil
	}
	if branches := schemaBranches(m["anyOf"]); len(branches) > 0 {
		return validateAnySchemaBranch(value, branches, path, "anyOf")
	}
	if branches := schemaBranches(m["oneOf"]); len(branches) > 0 {
		return validateAnySchemaBranch(value, branches, path, "oneOf")
	}
	for _, branch := range schemaBranches(m["allOf"]) {
		if err := validateJSONSchemaSubset(value, branch, path); err != nil {
			return err
		}
	}
	if values := schemaEnumValues(m["enum"]); len(values) > 0 && !valueInJSONEnum(value, values) {
		return fmt.Errorf("%s value %s is not in enum", path, jsonValuePreview(value))
	}
	types := schemaTypeNames(m["type"])
	if len(types) > 0 && !jsonValueMatchesAnyType(value, types) {
		return fmt.Errorf("%s type %s does not match schema type %s", path, jsonValueTypeName(value), strings.Join(types, "|"))
	}
	if shouldValidateObjectSubset(value, m, types) {
		obj, ok := value.(map[string]any)
		if !ok {
			return fmt.Errorf("%s type %s does not match schema type object", path, jsonValueTypeName(value))
		}
		if err := validateJSONObjectSubset(obj, m, path); err != nil {
			return err
		}
	}
	if shouldValidateArraySubset(value, m, types) {
		items, ok := value.([]any)
		if !ok {
			return fmt.Errorf("%s type %s does not match schema type array", path, jsonValueTypeName(value))
		}
		if itemSchema := schemaArrayItemSchema(m); itemSchema != nil {
			for i, item := range items {
				if err := validateJSONSchemaSubset(item, itemSchema, fmt.Sprintf("%s[%d]", path, i)); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func validateAnySchemaBranch(value any, branches []any, path, branchType string) error {
	var reasons []string
	for _, branch := range branches {
		if err := validateJSONSchemaSubset(value, branch, path); err == nil {
			return nil
		} else {
			reasons = append(reasons, err.Error())
		}
	}
	sort.Strings(reasons)
	if len(reasons) == 0 {
		return fmt.Errorf("%s does not match %s", path, branchType)
	}
	return fmt.Errorf("%s does not match %s: %s", path, branchType, reasons[0])
}

func shouldValidateObjectSubset(value any, schema map[string]any, types []string) bool {
	return containsString(types, "object") || len(schemaObjectProperties(schema)) > 0 || len(schemaStringList(schema["required"])) > 0 || isAdditionalPropertiesFalse(schema)
}

func validateJSONObjectSubset(obj map[string]any, schema map[string]any, path string) error {
	props := schemaObjectProperties(schema)
	for _, key := range schemaStringList(schema["required"]) {
		if _, ok := obj[key]; !ok {
			return fmt.Errorf("%s.%s is required", path, key)
		}
	}
	if isAdditionalPropertiesFalse(schema) {
		for key := range obj {
			if _, ok := props[key]; !ok {
				return fmt.Errorf("%s.%s is not allowed by additionalProperties=false", path, key)
			}
		}
	}
	for key, propSchema := range props {
		child, ok := obj[key]
		if !ok {
			continue
		}
		if err := validateJSONSchemaSubset(child, propSchema, path+"."+key); err != nil {
			return err
		}
	}
	return nil
}

func shouldValidateArraySubset(value any, schema map[string]any, types []string) bool {
	return containsString(types, "array") || schemaArrayItemSchema(schema) != nil
}

func isAdditionalPropertiesFalse(schema map[string]any) bool {
	allowed, ok := schema["additionalProperties"].(bool)
	return ok && !allowed
}

func schemaTypeNames(value any) []string {
	switch v := value.(type) {
	case string:
		if strings.TrimSpace(v) == "" {
			return nil
		}
		return []string{strings.TrimSpace(v)}
	case []any:
		out := make([]string, 0, len(v))
		for _, item := range v {
			if s, ok := item.(string); ok && strings.TrimSpace(s) != "" {
				out = append(out, strings.TrimSpace(s))
			}
		}
		return out
	default:
		return nil
	}
}

func schemaStringList(value any) []string {
	raw, ok := value.([]any)
	if !ok {
		return nil
	}
	out := make([]string, 0, len(raw))
	for _, item := range raw {
		if s, ok := item.(string); ok && strings.TrimSpace(s) != "" {
			out = append(out, strings.TrimSpace(s))
		}
	}
	return out
}

func schemaEnumValues(value any) []any {
	raw, ok := value.([]any)
	if !ok {
		return nil
	}
	return raw
}

func valueInJSONEnum(value any, allowed []any) bool {
	for _, item := range allowed {
		if jsonValuesEqual(value, item) {
			return true
		}
	}
	return false
}

func jsonValuesEqual(a, b any) bool {
	encodedA, errA := json.Marshal(a)
	encodedB, errB := json.Marshal(b)
	return errA == nil && errB == nil && string(encodedA) == string(encodedB)
}

func jsonValueMatchesAnyType(value any, types []string) bool {
	for _, t := range types {
		if jsonValueMatchesType(value, t) {
			return true
		}
	}
	return false
}

func jsonValueMatchesType(value any, typ string) bool {
	switch typ {
	case "object":
		_, ok := value.(map[string]any)
		return ok
	case "array":
		_, ok := value.([]any)
		return ok
	case "string":
		_, ok := value.(string)
		return ok
	case "boolean":
		_, ok := value.(bool)
		return ok
	case "number":
		return isJSONNumber(value)
	case "integer":
		return isJSONInteger(value)
	case "null":
		return value == nil
	default:
		return true
	}
}

func isJSONNumber(value any) bool {
	switch value.(type) {
	case float64, float32, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
		return true
	default:
		return false
	}
}

func isJSONInteger(value any) bool {
	switch v := value.(type) {
	case float64:
		return math.Trunc(v) == v
	case float32:
		return math.Trunc(float64(v)) == float64(v)
	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
		return true
	default:
		return false
	}
}

func jsonValueTypeName(value any) string {
	switch value.(type) {
	case nil:
		return "null"
	case map[string]any:
		return "object"
	case []any:
		return "array"
	case string:
		return "string"
	case bool:
		return "boolean"
	default:
		if isJSONInteger(value) {
			return "integer"
		}
		if isJSONNumber(value) {
			return "number"
		}
		return fmt.Sprintf("%T", value)
	}
}

func jsonValuePreview(value any) string {
	encoded, err := json.Marshal(value)
	if err != nil {
		return fmt.Sprintf("%v", value)
	}
	return string(encoded)
}

func containsString(values []string, want string) bool {
	for _, value := range values {
		if value == want {
			return true
		}
	}
	return false
}

// --- runtime-enabled tool validation (Core single owner) ---------------------

const (
	// openAIToolValidationFilterID is the production Edge filter that carries the
	// existing tool-validation semantics into the Core terminal gate. The Core
	// owns mutation, rebuild, dispatch, and the recovery budget; the filter only
	// judges the attempt result and asks for an exact replay.
	openAIToolValidationFilterID = "openai.tool_validation"
	openAIToolValidationRuleID   = "openai.tool_validation.contract"
	// openAIToolValidationPriority must match the registration priority: the
	// arbiter requires an exact match between a violation intent priority and the
	// filter's effective registration priority.
	openAIToolValidationPriority = 20
)

// openAIBufferedAttemptResult is one buffered attempt's assembled OpenAI output
// plus the validation verdict derived from the existing
// collectChatCompletionOutput / validateToolCallResponse semantics.
type openAIBufferedAttemptResult struct {
	output     chatCompletionOutput
	dispatch   edgeservice.RunDispatch
	collectErr error
	validErr   error
}

// openAIBufferedResultHolder is the request-local handoff between a buffered
// attempt's event source (which assembles exactly once per attempt), the
// terminal-gate filter (which judges it), and the release sink (which renders
// the passing attempt). Each new attempt resets it, so a discarded attempt's
// output can never be rendered or re-judged.
type openAIBufferedResultHolder struct {
	mu       sync.Mutex
	current  openAIBufferedAttemptResult
	hasValue bool
	attempts int
	lastFail string
	lastRun  string
}

func newOpenAIBufferedResultHolder() *openAIBufferedResultHolder {
	return &openAIBufferedResultHolder{attempts: 1}
}

// beginAttempt clears the previous attempt's result and returns the 1-based
// attempt number of the attempt that is about to produce output.
func (h *openAIBufferedResultHolder) beginAttempt() int {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.current = openAIBufferedAttemptResult{}
	h.hasValue = false
	return h.attempts
}

// nextAttemptMetadata records that a recovery attempt is being admitted and
// returns the legacy tool-validation attempt metadata inputs (attempt number,
// failed run id, failure reason) so the provider still sees the same
// request-level retry annotation it saw on the legacy path.
func (h *openAIBufferedResultHolder) nextAttemptMetadata() (attempt int, retryOf, reason string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.attempts++
	return h.attempts, h.lastRun, h.lastFail
}

func (h *openAIBufferedResultHolder) set(result openAIBufferedAttemptResult) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.current = result
	h.hasValue = true
	h.lastRun = result.dispatch.RunID
	switch {
	case result.collectErr != nil:
		h.lastFail = result.collectErr.Error()
	case result.validErr != nil:
		h.lastFail = result.validErr.Error()
	default:
		h.lastFail = ""
	}
}

func (h *openAIBufferedResultHolder) get() (openAIBufferedAttemptResult, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.current, h.hasValue
}

// validationFailure returns the tool-validation error of the current attempt,
// or nil when the attempt passed, failed for a non-validation reason, or has
// not produced a result yet.
func (h *openAIBufferedResultHolder) validationFailure() error {
	h.mu.Lock()
	defer h.mu.Unlock()
	if !h.hasValue || h.current.collectErr != nil {
		return nil
	}
	return h.current.validErr
}

// openAIToolValidationFilter is the production terminal-gate consumer for the
// runtime-enabled buffered and non-stream chat paths. It reuses the existing
// Edge semantics only: the attempt result holder already carries the outcome of
// collectChatCompletionOutput and validateToolCallResponse, so this filter adds
// no new consumer policy (no repeat, missing-tool, or schema-repair rules).
type openAIToolValidationFilter struct {
	streamgate.FilterBase
	holder     *openAIBufferedResultHolder
	requestRef string
}

func newOpenAIToolValidationFilter(holder *openAIBufferedResultHolder, requestRef string) (*openAIToolValidationFilter, error) {
	if holder == nil {
		return nil, fmt.Errorf("openai tool validation filter requires a result holder")
	}
	base, err := streamgate.NewFilterBase(openAIToolValidationFilterID)
	if err != nil {
		return nil, err
	}
	return &openAIToolValidationFilter{FilterBase: base, holder: holder, requestRef: requestRef}, nil
}

// Applies re-resolves per attempt: a provider-tunnel attempt relays raw provider
// bytes, so the normalized tool-validation contract does not apply to it. A
// pool recovery that switches the actual execution path therefore drops or
// regains this filter with the new binding.
func (f *openAIToolValidationFilter) Applies(fctx streamgate.FilterContext) bool {
	return fctx.ExecutionPath() != string(edgeservice.ProviderPoolPathTunnel)
}

// HoldRequirement is a terminal gate: the verdict is made exactly once, on the
// terminal event, before any byte is committed. The terminal kind is part of
// the subscription so the gate stays applicable even though the buffered event
// source publishes no content events at all — the assembled output lives in the
// result holder, which keeps an arbitrarily large response out of the Core
// evidence buffer while still blocking release until the verdict exists. Hold
// mode none would not work here: with no blocking requirement on the channel,
// Core only makes a none-mode filter trigger-ready on provider_error.
func (f *openAIToolValidationFilter) HoldRequirement(streamgate.FilterContext) streamgate.FilterHoldRequirement {
	req, err := streamgate.NewFilterHoldRequirementTerminalGate(
		streamGateChannelDefault,
		[]streamgate.EventKind{streamgate.EventKindTextDelta, streamgate.EventKindTerminal},
		streamgate.EventKindTerminal,
	)
	if err != nil {
		// The kind set is a compile-time constant; a failure here would mean a
		// Core contract change, so fall back to the text-delta subscription
		// rather than panicking in a request path.
		req, _ = streamgate.NewFilterHoldRequirementTerminalGate(
			streamGateChannelDefault,
			[]streamgate.EventKind{streamgate.EventKindTextDelta},
			streamgate.EventKindTerminal,
		)
	}
	return req
}

func (f *openAIToolValidationFilter) Evaluate(ctx context.Context, fctx streamgate.FilterContext, batch streamgate.EvidenceBatch) (streamgate.FilterDecision, error) {
	verr := f.holder.validationFailure()
	kind := streamgate.EventKindTerminal
	if events := batch.Events(); len(events) > 0 {
		kind = events[len(events)-1].Kind()
	}
	code := "tool_validation_pass"
	if verr != nil {
		code = "tool_validation_failed"
	}
	evidence, err := streamgate.NewSanitizedEvidence(
		kind, streamGateChannelDefault, openAIToolValidationRuleID, code,
		openAIToolValidationFingerprint(verr), 1, 0,
		streamgate.FilterOutcomeKindEvaluated, openAIToolValidationEvaluatedAt(batch.CapturedAt()),
	)
	if err != nil {
		return streamgate.FilterDecision{}, err
	}
	if verr == nil {
		return streamgate.NewFilterDecision(streamgate.FilterDecisionKindPass, openAIRebuildFamily, f.ID(), openAIToolValidationRuleID, evidence, nil)
	}
	directive, err := streamgate.NewRecoveryDirectiveExact(f.requestRef)
	if err != nil {
		return streamgate.FilterDecision{}, err
	}
	intent, err := streamgate.NewRecoveryIntent(streamgate.RecoveryStrategyExactReplay, directive, code, openAIToolValidationPriority)
	if err != nil {
		return streamgate.FilterDecision{}, err
	}
	return streamgate.NewFilterDecision(streamgate.FilterDecisionKindViolation, openAIRebuildFamily, f.ID(), openAIToolValidationRuleID, evidence, &intent)
}

// openAIToolValidationFingerprint derives a stable, raw-free fingerprint from
// the validation reason so repeated identical failures are comparable without
// carrying provider text into the Core evidence contract.
func openAIToolValidationFingerprint(verr error) streamgate.FixedFingerprint {
	reason := "pass"
	if verr != nil {
		reason = verr.Error()
	}
	return streamgate.FixedFingerprint(sha256.Sum256([]byte(openAIToolValidationRuleID + "\x00" + reason)))
}

var _ streamgate.Filter = (*openAIToolValidationFilter)(nil)

// openAIToolValidationEvaluatedAt is the batch capture fallback used when a
// batch carries no timestamp; Core requires a non-zero evidence timestamp.
func openAIToolValidationEvaluatedAt(ts time.Time) time.Time {
	if ts.IsZero() {
		return time.Now()
	}
	return ts
}