Skip to content

Commit 3818858

Browse files
committed
internal/mcp/protocol: make type name more convenient
You can now write "*" for a name in a config to use the original name. Also, fix a bug in dealing with initialisms. Change-Id: I4d5d0acb1200fff6500644457ec34c5026d63f9f Reviewed-on: https://go-review.googlesource.com/c/tools/+/671361 LUCI-TryBot-Result: Go LUCI <[email protected]> Reviewed-by: Robert Findley <[email protected]> Reviewed-by: Sam Thanawalla <[email protected]>
1 parent bbef3e4 commit 3818858

File tree

1 file changed

+56
-36
lines changed

1 file changed

+56
-36
lines changed

internal/mcp/protocol/generate.go

Lines changed: 56 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -46,64 +46,71 @@ type config map[string]*typeConfig
4646

4747
// declarations configures the set of declarations to write.
4848
//
49-
// Top level declarations are only created if they are configured with a
50-
// non-empty Name. Otherwise, they are discarded, though their fields may be
49+
// Top level declarations are created unless configured with Name=="-",
50+
// in which case they are discarded, though their fields may be
5151
// extracted to types if they have a nested field configuration.
52+
// If Name == "", the map key is used as the type name.
5253
var declarations = config{
53-
"Annotations": {Name: "Annotations"},
54+
"Annotations": {},
5455
"CallToolRequest": {
56+
Name: "-",
5557
Fields: config{"Params": {Name: "CallToolParams"}},
5658
},
57-
"CallToolResult": {Name: "CallToolResult"},
59+
"CallToolResult": {},
5860
"CancelledNotification": {
61+
Name: "-",
5962
Fields: config{"Params": {Name: "CancelledParams"}},
6063
},
61-
"ClientCapabilities": {Name: "ClientCapabilities"},
64+
"ClientCapabilities": {},
6265
"GetPromptRequest": {
66+
Name: "-",
6367
Fields: config{"Params": {Name: "GetPromptParams"}},
6468
},
65-
"GetPromptResult": {Name: "GetPromptResult"},
66-
"Implementation": {Name: "Implementation"},
69+
"GetPromptResult": {},
70+
"Implementation": {},
6771
"InitializeRequest": {
72+
Name: "-",
6873
Fields: config{"Params": {Name: "InitializeParams"}},
6974
},
70-
"InitializeResult": {Name: "InitializeResult"},
75+
"InitializeResult": {},
7176
"InitializedNotification": {
77+
Name: "-",
7278
Fields: config{"Params": {Name: "InitializedParams"}},
7379
},
7480
"ListPromptsRequest": {
81+
Name: "-",
7582
Fields: config{"Params": {Name: "ListPromptsParams"}},
7683
},
77-
"ListPromptsResult": {Name: "ListPromptsResult"},
84+
"ListPromptsResult": {},
7885
"ListRootsRequest": {
86+
Name: "-",
7987
Fields: config{"Params": {Name: "ListRootsParams"}},
8088
},
81-
"ListRootsResult": {Name: "ListRootsResult"},
89+
"ListRootsResult": {},
8290
"ListToolsRequest": {
91+
Name: "-",
8392
Fields: config{"Params": {Name: "ListToolsParams"}},
8493
},
85-
"ListToolsResult": {Name: "ListToolsResult"},
86-
"Prompt": {Name: "Prompt"},
87-
"PromptMessage": {Name: "PromptMessage"},
88-
"PromptArgument": {Name: "PromptArgument"},
89-
"ProgressToken": {Substitute: "any"}, // null|number|string
90-
"RequestId": {Substitute: "any"}, // null|number|string
91-
"Role": {Name: "Role"},
92-
"Root": {Name: "Root"},
94+
"ListToolsResult": {},
95+
"Prompt": {},
96+
"PromptMessage": {},
97+
"PromptArgument": {},
98+
"ProgressToken": {Name: "-", Substitute: "any"}, // null|number|string
99+
"RequestId": {Name: "-", Substitute: "any"}, // null|number|string
100+
"Role": {},
101+
"Root": {},
93102

94103
"ServerCapabilities": {
95-
Name: "ServerCapabilities",
96104
Fields: config{
97105
"Prompts": {Name: "PromptCapabilities"},
98106
"Resources": {Name: "ResourceCapabilities"},
99107
"Tools": {Name: "ToolCapabilities"},
100108
},
101109
},
102110
"Tool": {
103-
Name: "Tool",
104111
Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}},
105112
},
106-
"ToolAnnotations": {Name: "ToolAnnotations"},
113+
"ToolAnnotations": {},
107114
}
108115

