diff --git a/jsonschema/schema.go b/jsonschema/schema.go index 1a73cc6..26623f1 100644 --- a/jsonschema/schema.go +++ b/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. diff --git a/jsonschema/schema_test.go b/jsonschema/schema_test.go index 0d3ea09..4ceb1ee 100644 --- a/jsonschema/schema_test.go +++ b/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) } diff --git a/jsonschema/util.go b/jsonschema/util.go index f884aff..71c3443 100644 --- a/jsonschema/util.go +++ b/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 +} diff --git a/jsonschema/util_test.go b/jsonschema/util_test.go index 189cb26..03ccb4d 100644 --- a/jsonschema/util_test.go +++ b/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) + } + }) +}