st-pierre/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"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"
	"go.uber.org/zap"
)

var version = "0.1.0"

func main() {
	logger, err := zap.NewProduction()
	if err != nil {
		panic(fmt.Sprintf("failed to init logger: %v", err))
	}
	defer logger.Sync()

	qdrantURL := requireEnv(logger, "QDRANT_URL")
	openAIKey := requireEnv(logger, "OPENAI_API_KEY")
	token := requireEnv(logger, "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 {
		logger.Fatal("failed to create qdrant client", zap.Error(err))
	}
	defer qdrantClient.Close()

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

	sqlitePath := envOrDefault("SQLITE_PATH", "knowledge.db")
	sqliteStore, err := newSqliteStore(sqlitePath)
	if err != nil {
		logger.Fatal("failed to init sqlite", zap.Error(err))
	}
	defer sqliteStore.close()

	if err := qdrantStore.ensureCollection(context.Background()); err != nil {
		logger.Fatal("failed to ensure collection", zap.Error(err))
	}

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

	registerStoreKnowledge(mcpServer, sqliteStore, qdrantStore, logger)
	registerSearchKnowledge(mcpServer, sqliteStore, qdrantStore, logger)

	streamableServer := server.NewStreamableHTTPServer(mcpServer)

	authHandler := authMiddleware(token, streamableServer)

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

	logger.Info("starting", zap.String("port", port))

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

	if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		logger.Fatal("server stopped", zap.Error(err))
	}
}

func registerStoreKnowledge(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) {
	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("note",
				mcp.Required(),
				mcp.Description("The markdown body — keep it concise, one paragraph usually"),
			),
			mcp.WithArray("repos",
				mcp.Description("Optional list of repository names"),
			),
			mcp.WithString("host",
				mcp.Description("Optional hostname or null if universal"),
			),
			mcp.WithArray("tags",
				mcp.Description("Optional list of tags"),
			),
			mcp.WithString("user",
				mcp.Description("Optional user identifier"),
			),
			mcp.WithString("context",
				mcp.Description("Optional context — what triggered this"),
			),
		),
		handleStoreKnowledge(sqlite, qdrant, logger),
	)
}

func registerSearchKnowledge(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) {
	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(sqlite, qdrant),
	)
}

func handleStoreKnowledge(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) 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
		}

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

		repos := parseStringSlice(args["repos"])
		host := parseOptionalString(args["host"])
		tags := parseStringSlice(args["tags"])
		user := parseOptionalString(args["user"])
		context_ := parseOptionalString(args["context"])

		id, err := sqlite.save(ctx, note, host, user, context_, repos, tags)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to persist: %v", err)), nil
		}

		if err := qdrant.index(ctx, id, note); err != nil {
			logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err))
		}

		entry, err := sqlite.getByID(ctx, id)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to retrieve: %v", err)), nil
		}

		output, err := json.MarshalIndent(entry, "", "  ")
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to marshal result: %v", err)), nil
		}

		return mcp.NewToolResultText(string(output)), nil
	}
}

func handleSearchKnowledge(sqlite *sqliteStore, qdrant *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 := qdrant.search(ctx, query, limit)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
		}

		for i, r := range results {
			if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil {
				results[i].Entry = *cached
			}
		}

		output, err := json.MarshalIndent(results, "", "  ")
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to marshal results: %v", err)), nil
		}

		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 := strings.ToLower(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(logger *zap.Logger, key string) string {
	val := os.Getenv(key)
	if val == "" {
		logger.Fatal("missing required env var", zap.String("key", key))
	}
	return val
}

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