st-pierre/main.go

package main

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

	"github.com/joho/godotenv"
	"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() {
	godotenv.Load()

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

	openAIKey := requireEnv(logger, "OPENAI_API_KEY")
	port := envOrDefault("PORT", "8080")

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

	store.bootstrapAdmin(context.Background(), logger)

	hmacKeyFile := envOrDefault("HMAC_KEY_FILE", ".secret")
	hmacKey, err := loadHMACKey(hmacKeyFile, logger)
	if err != nil {
		logger.Fatal("failed to load HMAC key", zap.Error(err))
	}

	searchEngine := createSearchEngine(logger, store, openAIKey)

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

	registerStoreFlashcard(mcpServer, store, searchEngine, logger)
	registerSearchFlashcard(mcpServer, searchEngine)
	registerListRecentFlashcard(mcpServer, store, logger)
	registerUpdateFlashcard(mcpServer, store, searchEngine, logger)
	registerDeleteFlashcard(mcpServer, store, searchEngine, logger)
	registerVoteFlashcard(mcpServer, store, logger)

	streamableServer := server.NewStreamableHTTPServer(mcpServer)

	withAuth := func(next http.Handler) http.Handler {
		return authMiddleware(store, hmacKey, logger, next)
	}

	mux := http.NewServeMux()

	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{
			"build_date": os.Getenv("BUILD_DATE"),
			"git_commit": os.Getenv("GIT_COMMIT"),
		})
	})

	mux.Handle("/mcp", withAuth(streamableServer))

	mux.HandleFunc("/api/v1/auth/login", loginHandler(store, hmacKey, logger))
	mux.HandleFunc("/api/v1/auth/logout", logoutHandler())
	mux.HandleFunc("/api/v1/auth/me", withAuth(getUserHandler(store, logger)).ServeHTTP)
	mux.HandleFunc("/api/v1/auth/password", withAuth(changePasswordHandler(store, logger)).ServeHTTP)

	mux.Handle("/api/v1/sessions", withAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodPost:
			createSessionHandler(store, logger).ServeHTTP(w, r)
		case http.MethodGet:
			listSessionsHandler(store, logger).ServeHTTP(w, r)
		default:
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
		}
	})))
	mux.Handle("/api/v1/sessions/", withAuth(deleteSessionHandler(store, logger)))

	mux.Handle("/api/v1/users", withAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodPost:
			createUserHandler(store, logger).ServeHTTP(w, r)
		case http.MethodGet:
			listUsersHandler(store, logger).ServeHTTP(w, r)
		default:
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
		}
	})))

	mux.Handle("/api/v1/flashcards", withAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}
		summaries, err := store.listAll(r.Context())
		if err != nil {
			logger.Error("list flashcards", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}
		writeJSON(w, http.StatusOK, map[string][]flashcardSummary{"flashcards": summaries})
	})))
	mux.Handle("/api/v1/flashcards/", withAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}
		serveFlashcardDetail(w, r, store)
	})))

	mux.Handle("/", publicHandler())

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

	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))
		}
		store.close()
		os.Exit(0)
	}()

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

func createSearchEngine(logger *zap.Logger, sqliteStore *sqliteStore, openAIKey string) SearchEngine {
	qdrantURL := os.Getenv("QDRANT_URL")
	if qdrantURL == "" {
		logger.Warn("QDRANT_URL not set, using SQLite FTS5 for search")
		return newFTS5Engine(sqliteStore)
	}

	embeddingModel := os.Getenv("EMBEDDING_MODEL")
	embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL")

	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"

	client, 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))
	}

	emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL)
	store := newQdrantStore(client, emb, sqliteStore, logger)

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

	return store
}

func registerStoreFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("store_flashcard",
			mcp.WithDescription("Store a flashcard 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.WithNumber("public_id",
				mcp.Description("Optional public identifier. Auto-generated if omitted."),
			),
			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"),
			),
		),
		handleStoreFlashcard(sqlite, engine, logger),
	)
}

