Skip to content

fix(streamable_http): preserve error context for client responses #822

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.

Sign up for GitHub

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

Open
lawrence3699 wants to merge 1 commit into
base: main
Choose a base branch
from
+170 −8
Open

fix(streamable_http): preserve error context for client responses #822

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 65 additions & 8 deletions server/streamable_http.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -808,6 +808,13 @@ func (s *StreamableHTTPServer) handleSamplingResponse(w http.ResponseWriter, r *
requestID: requestID,
}

// Look up the pending request so we can preserve the original request context
responseContext, publicMessage, statusCode, err := s.pendingResponseContext(sessionID, requestID)
if err != nil {
http.Error(w, publicMessage, statusCode)
return err
}
Comment on lines +811 to +816

// Parse result or error
if responseMessage.Error != nil {
// Parse error
Expand All @@ -818,7 +825,7 @@ func (s *StreamableHTTPServer) handleSamplingResponse(w http.ResponseWriter, r *
if err := json.Unmarshal(responseMessage.Error, &jsonrpcError); err != nil {
response.err = fmt.Errorf("failed to parse error: %v", err)
} else {
response.err = fmt.Errorf("sampling error %d: %s", jsonrpcError.Code, jsonrpcError.Message)
response.err = fmt.Errorf("%s error %d: %s", pendingResponseErrorLabel(responseContext.method), jsonrpcError.Code, jsonrpcError.Message)
}
} else if responseMessage.Result != nil {
// Store the result to be unmarshaled later
Expand All @@ -839,6 +846,30 @@ func (s *StreamableHTTPServer) handleSamplingResponse(w http.ResponseWriter, r *
return nil
}

func (s *StreamableHTTPServer) pendingResponseContext(sessionID string, requestID int64) (pendingResponseItem, string, int, error) {
sessionInterface, ok := s.activeSessions.Load(sessionID)
if !ok {
return pendingResponseItem{}, "No active session found for the given session ID", http.StatusNotFound, fmt.Errorf("no active session found for session %s", sessionID)
}

session, ok := sessionInterface.(*streamableHttpSession)
if !ok {
return pendingResponseItem{}, "Invalid session type for the given session ID", http.StatusInternalServerError, fmt.Errorf("invalid session type for session %s", sessionID)
}

responseContextInterface, exists := session.samplingRequests.Load(requestID)
if !exists {
return pendingResponseItem{}, "No pending sampling request found for the given request ID", http.StatusBadRequest, fmt.Errorf("no pending request found for session %s, request %d", sessionID, requestID)
}
Comment on lines +860 to +863

responseContext, ok := responseContextInterface.(pendingResponseItem)
if !ok {
return pendingResponseItem{}, "Failed to deliver response", http.StatusInternalServerError, fmt.Errorf("invalid pending response type for session %s, request %d", sessionID, requestID)
}

return responseContext, "", http.StatusOK, nil
}

// deliverSamplingResponse delivers a sampling response to the appropriate session.
// On failure it writes the HTTP error status directly to w.
func (s *StreamableHTTPServer) deliverSamplingResponse(w http.ResponseWriter, sessionID string, response samplingResponseItem) error {
Expand All @@ -856,17 +887,18 @@ func (s *StreamableHTTPServer) deliverSamplingResponse(w http.ResponseWriter, se
}

// Look up the dedicated response channel for this specific request
responseChannelInterface, exists := session.samplingRequests.Load(response.requestID)
responseContextInterface, exists := session.samplingRequests.Load(response.requestID)
if !exists {
http.Error(w, "No pending sampling request found for the given request ID", http.StatusBadRequest)
return fmt.Errorf("no pending request found for session %s, request %d", sessionID, response.requestID)
Comment on lines 892 to 893
}

responseChan, ok := responseChannelInterface.(chan samplingResponseItem)
responseContext, ok := responseContextInterface.(pendingResponseItem)
if !ok {
http.Error(w, "Failed to deliver response", http.StatusInternalServerError)
return fmt.Errorf("invalid response channel type for session %s, request %d", sessionID, response.requestID)
return fmt.Errorf("invalid pending response type for session %s, request %d", sessionID, response.requestID)
}
responseChan := responseContext.ch

// Attempt to deliver the response with timeout to prevent indefinite blocking
select {
Expand All @@ -879,6 +911,17 @@ func (s *StreamableHTTPServer) deliverSamplingResponse(w http.ResponseWriter, se
}
}

func pendingResponseErrorLabel(method mcp.MCPMethod) string {
switch method {
case mcp.MethodElicitationCreate:
return "elicitation"
case mcp.MethodListRoots:
return "roots/list"
default:
return "sampling"
}
}

// writeJSONRPCError writes a JSON-RPC error response with the given error details.
func (s *StreamableHTTPServer) writeJSONRPCError(
w http.ResponseWriter,
Expand Down Expand Up @@ -1134,6 +1177,11 @@ type samplingResponseItem struct {
err error
}

type pendingResponseItem struct {
method mcp.MCPMethod
ch chan samplingResponseItem
}

// Elicitation support types for HTTP transport
type elicitationRequestItem struct {
requestID int64
Expand Down Expand Up @@ -1167,7 +1215,7 @@ type streamableHttpSession struct {
elicitationRequestChan chan elicitationRequestItem // server -> client elicitation requests
rootsRequestChan chan rootsRequestItem // server -> client list roots requests

samplingRequests sync.Map // requestID -> pending sampling request context
samplingRequests sync.Map // requestID -> pending response context
requestIDCounter atomic.Int64 // for generating unique request IDs
}

Expand Down Expand Up @@ -1294,7 +1342,10 @@ func (s *streamableHttpSession) RequestSampling(ctx context.Context, request mcp
}

// Store the pending request
s.samplingRequests.Store(requestID, responseChan)
s.samplingRequests.Store(requestID, pendingResponseItem{
method: mcp.MethodSamplingCreateMessage,
ch: responseChan,
})
defer s.samplingRequests.Delete(requestID)

// Send the sampling request via the channel (non-blocking)
Expand Down Expand Up @@ -1351,7 +1402,10 @@ func (s *streamableHttpSession) ListRoots(ctx context.Context, request mcp.ListR
}

// Store the pending request
s.samplingRequests.Store(requestID, responseChan)
s.samplingRequests.Store(requestID, pendingResponseItem{
method: mcp.MethodListRoots,
ch: responseChan,
})
defer s.samplingRequests.Delete(requestID)

// Send the list roots request via the channel (non-blocking)
Expand Down Expand Up @@ -1396,7 +1450,10 @@ func (s *streamableHttpSession) RequestElicitation(ctx context.Context, request
}

// Store the pending request
s.samplingRequests.Store(requestID, responseChan)
s.samplingRequests.Store(requestID, pendingResponseItem{
method: mcp.MethodElicitationCreate,
ch: responseChan,
})
defer s.samplingRequests.Delete(requestID)

// Send the sampling request via the channel (non-blocking)
Expand Down
105 changes: 105 additions & 0 deletions server/streamable_http_sampling_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,118 @@ import (
"encoding/json"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"

"github.com/mark3labs/mcp-go/mcp"
)

func TestStreamableHTTPServer_ResponseErrorsUsePendingRequestContext(t *testing.T) {
tests := []struct {
name string
expectedError string
startRequest func(context.Context, *streamableHttpSession) (int64, <-chan error)
}{
{
name: "sampling requests keep sampling label",
expectedError: "sampling error -32601: Method not found",
startRequest: func(ctx context.Context, session *streamableHttpSession) (int64, <-chan error) {
errCh := make(chan error, 1)
go func() {
_, err := session.RequestSampling(ctx, mcp.CreateMessageRequest{})
errCh <- err
}()

request := <-session.samplingRequestChan
return request.requestID, errCh
},
},
{
name: "elicitation requests use elicitation label",
expectedError: "elicitation error -32601: Method not found",
startRequest: func(ctx context.Context, session *streamableHttpSession) (int64, <-chan error) {
errCh := make(chan error, 1)
go func() {
_, err := session.RequestElicitation(ctx, mcp.ElicitationRequest{
Params: mcp.ElicitationParams{
Message: "Need input",
RequestedSchema: map[string]any{"type": "object"},
},
})
errCh <- err
}()

request := <-session.elicitationRequestChan
return request.requestID, errCh
},
},
{
name: "roots requests use roots label",
expectedError: "roots/list error -32601: Method not found",
startRequest: func(ctx context.Context, session *streamableHttpSession) (int64, <-chan error) {
errCh := make(chan error, 1)
go func() {
_, err := session.ListRoots(ctx, mcp.ListRootsRequest{})
errCh <- err
}()

request := <-session.rootsRequestChan
return request.requestID, errCh
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mcpServer := NewMCPServer("test-server", "1.0.0")
httpServer := NewStreamableHTTPServer(mcpServer, WithStateLess(true))

sessionID := "test-session"
session := newStreamableHttpSession(sessionID, nil, nil, nil, nil)
httpServer.activeSessions.Store(sessionID, session)

ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()

requestID, errCh := tt.startRequest(ctx, session)

recorder := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(nil))
req.Header.Set(HeaderKeySessionID, sessionID)

err := httpServer.handleSamplingResponse(recorder, req, struct {
ID json.RawMessage `json:"id"`
Result json.RawMessage `json:"result,omitempty"`
Error json.RawMessage `json:"error,omitempty"`
Method mcp.MCPMethod `json:"method,omitempty"`
}{
ID: []byte(strconv.FormatInt(requestID, 10)),
Error: []byte(`{"code":-32601,"message":"Method not found"}`),
})
if err != nil {
t.Fatalf("handleSamplingResponse returned error: %v", err)
}
if recorder.Code != http.StatusAccepted {
t.Fatalf("expected status %d, got %d", http.StatusAccepted, recorder.Code)
}

select {
case requestErr := <-errCh:
if requestErr == nil {
t.Fatal("expected request error, got nil")
}
if requestErr.Error() != tt.expectedError {
t.Fatalf("expected error %q, got %q", tt.expectedError, requestErr.Error())
}
case <-time.After(time.Second):
t.Fatal("timed out waiting for request error")
}
})
}
}

// TestStreamableHTTPServer_SamplingBasic tests basic sampling session functionality
func TestStreamableHTTPServer_SamplingBasic(t *testing.T) {
// Create MCP server with sampling enabled
Expand Down
Loading