From dbf57af24be7eb8f5aba2afc97427e5baefed79e Mon Sep 17 00:00:00 2001
From: Rob Findley <rfindley@google.com>
Date: Wed, 14 May 2025 14:59:59 +0000
Subject: [PATCH] internal/mcp: rename LocalTransport->InMemoryTransport, and
 document

Update the design to mention our in-memory transport, since it is
generally useful, and update the implementation to use distinct types
for InMemoryTransport and StdIOTransport (the fact that they are both
implemented as an io.ReadWriteCloser is an implementation detail).

Change-Id: I8228ccf2092a56a730c113c4db105fb492cce112
Reviewed-on: https://go-review.googlesource.com/c/tools/+/672715
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
Commit-Queue: Robert Findley <rfindley@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
TryBot-Bypass: Robert Findley <rfindley@google.com>
---
 design/design.md           |  8 ++++++++
 mcp/mcp_test.go            |  6 +++---
 mcp/server_example_test.go |  2 +-
 mcp/transport.go           | 30 +++++++++++++++++++++---------
 4 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/design/design.md b/design/design.md
index df86171..7d069f1 100644
--- a/design/design.md
+++ b/design/design.md
@@ -283,6 +283,14 @@ func NewStreamableClientTransport(url string) *StreamableClientTransport {
 func (*StreamableClientTransport) Connect(context.Context) (Stream, error)
 ```
 
+Finally, we also provide an in-memory transport, for scenarios such as testing,
+where the MCP client and server are in the same process.
+
+```go
+type InMemoryTransport struct { /* ... */ }
+func NewInMemoryTransport() (*InMemoryTransport, *InMemoryTransport)
+```
+
 **Differences from mcp-go**: The Go team has a battle-tested JSON-RPC
 implementation that we use for gopls, our Go LSP server. We are using the new
 version of this library as part of our MCP SDK. It handles all JSON-RPC 2.0
diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go
index 1165633..0a87f93 100644
--- a/mcp/mcp_test.go
+++ b/mcp/mcp_test.go
@@ -33,7 +33,7 @@ func sayHi(ctx context.Context, cc *ServerSession, v hiParams) ([]*Content, erro
 
 func TestEndToEnd(t *testing.T) {
 	ctx := context.Background()
-	ct, st := NewLocalTransport()
+	ct, st := NewInMemoryTransport()
 
 	s := NewServer("testServer", "v1.0.0", nil)
 
@@ -286,7 +286,7 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Clien
 	t.Helper()
 
 	ctx := context.Background()
-	ct, st := NewLocalTransport()
+	ct, st := NewInMemoryTransport()
 
 	s := NewServer("testServer", "v1.0.0", nil)
 
@@ -329,7 +329,7 @@ func TestServerClosing(t *testing.T) {
 
 func TestBatching(t *testing.T) {
 	ctx := context.Background()
-	ct, st := NewLocalTransport()
+	ct, st := NewInMemoryTransport()
 
 	s := NewServer("testServer", "v1.0.0", nil)
 	_, err := s.Connect(ctx, st, nil)
diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go
index 14a3248..85fccd1 100644
--- a/mcp/server_example_test.go
+++ b/mcp/server_example_test.go
@@ -24,7 +24,7 @@ func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*
 
 func ExampleServer() {
 	ctx := context.Background()
-	clientTransport, serverTransport := mcp.NewLocalTransport()
+	clientTransport, serverTransport := mcp.NewInMemoryTransport()
 
 	server := mcp.NewServer("greeter", "v0.0.1", nil)
 	server.AddTools(mcp.NewTool("greet", "say hi", SayHi))
diff --git a/mcp/transport.go b/mcp/transport.go
index 29a29c3..1cbc1d2 100644
--- a/mcp/transport.go
+++ b/mcp/transport.go
@@ -48,27 +48,39 @@ type ConnectionOptions struct {
 	Logger    io.Writer // if set, write RPC logs
 }
 
-// An IOTransport is a [Transport] that communicates using newline-delimited
+// A StdIOTransport is a [Transport] that communicates over stdin/stdout using
+// newline-delimited JSON.
+type StdIOTransport struct {
+	ioTransport
+}
+
+// An ioTransport is a [Transport] that communicates using newline-delimited
 // JSON over an io.ReadWriteCloser.
-type IOTransport struct {
+type ioTransport struct {
 	rwc io.ReadWriteCloser
 }
 
-func (t *IOTransport) Connect(context.Context) (Stream, error) {
+func (t *ioTransport) Connect(context.Context) (Stream, error) {
 	return newIOStream(t.rwc), nil
 }
 
 // NewStdIOTransport constructs a transport that communicates over
 // stdin/stdout.
-func NewStdIOTransport() *IOTransport {
-	return &IOTransport{rwc{os.Stdin, os.Stdout}}
+func NewStdIOTransport() *StdIOTransport {
+	return &StdIOTransport{ioTransport{rwc{os.Stdin, os.Stdout}}}
 }
 
-// NewLocalTransport returns two in-memory transports that connect to
-// each other, for testing purposes.
-func NewLocalTransport() (*IOTransport, *IOTransport) {
+// An InMemoryTransport is a [Transport] that communicates over an in-memory
+// network connection, using newline-delimited JSON.
+type InMemoryTransport struct {
+	ioTransport
+}
+
+// NewInMemoryTransport returns two InMemoryTransports that connect to each
+// other.
+func NewInMemoryTransport() (*InMemoryTransport, *InMemoryTransport) {
 	c1, c2 := net.Pipe()
-	return &IOTransport{c1}, &IOTransport{c2}
+	return &InMemoryTransport{ioTransport{c1}}, &InMemoryTransport{ioTransport{c2}}
 }
 
 // handler is an unexported version of jsonrpc2.Handler.