-
Notifications
You must be signed in to change notification settings - Fork 346
Implement MCP Streamable HTTP Transport Protocol #228
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
WalkthroughThis change introduces a complete implementation and documentation of the MCP Streamable HTTP transport in Go. It adds a new server-side transport ( Changes
Possibly related PRs
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the ✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 4
🧹 Nitpick comments (17)
server/streamable_http.go (5)
424-429
: Guardhttp.Flusher
access
w.(http.Flusher)
will panic if the underlying writer does not implementFlusher
(e.g. certain test doubles). Always assert the capability first.- w.(http.Flusher).Flush() + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + }
508-555
: Notification fan-out may silently drop messages
eventChan
is buffered with capacity 10. When the buffer is full the fallback is to block; however the select defaults to returning on<-notifDone>
.
If the client is slow, notifications will be lost without feedback.
Consider:
- Using a larger buffer or an unbounded queue.
- Back-pressure mechanisms (e.g. context cancellation).
- Logging or propagating the overflow error.
748-759
:event
argument is unused in callers – clarify or remove
writeSSEEvent
allows specifying anevent:
field but nothing in this file invokes it with a non-empty value. If custom event names are required, document and use them; otherwise drop the parameter to avoid confusion.
323-329
: Errors from notification handling are discarded
server.HandleMessage
may return an error for notifications (e.g. malformed data). Returning202 Accepted
unconditionally can hide failures.
At minimum, log the error and return400/500
when appropriate.
829-845
: Dead code –validateSession
is never usedThe helper looks correct but is not referenced anywhere. Remove it or wire it into
handlePost
/handleGet
to ensure consistent validation.examples/minimal_client/main.go (2)
24-26
: Shared context may cancel long-running requestsBoth the
initialize
andtools/call
requests reuse the same 30 s context. If the first call consumes most of the timeout, the second may fail prematurely. Prefer deriving a fresh context per operation.initCtx, cancel := context.WithTimeout(ctx, 30*time.Second) ... callCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
21-22
: IgnoreClose
error
trans.Close()
returns an error (according to transport interface) but it is discarded. Consider logging it to aid debugging of lingering connections.examples/minimal_server/main.go (1)
64-69
: Handlehttp.ErrServerClosed
gracefullyWhen
Shutdown
is invoked the listener returnshttp.ErrServerClosed
, which is expected and should not be treated as fatal.-if err := streamableServer.Start(":8080"); err != nil { - log.Fatalf("Failed to start server: %v", err) -} +if err := streamableServer.Start(":8080"); err != nil && err != http.ErrServerClosed { + log.Fatalf("Failed to start server: %v", err) +}examples/streamable_http_client/main.go (1)
1-56
: Well-structured example with clear flow and good practicesThis example demonstrates a streamable HTTP client with all essential components: transport creation, notification handling, session initialization, and response display. The code follows Go best practices with proper error handling and resource cleanup.
Consider adding more detailed comments explaining the expected behavior and outputs at key points, especially for the notification handler. For instance, explaining what kind of notifications might be received and how they're structured would help users understand the example better.
examples/streamable_http_client_complete/main.go (1)
1-131
: Comprehensive client example with complete MCP workflowThis example builds on the basic client by adding tool listing, tool invocation, and notification handling with signal management. The structure is logical and demonstrates a complete interaction flow with an MCP server.
Consider adding comments explaining the expected notification content from the echo tool, particularly in the notification handler setup (lines 27-31) or near the waiting section (lines 116-119). This would help users understand what to expect when running the example.
examples/streamable_http_server/main.go (1)
61-66
: Consider adding context check in notification goroutineThe notification goroutine doesn't check if the context is still valid before sending the notification, which could lead to errors if the server is shutting down.
go func() { time.Sleep(1 * time.Second) + // Check if context is still valid before sending notification + select { + case <-ctx.Done(): + return + default: mcpServer.SendNotificationToClient(ctx, "echo/notification", map[string]interface{}{ "message": "Echo notification: " + message, }) + } }()README-streamable-http.md (3)
22-25
: Fix bullet point formattingThe bullet points in this section have loose punctuation marks that should be fixed for better readability.
### Key Components - `StreamableHTTPServer`: The main server implementation that handles HTTP requests and responses - `streamableHTTPSession`: Represents an active session with a client - `EventStore`: Interface for storing and retrieving events for resumability - `InMemoryEventStore`: A simple in-memory implementation of the EventStore interface🧰 Tools
🪛 LanguageTool
[uncategorized] ~22-~22: Loose punctuation mark.
Context: ...Key Components -StreamableHTTPServer
: The main server implementation that han...(UNLIKELY_OPENING_PUNCTUATION)
34-37
: Consider expanding the client implementation sectionThe client implementation section is quite brief compared to the server section. Adding more details about the client's key components, options, and design considerations would make the documentation more balanced.
Consider expanding this section to include:
- Key client components
- Client options (similar to server options section)
- Design considerations specific to the client implementation
273-302
: Add a troubleshooting sectionThe documentation would benefit from a troubleshooting section that addresses common issues users might encounter, such as connection problems, session expiration, or event replay failures.
Consider adding a "Troubleshooting" section that covers:
- Common error scenarios and their solutions
- Debugging tips (e.g., enabling verbose logging)
- Best practices for error handling in both client and server implementations
server/streamable_http_test.go (3)
73-92
: Consider flattening nested conditionals for better readabilityThe nested conditionals for validating the response structure could be simplified for better readability and easier debugging when tests fail.
- if result, ok := response["result"].(map[string]interface{}); ok { - if serverInfo, ok := result["serverInfo"].(map[string]interface{}); ok { - if serverInfo["name"] != "test-server" { - t.Errorf("Expected server name test-server, got %v", serverInfo["name"]) - } - if serverInfo["version"] != "1.0.0" { - t.Errorf("Expected server version 1.0.0, got %v", serverInfo["version"]) - } - } else { - t.Errorf("Expected serverInfo in result, got none") - } - } else { - t.Errorf("Expected result in response, got none") - } + result, ok := response["result"].(map[string]interface{}) + if !ok { + t.Fatalf("Expected result in response, got none") + } + + serverInfo, ok := result["serverInfo"].(map[string]interface{}) + if !ok { + t.Fatalf("Expected serverInfo in result, got none") + } + + if serverInfo["name"] != "test-server" { + t.Errorf("Expected server name test-server, got %v", serverInfo["name"]) + } + + if serverInfo["version"] != "1.0.0" { + t.Errorf("Expected server version 1.0.0, got %v", serverInfo["version"]) + }
257-257
: Replace hard-coded sleep with a more reliable mechanismUsing a fixed sleep duration can lead to flaky tests on different environments or under different load conditions.
Consider using a more robust synchronization mechanism or at least making the sleep duration configurable:
- // Wait a bit for the stream to be established - time.Sleep(100 * time.Millisecond) + // Wait for the stream to be established + streamEstablishmentTimeout := 200 * time.Millisecond + time.Sleep(streamEstablishmentTimeout)For even better reliability, you could implement a ready signal from the server once the stream is established.
293-312
: Implement more robust SSE event parsingThe current SSE event parsing is simplified and assumes a specific format. A more robust implementation would handle multiple data lines, event IDs, and other SSE features.
Consider implementing a more comprehensive SSE parser or using the same timeout pattern recommended for the previous event reading loop.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
README-streamable-http.md
(1 hunks)examples/minimal_client/main.go
(1 hunks)examples/minimal_server/main.go
(1 hunks)examples/streamable_http_client/main.go
(1 hunks)examples/streamable_http_client_complete/main.go
(1 hunks)examples/streamable_http_server/main.go
(1 hunks)server/streamable_http.go
(1 hunks)server/streamable_http_test.go
(1 hunks)
🧰 Additional context used
🪛 LanguageTool
README-streamable-http.md
[uncategorized] ~22-~22: Loose punctuation mark.
Context: ...Key Components - StreamableHTTPServer
: The main server implementation that han...
(UNLIKELY_OPENING_PUNCTUATION)
[uncategorized] ~29-~29: Loose punctuation mark.
Context: ...rver Options - WithSessionIDGenerator
: Sets a custom session ID generator - `W...
(UNLIKELY_OPENING_PUNCTUATION)
🔇 Additional comments (6)
server/streamable_http.go (1)
448-457
: 🛠️ Refactor suggestionDeferred close order can panic – potential send-on-closed-channel
defer close(eventChan)
is registered beforedefer close(notifDone)
.
If the goroutine is still attempting to send oneventChan
whenhandleSSEResponse
returns, closingeventChan
first can trigger a panic.
Reverse the order or signal the goroutine before the channel is closed.- defer close(eventChan) - ... - notifDone := make(chan struct{}) - defer close(notifDone) + notifDone := make(chan struct{}) + defer close(notifDone) // signal first + defer close(eventChan) // close after goroutine stopsLikely an incorrect or invalid review comment.
examples/streamable_http_client_complete/main.go (1)
120-130
: Good signal handling pattern for graceful terminationThe signal handling implementation with a select statement provides a clean way to wait for either notifications or termination signals.
examples/streamable_http_server/main.go (1)
1-97
: Well-structured server example with proper lifecycle managementThe example demonstrates a complete MCP server with streamable HTTP transport, including tool registration, async notification handling, and graceful shutdown. The code follows good practices for server implementation in Go.
README-streamable-http.md (1)
1-302
: Comprehensive and well-structured documentationThe README provides thorough documentation of the MCP Streamable HTTP implementation, covering both server and client sides, with clear examples and protocol details. It aligns well with the provided example code and offers valuable context for users.
🧰 Tools
🪛 LanguageTool
[uncategorized] ~22-~22: Loose punctuation mark.
Context: ...Key Components -StreamableHTTPServer
: The main server implementation that han...(UNLIKELY_OPENING_PUNCTUATION)
[uncategorized] ~29-~29: Loose punctuation mark.
Context: ...rver Options -WithSessionIDGenerator
: Sets a custom session ID generator - `W...(UNLIKELY_OPENING_PUNCTUATION)
server/streamable_http_test.go (2)
16-33
: LGTM - Test setup looks goodThe initial setup is clear and concise, properly initializing the MCP server with required capabilities, configuring the StreamableHTTPServer, and creating a test HTTP server.
336-401
: Session termination test is well structuredThe session termination test effectively verifies both the successful termination of a session and the expected behavior when attempting to use a terminated session.
if isInitialize && response != nil { | ||
// Only create a session if we're not in stateless mode | ||
if s.sessionIDGenerator() != "" { | ||
newSessionID := s.sessionIDGenerator() | ||
newSession := &streamableHTTPSession{ | ||
sessionID: newSessionID, | ||
notificationChannel: make(chan mcp.JSONRPCNotification, 100), | ||
eventStore: s.eventStore, | ||
} | ||
|
||
// Register the session | ||
s.sessions.Store(newSessionID, newSession) | ||
if err := s.server.RegisterSession(ctx, newSession); err != nil { | ||
http.Error(w, fmt.Sprintf("Failed to register session: %v", err), http.StatusInternalServerError) | ||
return | ||
} | ||
|
||
// Set the session ID in the response header | ||
w.Header().Set("Mcp-Session-Id", newSessionID) | ||
|
||
// Update the session reference for further processing | ||
session = newSession | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
New sessions are never marked as initialised
streamableHTTPSession
has an initialized
atomic flag and a helper Initialize()
, yet it is never set to true
after a successful "initialize"
request.
Down-stream code that relies on session.Initialized()
(e.g. validateSession
) will therefore misbehave.
+ // After the server accepts the initialise request and BEFORE returning the
+ // first response:
+ newSession.Initialize()
if !isInitialize && session == nil && s.sessionIDGenerator() != "" { | ||
http.Error(w, "Bad Request: Server not initialized", http.StatusBadRequest) | ||
return | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Session-presence check is ineffective – sessionIDGenerator()
is always non-empty
sessionIDGenerator()
is invoked to decide whether the server runs in stateless mode, but every call produces a new UUID, hence the comparison against ""
is always false and the "Server not initialized"
branch will never be taken.
Either store a boolean flag (e.g. stateless bool
) at construction time, or treat a nil/empty generator explicitly.
- if !isInitialize && session == nil && s.sessionIDGenerator() != "" {
+ if !isInitialize && session == nil && !s.stateless {
Committable suggestion skipped: line range outside the PR's diff.
// Send a notification to the session | ||
go func() { | ||
// Wait a bit for the stream to be established | ||
time.Sleep(100 * time.Millisecond) | ||
|
||
// Create a notification | ||
notification := mcp.JSONRPCNotification{ | ||
JSONRPC: "2.0", | ||
Notification: mcp.Notification{ | ||
Method: "test/notification", | ||
Params: mcp.NotificationParams{ | ||
AdditionalFields: map[string]interface{}{ | ||
"message": "Hello, world!", | ||
}, | ||
}, | ||
}, | ||
} | ||
|
||
// Find the session | ||
sessionValue, ok := streamableServer.sessions.Load(sessionID) | ||
if !ok { | ||
t.Errorf("Session not found: %s", sessionID) | ||
return | ||
} | ||
|
||
// Send the notification | ||
session, ok := sessionValue.(*streamableHTTPSession) | ||
if !ok { | ||
t.Errorf("Invalid session type") | ||
return | ||
} | ||
|
||
session.notificationChannel <- notification | ||
}() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Avoid directly accessing internal implementation details
The test directly accesses internal implementation details of the StreamableHTTPServer (sessions map and streamableHTTPSession type). This creates tight coupling between the test and the implementation.
Consider adding a public method to StreamableHTTPServer for testing purposes:
// SendNotificationToSession sends a notification to a specific session (for testing)
func (s *StreamableHTTPServer) SendNotificationToSession(
sessionID string,
notification mcp.JSONRPCNotification) error {
sessionValue, ok := s.sessions.Load(sessionID)
if !ok {
return fmt.Errorf("session not found: %s", sessionID)
}
session, ok := sessionValue.(*streamableHTTPSession)
if !ok {
return fmt.Errorf("invalid session type")
}
session.notificationChannel <- notification
return nil
}
Then update the test:
- // Find the session
- sessionValue, ok := streamableServer.sessions.Load(sessionID)
- if !ok {
- t.Errorf("Session not found: %s", sessionID)
- return
- }
-
- // Send the notification
- session, ok := sessionValue.(*streamableHTTPSession)
- if !ok {
- t.Errorf("Invalid session type")
- return
- }
-
- session.notificationChannel <- notification
+ err := streamableServer.SendNotificationToSession(sessionID, notification)
+ if err != nil {
+ t.Errorf("Failed to send notification: %v", err)
+ return
+ }
for { | ||
line, err := reader.ReadString('\n') | ||
if err != nil { | ||
if err == io.EOF { | ||
break | ||
} | ||
t.Fatalf("Failed to read line: %v", err) | ||
} | ||
|
||
line = strings.TrimRight(line, "\r\n") | ||
if line == "" { | ||
// End of event | ||
break | ||
} | ||
|
||
if strings.HasPrefix(line, "data:") { | ||
eventData = strings.TrimSpace(strings.TrimPrefix(line, "data:")) | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Add timeout to prevent test hanging
The event reading loop lacks a timeout mechanism which could cause the test to hang indefinitely if there's an issue with the event stream.
+ timeout := time.NewTimer(5 * time.Second)
+ defer timeout.Stop()
+
for {
+ select {
+ case <-timeout.C:
+ t.Fatalf("Timeout waiting for SSE event")
+ return
+ default:
line, err := reader.ReadString('\n')
if err != nil {
if err == io.EOF {
break
}
t.Fatalf("Failed to read line: %v", err)
}
line = strings.TrimRight(line, "\r\n")
if line == "" {
// End of event
break
}
if strings.HasPrefix(line, "data:") {
eventData = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
}
+ }
}
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
for { | |
line, err := reader.ReadString('\n') | |
if err != nil { | |
if err == io.EOF { | |
break | |
} | |
t.Fatalf("Failed to read line: %v", err) | |
} | |
line = strings.TrimRight(line, "\r\n") | |
if line == "" { | |
// End of event | |
break | |
} | |
if strings.HasPrefix(line, "data:") { | |
eventData = strings.TrimSpace(strings.TrimPrefix(line, "data:")) | |
} | |
} | |
// Add a timeout so the test can’t hang forever waiting for an event | |
timeout := time.NewTimer(5 * time.Second) | |
defer timeout.Stop() | |
for { | |
select { | |
case <-timeout.C: | |
t.Fatalf("Timeout waiting for SSE event") | |
return | |
default: | |
line, err := reader.ReadString('\n') | |
if err != nil { | |
if err == io.EOF { | |
break | |
} | |
t.Fatalf("Failed to read line: %v", err) | |
} | |
line = strings.TrimRight(line, "\r\n") | |
if line == "" { | |
// End of event | |
break | |
} | |
if strings.HasPrefix(line, "data:") { | |
eventData = strings.TrimSpace(strings.TrimPrefix(line, "data:")) | |
} | |
} | |
} |
Key Features
Implementation Details
Summary by CodeRabbit
New Features
Examples
Documentation
Tests