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 <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>

mcp/protocol/generate.go

@@ -46,53 +46,61 @@ type config map[string]*typeConfig
 
 // declarations configures the set of declarations to write.
 //
-// Top level declarations are only created if they are configured with a
+// Top level declarations are created unless configured with Name=="-",
-// non-empty Name. Otherwise, they are discarded, though their fields may be
+// in which case they are discarded, though their fields may be
 // extracted to types if they have a nested field configuration.
+// If Name == "", the map key is used as the type name.
 var declarations = config{
-	"Annotations": {Name: "Annotations"},
+	"Annotations": {},
 	"CallToolRequest": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "CallToolParams"}},
 	},
-	"CallToolResult": {Name: "CallToolResult"},
+	"CallToolResult": {},
 	"CancelledNotification": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "CancelledParams"}},
 	},
-	"ClientCapabilities": {Name: "ClientCapabilities"},
+	"ClientCapabilities": {},
 	"GetPromptRequest": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "GetPromptParams"}},
 	},
-	"GetPromptResult": {Name: "GetPromptResult"},
+	"GetPromptResult": {},
-	"Implementation":  {Name: "Implementation"},
+	"Implementation":  {},
 	"InitializeRequest": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "InitializeParams"}},
 	},
-	"InitializeResult": {Name: "InitializeResult"},
+	"InitializeResult": {},
 	"InitializedNotification": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "InitializedParams"}},
 	},
 	"ListPromptsRequest": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "ListPromptsParams"}},
 	},
-	"ListPromptsResult": {Name: "ListPromptsResult"},
+	"ListPromptsResult": {},
 	"ListRootsRequest": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "ListRootsParams"}},
 	},
-	"ListRootsResult": {Name: "ListRootsResult"},
+	"ListRootsResult": {},
 	"ListToolsRequest": {
+		Name:   "-",
 		Fields: config{"Params": {Name: "ListToolsParams"}},
 	},
-	"ListToolsResult": {Name: "ListToolsResult"},
+	"ListToolsResult": {},
-	"Prompt":          {Name: "Prompt"},
+	"Prompt":          {},
-	"PromptMessage":   {Name: "PromptMessage"},
+	"PromptMessage":   {},
-	"PromptArgument":  {Name: "PromptArgument"},
+	"PromptArgument":  {},
-	"ProgressToken":   {Substitute: "any"}, // null|number|string
+	"ProgressToken":   {Name: "-", Substitute: "any"}, // null|number|string
-	"RequestId":       {Substitute: "any"}, // null|number|string
+	"RequestId":       {Name: "-", Substitute: "any"}, // null|number|string
-	"Role":            {Name: "Role"},
+	"Role":            {},
-	"Root":            {Name: "Root"},
+	"Root":            {},
 
 	"ServerCapabilities": {
-		Name: "ServerCapabilities",
 		Fields: config{
 			"Prompts":   {Name: "PromptCapabilities"},
 			"Resources": {Name: "ResourceCapabilities"},
@@ -100,10 +108,9 @@ var declarations = config{
 		},
 	},
 	"Tool": {
-		Name:   "Tool",
 		Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}},
 	},
-	"ToolAnnotations": {Name: "ToolAnnotations"},
+	"ToolAnnotations": {},
 }
 
 func main() {
@@ -128,7 +135,7 @@ func main() {
 		if config == nil {
 			continue
 		}
-		if err := writeDecl(*config, def, named); err != nil {
+		if err := writeDecl(name, *config, def, named); err != nil {
 			log.Fatal(err)
 		}
 	}
@@ -199,19 +206,22 @@ func loadSchema(schemaFile string) (data []byte, err error) {
 	return data, nil
 }
 
-func writeDecl(config typeConfig, def *jsonschema.Schema, named map[string]*bytes.Buffer) error {
+func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, named map[string]*bytes.Buffer) error {
 	var w io.Writer = io.Discard
-	if name := config.Name; name != "" {
+	if typeName := config.Name; typeName != "-" {
-		if _, ok := named[name]; ok {
+		if typeName == "" {
+			typeName = configName
+		}
+		if _, ok := named[typeName]; ok {
 			return nil
 		}
 		buf := new(bytes.Buffer)
 		w = buf
-		named[name] = buf
+		named[typeName] = buf
 		if def.Description != "" {
 			fmt.Fprintf(buf, "%s\n", toComment(def.Description))
 		}
-		fmt.Fprintf(buf, "type %s ", name)
+		fmt.Fprintf(buf, "type %s ", typeName)
 	}
 	if err := writeType(w, &config, def, named); err != nil {
 		return err // Better error here?
@@ -234,12 +244,12 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
 		// definition is missing, *but only if w is not io.Discard*. That's not a
 		// great API: see if we can do something more explicit than io.Discard.
 		if cfg, ok := declarations[name]; ok {
-			if cfg.Name == "" && cfg.Substitute == "" {
+			if cfg.Name == "-" && cfg.Substitute == "" {
 				panic(fmt.Sprintf("referenced type %q cannot be referred to (no name or substitution)", name))
 			}
 			if cfg.Substitute != "" {
 				name = cfg.Substitute
-			} else {
+			} else if cfg.Name != "" {
 				name = cfg.Name
 			}
 		}
@@ -311,14 +321,18 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
 					if r.Substitute != "" {
 						fmt.Fprintf(w, r.Substitute)
 					} else {
-						assert(r.Name != "", "missing ExtractTo")
+						assert(r.Name != "-", "missing ExtractTo")
-						if err := writeDecl(*r, fieldDef, named); err != nil {
+						typename := export
+						if r.Name != "" {
+							typename = r.Name
+						}
+						if err := writeDecl(typename, *r, fieldDef, named); err != nil {
 							return err
 						}
 						if needPointer {
 							fmt.Fprintf(w, "*")
 						}
-						fmt.Fprintf(w, r.Name)
+						fmt.Fprintf(w, typename)
 					}
 				} else {
 					if needPointer {
@@ -410,7 +424,12 @@ func exportName(s string) string {
 		// at once, because the replacement will change the indices.)
 		for {
 			if loc := re.FindStringIndex(s); loc != nil {
-				s = s[:loc[0]] + replacement + s[loc[1]:]
+				// Don't replace the rune after the initialism, if any.
+				end := loc[1]
+				if end < len(s) {
+					end--
+				}
+				s = s[:loc[0]] + replacement + s[end:]
 			} else {
 				break
 			}
@@ -421,9 +440,10 @@ func exportName(s string) string {
 
 // Map from initialism to the regexp that matches it.
 var initialisms = map[string]*regexp.Regexp{
-	"Id":  nil,
+	"Id":   nil,
-	"Url": nil,
+	"Url":  nil,
-	"Uri": nil,
+	"Uri":  nil,
+	"Mime": nil,
 }
 
 func init() {