executus/run/result_schema.go

package run

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"log/slog"
	"sort"
	"strings"

	"github.com/santhosh-tekuri/jsonschema/v6"

	"gitea.stevedudenhoeffer.com/steve/majordomo/agent"
	llm "gitea.stevedudenhoeffer.com/steve/majordomo/llm"

	"gitea.stevedudenhoeffer.com/steve/executus/tool"
)

// result_schema.go — kernel-native structured output for single-loop runs.
//
// When tool.Invocation.ResultSchema is set, the executor validates the loop's
// final answer against the schema and, on violation, sends the loop back for
// bounded corrective rounds ON THE SAME CONVERSATION — the model repairs its
// own answer with full task context, which beats any post-hoc host-side
// reformat (the host can't know what the model meant). Mirrors the FinalGuard
// extra-round mechanics (same model, same observers, small step cap), with one
// deliberate contrast: FinalGuard discards the extra round's text, while a
// schema round's text IS the new candidate answer.
//
// Fail-open posture: a schema that doesn't compile logs a WARN and skips
// validation entirely — the host's pre-dispatch compile is the authoritative
// gate, and the kernel must not brick runs on a host bug. Persistent
// violations keep the last candidate output; hosts layer their own fallback
// validator/envelope on top.

const (
	// resultSchemaMaxBytes caps the schema the kernel will compile.
	resultSchemaMaxBytes = 64 << 10
	// resultSchemaMaxRounds caps corrective rounds per run.
	resultSchemaMaxRounds = 2
	// resultSchemaRoundMaxSteps caps each corrective round. Producing a
	// corrected JSON document needs no tools; 2 steps tolerates one stray
	// tool call before the reply.
	resultSchemaRoundMaxSteps = 2
	// resultSchemaExtractCap bounds the output size scanned for a JSON
	// document (the balanced-span scan is quadratic in the worst case).
	resultSchemaExtractCap = 64 << 10
)

// denyLoader fails every $ref resolution. jsonschema/v6's default compiler
// resolves file:// refs from local disk (and can be configured for http) —
// a caller-supplied schema must never reach either.
type denyLoader struct{}

func (denyLoader) Load(url string) (any, error) {
	return nil, fmt.Errorf("external $ref %q refused: remote/file schema loading is disabled: %w", url, fs.ErrNotExist)
}

// compileResultSchema compiles a raw schema with external loading disabled.
func compileResultSchema(raw json.RawMessage) (*jsonschema.Schema, error) {
	if len(raw) > resultSchemaMaxBytes {
		return nil, fmt.Errorf("schema is %d bytes (cap %d)", len(raw), resultSchemaMaxBytes)
	}
	doc, err := jsonschema.UnmarshalJSON(strings.NewReader(string(raw)))
	if err != nil {
		return nil, fmt.Errorf("schema is not valid JSON: %w", err)
	}
	c := jsonschema.NewCompiler()
	c.UseLoader(denyLoader{})
	if err := c.AddResource("inline://result-schema", doc); err != nil {
		return nil, err
	}
	return c.Compile("inline://result-schema")
}

// extractResultJSON finds the answer's JSON document: the whole string, a
// fenced block, or the first parseable balanced {...}/[...] span scanning
// left to right. Input is capped at resultSchemaExtractCap.
func extractResultJSON(s string) (json.RawMessage, bool) {
	s = strings.TrimSpace(s)
	if len(s) > resultSchemaExtractCap {
		s = s[:resultSchemaExtractCap]
	}
	if json.Valid([]byte(s)) {
		return json.RawMessage(s), true
	}
	// Fenced block: ```json ... ``` or bare ``` ... ```.
	for _, fence := range []string{"```json", "```"} {
		if i := strings.Index(s, fence); i >= 0 {
			rest := s[i+len(fence):]
			if j := strings.Index(rest, "```"); j >= 0 {
				candidate := strings.TrimSpace(rest[:j])
				if json.Valid([]byte(candidate)) && candidate != "" {
					return json.RawMessage(candidate), true
				}
			}
		}
	}
	// First parseable balanced span.
	for i := 0; i < len(s); i++ {
		if s[i] != '{' && s[i] != '[' {
			continue
		}
		if end, ok := balancedSpanEnd(s, i); ok {
			candidate := s[i : end+1]
			if json.Valid([]byte(candidate)) {
				return json.RawMessage(candidate), true
			}
		}
	}
	return nil, false
}

// balancedSpanEnd returns the index of the bracket closing the span opened
// at start, honoring JSON string/escape rules.
func balancedSpanEnd(s string, start int) (int, bool) {
	depth := 0
	inStr := false
	esc := false
	for i := start; i < len(s); i++ {
		c := s[i]
		if inStr {
			switch {
			case esc:
				esc = false
			case c == '\\':
				esc = true
			case c == '"':
				inStr = false
			}
			continue
		}
		switch c {
		case '"':
			inStr = true
		case '{', '[':
			depth++
		case '}', ']':
			depth--
			if depth == 0 {
				return i, true
			}
		}
	}
	return 0, false
}

