mcp: add Session.ID (#67)

Support retrieving the session ID from client and server sessions.

Fixes #65.
This commit is contained in:
Jonathan Amsterdam 2025-07-01 13:49:23 -04:00 • committed by GitHub
commit 9b6327b3b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -161,6 +161,10 @@ func NewStreamableServerTransport(sessionID string) *StreamableServerTransport {
}
}
func (t *StreamableServerTransport) SessionID() string {
return t.id
}
// A StreamableServerTransport implements the [Transport] interface for a
// single session.
type StreamableServerTransport struct {
@ -331,7 +335,7 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R
http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
return
}
var requests = make(map[JSONRPCID]struct{})
requests := make(map[JSONRPCID]struct{})
for _, msg := range incoming {
if req, ok := msg.(*JSONRPCRequest); ok && req.ID.IsValid() {
requests[req.ID] = struct{}{}
@ -624,20 +628,26 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, er
}
type streamableClientConn struct {
url string
sessionID string
client *http.Client
incoming chan []byte
done chan struct{}
url string
client *http.Client
incoming chan []byte
done chan struct{}
closeOnce sync.Once
closeErr error
mu sync.Mutex
mu sync.Mutex
_sessionID string
// bodies map[*http.Response]io.Closer
err error
}
func (c *streamableClientConn) SessionID() string {
c.mu.Lock()
defer c.mu.Unlock()
return c._sessionID
}
// Read implements the [Connection] interface.
func (s *streamableClientConn) Read(ctx context.Context) (JSONRPCMessage, error) {
select {
@ -658,7 +668,7 @@ func (s *streamableClientConn) Write(ctx context.Context, msg JSONRPCMessage) er
return s.err
}
sessionID := s.sessionID
sessionID := s._sessionID
if sessionID == "" {
// Hold lock for the first request.
defer s.mu.Unlock()
@ -681,7 +691,7 @@ func (s *streamableClientConn) Write(ctx context.Context, msg JSONRPCMessage) er
if sessionID == "" {
// locked
s.sessionID = gotSessionID
s._sessionID = gotSessionID
}
return nil
@ -753,7 +763,7 @@ func (s *streamableClientConn) Close() error {
if err != nil {
s.closeErr = err
} else {
req.Header.Set("Mcp-Session-Id", s.sessionID)
req.Header.Set("Mcp-Session-Id", s._sessionID)
if _, err := s.client.Do(req); err != nil {
s.closeErr = err
}