Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
98 changes: 79 additions & 19 deletions mcp/streamable.go
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,25 @@ type sessionInfo struct {
timer *time.Timer
}

// StreamableHTTPRequestSummary contains redacted metadata about a single
// JSON-RPC message decoded from a streamable HTTP POST body.
type StreamableHTTPRequestSummary struct {
// Method is the method of a decoded JSON-RPC request. It is empty for a
// response. The method has not yet been validated and may contain an
// arbitrary attacker-controlled string.
Method string

// RequestID is valid only for a JSON-RPC call. IDs use the same coercion
// rules as [jsonrpc.DecodeMessage].
RequestID jsonrpc.ID

// IsNotification reports whether the message is a JSON-RPC notification.
IsNotification bool

// IsResponse reports whether the message is a JSON-RPC response.
IsResponse bool
}

// startPOST signals that a POST request for this session is starting (which
// carries a client->server message), pausing the session timeout if it was
// running.
Expand Down Expand Up @@ -218,6 +237,23 @@ type StreamableHTTPOptions struct {
// Requests using older protocol versions (including those routed through
// the allowsessionsinstateless compatibility path) are unaffected.
PropagateRequestCancellation bool

// OnRequestSummary, when non-nil, observes redacted metadata for a single
// JSON-RPC message decoded from a streamable HTTP POST body. It is not called
// for JSON-RPC batches. The callback receives the HTTP request's context and
// runs synchronously before validation and dispatch of the decoded message.
// It may be called concurrently for different requests and should return
// promptly; in particular, it must not wait for processing of the same
// request. Panics are not recovered. Only the summary is redacted; the context
// may contain values added by authentication or other middleware.
//
// The callback is not invoked for requests rejected before this point,
// including HTTP, authorization, session-routing, and connection failures,
// even if connection setup previously inspected the body. Use HTTP middleware
// to observe those failures. Use [Server.AddReceivingMiddleware] to observe
// dispatched messages; this callback additionally observes decoded messages
// that are rejected before dispatch without exposing their parameters.
OnRequestSummary func(context.Context, StreamableHTTPRequestSummary)
}

// DefaultMaxRequestBodyBytes is the default value used for
Expand Down Expand Up @@ -422,14 +458,8 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.
}
}

transport := &StreamableServerTransport{
SessionID: sessionID,
Stateless: true,
EventStore: h.opts.EventStore,
jsonResponse: h.opts.JSONResponse,
logger: h.opts.Logger,
shouldPropagateCancellation: info.usesNewProtocol && (info.isSubscriptionsListen || h.opts.PropagateRequestCancellation),
}
transport := h.newStreamableServerTransport(sessionID, true)
transport.shouldPropagateCancellation = info.usesNewProtocol && (info.isSubscriptionsListen || h.opts.PropagateRequestCancellation)

session, err := connectStreamable(req.Context(), server, transport, info.opts)
if err != nil {
Expand Down Expand Up @@ -534,6 +564,17 @@ func connectStreamable(ctx context.Context, server *Server, transport *Streamabl
return s, nil
}

func (h *StreamableHTTPHandler) newStreamableServerTransport(sessionID string, stateless bool) *StreamableServerTransport {
return &StreamableServerTransport{
SessionID: sessionID,
Stateless: stateless,
EventStore: h.opts.EventStore,
jsonResponse: h.opts.JSONResponse,
logger: h.opts.Logger,
onRequestSummary: h.opts.OnRequestSummary,
}
}

// serveStateful handles requests for stateful servers.
// Stateful servers support GET, POST, and DELETE, and maintain persistent
// sessions keyed by session ID.
Expand Down Expand Up @@ -652,13 +693,7 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht
}
sessionID = server.opts.GetSessionID()

transport := &StreamableServerTransport{
SessionID: sessionID,
Stateless: false,
EventStore: h.opts.EventStore,
jsonResponse: h.opts.JSONResponse,
logger: h.opts.Logger,
}
transport := h.newStreamableServerTransport(sessionID, false)

// Sessions without a session ID (GetSessionID returned "") are ephemeral:
// there's no way to address them, so they are closed after the request.
Expand Down Expand Up @@ -832,6 +867,10 @@ type StreamableServerTransport struct {
// [streamableServerConn]. See its docstring.
shouldPropagateCancellation bool

// onRequestSummary is forwarded from StreamableHTTPOptions for transports
// created by StreamableHTTPHandler.
onRequestSummary func(context.Context, StreamableHTTPRequestSummary)

// connection is non-nil if and only if the transport has been connected.
connection *streamableServerConn
}
Expand All @@ -848,6 +887,7 @@ func (t *StreamableServerTransport) Connect(ctx context.Context) (Connection, er
jsonResponse: t.jsonResponse,
logger: ensureLogger(t.logger), // see #556: must be non-nil
shouldPropagateCancellation: t.shouldPropagateCancellation,
onRequestSummary: t.onRequestSummary,
incoming: make(chan jsonrpc.Message, 10),
done: make(chan struct{}),
streams: make(map[string]*stream),
Expand Down Expand Up @@ -876,10 +916,11 @@ func (t *StreamableServerTransport) SupportsProtocolVersion(version string) bool
}

type streamableServerConn struct {
sessionID string
stateless bool
jsonResponse bool
eventStore EventStore
sessionID string
stateless bool
jsonResponse bool
eventStore EventStore
onRequestSummary func(context.Context, StreamableHTTPRequestSummary)

// shouldPropagateCancellation is true when the underlying HTTP request's
// lifetime IS the connection's cancellation signal (e.g., a stateless
Expand Down Expand Up @@ -1437,6 +1478,9 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
return
}
if c.onRequestSummary != nil && !isBatch && len(incoming) == 1 {
c.onRequestSummary(req.Context(), summarizeStreamableHTTPRequest(incoming[0]))
}

protocolVersion := protocolVersionFromContext(req.Context())
if protocolVersion == "" {
Expand Down Expand Up @@ -1737,6 +1781,22 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
c.hangResponse(req.Context(), done)
}

func summarizeStreamableHTTPRequest(msg jsonrpc.Message) StreamableHTTPRequestSummary {
var summary StreamableHTTPRequestSummary
switch msg := msg.(type) {
case *jsonrpc.Request:
summary.Method = msg.Method
if msg.IsCall() {
summary.RequestID = msg.ID
} else {
summary.IsNotification = true
}
case *jsonrpc.Response:
summary.IsResponse = true
}
return summary
}

// Event IDs: encode both the logical connection ID and the index, as
// <streamID>_<idx>, to be consistent with the typescript implementation.

Expand Down
Loading