// validateResultJSON returns flattened error paths, nil when the doc conforms.
func validateResultJSON(sch *jsonschema.Schema, doc json.RawMessage) []string {
	inst, err := jsonschema.UnmarshalJSON(strings.NewReader(string(doc)))
	if err != nil {
		return []string{"document is not valid JSON: " + err.Error()}
	}
	err = sch.Validate(inst)
	if err == nil {
		return nil
	}
	var ve *jsonschema.ValidationError
	if !errors.As(err, &ve) {
		return []string{err.Error()}
	}
	var out []string
	flattenValidationError(ve, &out)
	sort.Strings(out)
	if len(out) == 0 {
		out = []string{ve.Error()}
	}
	return out
}

// flattenValidationError collects leaf causes as "/path: message" lines.
func flattenValidationError(ve *jsonschema.ValidationError, out *[]string) {
	if len(ve.Causes) == 0 {
		loc := "/" + strings.Join(ve.InstanceLocation, "/")
		if len(ve.InstanceLocation) == 0 {
			loc = "/"
		}
		*out = append(*out, fmt.Sprintf("%s: %v", loc, ve.ErrorKind))
		return
	}
	for _, c := range ve.Causes {
		flattenValidationError(c, out)
	}
}

// resultSchemaDeps bundles the per-run state applyResultSchema needs from
// the executor's single-loop path.
type resultSchemaDeps struct {
	inv            tool.Invocation
	ra             RunnableAgent
	runRes         *agent.Result
	rec            RunRecorder
	model          llm.Model
	toolbox        *llm.Toolbox
	sharedOpts     []agent.Option
	obs            func(agent.Step)
	steer          func() []llm.Message
	noteCheckpoint func(llm.Message)
}

// applyResultSchema runs the validate + corrective-round loop, mutating
// d.runRes in place. Panic-isolated like safeFinalNudge: a bug in the
// validator (or a pathological schema) must never convert a successful
// run into a panic error.
func (e *Executor) applyResultSchema(runCtx context.Context, d resultSchemaDeps) {
	defer func() {
		if r := recover(); r != nil {
			slog.Error("run: result_schema validation panicked; keeping run result",
				"run_id", d.inv.RunID, "panic", r)
		}
	}()
	sch, cerr := compileResultSchema(d.inv.ResultSchema)
	if cerr != nil {
		slog.Warn("run: result_schema failed to compile; skipping kernel validation",
			"run_id", d.inv.RunID, "error", cerr)
		if d.rec != nil {
			d.rec.LogEvent("result_schema_skipped", map[string]any{"error": cerr.Error()})
		}
		return
	}
	for round := 0; ; round++ {
		var errs []string
		doc, ok := extractResultJSON(d.runRes.Output)
		if ok {
			errs = validateResultJSON(sch, doc)
			if errs == nil {
				// Conforming: the bare validated document IS the answer.
				d.runRes.Output = string(doc)
				if d.rec != nil {
					d.rec.LogEvent("result_schema_valid", map[string]any{"rounds": round})
				}
				return
			}
		} else {
			errs = []string{"no JSON document found in output"}
		}
		if round >= resultSchemaMaxRounds {
			// Best effort exhausted — keep the last candidate; the host's
			// own validator decides what to do with it.
			if d.rec != nil {
				d.rec.LogEvent("result_schema_unmet", map[string]any{
					"rounds": round, "validation_errors": errs,
				})
			}
			return
		}
		correction := resultSchemaCorrection(d.inv.ResultSchema, errs)
		if d.rec != nil {
			d.rec.LogEvent("result_schema_round", map[string]any{
				"round": round + 1, "validation_errors": errs,
			})
		}
		if d.noteCheckpoint != nil {
			d.noteCheckpoint(llm.UserText(correction))
		}
		roundOpts := append([]agent.Option{
			agent.WithToolbox(d.toolbox),
			agent.WithMaxSteps(resultSchemaRoundMaxSteps),
			agent.WithStepObserver(d.obs),
		}, d.sharedOpts...)
		roundAg := agent.New(d.model, e.systemPrompt(d.ra), roundOpts...)
		roundRes, roundErr := roundAg.Run(runCtx, correction,
			agent.WithSteer(d.steer), agent.WithHistory(d.runRes.Messages))
		if roundRes != nil {
			roundUsage := roundRes.Usage
			d.runRes.Usage.Add(roundUsage)
		}
		if roundErr != nil || roundRes == nil {
			slog.Warn("run: result_schema corrective round failed; keeping last output",
				"run_id", d.inv.RunID, "error", roundErr)
			return
		}
		// The round's text is the NEW candidate (contrast FinalGuard, whose
		// extra-round text is discarded meta-commentary) — but a BLANK round
		// reply must not clobber the prior non-empty candidate.
		if strings.TrimSpace(roundRes.Output) != "" {
			d.runRes.Output = roundRes.Output
		}
		d.runRes.Messages = roundRes.Messages
	}
}

// resultSchemaCorrection builds the corrective user turn for one round.
func resultSchemaCorrection(raw json.RawMessage, errs []string) string {
	return fmt.Sprintf(`Your final answer must be a SINGLE JSON document conforming to this JSON Schema (no prose, no code fences):
%s

Your previous answer failed validation:
- %s

Reply now with ONLY the corrected JSON document. Preserve your answer's content; do not invent data; use null or empty values for anything missing.`,
		string(raw), strings.Join(errs, "\n- "))
}