jsonschema: pre-compile regexps

This is the first CL that deals with the process of preparing a schema for validation.

Perform some basic checks on the scheme. Along the way, compile regexps and store
them in the schema for use during validation.

Change-Id: I0e5e6ce28656dcecb6d2d4b2fdc98998fa05b6f1
Reviewed-on: https://go-review.googlesource.com/c/tools/+/669696
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>

internal/jsonschema/infer_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
 )
 
@@ -63,7 +64,7 @@ func TestForType(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			if diff := cmp.Diff(test.want, test.got); diff != "" {
+			if diff := cmp.Diff(test.want, test.got, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
 				t.Errorf("ForType mismatch (-want +got):\n%s", diff)
 			}
 		})

internal/jsonschema/resolve.go

@@ -0,0 +1,93 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+// This file deals with preparing a schema for validation, including various checks,
+// optimizations, and the resolution of cross-schema references.
+
+package jsonschema
+
+import (
+	"errors"
+	"fmt"
+	"regexp"
+)
+
+// A Resolved consists of a [Schema] along with associated information needed to
+// validate documents against it.
+// A Resolved has been validated against its meta-schema, and all its references
+// (the $ref and $dynamicRef keywords) have been resolved to their referenced Schemas.
+// Call [Schema.Resolve] to obtain a Resolved from a Schema.
+type Resolved struct {
+	root *Schema
+}
+
+// Resolve resolves all references within the schema and performs other tasks that
+// prepare the schema for validation.
+func (root *Schema) Resolve() (*Resolved, error) {
+	// There are three steps involved in preparing a schema to validate.
+	// 1. Check: validate the schema against a meta-schema, and perform other well-formedness
+	//    checks. Precompute some values along the way.
+	// 2. Resolve URIs: TODO.
+	// 3. Resolve references: TODO.
+	if err := root.check(); err != nil {
+		return nil, err
+	}
+	return &Resolved{root: root}, nil
+}
+
+func (s *Schema) check() error {
+	if s == nil {
+		return errors.New("nil schema")
+	}
+	var errs []error
+	report := func(err error) { errs = append(errs, err) }
+
+	for ss := range s.all() {
+		ss.checkLocal(report)
+	}
+	return errors.Join(errs...)
+}
+
+// checkLocal checks s for validity, independently of other schemas it may refer to.
+// Since checking a regexp involves compiling it, checkLocal saves those compiled regexps
+// in the schema for later use.
+// It appends the errors it finds to errs.
+func (s *Schema) checkLocal(report func(error)) {
+	addf := func(format string, args ...any) {
+		report(fmt.Errorf("jsonschema.Schema: "+format, args...))
+	}
+
+	if s == nil {
+		addf("nil subschema")
+		return
+	}
+	if err := s.basicChecks(); err != nil {
+		report(err)
+		return
+	}
+
+	// TODO: validate the schema's properties,
+	// ideally by jsonschema-validating it against the meta-schema.
+
+	// Check and compile regexps.
+	if s.Pattern != "" {
+		re, err := regexp.Compile(s.Pattern)
+		if err != nil {
+			addf("pattern: %w", err)
+		} else {
+			s.pattern = re
+		}
+	}
+	if len(s.PatternProperties) > 0 {
+		s.patternProperties = map[*regexp.Regexp]*Schema{}
+		for reString, subschema := range s.PatternProperties {
+			re, err := regexp.Compile(reString)
+			if err != nil {
+				addf("patternProperties[%q]: %w", reString, err)
+				continue
+			}
+			s.patternProperties[re] = subschema
+		}
+	}
+}

internal/jsonschema/resolve_test.go

@@ -0,0 +1,37 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package jsonschema
+
+import (
+	"regexp"
+	"testing"
+)
+
+func TestCheckLocal(t *testing.T) {
+	for _, tt := range []struct {
+		s    *Schema
+		want string // error must be non-nil and match this regexp
+	}{
+		{nil, "nil"},
+		{
+			&Schema{Pattern: "]["},
+			"regexp",
+		},
+		{
+			&Schema{PatternProperties: map[string]*Schema{"*": nil}},
+			"regexp",
+		},
+	} {
+		_, err := tt.s.Resolve()
+		if err == nil {
+			t.Errorf("%s: unexpectedly passed", tt.s.json())
+			continue
+		}
+		if !regexp.MustCompile(tt.want).MatchString(err.Error()) {
+			t.Errorf("%s: did not match\nerror: %s\nregexp: %s",
+				tt.s.json(), err, tt.want)
+		}
+	}
+}

internal/jsonschema/schema.go

@@ -13,6 +13,7 @@ import (
 	"fmt"
 	"iter"
 	"math"
+	"regexp"
 )
 
 // A Schema is a JSON schema object.