func registerSearchFlashcard(mcpServer *server.MCPServer, engine SearchEngine) {
	mcpServer.AddTool(
		mcp.NewTool("search_flashcard",
			mcp.WithDescription("Search previously stored flashcards using semantic search. Returns the most relevant 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)"),
			),
		),
		handleSearchFlashcard(engine),
	)
}

func registerListRecentFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("list_recent_flashcards",
			mcp.WithDescription("List recently created flashcards, optionally filtered by tag, repo, or host. Use this at the start of a session to recall what was learned recently."),
			mcp.WithNumber("limit",
				mcp.Description("Maximum number of results to return (default: 20)"),
			),
			mcp.WithString("tag",
				mcp.Description("Filter by tag (e.g. 'preference', 'convention')"),
			),
			mcp.WithString("repo",
				mcp.Description("Filter by repository name (e.g. 'st-pierre', 'loopagent')"),
			),
			mcp.WithString("host",
				mcp.Description("Filter by hostname (e.g. 'shulgi')"),
			),
		),
		handleListRecentFlashcard(sqlite, logger),
	)
}

func handleListRecentFlashcard(sqlite *sqliteStore, 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
		}

		tag := parseOptionalString(args["tag"])
		repo := parseOptionalString(args["repo"])
		host := parseOptionalString(args["host"])

		limit := uint64(20)
		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)
			}
		}

		entries, err := sqlite.listRecent(ctx, tag, repo, host, limit)
		if err != nil {
			logger.Error("list recent failed", zap.Error(err))
			return mcp.NewToolResultError(fmt.Sprintf("list recent failed: %v", err)), nil
		}

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

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

func registerUpdateFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("update_flashcard",
			mcp.WithDescription("Update an existing flashcard. Creates a new version while preserving history."),
			mcp.WithNumber("public_id",
				mcp.Required(),
				mcp.Description("Public identifier of the flashcard to update"),
			),
			mcp.WithString("note",
				mcp.Description("The markdown body"),
			),
			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"),
			),
		),
		handleUpdateFlashcard(sqlite, engine, logger),
	)
}

func handleUpdateFlashcard(sqlite *sqliteStore, engine SearchEngine, 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
		}

		publicID := parseInt64(args["public_id"])
		if publicID == 0 {
			return mcp.NewToolResultError("public_id is required and must be a positive integer"), nil
		}

		latest, err := sqlite.getLatestByPublicID(ctx, publicID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to query: %v", err)), nil
		}
		if latest == nil {
			return mcp.NewToolResultError(fmt.Sprintf("no flashcard found with public_id %d", publicID)), nil
		}

		note := latest.Note
		repos := latest.Repos
		host := latest.Host
		tags := latest.Tags
		user := latest.User
		context_ := latest.Context

		if v := parseOptionalString(args["note"]); v != "" {
			note = v
		}
		if v := parseStringSlice(args["repos"]); v != nil {
			repos = v
		}
		if v := parseOptionalString(args["host"]); v != "" {
			host = v
		}
		if v := parseStringSlice(args["tags"]); v != nil {
			tags = v
		}
		if v := parseOptionalString(args["user"]); v != "" {
			user = v
		}
		if v := parseOptionalString(args["context"]); v != "" {
			context_ = v
		}

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

		if err := engine.Delete(ctx, latest.ID, publicID); err != nil {
			logger.Warn("search engine delete failed", zap.Int64("id", latest.ID), zap.Error(err))
		}
		if err := engine.Index(ctx, id, publicID, version, note); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("search engine index failed: %v", err)), nil
		}

		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 registerDeleteFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, engine SearchEngine, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("delete_flashcard",
			mcp.WithDescription("Delete a flashcard. Creates a tombstone version and removes it from the search index."),
			mcp.WithNumber("public_id",
				mcp.Required(),
				mcp.Description("Public identifier of the flashcard to delete"),
			),
		),
		handleDeleteFlashcard(sqlite, engine, logger),
	)
}

