Cyrene_v2/internal/capability/mcp/http.go

package mcp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

// mcpSessionIDHeader carries the Streamable HTTP session identifier returned by a
// server that maintains a session (MCP 2025-11-25 Streamable HTTP).
const mcpSessionIDHeader = "Mcp-Session-Id"

// maxHTTPResponse bounds how much of a single POST reply body is read before the
// response is considered non-matching.
const maxHTTPResponse = 1 << 20

// streamableHTTPTransport speaks JSON-RPC 2.0 to an MCP server over Streamable
// HTTP (MCP 2025-11-25). Client→server messages are POSTed to the endpoint, and
// the server answers either with a single application/json body or an SSE stream.
// Server-initiated notifications (e.g. notifications/tools/list_changed) are
// delivered over a standalone GET SSE stream opened at Start.
type streamableHTTPTransport struct {
	cfg  MCPConfig
	log  *slog.Logger
	url  string
	http *http.Client

	mu       sync.Mutex
	nextID   int
	started  bool
	session  string
	protocol func() string

	onNotification func(method string)

	closeCh   chan struct{}
	closeOnce sync.Once
	getCancel context.CancelFunc
}

// newHTTPTransport builds a Streamable HTTP transport for cfg.
func newHTTPTransport(cfg MCPConfig, log *slog.Logger) *streamableHTTPTransport {
	return &streamableHTTPTransport{
		cfg:     cfg,
		log:     log,
		url:     cfg.URL,
		http:    http.DefaultClient,
		closeCh: make(chan struct{}),
	}
}

// Start validates the endpoint and opens the standalone GET SSE stream for
// server-initiated messages. The GET stream is best-effort: a server that only
// answers via POST replies 4xx, which we tolerate because request/response still
// works over POST.
func (s *streamableHTTPTransport) Start(ctx context.Context) error {
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("mcp http start: %w", err)
	}
	if s.cfg.URL == "" {
		return fmt.Errorf("mcp http: %w", ErrNoTransport)
	}
	if !strings.HasPrefix(s.cfg.URL, "http://") && !strings.HasPrefix(s.cfg.URL, "https://") {
		return fmt.Errorf("mcp http: invalid url %q", s.cfg.URL)
	}
	s.mu.Lock()
	s.started = true
	s.mu.Unlock()

	getCtx, cancel := context.WithCancel(context.Background())
	s.mu.Lock()
	s.getCancel = cancel
	s.mu.Unlock()
	go s.readStandalone(getCtx)
	return nil
}

// setOnNotification records the handler invoked when the server sends a
// server-initiated notification over the standalone SSE stream.
func (s *streamableHTTPTransport) setOnNotification(fn func(method string)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onNotification = fn
}

// setProtocolProvider records the getter for the negotiated MCP protocol version,
// echoing it back in the Mcp-Protocol-Version header on every request.
func (s *streamableHTTPTransport) setProtocolProvider(fn func() string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.protocol = fn
}

// Call POSTs a request and decodes the matching response, from either a JSON body
// or an SSE stream returned by the server.
func (s *streamableHTTPTransport) Call(ctx context.Context, method string, params any, result any) error {
	s.mu.Lock()
	if !s.started {
		s.mu.Unlock()
		return fmt.Errorf("mcp http %s: %w", method, ErrNotConnected)
	}
	s.nextID++
	id := strconv.Itoa(s.nextID)
	s.mu.Unlock()

	reqBody, err := encodeRequest(id, method, params)
	if err != nil {
		return fmt.Errorf("mcp http %s: %w", method, err)
	}
	resp, err := s.post(ctx, reqBody)
	if err != nil {
		return fmt.Errorf("mcp http %s: %w", method, err)
	}
	defer resp.Body.Close()
	s.captureSession(resp)
	// Close the response body if the caller cancels, so a blocked streamed read
	// unblocks rather than leaking a goroutine.
	stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
	defer stopClose()

	return s.decodeResponse(resp, id, method, result)
}

