🔌

Model Context Protocol サーバーをGoで実装する

2024/11/27に公開

Model Context ProtocolはAIが利用するツールやリソースを提供する仕組みですが、シンプルだと主張する割には意外と実装が複雑なのでGoで実装してみました。

https://zenn.dev/masacento/scraps/fc1e6e063bd732

参考にしたのはQuickstartもあるSQLiteを読み書きするサーバーです。

概要

Quickstartにも書いてあるとおり、MCP Serverはリソースやツールを提供し、クライアントの要求に従うことで、クライアントがモデルを利用する際にコンテキストを渡したりツールを使えるようにするものです。

このような仕組みはモデルをAPIで利用する場合にはFunction CallやTool Callで実装されていますが、それを双方向にした上でプロトコル化しようというものです。

フローとしてはクライアント(Claude Desktop)がサーバーを起動したのちにInitializeして接続し、リクエストやレスポンスをJSONRPCを通じて行うというものです。

一見単純ですが主にこれを標準入力・出力を通じで行います。つまり、Claude Desktopは登録されたこのサーバーを基本的に立ち上げっぱなしにして、更新がないかなどの通信を常に行うことになります。
(http/SSEによる通信方法もあります)

stdioとは何か

標準入力をインプットに、標準出力をアウトプットとしてやり取りする、コマンドラインパイプの時代からある古風な仕組みです。これなら様々な言語で実装できると考えたのでしょうが、ログや状態の把握が難しく開発はしづらいです。(stderrもエラー扱いになるので使えない)

それを考慮してかInspectorというデバッグクライアントがあり、これに接続することでかなり動作を理解することができます。

実装

Goの実装としてはchannelでRequestとResponseを扱い、stdin, stdoutとそれを保持したサーバーがそれぞれread, writeでloopすることで双方向通信をしてみました。

type StdioServer struct {
	inReader  *bufio.Scanner
	outWriter io.Writer
	readCh    chan JSONRPCMessage
	writeCh   chan JSONRPCResponse
	wg        sync.WaitGroup
}

func (s *StdioServer) readLoop(ctx context.Context) {
	defer s.wg.Done()
	for {
		select {
		case <-ctx.Done():
			return
		default:
			var msg JSONRPCMessage
			if err := json.Unmarshal([]byte(s.inReader.Text()), &msg); err != nil {
				slog.Error("Failed to parse JSON", "error", err)
				continue
			}

			s.readCh <- msg
		}
	}
}

func (s *StdioServer) writeLoop(ctx context.Context) {
	defer s.wg.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case msg, ok := <-s.writeCh:
			if !ok {
				return
			}

			jsonData, err := json.Marshal(msg)
			if err != nil {
				slog.Error("Failed to encode JSON", "error", err)
				continue
			}

			if _, err := fmt.Fprintln(s.outWriter, string(jsonData)); err != nil {
				slog.Error("Failed to write to stdout", "error", err)
				return
			}
		}
	}
}

これを利用し、stdin->ReadChannelに届いたRequestのメッセージに応じたResponseを作成し、それをWriteChannel->stdoutを通じて出力して通信します。

server = NewStdioServer(os.Stdin, os.Stdout)
server.Start(ctx)
for msg := range server.ReadChannel() {
    slog.Info("received message", "method", msg.Method, "id", msg.ID)

    switch msg.Method {
    case "ping":
        resp := JSONRPCResponse{
            JSONRPC: "2.0",
            ID:      msg.ID,
            Result:  map[string]any{},
        }
        slog.Debug("sending ping response", "msg", msg, "resp", resp)
        server.WriteChannel() <- resp

これを利用することでInspectorからの接続も確認できました。あとはMethodに応じたSQLiteの操作などを実装(AIに書いてもらった)しました。

実行

go build .
npx -y @modelcontextprotocol/inspector ./mcp-go-example 

リポジトリ

https://github.com/masacento/mcp-go-example

近頃はPythonのサンプルも使っているuvの登場によりPythonの実行は楽になったとはいえ、依存のないシングルバイナリをビルドでき、非同期通信も得意なGoはMCP Serverの実装に役に立つでしょう。

Discussion