tools: add WithObject and WithArray to supplement type definition utility (#47)

Author: xhd2015

mcp/tools.go

@@ -367,3 +367,120 @@ func WithString(name string, opts ...PropertyOption) ToolOption {
 		t.InputSchema.Properties[name] = schema
 	}
 }
+
+// WithObject adds an object property to the tool schema.
+// It accepts property options to configure the object property's behavior and constraints.
+func WithObject(name string, opts ...PropertyOption) ToolOption {
+	return func(t *Tool) {
+		schema := map[string]interface{}{
+			"type":       "object",
+			"properties": map[string]interface{}{},
+		}
+
+		for _, opt := range opts {
+			opt(schema)
+		}
+
+		// Remove required from property schema and add to InputSchema.required
+		if required, ok := schema["required"].(bool); ok && required {
+			delete(schema, "required")
+			if t.InputSchema.Required == nil {
+				t.InputSchema.Required = []string{name}
+			} else {
+				t.InputSchema.Required = append(t.InputSchema.Required, name)
+			}
+		}
+
+		t.InputSchema.Properties[name] = schema
+	}
+}
+
+// WithArray adds an array property to the tool schema.
+// It accepts property options to configure the array property's behavior and constraints.
+func WithArray(name string, opts ...PropertyOption) ToolOption {
+	return func(t *Tool) {
+		schema := map[string]interface{}{
+			"type": "array",
+		}
+
+		for _, opt := range opts {
+			opt(schema)
+		}
+
+		// Remove required from property schema and add to InputSchema.required
+		if required, ok := schema["required"].(bool); ok && required {
+			delete(schema, "required")
+			if t.InputSchema.Required == nil {
+				t.InputSchema.Required = []string{name}
+			} else {
+				t.InputSchema.Required = append(t.InputSchema.Required, name)
+			}
+		}
+
+		t.InputSchema.Properties[name] = schema
+	}
+}
+
+// Properties defines the properties for an object schema
+func Properties(props map[string]interface{}) PropertyOption {
+	return func(schema map[string]interface{}) {
+		schema["properties"] = props
+	}
+}
+
+// AdditionalProperties specifies whether additional properties are allowed in the object
+// or defines a schema for additional properties
+func AdditionalProperties(schema interface{}) PropertyOption {
+	return func(schemaMap map[string]interface{}) {
+		schemaMap["additionalProperties"] = schema
+	}
+}
+
+// MinProperties sets the minimum number of properties for an object
+func MinProperties(min int) PropertyOption {
+	return func(schema map[string]interface{}) {
+		schema["minProperties"] = min
+	}
+}
+
+// MaxProperties sets the maximum number of properties for an object
+func MaxProperties(max int) PropertyOption {
+	return func(schema map[string]interface{}) {
+		schema["maxProperties"] = max
+	}
+}
+
+// PropertyNames defines a schema for property names in an object
+func PropertyNames(schema map[string]interface{}) PropertyOption {
+	return func(schemaMap map[string]interface{}) {
+		schemaMap["propertyNames"] = schema
+	}
+}
+
+// Items defines the schema for array items
+func Items(schema interface{}) PropertyOption {
+	return func(schemaMap map[string]interface{}) {
+		schemaMap["items"] = schema
+	}
+}
+
+// MinItems sets the minimum number of items for an array
+func MinItems(min int) PropertyOption {
+	return func(schema map[string]interface{}) {
+		schema["minItems"] = min
+	}
+}
+
+// MaxItems sets the maximum number of items for an array
+func MaxItems(max int) PropertyOption {
+	return func(schema map[string]interface{}) {
+		schema["maxItems"] = max
+	}
+}
+
+// UniqueItems specifies whether array items must be unique
+func UniqueItems(unique bool) PropertyOption {
+	return func(schema map[string]interface{}) {
+		schema["uniqueItems"] = unique
+	}
+}

mcp/tools_test.go

@@ -141,4 +141,102 @@ func TestUnmarshalToolWithoutRawSchema(t *testing.T) {
 	})
 	assert.Empty(t, toolUnmarshalled.InputSchema.Required)
 	assert.Empty(t, toolUnmarshalled.RawInputSchema)
-}
+}
+
+func TestToolWithObjectAndArray(t *testing.T) {
+	// Create a tool with both object and array properties
+	tool := NewTool("reading-list",
+		WithDescription("A tool for managing reading lists"),
+		WithObject("preferences",
+			Description("User preferences for the reading list"),
+			Properties(map[string]interface{}{
+				"theme": map[string]interface{}{
+					"type":        "string",
+					"description": "UI theme preference",
+					"enum":        []string{"light", "dark"},
+				},
+				"maxItems": map[string]interface{}{
+					"type":        "number",
+					"description": "Maximum number of items in the list",
+					"minimum":     1,
+					"maximum":     100,
+				},
+			})),
+		WithArray("books",
+			Description("List of books to read"),
+			Required(),
+			Items(map[string]interface{}{
+				"type": "object",
+				"properties": map[string]interface{}{
+					"title": map[string]interface{}{
+						"type":        "string",
+						"description": "Book title",
+						"required":    true,
+					},
+					"author": map[string]interface{}{
+						"type":        "string",
+						"description": "Book author",
+					},
+					"year": map[string]interface{}{
+						"type":        "number",
+						"description": "Publication year",
+						"minimum":     1000,
+					},
+				},
+			})))
+
+	// Marshal to JSON
+	data, err := json.Marshal(tool)
+	assert.NoError(t, err)
+
+	// Unmarshal to verify the structure
+	var result map[string]interface{}
+	err = json.Unmarshal(data, &result)
+	assert.NoError(t, err)
+
+	// Verify tool properties
+	assert.Equal(t, "reading-list", result["name"])
+	assert.Equal(t, "A tool for managing reading lists", result["description"])
+
+	// Verify schema was properly included
+	schema, ok := result["inputSchema"].(map[string]interface{})
+	assert.True(t, ok)
+	assert.Equal(t, "object", schema["type"])
+
+	// Verify properties
+	properties, ok := schema["properties"].(map[string]interface{})
+	assert.True(t, ok)
+
+	// Verify preferences object
+	preferences, ok := properties["preferences"].(map[string]interface{})
+	assert.True(t, ok)
+	assert.Equal(t, "object", preferences["type"])
+	assert.Equal(t, "User preferences for the reading list", preferences["description"])
+
+	prefProps, ok := preferences["properties"].(map[string]interface{})
+	assert.True(t, ok)
+	assert.Contains(t, prefProps, "theme")
+	assert.Contains(t, prefProps, "maxItems")
+
+	// Verify books array
+	books, ok := properties["books"].(map[string]interface{})
+	assert.True(t, ok)
+	assert.Equal(t, "array", books["type"])
+	assert.Equal(t, "List of books to read", books["description"])
+
+	// Verify array items schema
+	items, ok := books["items"].(map[string]interface{})
+	assert.True(t, ok)
+	assert.Equal(t, "object", items["type"])
+
+	itemProps, ok := items["properties"].(map[string]interface{})
+	assert.True(t, ok)
+	assert.Contains(t, itemProps, "title")
+	assert.Contains(t, itemProps, "author")
+	assert.Contains(t, itemProps, "year")
+
+	// Verify required fields
+	required, ok := schema["required"].([]interface{})
+	assert.True(t, ok)
+	assert.Contains(t, required, "books")
+}