jsonchema: support additional keywords

Support additional (extra) keywords beyond those specified in the
`Schema` type. The current behavior is additional keywords are dropped.

To follow up:

- There are similar util functions in `mcp/util.go` which don't seem to
be used. Maybe we should move it to `internal/util` instead of
`jsonchema/util.go` package?
- The method `[Un]MarshalJSON` is only on `pointer receiver`, is this
intended? (that means it will only work if we [un]marshal pointer to the
`Schema` instance)
- The fix is using the [un]marshalStructWithMap from `mcp/util.go` and
there are requirements for the function to work properly but these are
not strictly checked e.g. need `json:"-"` tag for the Extra field.

Fixes modelcontextprotocol#69.

Author: albertsundjaja

jsonschema/schema.go

@@ -127,6 +127,9 @@ type Schema struct {
 	// https://json-schema.org/draft/2020-12/draft-bhutton-json-schema-validation-00#rfc.section.7
 	Format string `json:"format,omitempty"`
 
+	// Extra allows for additional keywords beyond those specified.
+	Extra map[string]any `json:"-"`
+
 	// computed fields
 
 	// This schema's base schema.
@@ -236,7 +239,7 @@ func (s *Schema) MarshalJSON() ([]byte, error) {
 		Type:                 typ,
 		schemaWithoutMethods: (*schemaWithoutMethods)(s),
 	}
-	return json.Marshal(ms)
+	return marshalStructWithMap(&ms, "Extra")
 }
 
 func (s *Schema) UnmarshalJSON(data []byte) error {
@@ -269,7 +272,7 @@ func (s *Schema) UnmarshalJSON(data []byte) error {
 	}{
 		schemaWithoutMethods: (*schemaWithoutMethods)(s),
 	}
-	if err := json.Unmarshal(data, &ms); err != nil {
+	if err := unmarshalStructWithMap(data, &ms, "Extra"); err != nil {
 		return err
 	}
 	// Unmarshal "type" as either Type or Types.

jsonschema/schema_test.go

@@ -26,6 +26,7 @@ func TestGoRoundTrip(t *testing.T) {
 		{Const: Ptr(any(map[string]any{}))},
 		{Default: mustMarshal(1)},
 		{Default: mustMarshal(nil)},
+		{Extra: map[string]any{"test": "value"}},
 	} {
 		data, err := json.Marshal(s)
 		if err != nil {
@@ -64,11 +65,19 @@ func TestJSONRoundTrip(t *testing.T) {
 			`{"$vocabulary":{"b":true, "a":false}}`,
 			`{"$vocabulary":{"a":false,"b":true}}`,
 		},
-		{`{"unk":0}`, `{}`}, // unknown fields are dropped, unfortunately
+		{`{"unk":0}`, `{"unk":0}`}, // unknown fields are not dropped
+		{
+			// known and unknown fields are not dropped
+			// note that the order will be by the declaration order in the anonymous struct inside MarshalJSON
+			`{"comment":"test","type":"example","unk":0}`,
+			`{"type":"example","comment":"test","unk":0}`,
+		},
+		{`{"extra":0}`, `{"extra":0}`}, // extra is not a special keyword and should not be dropped
+		{`{"Extra":0}`, `{"Extra":0}`}, // Extra is not a special keyword and should not be dropped
 	} {
 		var s Schema
 		mustUnmarshal(t, []byte(tt.in), &s)
-		data, err := json.Marshal(s)
+		data, err := json.Marshal(&s)
 		if err != nil {
 			t.Fatal(err)
 		}

jsonschema/util.go

@@ -15,6 +15,9 @@ import (
 	"math/big"
 	"reflect"
 	"slices"
+	"sync"
+
+	"github.com/modelcontextprotocol/go-sdk/internal/util"
 )
 
 // Equal reports whether two Go values representing JSON values are equal according
@@ -282,3 +285,136 @@ func assert(cond bool, msg string) {
 		panic("assertion failed: " + msg)
 	}
 }
+
+// marshalStructWithMap marshals its first argument to JSON, treating the field named
+// mapField as an embedded map. The first argument must be a pointer to
+// a struct. The underlying type of mapField must be a map[string]any, and it must have
+// a "-" json tag, meaning it will not be marshaled.
+//
+// For example, given this struct:
+//
+//	type S struct {
+//	   A int
+//	   Extra map[string] any `json:"-"`
+//	}
+//
+// and this value:
+//
+//	s := S{A: 1, Extra: map[string]any{"B": 2}}
+//
+// the call marshalJSONWithMap(s, "Extra") would return
+//
+//	{"A": 1, "B": 2}
+//
+// It is an error if the map contains the same key as another struct field's
+// JSON name.
+//
+// marshalStructWithMap calls json.Marshal on a value of type T, so T must not
+// have a MarshalJSON method that calls this function, on pain of infinite regress.
+//
+// Note that there is a similar function in mcp/util.go, but they are not the same.
+// Here the function requires `-` json tag, does not clear the mapField map,
+// and handles embedded struct due to the implementation of jsonNames in this package.
+//
+// TODO: avoid this restriction on T by forcing it to marshal in a default way.
+// See https://go.dev/play/p/EgXKJHxEx_R.
+func marshalStructWithMap[T any](s *T, mapField string) ([]byte, error) {
+	// Marshal the struct and the map separately, and concatenate the bytes.
+	// This strategy is dramatically less complicated than
+	// constructing a synthetic struct or map with the combined keys.
+	if s == nil {
+		return []byte("null"), nil
+	}
+	s2 := *s
+	vMapField := reflect.ValueOf(&s2).Elem().FieldByName(mapField)
+	mapVal := vMapField.Interface().(map[string]any)
+
+	// Check for duplicates.
+	names := jsonNames(reflect.TypeFor[T]())
+	for key := range mapVal {
+		if names[key] {
+			return nil, fmt.Errorf("map key %q duplicates struct field", key)
+		}
+	}
+
+	structBytes, err := json.Marshal(s2)
+	if err != nil {
+		return nil, fmt.Errorf("marshalStructWithMap(%+v): %w", s, err)
+	}
+	if len(mapVal) == 0 {
+		return structBytes, nil
+	}
+	mapBytes, err := json.Marshal(mapVal)
+	if err != nil {
+		return nil, err
+	}
+	if len(structBytes) == 2 { // must be "{}"
+		return mapBytes, nil
+	}
+	// "{X}" + "{Y}" => "{X,Y}"
+	res := append(structBytes[:len(structBytes)-1], ',')
+	res = append(res, mapBytes[1:]...)
+	return res, nil
+}
+
+// unmarshalStructWithMap is the inverse of marshalStructWithMap.
+// T has the same restrictions as in that function.
+//
+// Note that there is a similar function in mcp/util.go, but they are not the same.
+// Here jsonNames also returns fields from embedded structs, hence this function
+// handles embedded structs as well.
+func unmarshalStructWithMap[T any](data []byte, v *T, mapField string) error {
+	// Unmarshal into the struct, ignoring unknown fields.
+	if err := json.Unmarshal(data, v); err != nil {
+		return err
+	}
+	// Unmarshal into the map.
+	m := map[string]any{}
+	if err := json.Unmarshal(data, &m); err != nil {
+		return err
+	}
+	// Delete from the map the fields of the struct.
+	for n := range jsonNames(reflect.TypeFor[T]()) {
+		delete(m, n)
+	}
+	if len(m) != 0 {
+		reflect.ValueOf(v).Elem().FieldByName(mapField).Set(reflect.ValueOf(m))
+	}
+	return nil
+}
+
+var jsonNamesMap sync.Map // from reflect.Type to map[string]bool
+
+// jsonNames returns the set of JSON object keys that t will marshal into,
+// including fields from embedded structs in t.
+// t must be a struct type.
+//
+// Note that there is a similar function in mcp/util.go, but they are not the same
+// Here the function recurses over embedded structs and includes fields from them.
+func jsonNames(t reflect.Type) map[string]bool {
+	// Lock not necessary: at worst we'll duplicate work.
+	if val, ok := jsonNamesMap.Load(t); ok {
+		return val.(map[string]bool)
+	}
+	m := map[string]bool{}
+	for i := range t.NumField() {
+		field := t.Field(i)
+		// handle embedded structs
+		if field.Anonymous {
+			fieldType := field.Type
+			if fieldType.Kind() == reflect.Ptr {
+				fieldType = fieldType.Elem()
+			}
+			for n := range jsonNames(fieldType) {
+				m[n] = true
+			}
+			continue
+		}
+		info := util.FieldJSONInfo(field)
+		if !info.Omit {
+			m[info.Name] = true
+		}
+	}
+	jsonNamesMap.Store(t, m)
+	return m
+}

jsonschema/util_test.go

@@ -8,7 +8,11 @@ import (
 	"encoding/json"
 	"hash/maphash"
 	"reflect"
+	"strings"
 	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
 func TestEqual(t *testing.T) {
@@ -125,3 +129,58 @@ func TestHash(t *testing.T) {
 	}
 	_ = hash(null)
 }
+
+func TestMarshalStructWithMap(t *testing.T) {
+	type S struct {
+		A int
+		B string `json:"b,omitempty"`
+		u bool
+		M map[string]any `json:"-"`
+	}
+	t.Run("basic", func(t *testing.T) {
+		s := S{A: 1, B: "two", M: map[string]any{"!@#": true}}
+		got, err := marshalStructWithMap(&s, "M")
+		if err != nil {
+			t.Fatal(err)
+		}
+		want := `{"A":1,"b":"two","!@#":true}`
+		if g := string(got); g != want {
+			t.Errorf("\ngot  %s\nwant %s", g, want)
+		}
+
+		var un S
+		if err := unmarshalStructWithMap(got, &un, "M"); err != nil {
+			t.Fatal(err)
+		}
+		if diff := cmp.Diff(s, un, cmpopts.IgnoreUnexported(S{})); diff != "" {
+			t.Errorf("mismatch (-want, +got):\n%s", diff)
+		}
+	})
+	t.Run("duplicate", func(t *testing.T) {
+		s := S{A: 1, B: "two", M: map[string]any{"b": "dup"}}
+		_, err := marshalStructWithMap(&s, "M")
+		if err == nil || !strings.Contains(err.Error(), "duplicate") {
+			t.Errorf("got %v, want error with 'duplicate'", err)
+		}
+	})
+	t.Run("embedded", func(t *testing.T) {
+		type Embedded struct {
+			A     int
+			B     int
+			Extra map[string]any `json:"-"`
+		}
+		type S struct {
+			C int
+			Embedded
+		}
+		s := S{C: 1, Embedded: Embedded{A: 2, B: 3, Extra: map[string]any{"d": 4, "e": 5}}}
+		got, err := marshalStructWithMap(&s, "Extra")
+		if err != nil {
+			t.Fatal(err)
+		}
+		want := `{"C":1,"A":2,"B":3,"d":4,"e":5}`
+		if g := string(got); g != want {
+			t.Errorf("got %v, want %v", g, want)
+		}
+	})
+}