@@ -106,6 +107,10 @@ type Schema struct {
 	Then             *Schema            `json:"then,omitempty"`
 	Else             *Schema            `json:"else,omitempty"`
 	DependentSchemas map[string]*Schema `json:"dependentSchemas,omitempty"`
+
+	// computed fields
+	pattern           *regexp.Regexp
+	patternProperties map[*regexp.Regexp]*Schema
 }
 
 // String returns a short description of the schema.

internal/jsonschema/validate.go

@@ -10,7 +10,6 @@ import (
 	"math"
 	"math/big"
 	"reflect"
-	"regexp"
 	"slices"
 	"strings"
 	"unicode/utf8"
@@ -19,16 +18,9 @@ import (
 // The value of the "$schema" keyword for the version that we can validate.
 const draft202012 = "https://json-schema.org/draft/2020-12/schema"
 
-// Temporary definition of ResolvedSchema.
-// The full definition deals with references between schemas, specifically the $id, $anchor and $ref keywords.
-// We'll ignore that for now.
-type ResolvedSchema struct {
-	root *Schema
-}
-
 // Validate validates the instance, which must be a JSON value, against the schema.
 // It returns nil if validation is successful or an error if it is not.
-func (rs *ResolvedSchema) Validate(instance any) error {
+func (rs *Resolved) Validate(instance any) error {
 	if s := rs.root.Schema; s != "" && s != draft202012 {
 		return fmt.Errorf("cannot validate version %s, only %s", s, draft202012)
 	}
@@ -39,7 +31,7 @@ func (rs *ResolvedSchema) Validate(instance any) error {
 
 // state is the state of single call to ResolvedSchema.Validate.
 type state struct {
-	rs    *ResolvedSchema
+	rs    *Resolved
 	depth int
 }
 
@@ -60,10 +52,8 @@ func (st *state) validate(instance reflect.Value, schema *Schema, callerAnns *an
 		return fmt.Errorf("max recursion depth of %d reached", st.depth)
 	}
 
-	// Treat the nil schema like the empty schema, as accepting everything.
-	if schema == nil {
-		return nil
-	}
+	// We checked for nil schemas in [Schema.Resolve].
+	assert(schema != nil, "nil schema")
 
 	// Step through interfaces.
 	if instance.IsValid() && instance.Kind() == reflect.Interface {
@@ -156,15 +146,8 @@ func (st *state) validate(instance reflect.Value, schema *Schema, callerAnns *an
 			}
 		}
 
-		if schema.Pattern != "" {
-			// TODO(jba): compile regexps during schema validation.
-			m, err := regexp.MatchString(schema.Pattern, str)
-			if err != nil {
-				return err
-			}
-			if !m {
-				return fmt.Errorf("pattern: %q does not match pattern %q", str, schema.Pattern)
-			}
+		if schema.Pattern != "" && !schema.pattern.MatchString(str) {
+			return fmt.Errorf("pattern: %q does not match regular expression %q", str, schema.Pattern)
 		}
 	}
 
@@ -364,13 +347,8 @@ func (st *state) validate(instance reflect.Value, schema *Schema, callerAnns *an
 			for vprop, val := range instance.Seq2() {
 				prop := vprop.String()
 				// Check every matching pattern.
-				for pattern, schema := range schema.PatternProperties {
-					// TODO(jba): pre-compile regexps
-					m, err := regexp.MatchString(pattern, prop)
-					if err != nil {
-						return err
-					}
-					if m {
+				for re, schema := range schema.patternProperties {
+					if re.MatchString(prop) {
 						if err := st.validate(val, schema, nil, append(path, prop)); err != nil {
 							return err
 						}

internal/jsonschema/validate_test.go

@@ -51,7 +51,10 @@ func TestValidate(t *testing.T) {
 			}
 			for _, g := range groups {
 				t.Run(g.Description, func(t *testing.T) {
-					rs := &ResolvedSchema{root: g.Schema}
+					rs, err := g.Schema.Resolve()
+					if err != nil {
+						t.Fatal(err)
+					}
 					for s := range g.Schema.all() {
 						if s.Defs != nil || s.Ref != "" {
 							t.Skip("schema or subschema has unimplemented keywords")

mcp/mcp_test.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
 	"github.com/modelcontextprotocol/go-sdk/internal/protocol"
 )
@@ -149,7 +150,7 @@ func TestEndToEnd(t *testing.T) {
 			AdditionalProperties: falseSchema,
 		},
 	}}
-	if diff := cmp.Diff(wantTools, gotTools); diff != "" {
+	if diff := cmp.Diff(wantTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
 		t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
 	}
 

mcp/tool_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
 )
@@ -82,7 +83,7 @@ func TestMakeTool(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		if diff := cmp.Diff(test.want, test.tool.Definition.InputSchema); diff != "" {
+		if diff := cmp.Diff(test.want, test.tool.Definition.InputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
 			t.Errorf("MakeTool(%v) mismatch (-want +got):\n%s", test.tool.Definition.Name, diff)
 		}
 	}