From 5bcbf820f2c4f32f214aec0c62498d57a24e338e Mon Sep 17 00:00:00 2001
From: Jonathan Amsterdam <jba@google.com>
Date: Thu, 8 May 2025 07:09:48 -0400
Subject: [PATCH] internal/mcp/protocol: spell "Id" idiomatically

Enhance the generator so that it recognizes initialisms like "Id" and
writes them in all caps, as is idiomatic for Go.

Also, fix the import path of jsonschema in the generated code.

Change-Id: Ib12e9fbd543d360035d5e324f0aefca584ac7c85
Reviewed-on: https://go-review.googlesource.com/c/tools/+/670916
Auto-Submit: Jonathan Amsterdam <jba@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
---
 mcp/protocol/generate.go | 40 +++++++++++++++++++++++++++++++++++++---
 mcp/protocol/protocol.go |  2 +-
 mcp/transport.go         |  4 ++--
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/mcp/protocol/generate.go b/mcp/protocol/generate.go
index 60b0f0a..e082502 100644
--- a/mcp/protocol/generate.go
+++ b/mcp/protocol/generate.go
@@ -22,11 +22,12 @@ import (
 	"net/http"
 	"os"
 	"reflect"
+	"regexp"
 	"slices"
 	"strings"
 
-	"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
 	"github.com/modelcontextprotocol/go-sdk/internal/util"
+	"github.com/modelcontextprotocol/go-sdk/jsonschema"
 )
 
 var schemaFile = flag.String("schema_file", "", "if set, use this file as the persistent schema file")
@@ -138,7 +139,7 @@ package protocol
 import (
 	"encoding/json"
 
-	"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
+	"github.com/modelcontextprotocol/go-sdk/jsonschema"
 )
 `)
 
@@ -381,11 +382,44 @@ func canHaveAdditionalProperties(s *jsonschema.Schema) bool {
 
 // exportName returns an exported name for a Go symbol, based on the given name
 // in the JSON schema, removing leading underscores and capitalizing.
+// It also rewrites initialisms.
 func exportName(s string) string {
 	if strings.HasPrefix(s, "_") {
 		s = s[1:]
 	}
-	return strings.ToUpper(s[:1]) + s[1:]
+	s = strings.ToUpper(s[:1]) + s[1:]
+	// Replace an initialism if it is its own "word": see the init function below for
+	// a definition.
+	// There is probably a clever way to write this whole thing with one regexp and
+	// a Replace method, but it would be quite obscure.
+	// This doesn't have to be fast, because the first match will rarely succeed.
+	for ism, re := range initialisms {
+		replacement := strings.ToUpper(ism)
+		// Find the index of one match at a time, and replace. (We can't find all
+		// 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]:]
+			} else {
+				break
+			}
+		}
+	}
+	return s
+}
+
+// Map from initialism to the regexp that matches it.
+var initialisms = map[string]*regexp.Regexp{
+	"Id":  nil,
+	"Url": nil,
+	"Uri": nil,
+}
+
+func init() {
+	for ism := range initialisms {
+		// Match ism if it is at the end, or followed by an uppercase letter or a number.
+		initialisms[ism] = regexp.MustCompile(ism + `($|[A-Z0-9])`)
+	}
 }
 
 func assert(cond bool, msg string) {
diff --git a/mcp/protocol/protocol.go b/mcp/protocol/protocol.go
index 800ed88..a993a62 100644
--- a/mcp/protocol/protocol.go
+++ b/mcp/protocol/protocol.go
@@ -62,7 +62,7 @@ type CancelledParams struct {
 	//
 	// This MUST correspond to the ID of a request previously issued in the same
 	// direction.
-	RequestId any `json:"requestId"`
+	RequestID any `json:"requestId"`
 }
 
 // Capabilities a client may support. Known capabilities are defined here, in
diff --git a/mcp/transport.go b/mcp/transport.go
index 6bbe5a5..0b9aa0b 100644
--- a/mcp/transport.go
+++ b/mcp/transport.go
@@ -135,7 +135,7 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc2.Request) (result
 		if err := json.Unmarshal(req.Params, &params); err != nil {
 			return nil, err
 		}
-		id, err := jsonrpc2.MakeID(params.RequestId)
+		id, err := jsonrpc2.MakeID(params.RequestID)
 		if err != nil {
 			return nil, err
 		}
@@ -156,7 +156,7 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params,
 		// Notify the peer of cancellation.
 		err := conn.Notify(xcontext.Detach(ctx), "notifications/cancelled", &protocol.CancelledParams{
 			Reason:    ctx.Err().Error(),
-			RequestId: call.ID().Raw(),
+			RequestID: call.ID().Raw(),
 		})
 		return errors.Join(ctx.Err(), err)
 	case err != nil: