byre/internal/mcp/server.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"time"
)

const protocolVersion = "2025-06-18"

type Tool struct {
	Name        string
	Description string
	InputSchema json.RawMessage
	Handler     func(ctx context.Context, args json.RawMessage) (string, error)
	Rich        func(ctx context.Context, args json.RawMessage) Result
}

type Content struct {
	Type     string `json:"type"`
	Text     string `json:"text"`
	Data     string `json:"data"`
	MimeType string `json:"mimeType"`
}

func (c Content) MarshalJSON() ([]byte, error) {
	if c.Type == "image" {
		return json.Marshal(map[string]string{"type": c.Type, "data": c.Data, "mimeType": c.MimeType})
	}
	return json.Marshal(map[string]string{"type": c.Type, "text": c.Text})
}

type Result struct {
	Content []Content `json:"content"`
	IsError bool      `json:"isError,omitempty"`
}

func TextResult(text string, isError bool) Result {
	return Result{Content: []Content{{Type: "text", Text: text}}, IsError: isError}
}

const (
	MaxTimeout = 10 * time.Minute
	Grace      = 5 * time.Second
)

type timeoutKey struct{}

func Timeout(ctx context.Context) time.Duration {
	return ctx.Value(timeoutKey{}).(time.Duration)
}

func (t Tool) Schema() json.RawMessage {
	s := map[string]any{}
	err := json.Unmarshal(t.InputSchema, &s)
	if err != nil {
		panic(fmt.Sprintf("tool %s: input schema: %s", t.Name, err))
	}
	props, _ := s["properties"].(map[string]any)
	if props == nil {
		props = map[string]any{}
		s["properties"] = props
	}
	if props["timeout_ms"] == nil {
		props["timeout_ms"] = map[string]any{"type": "integer", "description": fmt.Sprintf("Give up on the call after this long (at most %d).", MaxTimeout.Milliseconds())}
	}
	required, _ := s["required"].([]any)
	s["required"] = append(required, "timeout_ms")
	out, err := json.Marshal(s)
	if err != nil {
		panic(fmt.Sprintf("tool %s: input schema: %s", t.Name, err))
	}
	return out
}

func (t Tool) Call(ctx context.Context, args json.RawMessage) Result {
	p := struct {
		TimeoutMs *int `json:"timeout_ms"`
	}{}
	err := json.Unmarshal(args, &p)
	if err != nil {
		return TextResult(fmt.Sprintf("malformed %s call: %s", t.Name, err), true)
	}
	if p.TimeoutMs == nil || *p.TimeoutMs <= 0 || *p.TimeoutMs > int(MaxTimeout.Milliseconds()) {
		return TextResult(fmt.Sprintf("timeout_ms is required, between 1 and %d", MaxTimeout.Milliseconds()), true)
	}
	timeout := time.Duration(*p.TimeoutMs) * time.Millisecond

	cctx, cancel := context.WithTimeout(context.WithValue(ctx, timeoutKey{}, timeout), timeout)
	defer cancel()
	done := make(chan Result, 1)
	go func() {
		done <- t.run(cctx, args)
	}()
	late := time.NewTimer(timeout + Grace)
	defer late.Stop()
	select {
	case res := <-done:
		return res
	case <-ctx.Done():
		return TextResult(ctx.Err().Error(), true)
	case <-late.C:
		return TextResult(fmt.Sprintf("%s gave no result within %s", t.Name, timeout), true)
	}
}

func (t Tool) run(ctx context.Context, args json.RawMessage) Result {
	if t.Rich != nil {
		return t.Rich(ctx, args)
	}
	text, err := t.Handler(ctx, args)
	if err != nil {
		return TextResult(err.Error(), true)
	}
	return TextResult(text, false)
}

type Server struct {
	Name    string
	Version string
	Tools   []Tool
}

type Stdio struct {
	Name    string
	Command string
	Args    []string
}

type request struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type response struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Result  any             `json:"result,omitempty"`
	Error   *rpcError       `json:"error,omitempty"`
}

type toolDesc struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"inputSchema"`
}

func (s *Server) Tool(name string) (Tool, bool) {
	for _, t := range s.Tools {
		if t.Name == name {
			return t, true
		}
	}
	return Tool{}, false
}

func (s *Server) Handle(ctx context.Context, raw json.RawMessage) (json.RawMessage, error) {
	req := request{}
	err := json.Unmarshal(raw, &req)
	if err != nil {
		return nil, fmt.Errorf("parse jsonrpc: %w", err)
	}

	if req.ID == nil {
		return json.Marshal(response{JSONRPC: "2.0", Result: map[string]any{}})
	}

	result, rerr := s.dispatch(ctx, req)
	resp := response{JSONRPC: "2.0", ID: req.ID, Result: result, Error: rerr}
	return json.Marshal(resp)
}

func (s *Server) dispatch(ctx context.Context, req request) (any, *rpcError) {
	switch req.Method {
	case "initialize":
		return map[string]any{
			"protocolVersion": protocolVersion,
			"capabilities":    map[string]any{"tools": map[string]any{}},
			"serverInfo":      map[string]any{"name": s.Name, "version": s.Version},
		}, nil

	case "ping":
		return map[string]any{}, nil

	case "tools/list":
		tools := []toolDesc{}
		for _, t := range s.Tools {
			tools = append(tools, toolDesc{Name: t.Name, Description: t.Description, InputSchema: t.Schema()})
		}
		return map[string]any{"tools": tools}, nil

	case "tools/call":
		return s.call(ctx, req.Params)
	}

	return nil, &rpcError{Code: -32601, Message: fmt.Sprintf("method not found: %s", req.Method)}
}

func (s *Server) call(ctx context.Context, params json.RawMessage) (any, *rpcError) {
	p := struct {
		Name      string          `json:"name"`
		Arguments json.RawMessage `json:"arguments"`
	}{}
	err := json.Unmarshal(params, &p)
	if err != nil {
		return nil, &rpcError{Code: -32602, Message: err.Error()}
	}

	t, ok := s.Tool(p.Name)
	if !ok {
		return nil, &rpcError{Code: -32602, Message: fmt.Sprintf("unknown tool: %s", p.Name)}
	}

	return t.Call(ctx, p.Arguments), nil
}