st-pierre/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"github.com/qdrant/go-client/qdrant"
)

var version = "0.1.0"

func main() {
	qdrantURL := requireEnv("QDRANT_URL")
	openAIKey := requireEnv("OPENAI_API_KEY")
	token := requireEnv("MCP_TOKEN")
	embeddingModel := os.Getenv("EMBEDDING_MODEL")
	embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL")
	port := envOrDefault("PORT", "8080")

	qdrantPort := 6334
	if p, err := strconv.Atoi(envOrDefault("QDRANT_PORT", "6334")); err == nil {
		qdrantPort = p
	}
	qdrantAPIKey := os.Getenv("QDRANT_API_KEY")
	qdrantTLS := os.Getenv("QDRANT_TLS") == "true"

	qdrantClient, err := qdrant.NewClient(&qdrant.Config{
		Host:   qdrantURL,
		Port:   qdrantPort,
		APIKey: qdrantAPIKey,
		UseTLS: qdrantTLS,
	})
	if err != nil {
		log.Fatalf("failed to create qdrant client: %v", err)
	}
	defer qdrantClient.Close()

	emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL)
	store := newQdrantStore(qdrantClient, emb)

	if err := store.ensureCollection(context.Background()); err != nil {
		log.Fatalf("failed to ensure collection: %v", err)
	}

	mcpServer := server.NewMCPServer(
		"knowledge-mcp",
		version,
		server.WithToolCapabilities(true),
		server.WithLogging(),
	)

	registerTools(mcpServer, store)

	streamableServer := server.NewStreamableHTTPServer(mcpServer)

	authHandler := authMiddleware(token, streamableServer)

	httpServer := &http.Server{
		Addr:    ":" + port,
		Handler: corsMiddleware(authHandler),
	}

	log.Printf("MCP Knowledge Server starting on http://0.0.0.0:%s", port)

	go func() {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
		sig := <-sigChan
		log.Printf("received signal %v, shutting down...", sig)
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if err := streamableServer.Shutdown(shutdownCtx); err != nil {
			log.Printf("shutdown error: %v", err)
		}
		qdrantClient.Close()
		os.Exit(0)
	}()

	if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("server stopped: %v", err)
	}
}

func registerTools(mcpServer *server.MCPServer, store *qdrantStore) {
	mcpServer.AddTool(
		mcp.NewTool("store_knowledge",
			mcp.WithDescription("Store a piece of knowledge for later retrieval. Use this to persist useful information discovered during a session."),
			mcp.WithString("content",
				mcp.Required(),
				mcp.Description("The text content to store"),
			),
			mcp.WithObject("metadata",
				mcp.Description("Optional metadata as key-value pairs (e.g. {\"source\": \"docs\", \"category\": \"golang\"})"),
			),
		),
		handleStoreKnowledge(store),
	)

	mcpServer.AddTool(
		mcp.NewTool("search_knowledge",
			mcp.WithDescription("Search previously stored knowledge using semantic search. Returns the most relevant knowledge entries for the given query."),
			mcp.WithString("query",
				mcp.Required(),
				mcp.Description("The search query"),
			),
			mcp.WithNumber("limit",
				mcp.Description("Maximum number of results to return (default: 5)"),
			),
		),
		handleSearchKnowledge(store),
	)
}

func handleStoreKnowledge(store *qdrantStore) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args, ok := request.Params.Arguments.(map[string]any)
		if !ok {
			return mcp.NewToolResultError("invalid arguments"), nil
		}

		content, ok := args["content"].(string)
		if !ok || content == "" {
			return mcp.NewToolResultError("content is required and must be a non-empty string"), nil
		}

		var metadata map[string]any
		if rawMeta, ok := args["metadata"]; ok && rawMeta != nil {
			metaMap, ok := rawMeta.(map[string]any)
			if !ok {
				return mcp.NewToolResultError("metadata must be a key-value object"), nil
			}
			metadata = metaMap
		}

		result, err := store.store(ctx, content, metadata)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to store knowledge: %v", err)), nil
		}

		output, _ := json.MarshalIndent(result, "", "  ")
		return mcp.NewToolResultText(string(output)), nil
	}
}

func handleSearchKnowledge(store *qdrantStore) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args, ok := request.Params.Arguments.(map[string]any)
		if !ok {
			return mcp.NewToolResultError("invalid arguments"), nil
		}

		query, ok := args["query"].(string)
		if !ok || query == "" {
			return mcp.NewToolResultError("query is required and must be a non-empty string"), nil
		}

		limit := uint64(5)
		if rawLimit := args["limit"]; rawLimit != nil {
			switch v := rawLimit.(type) {
			case float64:
				limit = uint64(v)
			case int:
				limit = uint64(v)
			case int64:
				limit = uint64(v)
			}
		}

		results, err := store.search(ctx, query, limit)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
		}

		output, _ := json.MarshalIndent(results, "", "  ")
		return mcp.NewToolResultText(string(output)), nil
	}
}

func authMiddleware(token string, next http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/health" {
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("ok"))
			return
		}

		authHeader := r.Header.Get("Authorization")
		if !strings.HasPrefix(authHeader, "Bearer ") || strings.TrimPrefix(authHeader, "Bearer ") != token {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte(`{"error":"unauthorized"}`))
			return
		}

		next.ServeHTTP(w, r)
	}
}

func corsMiddleware(next http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")

		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusOK)
			return
		}

		next.ServeHTTP(w, r)
	}
}

func requireEnv(key string) string {
	val := os.Getenv(key)
	if val == "" {
		log.Fatalf("environment variable %s is required", key)
	}
	return val
}

func envOrDefault(key, defaultVal string) string {
	val := os.Getenv(key)
	if val == "" {
		return defaultVal
	}
	return val
}