feat(capability/mcp): MCP 客户端（stdlib JSON-RPC over stdio/SSE）+ Registry 自动注册/Unregister（registry 增 Remove）
This commit is contained in:
Aska committed
1 parent
08223c06a5
commit
488cf9c79c
14 files changed
+2212
No files matched your search
@@ -0,0 +1,250 @@
|
||||
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
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
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
|
||||
}
|
||||
Reference in new issue
Block a user