109116
func main() {
@@ -128,7 +135,7 @@ func main() {
128135
if config == nil {
129136
continue
130137
}
131-
if err := writeDecl(*config, def, named); err != nil {
138+
if err := writeDecl(name, *config, def, named); err != nil {
132139
log.Fatal(err)
133140
}
134141
}
@@ -199,19 +206,22 @@ func loadSchema(schemaFile string) (data []byte, err error) {
199206
return data, nil
200207
}
201208

202-
func writeDecl(config typeConfig, def *jsonschema.Schema, named map[string]*bytes.Buffer) error {
209+
func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, named map[string]*bytes.Buffer) error {
203210
var w io.Writer = io.Discard
204-
if name := config.Name; name != "" {
205-
if _, ok := named[name]; ok {
211+
if typeName := config.Name; typeName != "-" {
212+
if typeName == "" {
213+
typeName = configName
214+
}
215+
if _, ok := named[typeName]; ok {
206216
return nil
207217
}
208218
buf := new(bytes.Buffer)
209219
w = buf
210-
named[name] = buf
220+
named[typeName] = buf
211221
if def.Description != "" {
212222
fmt.Fprintf(buf, "%s\n", toComment(def.Description))
213223
}
214-
fmt.Fprintf(buf, "type %s ", name)
224+
fmt.Fprintf(buf, "type %s ", typeName)
215225
}
216226
if err := writeType(w, &config, def, named); err != nil {
217227
return err // Better error here?
@@ -234,12 +244,12 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
234244
// definition is missing, *but only if w is not io.Discard*. That's not a
235245
// great API: see if we can do something more explicit than io.Discard.
236246
if cfg, ok := declarations[name]; ok {
237-
if cfg.Name == "" && cfg.Substitute == "" {
247+
if cfg.Name == "-" && cfg.Substitute == "" {
238248
panic(fmt.Sprintf("referenced type %q cannot be referred to (no name or substitution)", name))
239249
}
240250
if cfg.Substitute != "" {
241251
name = cfg.Substitute
242-
} else {
252+
} else if cfg.Name != "" {
243253
name = cfg.Name
244254
}
245255
}
@@ -311,14 +321,18 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
311321
if r.Substitute != "" {
312322
fmt.Fprintf(w, r.Substitute)
313323
} else {
314-
assert(r.Name != "", "missing ExtractTo")
315-
if err := writeDecl(*r, fieldDef, named); err != nil {
324+
assert(r.Name != "-", "missing ExtractTo")
325+
typename := export
326+
if r.Name != "" {
327+
typename = r.Name
328+
}
329+
if err := writeDecl(typename, *r, fieldDef, named); err != nil {
316330
return err
317331
}
318332
if needPointer {
319333
fmt.Fprintf(w, "*")
320334
}
321-
fmt.Fprintf(w, r.Name)
335+
fmt.Fprintf(w, typename)
322336
}
323337
} else {
324338
if needPointer {
@@ -410,7 +424,12 @@ func exportName(s string) string {
410424
// at once, because the replacement will change the indices.)
411425
for {
412426
if loc := re.FindStringIndex(s); loc != nil {
413-
s = s[:loc[0]] + replacement + s[loc[1]:]
427+
// Don't replace the rune after the initialism, if any.
428+
end := loc[1]
429+
if end < len(s) {
430+
end--
431+
}
432+
s = s[:loc[0]] + replacement + s[end:]
414433
} else {
415434
break
416435
}
@@ -421,9 +440,10 @@ func exportName(s string) string {
421440

422441
// Map from initialism to the regexp that matches it.
423442
var initialisms = map[string]*regexp.Regexp{
424-
"Id": nil,
425-
"Url": nil,
426-
"Uri": nil,
443+
"Id": nil,
444+
"Url": nil,
445+
"Uri": nil,
446+
"Mime": nil,
427447
}
428448

429449
func init() {

0 commit comments

Comments
 (0)