diff --git a/mcp/tools.go b/mcp/tools.go index 3e3931b09..88d0e0e28 100644 --- a/mcp/tools.go +++ b/mcp/tools.go @@ -740,7 +740,7 @@ func Pattern(pattern string) PropertyOption { // DefaultNumber sets the default value for a number property. // This value will be used if the property is not explicitly provided. -func DefaultNumber(value float64) PropertyOption { +func DefaultNumber[T int |int64 | float64](value T) PropertyOption { return func(schema map[string]any) { schema["default"] = value } @@ -748,7 +748,7 @@ func DefaultNumber(value float64) PropertyOption { // Max sets the maximum value for a number property. // The number value must not exceed this maximum. -func Max(max float64) PropertyOption { +func Max[T int | int64 | float64](max T) PropertyOption { return func(schema map[string]any) { schema["maximum"] = max } @@ -756,7 +756,7 @@ func Max(max float64) PropertyOption { // Min sets the minimum value for a number property. // The number value must not be less than this minimum. -func Min(min float64) PropertyOption { +func Min[T int | int64 | float64](min T) PropertyOption { return func(schema map[string]any) { schema["minimum"] = min } @@ -764,12 +764,11 @@ func Min(min float64) PropertyOption { // MultipleOf specifies that a number must be a multiple of the given value. // The number value must be divisible by this value. -func MultipleOf(value float64) PropertyOption { +func MultipleOf[T int | int64 | float64](value T) PropertyOption { return func(schema map[string]any) { schema["multipleOf"] = value } } - // // Boolean Property Options // @@ -820,6 +819,28 @@ func WithBoolean(name string, opts ...PropertyOption) ToolOption { } } +// WithInteger adds an integer property to the tool schema. +// It accepts propety operations to configure the integer property's behavior and constraints. +func WithInteger(name string, opts ...PropertyOption) ToolOption { + return func(t *Tool) { + schema := map[string]any{ + "type": "integer", + } + + 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") + t.InputSchema.Required = append(t.InputSchema.Required, name) + } + + t.InputSchema.Properties[name] = schema + } +} + // WithNumber adds a number property to the tool schema. // It accepts property options to configure the number property's behavior and constraints. func WithNumber(name string, opts ...PropertyOption) ToolOption { @@ -1052,6 +1073,31 @@ func WithNumberItems(opts ...PropertyOption) PropertyOption { } } +// WithIntegerItems configures an array's items to be of type integer. +// +// Supported options: Description(), DefaultNumber(), Min(), Max(), MultipleOf() +// Note: Options like Required() are not valid for item schemas and will be ignored. +// +// Examples: +// +// mcp.WithArray("scores", mcp.WithIntegerItems(mcp.Min(0), mcp.Max(100))) +// mcp.WithArray("prices", mcp.WithIntegerItems(mcp.Min(0))) +// +// Limitations: Only supports simple integer arrays. Use Items() for complex objects. +func WithIntegerItems(opts ...PropertyOption) PropertyOption { + return func(schema map[string]any) { + itemSchema := map[string]any{ + "type": "integer", + } + + for _, opt := range opts { + opt(itemSchema) + } + + schema["items"] = itemSchema + } +} + // WithBooleanItems configures an array's items to be of type boolean. // // Supported options: Description(), DefaultBool() diff --git a/mcp/tools_test.go b/mcp/tools_test.go index 0cd71230e..69b6a358f 100644 --- a/mcp/tools_test.go +++ b/mcp/tools_test.go @@ -144,6 +144,54 @@ func TestUnmarshalToolWithoutRawSchema(t *testing.T) { assert.Empty(t, toolUnmarshalled.RawInputSchema) } +func TestToolWithStringInteger(t *testing.T) { + tool := NewTool("buy-item", + WithDescription("A tool for buying items"), + WithString("itemName", + Description("Name of the item to purchase"), + Required(), + ), + WithInteger("itemCount", + Description("Number of copies of the item to purchase"), + Required(), + ), + ) + + // Marshal to JSON + data, err := json.Marshal(tool) + assert.NoError(t, err) + + // Unmarshal to verify the structure + var result map[string]any + err = json.Unmarshal(data, &result) + assert.NoError(t, err) + + // Verify tool properties + assert.Equal(t, "buy-item", result["name"]) + assert.Equal(t, "A tool for buying items", result["description"]) + + // Verify schema was properly included + schema, ok := result["inputSchema"].(map[string]any) + assert.True(t, ok) + assert.Equal(t, "object", schema["type"]) + + // Verify properties + properties, ok := schema["properties"].(map[string]any) + assert.True(t, ok) + + // Verify itemName object + itemName, ok := properties["itemName"].(map[string]any) + assert.True(t, ok) + assert.Equal(t, "string", itemName["type"]) + assert.Equal(t, "Name of the item to purchase", itemName["description"]) + + // Verify itemName object + itemCount, ok := properties["itemCount"].(map[string]any) + assert.True(t, ok) + assert.Equal(t, "integer", itemCount["type"]) + assert.Equal(t, "Number of copies of the item to purchase", itemCount["description"]) +} + func TestToolWithObjectAndArray(t *testing.T) { // Create a tool with both object and array properties tool := NewTool("reading-list", @@ -656,6 +704,46 @@ func TestNewItemsAPICompatibility(t *testing.T) { ), ), }, + { + name: "WithIntegerItems basic", + oldTool: NewTool("old-integer-basic", + WithDescription("tool with integer array using old API"), + WithArray("scores", + Description("List of scores"), + Items(map[string]any{ + "type": "integer", + }), + ), + ), + newTool: NewTool("new-integer-basic", + WithDescription("tool with integer array using new API"), + WithArray("scores", + Description("List of scores"), + WithIntegerItems(), + ), + ), + }, + { + name: "WithIntegerItems with constraints", + oldTool: NewTool("old-integer-with-constraints", + WithDescription("Tool with constrained integer array using old API"), + WithArray("ratings", + Description("List of ratings"), + Items(map[string]any{ + "type": "integer", + "minimum": 0, + "maximum": 10, + }), + ), + ), + newTool: NewTool("new-number-with-constraints", + WithDescription("Tool with constrained number array using new API"), + WithArray("ratings", + Description("List of ratings"), + WithIntegerItems(Min(0), Max(10)), + ), + ), + }, } for _, tt := range tests {