mcp_server.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

type SearchArgs struct {
	Query    string `json:"query"`
	Provider string `json:"provider"`
}

func mcpHandler(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SearchArgs]) (*mcp.CallToolResultFor[any], error) {
	provider := params.Arguments.Provider
	if provider == "" {
		if len(config.EnabledProviders) > 0 {
			provider = config.EnabledProviders[0]
		} else {
			return nil, fmt.Errorf("no provider specified and no enabled providers configured")
		}
	}

	results := runQuery(params.Arguments.Query, provider)

	jsonResults, err := json.Marshal(results)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal search results: %w", err)
	}

	return &mcp.CallToolResultFor[any]{
		Content: []mcp.Content{
			&mcp.TextContent{
				Text: string(jsonResults),
			},
		},
	}, nil
}

func setupMCPServer(mux *http.ServeMux) {
	mcpServer := mcp.NewServer("seargo", "0.1.0", nil)

	providerNames := make([]any, 0, len(providersMap))
	for k := range providersMap {
		providerNames = append(providerNames, k)
	}

	tool := mcp.NewServerTool(
		"search",
		"Performs a search query across multiple providers.",
		mcpHandler,
		mcp.Input(
			mcp.Property(
				"query",
				mcp.Description("The search query string"),
				mcp.Required(true),
			),
			mcp.Property(
				"provider",
				mcp.Description("The single provider to query (e.g., 'duckduckgo', 'brave', 'bing). Defaults to Bing."),
				mcp.Enum(providerNames...),
				mcp.Required(false),
			),
		),
	)
	mcpServer.AddTools(tool)

	sseServer := mcp.NewSSEHandler(func(request *http.Request) *mcp.Server {
		return mcpServer
	})
	mux.Handle("/sse", sseServer)
	mux.Handle("/message", sseServer)

	streamableHttpServer := mcp.NewStreamableHTTPHandler(func(request *http.Request) *mcp.Server {
		return mcpServer
	}, nil)
	mux.Handle("/mcp", streamableHttpServer)
}