Cyrene_v2/internal/capability/mcp/sse.go

package mcp

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

// sseTransport speaks JSON-RPC 2.0 to an MCP server over SSE. It opens a
// long-lived GET stream for server-to-client messages and POSTs client-to-server
// requests to the message endpoint (discovered via an endpoint event, falling
// back to the configured URL).
type sseTransport struct {
	cfg            MCPConfig
	log            *slog.Logger
	url            string
	endpoint       string
	mu             sync.Mutex
	nextID         int
	pending        map[string]chan []byte
	resp           *http.Response
	body           io.ReadCloser
	closeCh        chan struct{}
	closeOnce      sync.Once
	started        bool
	onNotification func(method string)
}

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

// maxSSEResponse bounds how much of a POST reply body is read before falling
// back to the SSE stream.
const maxSSEResponse = 1 << 20

// newSSETransport builds an SSE transport for cfg.
func newSSETransport(cfg MCPConfig, log *slog.Logger) *sseTransport {
	return &sseTransport{
		cfg:      cfg,
		log:      log,
		url:      cfg.URL,
		endpoint: cfg.URL,
		pending:  make(map[string]chan []byte),
		closeCh:  make(chan struct{}),
	}
}

// Start opens the SSE stream and launches the reader goroutine.
func (s *sseTransport) Start(ctx context.Context) error {
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("mcp sse start: %w", err)
	}
	if s.cfg.URL == "" {
		return fmt.Errorf("mcp sse: %w", ErrNoTransport)
	}
	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, s.cfg.URL, nil)
	if err != nil {
		return fmt.Errorf("mcp sse get: %w", err)
	}
	req.Header.Set("Accept", "text/event-stream")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("mcp sse connect %s: %w", s.cfg.URL, err)
	}
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return fmt.Errorf("mcp sse connect %s: status %d", s.cfg.URL, resp.StatusCode)
	}
	s.resp = resp
	s.body = resp.Body
	s.started = true
	go s.readLoop()
	return nil
}

// readLoop parses SSE events and dispatches message events to pending calls.
func (s *sseTransport) readLoop() {
	r := bufio.NewReader(s.body)
	var event, data string
	for {
		line, err := r.ReadString('\n')
		if err != nil {
			return
		}
		line = strings.TrimRight(line, "\r\n")
		switch {
		case line == "":
			if data != "" {
				s.dispatch(event, data)
			}
			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
			}
		}
	}
}

// dispatch handles a single SSE event.
func (s *sseTransport) dispatch(event, data string) {
	switch event {
	case "endpoint":
		if e := strings.TrimSpace(data); e != "" {
			s.mu.Lock()
			s.endpoint = e
			s.mu.Unlock()
		}
	case "message":
		var resp rpcResponse
		if err := json.Unmarshal([]byte(data), &resp); err != nil {
			return
		}
		if resp.ID == nil {
			// A server-initiated notification; hand it to the client's handler.
			// The handler must not re-enter this transport synchronously.
			if resp.Method != "" {
				s.mu.Lock()
				fn := s.onNotification
				s.mu.Unlock()
				if fn != nil {
					fn(resp.Method)
				}
			}
			return
		}
		id, ok := resp.ID.(string)
		if !ok {
			return
		}
		s.mu.Lock()
		ch, has := s.pending[id]
		delete(s.pending, id)
		s.mu.Unlock()
		if has {
			ch <- []byte(data)
		}
	}
}

// Call POSTs a request and waits for a response, either from the POST reply body
// or from the SSE stream.
func (s *sseTransport) Call(ctx context.Context, method string, params any, result any) error {
	s.mu.Lock()
	if !s.started {
		s.mu.Unlock()
		return fmt.Errorf("mcp sse %s: %w", method, ErrNotConnected)
	}
	s.nextID++
	id := strconv.Itoa(s.nextID)
	ch := make(chan []byte, 1)
	s.pending[id] = ch
	endpoint := s.endpoint
	s.mu.Unlock()

	defer func() {
		s.mu.Lock()
		delete(s.pending, id)
		s.mu.Unlock()
	}()

	reqBody, err := encodeRequest(id, method, params)
	if err != nil {
		return fmt.Errorf("mcp sse %s: %w", method, err)
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody))
	if err != nil {
		return fmt.Errorf("mcp sse %s: %w", method, err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "application/json")
	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return fmt.Errorf("mcp sse %s post: %w", method, err)
	}
	body, _ := io.ReadAll(io.LimitReader(resp.Body, maxSSEResponse))
	resp.Body.Close()

	var direct rpcResponse
	if json.Unmarshal(bytes.TrimSpace(body), &direct) == nil && decodeIDString(direct, id) {
		return s.applyResponse(direct, method, result)
	}

	select {
	case data := <-ch:
		var rr rpcResponse
		if err := json.Unmarshal(data, &rr); err != nil {
			return fmt.Errorf("mcp sse %s decode: %w", method, err)
		}
		return s.applyResponse(rr, method, result)
	case <-ctx.Done():
		return fmt.Errorf("mcp sse %s: %w", method, ctx.Err())
	case <-s.closeCh:
		return fmt.Errorf("mcp sse %s: %w", method, ErrClosed)
	}
}

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

// notify POSTs a fire-and-forget JSON-RPC notification.
func (s *sseTransport) notify(ctx context.Context, method string, params any) error {
	s.mu.Lock()
	started := s.started
	endpoint := s.endpoint
	s.mu.Unlock()
	if !started {
		return fmt.Errorf("mcp sse %s: %w", method, ErrNotConnected)
	}
	reqBody, err := encodeRequest(nil, method, params)
	if err != nil {
		return fmt.Errorf("mcp notify %s: %w", method, err)
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody))
	if err != nil {
		return fmt.Errorf("mcp notify %s: %w", method, err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return fmt.Errorf("mcp notify %s: %w", method, err)
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)
	return nil
}

// Close closes the SSE stream and any pending responses.
func (s *sseTransport) Close() error {
	s.closeOnce.Do(func() {
		close(s.closeCh)
		if s.body != nil {
			_ = s.body.Close()
		}
		if s.resp != nil && s.resp.Body != nil {
			_ = s.resp.Body.Close()
		}
		s.mu.Lock()
		s.pending = make(map[string]chan []byte)
		s.mu.Unlock()
	})
	return nil
}