Preserve unmodeled top-level tool schema keywords

mcpcompat/mcp/tools.go

@@ -369,6 +369,14 @@ type ToolArgumentsSchema struct {
 	Properties           map[string]any `json:"properties"`
 	Required             []string       `json:"required,omitempty"`
 	AdditionalProperties any            `json:"additionalProperties,omitempty"`
+	// Extra preserves top-level JSON Schema keywords that are not modeled by
+	// the fields above — e.g. oneOf, anyOf, allOf, $ref, enum, const,
+	// patternProperties. Without it, such keywords are silently dropped on an
+	// unmarshal -> marshal round-trip (a schema like {"oneOf": [...]} would be
+	// gutted). Populated by UnmarshalJSON and re-emitted by MarshalJSON; keys
+	// here never overlap the modeled fields. Not a JSON field itself (json:"-");
+	// its contents are inlined at the top level.
+	Extra map[string]json.RawMessage `json:"-"`
 }
 
 // ToolInputSchema remains a named type for retro-compatibility, so its JSON
@@ -390,24 +398,41 @@ func (tos ToolOutputSchema) MarshalJSON() ([]byte, error) {
 
 // MarshalJSON implements the json.Marshaler interface for ToolArgumentsSchema.
 func (tas ToolArgumentsSchema) MarshalJSON() ([]byte, error) {
-	m := make(map[string]any)
-	m["type"] = tas.Type
+	m := make(map[string]any, len(tas.Extra)+5)
+
+	// Re-emit preserved unmodeled keywords first; the modeled fields below take
+	// precedence on the (spec-wise impossible) chance of a key collision.
+	for k, v := range tas.Extra {
+		m[k] = v
+	}
+
+	// Emit "type" only when set. Emitting it unconditionally fabricated
+	// "type":"" for schemas that legitimately omit a top-level type (e.g. a
+	// top-level oneOf), which is not valid JSON Schema and misleads consumers.
+	if tas.Type != "" {
+		m["type"] = tas.Type
+	}
 
 	if tas.Defs != nil {
 		m["$defs"] = tas.Defs
 	}
 
-	// Marshal Properties to '{}' rather than `nil` when its length equals zero
-	if tas.Properties != nil {
+	// For object schemas keep the historical behavior of always emitting an
+	// explicit (possibly empty) properties/required, which clients rely on for
+	// no-argument object tools. A schema that is not object-typed and carries
+	// no properties/required of its own (e.g. a top-level oneOf/anyOf) is left
+	// alone rather than being polluted with a spurious empty properties/required.
+	isObjectSchema := tas.Type == "object"
+	switch {
+	case tas.Properties != nil:
 		m["properties"] = tas.Properties
-	} else {
+	case isObjectSchema:
 		m["properties"] = map[string]any{}
 	}
-
-	// Marshal Required to '[]' rather than `nil` when its length equals zero
-	if len(tas.Required) > 0 {
+	switch {
+	case len(tas.Required) > 0:
 		m["required"] = tas.Required
-	} else {
+	case isObjectSchema:
 		m["required"] = []string{}
 	}
 
@@ -448,6 +473,20 @@ func (tas *ToolArgumentsSchema) UnmarshalJSON(data []byte) error {
 		tas.Defs = aux.Definitions
 	}
 
+	// Preserve any top-level keywords not modeled by the struct fields (oneOf,
+	// anyOf, allOf, $ref, enum, const, patternProperties, ...) so they survive
+	// an unmarshal -> marshal round-trip instead of being silently dropped.
+	var all map[string]json.RawMessage
+	if err := json.Unmarshal(data, &all); err != nil {
+		return err
+	}
+	for _, modeled := range []string{"$defs", "definitions", "type", "properties", "required", "additionalProperties"} {
+		delete(all, modeled)
+	}
+	if len(all) > 0 {
+		tas.Extra = all
+	}
+
 	return nil
 }
 

mcpcompat/mcp/tools_schema_fidelity_test.go

@@ -0,0 +1,128 @@
+// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package mcp_test
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp"
+)
+
+const (
+	schemaTypeObject = "object"
+	keyProperties    = "properties"
+)
+
+// TestToolInputSchema_RoundTripPreservesCompositors verifies that a tool input
+// schema using top-level JSON Schema keywords not modeled by ToolArgumentsSchema
+// (oneOf/anyOf/allOf/$ref/enum/...) survives an unmarshal -> marshal round-trip,
+// and that a schema without a top-level "type" does not gain a fabricated
+// "type":"". Regression guard for stacklok/toolhive#5976.
+func TestToolInputSchema_RoundTripPreservesCompositors(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name          string
+		inputSchema   string
+		wantContains  []string // substrings that must survive the round-trip
+		wantNoType    bool     // true => marshaled output must not contain a "type" key
+		wantTypeValue string   // when non-empty, the "type" value that must be preserved
+	}{
+		{
+			name: "top-level oneOf with no type",
+			inputSchema: `{"oneOf":[` +
+				`{"type":"object","properties":{"a":{"type":"string"}},"required":["a"]},` +
+				`{"type":"object","properties":{"b":{"type":"string"}},"required":["b"]}]}`,
+			wantContains: []string{"oneOf"},
+			wantNoType:   true,
+		},
+		{
+			name:         "anyOf with no type",
+			inputSchema:  `{"anyOf":[{"type":"string"},{"type":"number"}]}`,
+			wantContains: []string{"anyOf"},
+			wantNoType:   true,
+		},
+		{
+			name:         "properties without top-level type",
+			inputSchema:  `{"properties":{"c":{"type":"string"}},"required":["c"]}`,
+			wantContains: []string{keyProperties, "\"c\""},
+			wantNoType:   true,
+		},
+		{
+			name:          "ordinary object schema is unchanged",
+			inputSchema:   `{"type":"object","properties":{"x":{"type":"string"}},"required":["x"]}`,
+			wantContains:  []string{keyProperties, "\"x\""},
+			wantTypeValue: schemaTypeObject,
+		},
+		{
+			name:          "object schema with extra compositor keyword",
+			inputSchema:   `{"type":"object","properties":{"x":{"type":"string"}},"allOf":[{"required":["x"]}]}`,
+			wantContains:  []string{"allOf", keyProperties},
+			wantTypeValue: schemaTypeObject,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			var schema mcp.ToolInputSchema
+			require.NoError(t, json.Unmarshal([]byte(tt.inputSchema), &schema),
+				"unmarshal must succeed")
+
+			out, err := json.Marshal(schema)
+			require.NoError(t, err, "marshal must succeed")
+
+			var got map[string]any
+			require.NoError(t, json.Unmarshal(out, &got), "remarshaled output must be valid JSON")
+
+			for _, want := range tt.wantContains {
+				assert.Contains(t, string(out), want,
+					"round-trip must preserve %q; got %s", want, string(out))
+			}
+
+			if tt.wantNoType {
+				if typ, ok := got["type"]; ok {
+					assert.NotEqual(t, "", typ,
+						"a schema without a top-level type must not gain a fabricated empty type; got %s", string(out))
+				}
+			}
+			if tt.wantTypeValue != "" {
+				assert.Equal(t, tt.wantTypeValue, got["type"],
+					"top-level type must be preserved; got %s", string(out))
+			}
+		})
+	}
+}
+
+// TestToolInputSchema_RoundTripThroughTool verifies the fidelity holds when the
+// schema is nested inside a Tool decoded from a tools/list-style payload — the
+// exact path a client takes when ingesting a backend's advertised tools.
+func TestToolInputSchema_RoundTripThroughTool(t *testing.T) {
+	t.Parallel()
+
+	raw := `{
+		"name": "compose",
+		"description": "oneOf tool",
+		"inputSchema": {"oneOf":[{"type":"object"},{"type":"string"}]}
+	}`
+
+	var tool mcp.Tool
+	require.NoError(t, json.Unmarshal([]byte(raw), &tool))
+
+	out, err := json.Marshal(tool.InputSchema)
+	require.NoError(t, err)
+	assert.Contains(t, string(out), "oneOf",
+		"the oneOf compositor must survive ingestion into mcp.Tool; got %s", string(out))
+
+	var got map[string]any
+	require.NoError(t, json.Unmarshal(out, &got))
+	if typ, ok := got["type"]; ok {
+		assert.NotEqual(t, "", typ, "must not fabricate an empty top-level type; got %s", string(out))
+	}
+}