diff --git a/mcp/event.go b/mcp/event.go new file mode 100644 index 0000000..4562b81 --- /dev/null +++ b/mcp/event.go @@ -0,0 +1,138 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +// This file is for SSE events. +// See https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events. + +package mcp + +import ( + "bufio" + "bytes" + "errors" + "fmt" + "io" + "iter" + "net/http" + "strings" +) + +// An Event is a server-sent event. +// See https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#fields. +type Event struct { + Name string // the "event" field + ID string // the "id" field + Data []byte // the "data" field +} + +// Empty reports whether the Event is empty. +func (e Event) Empty() bool { + return e.Name == "" && e.ID == "" && len(e.Data) == 0 +} + +// writeEvent writes the event to w, and flushes. +func writeEvent(w io.Writer, evt Event) (int, error) { + var b bytes.Buffer + if evt.Name != "" { + fmt.Fprintf(&b, "event: %s\n", evt.Name) + } + if evt.ID != "" { + fmt.Fprintf(&b, "id: %s\n", evt.ID) + } + fmt.Fprintf(&b, "data: %s\n\n", string(evt.Data)) + n, err := w.Write(b.Bytes()) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + return n, err +} + +// scanEvents iterates SSE events in the given scanner. The iterated error is +// terminal: if encountered, the stream is corrupt or broken and should no +// longer be used. +// +// TODO(rfindley): consider a different API here that makes failure modes more +// apparent. +func scanEvents(r io.Reader) iter.Seq2[Event, error] { + scanner := bufio.NewScanner(r) + const maxTokenSize = 1 * 1024 * 1024 // 1 MiB max line size + scanner.Buffer(nil, maxTokenSize) + + // TODO: investigate proper behavior when events are out of order, or have + // non-standard names. + var ( + eventKey = []byte("event") + idKey = []byte("id") + dataKey = []byte("data") + ) + + return func(yield func(Event, error) bool) { + // iterate event from the wire. + // https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#examples + // + // - `key: value` line records. + // - Consecutive `data: ...` fields are joined with newlines. + // - Unrecognized fields are ignored. Since we only care about 'event', 'id', and + // 'data', these are the only three we consider. + // - Lines starting with ":" are ignored. + // - Records are terminated with two consecutive newlines. + var ( + evt Event + dataBuf *bytes.Buffer // if non-nil, preceding field was also data + ) + flushData := func() { + if dataBuf != nil { + evt.Data = dataBuf.Bytes() + dataBuf = nil + } + } + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + flushData() + // \n\n is the record delimiter + if !evt.Empty() && !yield(evt, nil) { + return + } + evt = Event{} + continue + } + before, after, found := bytes.Cut(line, []byte{':'}) + if !found { + yield(Event{}, fmt.Errorf("malformed line in SSE stream: %q", string(line))) + return + } + if !bytes.Equal(before, dataKey) { + flushData() + } + switch { + case bytes.Equal(before, eventKey): + evt.Name = strings.TrimSpace(string(after)) + case bytes.Equal(before, idKey): + evt.ID = strings.TrimSpace(string(after)) + case bytes.Equal(before, dataKey): + data := bytes.TrimSpace(after) + if dataBuf != nil { + dataBuf.WriteByte('\n') + dataBuf.Write(data) + } else { + dataBuf = new(bytes.Buffer) + dataBuf.Write(data) + } + } + } + if err := scanner.Err(); err != nil { + if errors.Is(err, bufio.ErrTooLong) { + err = fmt.Errorf("event exceeded max line length of %d", maxTokenSize) + } + if !yield(Event{}, err) { + return + } + } + flushData() + if !evt.Empty() { + yield(evt, nil) + } + } +} diff --git a/mcp/event_test.go b/mcp/event_test.go new file mode 100644 index 0000000..2e2e3d2 --- /dev/null +++ b/mcp/event_test.go @@ -0,0 +1,99 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "strings" + "testing" +) + +func TestScanEvents(t *testing.T) { + tests := []struct { + name string + input string + want []Event + wantErr string + }{ + { + name: "simple event", + input: "event: message\nid: 1\ndata: hello\n\n", + want: []Event{ + {Name: "message", ID: "1", Data: []byte("hello")}, + }, + }, + { + name: "multiple data lines", + input: "data: line 1\ndata: line 2\n\n", + want: []Event{ + {Data: []byte("line 1\nline 2")}, + }, + }, + { + name: "multiple events", + input: "data: first\n\nevent: second\ndata: second\n\n", + want: []Event{ + {Data: []byte("first")}, + {Name: "second", Data: []byte("second")}, + }, + }, + { + name: "no trailing newline", + input: "data: hello", + want: []Event{ + {Data: []byte("hello")}, + }, + }, + { + name: "malformed line", + input: "invalid line\n\n", + wantErr: "malformed line", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r := strings.NewReader(tt.input) + var got []Event + var err error + for e, err2 := range scanEvents(r) { + if err2 != nil { + err = err2 + break + } + got = append(got, e) + } + + if tt.wantErr != "" { + if err == nil { + t.Fatalf("scanEvents() got nil error, want error containing %q", tt.wantErr) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("scanEvents() error = %q, want containing %q", err, tt.wantErr) + } + return + } + + if err != nil { + t.Fatalf("scanEvents() returned unexpected error: %v", err) + } + + if len(got) != len(tt.want) { + t.Fatalf("scanEvents() got %d events, want %d", len(got), len(tt.want)) + } + + for i := range got { + if g, w := got[i].Name, tt.want[i].Name; g != w { + t.Errorf("event %d: name = %q, want %q", i, g, w) + } + if g, w := got[i].ID, tt.want[i].ID; g != w { + t.Errorf("event %d: id = %q, want %q", i, g, w) + } + if g, w := string(got[i].Data), string(tt.want[i].Data); g != w { + t.Errorf("event %d: data = %q, want %q", i, g, w) + } + } + }) + } +} diff --git a/mcp/sse.go b/mcp/sse.go index f0d7b34..4051f45 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -5,16 +5,12 @@ package mcp import ( - "bufio" "bytes" "context" - "errors" "fmt" "io" - "iter" "net/http" "net/url" - "strings" "sync" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" @@ -41,34 +37,6 @@ import ( // - Read reads off a message queue that is pushed to via POST requests. // - Close causes the hanging GET to exit. -// An event is a server-sent event. -type event struct { - name string - id string - data []byte -} - -func (e event) empty() bool { - return e.name == "" && e.id == "" && len(e.data) == 0 -} - -// writeEvent writes the event to w, and flushes. -func writeEvent(w io.Writer, evt event) (int, error) { - var b bytes.Buffer - if evt.name != "" { - fmt.Fprintf(&b, "event: %s\n", evt.name) - } - if evt.id != "" { - fmt.Fprintf(&b, "id: %s\n", evt.id) - } - fmt.Fprintf(&b, "data: %s\n\n", string(evt.data)) - n, err := w.Write(b.Bytes()) - if f, ok := w.(http.Flusher); ok { - f.Flush() - } - return n, err -} - // SSEHandler is an http.Handler that serves SSE-based MCP sessions as defined by // the [2024-11-05 version] of the MCP spec. // @@ -172,9 +140,9 @@ func (t *SSEServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) // See [SSEServerTransport] for more details on the [Connection] implementation. func (t *SSEServerTransport) Connect(context.Context) (Connection, error) { t.mu.Lock() - _, err := writeEvent(t.w, event{ - name: "endpoint", - data: []byte(t.endpoint), + _, err := writeEvent(t.w, Event{ + Name: "endpoint", + Data: []byte(t.endpoint), }) t.mu.Unlock() if err != nil { @@ -300,7 +268,7 @@ func (s sseServerConn) Write(ctx context.Context, msg jsonrpc.Message) error { return io.EOF } - _, err = writeEvent(s.t.w, event{name: "message", data: data}) + _, err = writeEvent(s.t.w, Event{Name: "message", Data: data}) return err } @@ -372,17 +340,17 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) { } msgEndpoint, err := func() (*url.URL, error) { - var evt event + var evt Event for evt, err = range scanEvents(resp.Body) { break } if err != nil { return nil, err } - if evt.name != "endpoint" { - return nil, fmt.Errorf("first event is %q, want %q", evt.name, "endpoint") + if evt.Name != "endpoint" { + return nil, fmt.Errorf("first event is %q, want %q", evt.Name, "endpoint") } - raw := string(evt.data) + raw := string(evt.Data) return c.sseEndpoint.Parse(raw) }() if err != nil { @@ -408,7 +376,7 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) { return } select { - case s.incoming <- evt.data: + case s.incoming <- evt.Data: case <-s.done: return } @@ -418,95 +386,6 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) { return s, nil } -// scanEvents iterates SSE events in the given scanner. The iterated error is -// terminal: if encountered, the stream is corrupt or broken and should no -// longer be used. -// -// TODO(rfindley): consider a different API here that makes failure modes more -// apparent. -func scanEvents(r io.Reader) iter.Seq2[event, error] { - scanner := bufio.NewScanner(r) - const maxTokenSize = 1 * 1024 * 1024 // 1 MiB max line size - scanner.Buffer(nil, maxTokenSize) - - // TODO: investigate proper behavior when events are out of order, or have - // non-standard names. - var ( - eventKey = []byte("event") - idKey = []byte("id") - dataKey = []byte("data") - ) - - return func(yield func(event, error) bool) { - // iterate event from the wire. - // https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#examples - // - // - `key: value` line records. - // - Consecutive `data: ...` fields are joined with newlines. - // - Unrecognized fields are ignored. Since we only care about 'event', 'id', and - // 'data', these are the only three we consider. - // - Lines starting with ":" are ignored. - // - Records are terminated with two consecutive newlines. - var ( - evt event - dataBuf *bytes.Buffer // if non-nil, preceding field was also data - ) - flushData := func() { - if dataBuf != nil { - evt.data = dataBuf.Bytes() - dataBuf = nil - } - } - for scanner.Scan() { - line := scanner.Bytes() - if len(line) == 0 { - flushData() - // \n\n is the record delimiter - if !evt.empty() && !yield(evt, nil) { - return - } - evt = event{} - continue - } - before, after, found := bytes.Cut(line, []byte{':'}) - if !found { - yield(event{}, fmt.Errorf("malformed line in SSE stream: %q", string(line))) - return - } - if !bytes.Equal(before, dataKey) { - flushData() - } - switch { - case bytes.Equal(before, eventKey): - evt.name = strings.TrimSpace(string(after)) - case bytes.Equal(before, idKey): - evt.id = strings.TrimSpace(string(after)) - case bytes.Equal(before, dataKey): - data := bytes.TrimSpace(after) - if dataBuf != nil { - dataBuf.WriteByte('\n') - dataBuf.Write(data) - } else { - dataBuf = new(bytes.Buffer) - dataBuf.Write(data) - } - } - } - if err := scanner.Err(); err != nil { - if errors.Is(err, bufio.ErrTooLong) { - err = fmt.Errorf("event exceeded max line length of %d", maxTokenSize) - } - if !yield(event{}, err) { - return - } - } - flushData() - if !evt.empty() { - yield(evt, nil) - } - } -} - // An sseClientConn is a logical jsonrpc2 connection that implements the client // half of the SSE protocol: // - Writes are POSTS to the session endpoint. diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 846d68c..b4e8eba 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -9,7 +9,6 @@ import ( "fmt" "net/http" "net/http/httptest" - "strings" "sync/atomic" "testing" @@ -90,95 +89,6 @@ func TestSSEServer(t *testing.T) { } } -func TestScanEvents(t *testing.T) { - tests := []struct { - name string - input string - want []event - wantErr string - }{ - { - name: "simple event", - input: "event: message\nid: 1\ndata: hello\n\n", - want: []event{ - {name: "message", id: "1", data: []byte("hello")}, - }, - }, - { - name: "multiple data lines", - input: "data: line 1\ndata: line 2\n\n", - want: []event{ - {data: []byte("line 1\nline 2")}, - }, - }, - { - name: "multiple events", - input: "data: first\n\nevent: second\ndata: second\n\n", - want: []event{ - {data: []byte("first")}, - {name: "second", data: []byte("second")}, - }, - }, - { - name: "no trailing newline", - input: "data: hello", - want: []event{ - {data: []byte("hello")}, - }, - }, - { - name: "malformed line", - input: "invalid line\n\n", - wantErr: "malformed line", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - r := strings.NewReader(tt.input) - var got []event - var err error - for e, err2 := range scanEvents(r) { - if err2 != nil { - err = err2 - break - } - got = append(got, e) - } - - if tt.wantErr != "" { - if err == nil { - t.Fatalf("scanEvents() got nil error, want error containing %q", tt.wantErr) - } - if !strings.Contains(err.Error(), tt.wantErr) { - t.Fatalf("scanEvents() error = %q, want containing %q", err, tt.wantErr) - } - return - } - - if err != nil { - t.Fatalf("scanEvents() returned unexpected error: %v", err) - } - - if len(got) != len(tt.want) { - t.Fatalf("scanEvents() got %d events, want %d", len(got), len(tt.want)) - } - - for i := range got { - if g, w := got[i].name, tt.want[i].name; g != w { - t.Errorf("event %d: name = %q, want %q", i, g, w) - } - if g, w := got[i].id, tt.want[i].id; g != w { - t.Errorf("event %d: id = %q, want %q", i, g, w) - } - if g, w := string(got[i].data), string(tt.want[i].data); g != w { - t.Errorf("event %d: data = %q, want %q", i, g, w) - } - } - }) - } -} - // roundTripperFunc is a helper to create a custom RoundTripper type roundTripperFunc func(*http.Request) (*http.Response, error) diff --git a/mcp/streamable.go b/mcp/streamable.go index d371c87..ef740d3 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -149,6 +149,9 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque // NewStreamableServerTransport returns a new [StreamableServerTransport] with // the given session ID. +// The session ID must be globally unique, that is, different from any other +// session ID anywhere, past and future. (We recommend using a crypto random number +// generator to produce one, as in [crypto/rand.Text].) // // A StreamableServerTransport implements the server-side of the streamable // transport. @@ -246,7 +249,7 @@ type streamID int64 // a streamableMsg is an SSE event with an index into its logical stream. type streamableMsg struct { idx int - event event + event Event } // Connect implements the [Transport] interface. @@ -549,10 +552,10 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg jsonrpc.Messa idx := len(t.outgoingMessages[forConn]) t.outgoingMessages[forConn] = append(t.outgoingMessages[forConn], &streamableMsg{ idx: idx, - event: event{ - name: "message", - id: formatEventID(forConn, idx), - data: data, + event: Event{ + Name: "message", + ID: formatEventID(forConn, idx), + Data: data, }, }) if replyTo.IsValid() { @@ -768,7 +771,7 @@ func (s *streamableClientConn) handleSSE(resp *http.Response) { select { case <-s.done: return - case s.incoming <- evt.data: + case s.incoming <- evt.Data: } } }() diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index e0e85a1..05bf59e 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -520,7 +520,7 @@ func streamingRequest(ctx context.Context, serverURL, sessionID, method string, } // TODO(rfindley): do we need to check evt.name? // Does the MCP spec say anything about this? - msg, err := jsonrpc2.DecodeMessage(evt.data) + msg, err := jsonrpc2.DecodeMessage(evt.Data) if err != nil { return newSessionID, resp.StatusCode, fmt.Errorf("decoding message: %w", err) }