mcp-gateway-nexus/internal/demomcp/server.go

// Package demomcp is a tiny, dependency-free MCP server that speaks the stdio
// transport. It lets you exercise the Nexus Gateway end-to-end with no external
// MCP servers installed, and doubles as a hermetic test fixture.
package demomcp

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"time"

	"gitea.thetempleofdoom.com/drjones/mcp-gateway-nexus/internal/mcp"
)

// ServerInfo is the demo server's advertised identity.
var ServerInfo = mcp.Implementation{Name: "nexus-demo-mcp", Version: "0.1.0"}

// nowFunc is overridable in tests.
var nowFunc = time.Now

// Run serves MCP over the given reader/writer (typically os.Stdin/os.Stdout)
// until the input is closed.
func Run(in io.Reader, out io.Writer) error {
	scanner := bufio.NewScanner(in)
	scanner.Buffer(make([]byte, 0, 64*1024), 8<<20)
	enc := json.NewEncoder(out)

	for scanner.Scan() {
		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}
		var msg mcp.Message
		if err := json.Unmarshal(line, &msg); err != nil {
			continue
		}
		if msg.IsNotification() {
			continue // e.g. notifications/initialized — nothing to reply
		}
		resp := dispatch(&msg)
		if resp != nil {
			if err := enc.Encode(resp); err != nil {
				return err
			}
		}
	}
	return scanner.Err()
}

func dispatch(msg *mcp.Message) *mcp.Message {
	switch msg.Method {
	case mcp.MethodInitialize:
		res := mcp.InitializeResult{
			ProtocolVersion: mcp.ProtocolVersion,
			Capabilities:    mcp.Capabilities{Tools: &mcp.ToolsCapability{}},
			ServerInfo:      ServerInfo,
		}
		m, _ := mcp.NewResult(msg.ID, res)
		return m
	case mcp.MethodPing:
		m, _ := mcp.NewResult(msg.ID, struct{}{})
		return m
	case mcp.MethodToolsList:
		m, _ := mcp.NewResult(msg.ID, mcp.ListToolsResult{Tools: toolDefs()})
		return m
	case mcp.MethodToolsCall:
		return callTool(msg)
	default:
		return mcp.NewError(msg.ID, mcp.CodeMethodNotFound, "unknown method: "+msg.Method)
	}
}

func toolDefs() []mcp.ToolDefinition {
	return []mcp.ToolDefinition{
		{
			Name:        "echo",
			Description: "Echo back the provided message.",
			InputSchema: map[string]any{
				"type":       "object",
				"properties": map[string]any{"message": map[string]any{"type": "string"}},
				"required":   []any{"message"},
			},
		},
		{
			Name:        "add",
			Description: "Add two numbers and return the sum.",
			InputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"a": map[string]any{"type": "number"},
					"b": map[string]any{"type": "number"},
				},
				"required": []any{"a", "b"},
			},
		},
		{
			Name:        "now",
			Description: "Return the current server time in RFC3339.",
			InputSchema: map[string]any{"type": "object", "properties": map[string]any{}},
		},
	}
}

func callTool(msg *mcp.Message) *mcp.Message {
	var params mcp.CallToolParams
	if err := msg.UnmarshalParams(&params); err != nil {
		return mcp.NewError(msg.ID, mcp.CodeInvalidParams, err.Error())
	}
	switch params.Name {
	case "echo":
		var args struct {
			Message string `json:"message"`
		}
		_ = json.Unmarshal(params.Arguments, &args)
		m, _ := mcp.NewResult(msg.ID, mcp.TextResult(args.Message, false))
		return m
	case "add":
		var args struct {
			A float64 `json:"a"`
			B float64 `json:"b"`
		}
		if err := json.Unmarshal(params.Arguments, &args); err != nil {
			m, _ := mcp.NewResult(msg.ID, mcp.TextResult("invalid arguments: "+err.Error(), true))
			return m
		}
		m, _ := mcp.NewResult(msg.ID, mcp.TextResult(fmt.Sprintf("%g", args.A+args.B), false))
		return m
	case "now":
		m, _ := mcp.NewResult(msg.ID, mcp.TextResult(nowFunc().UTC().Format(time.RFC3339), false))
		return m
	default:
		m, _ := mcp.NewResult(msg.ID, mcp.TextResult("unknown tool: "+params.Name, true))
		return m
	}
}