func handleDeleteFlashcard(sqlite *sqliteStore, engine SearchEngine, 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
		}

		publicID := parseInt64(args["public_id"])
		if publicID == 0 {
			return mcp.NewToolResultError("public_id is required and must be a positive integer"), nil
		}

		latest, err := sqlite.getLatestByPublicID(ctx, publicID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to query: %v", err)), nil
		}
		if latest == nil {
			return mcp.NewToolResultError(fmt.Sprintf("no flashcard found with public_id %d", publicID)), nil
		}
		if latest.DeleteTimestamp != "" {
			return mcp.NewToolResultError(fmt.Sprintf("flashcard %d is already deleted", publicID)), nil
		}

		if _, err := sqlite.tombstone(ctx, publicID, latest.Version+1); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to tombstone: %v", err)), nil
		}

		if err := engine.Delete(ctx, latest.ID, publicID); err != nil {
			logger.Warn("search engine delete failed", zap.Int64("id", latest.ID), zap.Error(err))
		}

		return mcp.NewToolResultText(fmt.Sprintf("flashcard %d deleted", publicID)), nil
	}
}

func registerVoteFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("vote_flashcard",
			mcp.WithDescription("Vote on a flashcard. Rate it from -3 (worst) to 3 (best). Requires detailed context about who is voting and why."),
			mcp.WithNumber("public_id",
				mcp.Required(),
				mcp.Description("Public identifier of the flashcard to vote on"),
			),
			mcp.WithNumber("value",
				mcp.Required(),
				mcp.Description("Vote value from -3 to 3"),
			),
			mcp.WithString("user",
				mcp.Required(),
				mcp.Description("Who is voting"),
			),
			mcp.WithString("repo",
				mcp.Required(),
				mcp.Description("Repository context for the vote"),
			),
			mcp.WithString("message",
				mcp.Required(),
				mcp.Description("Justification or context for the vote"),
			),
		),
		handleVoteFlashcard(sqlite, logger),
	)
}

func handleVoteFlashcard(sqlite *sqliteStore, 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
		}

		publicID := parseInt64(args["public_id"])
		if publicID == 0 {
			return mcp.NewToolResultError("public_id is required and must be a positive integer"), nil
		}

		value := int(parseInt64(args["value"]))
		if value < -3 || value > 3 {
			return mcp.NewToolResultError("value must be between -3 and 3"), nil
		}

		user := parseOptionalString(args["user"])
		if user == "" {
			return mcp.NewToolResultError("user is required"), nil
		}

		repo := parseOptionalString(args["repo"])
		if repo == "" {
			return mcp.NewToolResultError("repo is required"), nil
		}

		message := parseOptionalString(args["message"])
		if message == "" {
			return mcp.NewToolResultError("message is required"), nil
		}

		latest, err := sqlite.getLatestByPublicID(ctx, publicID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to query: %v", err)), nil
		}
		if latest == nil {
			return mcp.NewToolResultError(fmt.Sprintf("no flashcard found with public_id %d", publicID)), nil
		}
		if latest.DeleteTimestamp != "" {
			return mcp.NewToolResultError(fmt.Sprintf("flashcard %d is deleted", publicID)), nil
		}

		v, err := sqlite.saveVote(ctx, latest.ID, value, user, repo, message)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to save vote: %v", err)), nil
		}

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

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

func handleStoreFlashcard(sqlite *sqliteStore, engine SearchEngine, 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"])
		publicID := parseInt64(args["public_id"])

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

		if publicID == 0 {
			publicID = id
		}
		if err := engine.Index(ctx, id, publicID, 1, note); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("search engine index failed: %v", err)), nil
		}

		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 handleSearchFlashcard(engine SearchEngine) 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 := engine.Search(ctx, query, limit)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
		}

		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 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
}