diff --git a/go.mod b/go.mod index 8e3f4fe..58ebbb4 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index dacea4d..b5428e0 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/run/executor.go b/run/executor.go index 4176447..27eb4d6 100644 --- a/run/executor.go +++ b/run/executor.go @@ -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 diff --git a/run/result_schema.go b/run/result_schema.go new file mode 100644 index 0000000..5247043 --- /dev/null +++ b/run/result_schema.go @@ -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- ")) +} diff --git a/run/result_schema_test.go b/run/result_schema_test.go new file mode 100644 index 0000000..94e58eb --- /dev/null +++ b/run/result_schema_test.go @@ -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 +} diff --git a/tool/registry.go b/tool/registry.go index 008d20d..34b0fcd 100644 --- a/tool/registry.go +++ b/tool/registry.go @@ -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