@@ -7,6 +7,7 @@ package mcp
77import (
88 "context"
99 "encoding/json"
10+ "slices"
1011
1112 "golang.org/x/tools/internal/mcp/internal/jsonschema"
1213 "golang.org/x/tools/internal/mcp/internal/protocol"
@@ -23,13 +24,16 @@ type Tool struct {
2324
2425// MakeTool is a helper to make a tool using reflection on the given handler.
2526//
27+ // If provided, variadic [ToolOption] values may be used to customize the tool.
28+ //
2629// The input schema for the tool is extracted from the request type for the
27- // handler, and used to unmmarshal and validate requests to the handler.
30+ // handler, and used to unmmarshal and validate requests to the handler. This
31+ // schema may be customized using the [Input] option.
2832//
2933// It is the caller's responsibility that the handler request type can produce
3034// a valid schema, as documented by [jsonschema.ForType]; otherwise, MakeTool
3135// panics.
32- func MakeTool [TReq any ](name , description string , handler func (context.Context , * ClientConnection , TReq ) ([]Content , error )) * Tool {
36+ func MakeTool [TReq any ](name , description string , handler func (context.Context , * ClientConnection , TReq ) ([]Content , error ), opts ... ToolOption ) * Tool {
3337 schema , err := jsonschema .For [TReq ]()
3438 if err != nil {
3539 panic (err )
@@ -51,14 +55,18 @@ func MakeTool[TReq any](name, description string, handler func(context.Context,
5155 }
5256 return res , nil
5357 }
54- return & Tool {
58+ t := & Tool {
5559 Definition : protocol.Tool {
5660 Name : name ,
5761 Description : description ,
5862 InputSchema : schema ,
5963 },
6064 Handler : wrapped ,
6165 }
66+ for _ , opt := range opts {
67+ opt .set (t )
68+ }
69+ return t
6270}
6371
6472// unmarshalSchema unmarshals data into v and validates the result according to
@@ -68,3 +76,94 @@ func unmarshalSchema(data json.RawMessage, _ *jsonschema.Schema, v any) error {
6876 // Separate validation from assignment.
6977 return json .Unmarshal (data , v )
7078}
79+
80+ // A ToolOption configures the behavior of a Tool.
81+ type ToolOption interface {
82+ set (* Tool )
83+ }
84+
85+ type toolSetter func (* Tool )
86+
87+ func (s toolSetter ) set (t * Tool ) { s (t ) }
88+
89+ // Input applies the provided [SchemaOption] configuration to the tool's input
90+ // schema.
91+ func Input (opts ... SchemaOption ) ToolOption {
92+ return toolSetter (func (t * Tool ) {
93+ for _ , opt := range opts {
94+ opt .set (t .Definition .InputSchema )
95+ }
96+ })
97+ }
98+
99+ // A SchemaOption configures a jsonschema.Schema.
100+ type SchemaOption interface {
101+ set (s * jsonschema.Schema )
102+ }
103+
104+ type schemaSetter func (* jsonschema.Schema )
105+
106+ func (s schemaSetter ) set (schema * jsonschema.Schema ) { s (schema ) }
107+
108+ // Property configures the schema for the property of the given name.
109+ // If there is no such property in the schema, it is created.
110+ func Property (name string , opts ... SchemaOption ) SchemaOption {
111+ return schemaSetter (func (schema * jsonschema.Schema ) {
112+ propSchema , ok := schema .Properties [name ]
113+ if ! ok {
114+ propSchema = new (jsonschema.Schema )
115+ schema .Properties [name ] = propSchema
116+ }
117+ // Apply the options, with special handling for Required, as it needs to be
118+ // set on the parent schema.
119+ for _ , opt := range opts {
120+ if req , ok := opt .(required ); ok {
121+ if req {
122+ if ! slices .Contains (schema .Required , name ) {
123+ schema .Required = append (schema .Required , name )
124+ }
125+ } else {
126+ schema .Required = slices .DeleteFunc (schema .Required , func (s string ) bool {
127+ return s == name
128+ })
129+ }
130+ } else {
131+ opt .set (propSchema )
132+ }
133+ }
134+ })
135+ }
136+
137+ // Required sets whether the associated property is required. It is only valid
138+ // when used in a [Property] option: using Required outside of Property panics.
139+ func Required (v bool ) SchemaOption {
140+ return required (v )
141+ }
142+
143+ type required bool
144+
145+ func (required ) set (s * jsonschema.Schema ) {
146+ panic ("use of required outside of Property" )
147+ }
148+
149+ // Enum sets the provided values as the "enum" value of the schema.
150+ func Enum (values ... any ) SchemaOption {
151+ return schemaSetter (func (s * jsonschema.Schema ) {
152+ s .Enum = values
153+ })
154+ }
155+
156+ // Description sets the provided schema description.
157+ func Description (description string ) SchemaOption {
158+ return schemaSetter (func (schema * jsonschema.Schema ) {
159+ schema .Description = description
160+ })
161+ }
162+
163+ // Schema overrides the inferred schema with a shallow copy of the given
164+ // schema.
165+ func Schema (schema * jsonschema.Schema ) SchemaOption {
166+ return schemaSetter (func (s * jsonschema.Schema ) {
167+ * s = * schema
168+ })
169+ }
0 commit comments