// post sends a JSON-RPC request body to the endpoint with the streamable HTTP
// headers and returns the response.
func (s *streamableHTTPTransport) post(ctx context.Context, body []byte) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.url, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	s.applyHeaders(req)
	return s.http.Do(req)
}

// applyHeaders sets the Streamable HTTP request headers.
func (s *streamableHTTPTransport) applyHeaders(req *http.Request) {
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	s.mu.Lock()
	session := s.session
	proto := ""
	if s.protocol != nil {
		proto = s.protocol()
	}
	s.mu.Unlock()
	if session != "" {
		req.Header.Set(mcpSessionIDHeader, session)
	}
	if proto != "" {
		req.Header.Set("Mcp-Protocol-Version", proto)
	}
}

// captureSession stores a session id returned by the server.
func (s *streamableHTTPTransport) captureSession(resp *http.Response) {
	if sid := resp.Header.Get(mcpSessionIDHeader); sid != "" {
		s.mu.Lock()
		s.session = sid
		s.mu.Unlock()
	}
}

// decodeResponse reads a POST response and decodes the reply for id.
func (s *streamableHTTPTransport) decodeResponse(resp *http.Response, id, method string, result any) error {
	if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") {
		return s.readEventStream(resp.Body, resp, id, method, result)
	}
	// A single JSON body.
	body, err := io.ReadAll(io.LimitReader(resp.Body, maxHTTPResponse))
	if err != nil {
		return fmt.Errorf("mcp http %s read: %w", method, err)
	}
	var rr rpcResponse
	if err := json.Unmarshal(bytes.TrimSpace(body), &rr); err != nil {
		return fmt.Errorf("mcp http %s decode: %w", method, err)
	}
	if !decodeIDString(rr, id) {
		return fmt.Errorf("mcp http %s: %w", method, ErrClosed)
	}
	return s.applyResponse(rr, method, result)
}

// readEventStream parses an SSE stream and stops once the response matching id is
// received. Notifications (no id) are routed to the handler.
func (s *streamableHTTPTransport) readEventStream(r io.Reader, resp *http.Response, id, method string, result any) error {
	br := bufio.NewReader(r)
	var event, data string
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			if ctxErr := s.cancelErr(resp); ctxErr != nil && ctxErr != context.Canceled {
				return fmt.Errorf("mcp http %s: %w", method, ctxErr)
			}
			return fmt.Errorf("mcp http %s: %w", method, ErrClosed)
		}
		line = strings.TrimRight(line, "\r\n")
		switch {
		case line == "":
			if data != "" {
				stop, herr := s.handleEvent(event, data, id, method, result)
				if herr != nil {
					return herr
				}
				if stop {
					return nil
				}
			}
			event, data = "", ""
		case strings.HasPrefix(line, "event:"):
			event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
		case strings.HasPrefix(line, "data:"):
			d := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if data == "" {
				data = d
			} else {
				data += "\n" + d
			}
		}
	}
}

// handleEvent dispatches a single SSE data payload. It returns true when the
// event is the response being awaited, so the reader can stop.
func (s *streamableHTTPTransport) handleEvent(event, data, id, method string, result any) (bool, error) {
	if data == "" {
		return false, nil
	}
	var resp rpcResponse
	if err := json.Unmarshal([]byte(data), &resp); err != nil {
		return false, nil
	}
	if resp.ID == nil {
		s.routeNotification(resp.Method)
		return false, nil
	}
	if decodeIDString(resp, id) {
		if err := s.applyResponse(resp, method, result); err != nil {
			return false, err
		}
		return true, nil
	}
	return false, nil
}

// routeNotification hands a server-initiated notification to the client handler.
func (s *streamableHTTPTransport) routeNotification(method string) {
	if method == "" {
		return
	}
	s.mu.Lock()
	fn := s.onNotification
	s.mu.Unlock()
	if fn != nil {
		fn(method)
	}
}

// cancelErr reports whether the response's request context was cancelled.
func (s *streamableHTTPTransport) cancelErr(resp *http.Response) error {
	if resp.Request != nil && resp.Request.Context().Err() != nil {
		return resp.Request.Context().Err()
	}
	return nil
}

