didier
e700a81c25
Schema: public_id + version columns with UNIQUE constraint, delete_timestamp for soft-delete. Tools: store_flashcard accepts optional public_id, update_flashcard creates a new version preserving history, delete_flashcard inserts a tombstone row and removes the Qdrant point. Qdrant: point ID is public_id (only latest version indexed).
453 lines
14 KiB
Go
453 lines
14 KiB
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", "flashcard.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(
|
|
"flashcard-mcp",
|
|
version,
|
|
server.WithToolCapabilities(true),
|
|
server.WithLogging(),
|
|
)
|
|
|
|
registerStoreFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
|
registerSearchFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
|
registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
|
registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStore, logger)
|
|
|
|
streamableServer := server.NewStreamableHTTPServer(mcpServer)
|
|
|
|
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", authMiddleware(token, streamableServer))
|
|
|
|
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))
|
|
}
|
|
qdrantClient.Close()
|
|
sqliteStore.close()
|
|
os.Exit(0)
|
|
}()
|
|
|
|
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
logger.Fatal("server stopped", zap.Error(err))
|
|
}
|
|
}
|
|
|
|
func registerStoreFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, 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, qdrant, logger),
|
|
)
|
|
}
|
|
|
|
func registerSearchFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) {
|
|
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(sqlite, qdrant),
|
|
)
|
|
}
|
|
|
|
func registerUpdateFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, 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, qdrant, logger),
|
|
)
|
|
}
|
|
|
|
func handleUpdateFlashcard(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
|
|
}
|
|
|
|
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 := qdrant.index(ctx, id, publicID, version, 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 registerDeleteFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) {
|
|
mcpServer.AddTool(
|
|
mcp.NewTool("delete_flashcard",
|
|
mcp.WithDescription("Delete a flashcard. Creates a tombstone version and removes it from the vector index."),
|
|
mcp.WithNumber("public_id",
|
|
mcp.Required(),
|
|
mcp.Description("Public identifier of the flashcard to delete"),
|
|
),
|
|
),
|
|
handleDeleteFlashcard(sqlite, qdrant, logger),
|
|
)
|
|
}
|
|
|
|
func handleDeleteFlashcard(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
|
|
}
|
|
|
|
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 := qdrant.delete(ctx, publicID); err != nil {
|
|
logger.Warn("qdrant delete failed (tombstone in sqlite)", zap.Int64("public_id", publicID), zap.Error(err))
|
|
}
|
|
|
|
return mcp.NewToolResultText(fmt.Sprintf("flashcard %d deleted", publicID)), nil
|
|
}
|
|
}
|
|
|
|
func handleStoreFlashcard(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"])
|
|
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 := qdrant.index(ctx, id, publicID, 1, 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 handleSearchFlashcard(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) {
|
|
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
|
|
}
|