feat(run): kernel-native result_schema — validate + in-loop self-repair

tool.Invocation.ResultSchema (JSON Schema, draft 2020-12): the executor validates the single-loop final answer and, on violation, runs up to two bounded corrective rounds on the same conversation (FinalGuard-style extra rounds, but the round's text IS the new candidate); a conforming answer is replaced by the bare validated JSON. Fail-open on compile problems (host pre-dispatch compile is the authoritative gate) and best-effort on persistent violations (last output stands, hosts keep their fallback validator). External $refs fail closed — jsonschema/v6's default loader resolves file:// from local disk, overridden with a deny loader. Audit events: result_schema_round/_valid/_unmet/_skipped. Multi-phase runs unvalidated (FinalGuard's single-loop-only scope). Co-Authored-By: Claude Fable 5 <[email protected]>
2026-07-17 13:17:43 -04:00
parent 46e4eb3da6
commit b4080a8a6e
6 changed files with 496 additions and 0 deletions
@@ -19,6 +19,7 @@ require (
 	github.com/google/s2a-go v0.1.8 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
 	github.com/gorilla/websocket v1.5.3 // indirect
+	github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect
 	go.opencensus.io v0.24.0 // indirect
 	golang.org/x/net v0.55.0 // indirect
 	golang.org/x/sys v0.46.0 // indirect
@@ -52,6 +52,8 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
 github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
+github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ=
+github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
@@ -492,6 +492,78 @@ func (e *Executor) Run(ctx context.Context, ra RunnableAgent, inv tool.Invocatio
 			runRes, runErr = runAgent(runCtx, ag, input, inv.Images, agent.WithSteer(steer))
 		}

+		// ResultSchema (optional): validate the final answer and let the model
+		// self-repair on the same conversation before the host sees it. Runs
+		// BEFORE FinalGuard so the delivery check inspects the answer that will
+		// actually be returned. Fail-open on compile problems; best-effort on
+		// persistent violations (the host's fallback validator is authoritative).
+		if runErr == nil && runRes != nil && len(inv.ResultSchema) > 0 {
+			if sch, cerr := compileResultSchema(inv.ResultSchema); cerr != nil {
+				slog.Warn("run: result_schema failed to compile; skipping kernel validation",
+					"run_id", inv.RunID, "error", cerr)
+				if rec != nil {
+					rec.LogEvent("result_schema_skipped", map[string]any{"error": cerr.Error()})
+				}
+			} else {
+				for round := 0; ; round++ {
+					var errs []string
+					doc, ok := extractResultJSON(runRes.Output)
+					if ok {
+						errs = validateResultJSON(sch, doc)
+						if errs == nil {
+							// Conforming: the bare validated document IS the answer.
+							runRes.Output = string(doc)
+							if rec != nil && round > 0 {
+								rec.LogEvent("result_schema_valid", map[string]any{"rounds": round})
+							}
+							break
+						}
+					} 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 rec != nil {
+							rec.LogEvent("result_schema_unmet", map[string]any{
+								"rounds": round, "validation_errors": errs,
+							})
+						}
+						break
+					}
+					correction := resultSchemaCorrection(inv.ResultSchema, errs)
+					if rec != nil {
+						rec.LogEvent("result_schema_round", map[string]any{
+							"round": round + 1, "validation_errors": errs,
+						})
+					}
+					if noteCheckpoint != nil {
+						noteCheckpoint(llm.UserText(correction))
+					}
+					roundOpts := append([]agent.Option{
+						agent.WithToolbox(toolbox),
+						agent.WithMaxSteps(resultSchemaRoundMaxSteps),
+						agent.WithStepObserver(obs),
+					}, sharedOpts...)
+					roundAg := agent.New(model, e.systemPrompt(ra), roundOpts...)
+					roundRes, roundErr := roundAg.Run(runCtx, correction,
+						agent.WithSteer(steer), agent.WithHistory(runRes.Messages))
+					if roundRes != nil {
+						runRes.Usage.Add(roundRes.Usage)
+					}
+					if roundErr != nil || roundRes == nil {
+						slog.Warn("run: result_schema corrective round failed; keeping last output",
+							"run_id", inv.RunID, "error", roundErr)
+						break
+					}
+					// The round's text is the NEW candidate (contrast FinalGuard,
+					// whose extra-round text is discarded meta-commentary).
+					runRes.Output = roundRes.Output
+					runRes.Messages = roundRes.Messages
+				}
+			}
+		}
+
 		// FinalGuard (optional): the run finished successfully — before the result
 		// is finalized, let the host inspect it for undelivered work (the
 		// announce-then-stop failure: an artifact rendered into run storage, a
@@ -0,0 +1,198 @@
+package run
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/fs"
+	"sort"
+	"strings"
+
+	"github.com/santhosh-tekuri/jsonschema/v6"
+)
+
+// 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)) && len(s) > 0 {
+		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 ok := asValidationError(err, &ve); !ok {
+		return []string{err.Error()}
+	}
+	var out []string
+	flattenValidationError(ve, &out)
+	sort.Strings(out)
+	if len(out) == 0 {
+		out = []string{ve.Error()}
+	}
+	return out
+}
+
+func asValidationError(err error, target **jsonschema.ValidationError) bool {
+	ve, ok := err.(*jsonschema.ValidationError)
+	if ok {
+		*target = ve
+	}
+	return ok
+}
+
+// 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)
+	}
+}
+
+// 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- "))
+}
@@ -0,0 +1,208 @@
+package run
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"testing"
+
+	llm "gitea.stevedudenhoeffer.com/steve/majordomo/llm"
+	"gitea.stevedudenhoeffer.com/steve/majordomo/provider/fake"
+
+	"gitea.stevedudenhoeffer.com/steve/executus/tool"
+)
+
+func resultSchemaExecutor(t *testing.T, m llm.Model, rec RunRecorder) *Executor {
+	t.Helper()
+	ports := Ports{}
+	if rec != nil {
+		ports.Audit = auditFunc{start: func(RunInfo) RunRecorder { return rec }}
+	}
+	return New(Config{
+		Registry: tool.NewRegistry(),
+		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
+		Ports:    ports,
+	})
+}
+
+var personSchema = json.RawMessage(`{"type":"object","required":["name"],"properties":{"name":{"type":"string"}},"additionalProperties":false}`)
+
+// TestResultSchema_ConformingAnswerPassesThrough — a first-try conforming
+// answer (fence-wrapped, even) is replaced by the bare JSON with no extra
+// model rounds.
+func TestResultSchema_ConformingAnswerPassesThrough(t *testing.T) {
+	fp := fake.New("fake")
+	fp.Enqueue("m", fake.Reply("Here you go:\n```json\n{\"name\": \"mort\"}\n```"))
+	m, _ := fp.Model("m")
+
+	ex := resultSchemaExecutor(t, m, nil)
+	res := ex.Run(context.Background(),
+		RunnableAgent{Name: "g", ModelTier: "m"},
+		tool.Invocation{RunID: "r-rs1", CallerID: "c", ResultSchema: personSchema}, "who are you")
+
+	if res.Err != nil {
+		t.Fatalf("run error: %v", res.Err)
+	}
+	if res.Output != `{"name": "mort"}` {
+		t.Fatalf("output = %q, want the bare validated JSON", res.Output)
+	}
+	if n := fp.CallCount("m"); n != 1 {
+		t.Fatalf("model calls = %d, want 1 (no corrective round)", n)
+	}
+}
+
+// TestResultSchema_CorrectiveRoundRepairs — a non-conforming first answer
+// triggers one corrective round whose text becomes the run output.
+func TestResultSchema_CorrectiveRoundRepairs(t *testing.T) {
+	fp := fake.New("fake")
+	fp.Enqueue("m",
+		fake.Reply(`{"name": 42}`),      // wrong type
+		fake.Reply(`{"name": "fixed"}`), // corrective round conforms
+	)
+	m, _ := fp.Model("m")
+	rec := &eventRecorder{}
+
+	ex := resultSchemaExecutor(t, m, rec)
+	res := ex.Run(context.Background(),
+		RunnableAgent{Name: "g", ModelTier: "m"},
+		tool.Invocation{RunID: "r-rs2", CallerID: "c", ResultSchema: personSchema}, "go")
+
+	if res.Err != nil {
+		t.Fatalf("run error: %v", res.Err)
+	}
+	if res.Output != `{"name": "fixed"}` {
+		t.Fatalf("output = %q, want the corrected JSON", res.Output)
+	}
+	if n := fp.CallCount("m"); n != 2 {
+		t.Fatalf("model calls = %d, want 2", n)
+	}
+	if !recHas(rec, "result_schema_round") || !recHas(rec, "result_schema_valid") {
+		t.Fatalf("expected round + valid audit events, got %v", rec.events)
+	}
+}
+
+// TestResultSchema_PersistentViolationKeepsLastOutput — rounds exhaust, the
+// last candidate stands, and result_schema_unmet is recorded (the host's own
+// validator is the fallback authority).
+func TestResultSchema_PersistentViolationKeepsLastOutput(t *testing.T) {
+	fp := fake.New("fake")
+	fp.Enqueue("m",
+		fake.Reply("not json at all"),
+		fake.Reply("still prose"),
+		fake.Reply("prose forever"),
+	)
+	m, _ := fp.Model("m")
+	rec := &eventRecorder{}
+
+	ex := resultSchemaExecutor(t, m, rec)
+	res := ex.Run(context.Background(),
+		RunnableAgent{Name: "g", ModelTier: "m"},
+		tool.Invocation{RunID: "r-rs3", CallerID: "c", ResultSchema: personSchema}, "go")
+
+	if res.Err != nil {
+		t.Fatalf("run error: %v", res.Err)
+	}
+	if res.Output != "prose forever" {
+		t.Fatalf("output = %q, want the last candidate kept", res.Output)
+	}
+	if n := fp.CallCount("m"); n != 1+resultSchemaMaxRounds {
+		t.Fatalf("model calls = %d, want %d", n, 1+resultSchemaMaxRounds)
+	}
+	if !recHas(rec, "result_schema_unmet") {
+		t.Fatalf("expected result_schema_unmet, got %v", rec.events)
+	}
+}
+
+// TestResultSchema_NoSchemaUnchanged — anchor: without a schema the output
+// and call count are untouched.
+func TestResultSchema_NoSchemaUnchanged(t *testing.T) {
+	fp := fake.New("fake")
+	fp.Enqueue("m", fake.Reply("plain prose answer"))
+	m, _ := fp.Model("m")
+
+	ex := resultSchemaExecutor(t, m, nil)
+	res := ex.Run(context.Background(),
+		RunnableAgent{Name: "g", ModelTier: "m"},
+		tool.Invocation{RunID: "r-rs4", CallerID: "c"}, "go")
+
+	if res.Err != nil || res.Output != "plain prose answer" {
+		t.Fatalf("output = %q err = %v", res.Output, res.Err)
+	}
+	if n := fp.CallCount("m"); n != 1 {
+		t.Fatalf("model calls = %d, want 1", n)
+	}
+}
+
+// TestResultSchema_CompileFailureFailsOpen — an uncompilable schema skips
+// validation (WARN + audit event), never bricks the run.
+func TestResultSchema_CompileFailureFailsOpen(t *testing.T) {
+	fp := fake.New("fake")
+	fp.Enqueue("m", fake.Reply("prose"))
+	m, _ := fp.Model("m")
+	rec := &eventRecorder{}
+
+	ex := resultSchemaExecutor(t, m, rec)
+	res := ex.Run(context.Background(),
+		RunnableAgent{Name: "g", ModelTier: "m"},
+		tool.Invocation{RunID: "r-rs5", CallerID: "c",
+			ResultSchema: json.RawMessage(`{"type": "not-a-type"}`)}, "go")
+
+	if res.Err != nil {
+		t.Fatalf("run error: %v", res.Err)
+	}
+	if res.Output != "prose" {
+		t.Fatalf("output = %q, want untouched", res.Output)
+	}
+	if !recHas(rec, "result_schema_skipped") {
+		t.Fatalf("expected result_schema_skipped, got %v", rec.events)
+	}
+}
+
+// TestCompileResultSchema_ExternalRefsFailClosed — http(s) AND file $refs
+// must not resolve (jsonschema/v6's default loader reads local disk).
+func TestCompileResultSchema_ExternalRefsFailClosed(t *testing.T) {
+	for _, ref := range []string{
+		"https://example.com/schema.json",
+		"file:///etc/passwd",
+	} {
+		raw := json.RawMessage(`{"$ref": "` + ref + `"}`)
+		if _, err := compileResultSchema(raw); err == nil {
+			t.Errorf("$ref %q compiled — external loading must fail closed", ref)
+		} else if !strings.Contains(err.Error(), "refused") && !strings.Contains(err.Error(), "not exist") {
+			// Any error is acceptable as long as it FAILS; the refused
+			// wording just confirms the deny loader was consulted.
+			t.Logf("$ref %q failed with: %v", ref, err)
+		}
+	}
+}
+
+// TestExtractResultJSON_Shapes — extraction tolerance.
+func TestExtractResultJSON_Shapes(t *testing.T) {
+	cases := []struct {
+		in   string
+		want string
+		ok   bool
+	}{
+		{`{"a":1}`, `{"a":1}`, true},
+		{"prefix text {\"a\": {\"b\": \"}\"}} suffix", `{"a": {"b": "}"}}`, true},
+		{"```json\n[1,2]\n```", `[1,2]`, true},
+		{"no json here", "", false},
+		{"broken { \"a\": ", "", false},
+	}
+	for _, c := range cases {
+		got, ok := extractResultJSON(c.in)
+		if ok != c.ok || (ok && string(got) != c.want) {
+			t.Errorf("extract(%q) = (%q, %v), want (%q, %v)", c.in, got, ok, c.want, c.ok)
+		}
+	}
+}
+
+// recHas reports whether the recorder saw an event type.
+func recHas(r *eventRecorder, name string) bool {
+	for _, e := range r.events {
+		if e == name {
+			return true
+		}
+	}
+	return false
+}
@@ -19,6 +19,7 @@ package tool

 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"sync"
 	"time"
@@ -238,6 +239,20 @@ type Invocation struct {
 	// persisted agent row. Skill runs ignore this field.
 	SystemPromptPrepend string

+	// ResultSchema, when non-empty, is a JSON Schema (draft 2020-12) the
+	// run's final answer must conform to. The executor validates the
+	// single-loop final output against it and, on violation, runs up to
+	// two bounded corrective rounds on the same conversation so the
+	// model self-repairs with full task context; a conforming answer is
+	// replaced by the bare validated JSON document. The kernel is
+	// FAIL-OPEN on schema problems (an uncompilable schema logs a WARN
+	// and skips validation — the host's own pre-dispatch compile is the
+	// authoritative gate) and best-effort on persistent violations (the
+	// last output stands; hosts keep their own fallback validator /
+	// error envelope). Multi-phase runs are not validated, matching
+	// FinalGuard's single-loop-only scope.
+	ResultSchema json.RawMessage
+
 	// SuppressDelivery, when true, instructs the skill executor to
 	// SKIP its OutputTarget Delivery (Deliver / DeliverError) entirely.
 	// The run still produces an output string (returned from Run) and