// applyResponse decodes a JSON-RPC response into result, surfacing an error
// object if present.
func (s *streamableHTTPTransport) applyResponse(resp rpcResponse, method string, result any) error {
	if resp.Error != nil {
		return fmt.Errorf("mcp http %s: %w: %w", method, ErrRPC, resp.Error)
	}
	if result != nil {
		if err := json.Unmarshal(resp.Result, result); err != nil {
			return fmt.Errorf("mcp http %s decode: %w", method, err)
		}
	}
	return nil
}

// notify POSTs a fire-and-forget JSON-RPC notification (no id). The server
// typically answers with 202 Accepted and no body.
func (s *streamableHTTPTransport) notify(ctx context.Context, method string, params any) error {
	s.mu.Lock()
	if !s.started {
		s.mu.Unlock()
		return fmt.Errorf("mcp notify %s: %w", method, ErrNotConnected)
	}
	s.mu.Unlock()
	reqBody, err := encodeRequest(nil, method, params)
	if err != nil {
		return fmt.Errorf("mcp notify %s: %w", method, err)
	}
	resp, err := s.post(ctx, reqBody)
	if err != nil {
		return fmt.Errorf("mcp notify %s: %w", method, err)
	}
	defer resp.Body.Close()
	s.captureSession(resp)
	_, _ = io.Copy(io.Discard, resp.Body)
	return nil
}

// readStandalone reads the standalone GET SSE stream for server-initiated
// messages until the connection is closed or the stream terminates.
func (s *streamableHTTPTransport) readStandalone(ctx context.Context) {
	for {
		if ctx.Err() != nil {
			return
		}
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.url, nil)
		if err != nil {
			return
		}
		req.Header.Set("Accept", "text/event-stream")
		s.mu.Lock()
		session := s.session
		proto := ""
		if s.protocol != nil {
			proto = s.protocol()
		}
		s.mu.Unlock()
		if session != "" {
			req.Header.Set(mcpSessionIDHeader, session)
		}
		if proto != "" {
			req.Header.Set("Mcp-Protocol-Version", proto)
		}
		resp, err := s.http.Do(req)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			s.log.Debug("mcp http get stream error", "name", s.cfg.Name, "err", err)
			s.sleepBackoff(ctx)
			continue
		}
		if resp.StatusCode != http.StatusOK {
			_ = resp.Body.Close()
			s.log.Debug("mcp http get stream not accepted", "name", s.cfg.Name, "status", resp.StatusCode)
			return
		}
		s.readStandaloneEvents(ctx, resp)
		if ctx.Err() != nil {
			return
		}
		s.sleepBackoff(ctx)
	}
}

// readStandaloneEvents consumes one GET stream, routing notifications until it
// ends or the connection is closed.
func (s *streamableHTTPTransport) readStandaloneEvents(ctx context.Context, resp *http.Response) {
	stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
	defer stopClose()
	br := bufio.NewReader(resp.Body)
	var event, data string
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			return
		}
		line = strings.TrimRight(line, "\r\n")
		switch {
		case line == "":
			if data != "" {
				s.handleEvent(event, data, "", "", nil)
			}
			event, data = "", ""
		case strings.HasPrefix(line, "event:"):
			event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
		case strings.HasPrefix(line, "data:"):
			d := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if data == "" {
				data = d
			} else {
				data += "\n" + d
			}
		}
	}
}

// sleepBackoff waits between standalone stream reconnect attempts, exiting early
// when the connection is closed.
func (s *streamableHTTPTransport) sleepBackoff(ctx context.Context) {
	select {
	case <-time.After(time.Second):
	case <-ctx.Done():
	}
}

// Close closes the standalone stream and marks the transport closed. It is
// idempotent.
func (s *streamableHTTPTransport) Close() error {
	s.closeOnce.Do(func() {
		close(s.closeCh)
		s.mu.Lock()
		if s.getCancel != nil {
			s.getCancel()
		}
		s.mu.Unlock()
	})
	return nil
}