main.go

package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"

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

const (
	defaultHTTPAddr = ":8080"
	defaultHTTPPath = "/mcp"
)

func readEnv(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func toExecRelative(p string) (string, error) {
	if p == "" {
		return "", nil
	}
	if filepath.IsAbs(p) {
		return p, nil
	}
	exe, err := os.Executable()
	if err != nil {
		return "", err
	}
	return filepath.Join(filepath.Dir(exe), p), nil
}

func mustJSON(v any) string {
	b, _ := json.MarshalIndent(v, "", "  ")
	return string(b)
}

func main() {
	transport := strings.ToLower(readEnv("TRANSPORT", "stdio"))
	httpAddr := readEnv("HTTP_ADDR", defaultHTTPAddr)
	httpPath := readEnv("HTTP_PATH", defaultHTTPPath)
	memPathEnv := readEnv("MEMORY_FILE_PATH", "memory.db")
	boltPathEnv := readEnv("BOLT_DB_PATH", "")

	memPath, err := toExecRelative(memPathEnv)
	if err != nil {
		log.Fatalf("resolve memory path: %v", err)
	}
	boltPath := boltPathEnv
	if boltPath == "" {
		boltPath = memPath + ".bolt"
	}
	boltPath, err = toExecRelative(boltPath)
	if err != nil {
		log.Fatalf("resolve bolt path: %v", err)
	}

	mgr := NewGraphManager(memPath, boltPath)

	server := mcp.NewServer(&mcp.Implementation{Name: "memory-server", Version: "0.6.3"}, nil)

	mcp.AddTool(server, &mcp.Tool{
		Name:        "create_entities",
		Description: "Create multiple new entities in the knowledge graph",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		Entities []Entity `json:"entities"`
	}) (*mcp.CallToolResult, any, error) {
		created, err := mgr.CreateEntities(in.Entities)
		if err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: mustJSON(created)}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "create_relations",
		Description: "Create multiple new relations between entities in the knowledge graph. Relations should be in active voice",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		Relations []Relation `json:"relations"`
	}) (*mcp.CallToolResult, any, error) {
		created, err := mgr.CreateRelations(in.Relations)
		if err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: mustJSON(created)}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "add_observations",
		Description: "Add new observations to existing entities in the knowledge graph",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		Observations []struct {
			EntityName string   `json:"entityName"`
			Contents   []string `json:"contents"`
		} `json:"observations"`
	}) (*mcp.CallToolResult, any, error) {
		added, err := mgr.AddObservations(in.Observations)
		if err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: mustJSON(added)}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "delete_entities",
		Description: "Delete multiple entities and their associated relations from the knowledge graph",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		EntityNames []string `json:"entityNames"`
	}) (*mcp.CallToolResult, any, error) {
		if err := mgr.DeleteEntities(in.EntityNames); err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "Entities deleted successfully"}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "delete_observations",
		Description: "Delete specific observations from entities in the knowledge graph",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		Deletions []struct {
			EntityName   string   `json:"entityName"`
			Observations []string `json:"observations"`
		} `json:"deletions"`
	}) (*mcp.CallToolResult, any, error) {
		if err := mgr.DeleteObservations(in.Deletions); err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "Observations deleted successfully"}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "delete_relations",
		Description: "Delete multiple relations from the knowledge graph",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		Relations []Relation `json:"relations"`
	}) (*mcp.CallToolResult, any, error) {
		if err := mgr.DeleteRelations(in.Relations); err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "Relations deleted successfully"}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "read_graph",
		Description: "Read the entire knowledge graph",
	}, func(_ context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) {
		g, err := mgr.ReadGraph()
		if err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: mustJSON(g)}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "search_nodes",
		Description: "Search for nodes in the knowledge graph based on a query",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		Query string `json:"query"`
	}) (*mcp.CallToolResult, any, error) {
		g, err := mgr.SearchNodes(in.Query)
		if err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: mustJSON(g)}}}, nil, nil
	})

	mcp.AddTool(server, &mcp.Tool{
		Name:        "open_nodes",
		Description: "Open specific nodes in the knowledge graph by their names",
	}, func(_ context.Context, _ *mcp.CallToolRequest, in struct {
		Names []string `json:"names"`
	}) (*mcp.CallToolResult, any, error) {
		g, err := mgr.OpenNodes(in.Names)
		if err != nil {
			return nil, nil, err
		}
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: mustJSON(g)}}}, nil, nil
	})

	switch transport {
	case "stdio":
		if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil {
			log.Fatalf("stdio server failed: %v", err)
		}
	case "http":
		addr := httpAddr
		path := httpPath
		if path == "" {
			path = defaultHTTPPath
		}

		getServer := func(_ *http.Request) *mcp.Server { return server }
		h := mcp.NewStreamableHTTPHandler(getServer, &mcp.StreamableHTTPOptions{})

		var handler http.Handler = h

		user, pass, err := readBasicAuth()
		if err != nil {
			log.Fatalf("basic auth config error: %v", err)
		}
		if user != "" || pass != "" {
			handler = BasicAuthMiddleware(handler, user, pass)
		}
		handler = CORSMiddleware(handler)

		mux := http.NewServeMux()
		mux.Handle(path, handler)
		log.Printf("Knowledge Graph MCP Server running on HTTP %s%s", addr, path)
		if err := http.ListenAndServe(addr, mux); err != nil {
			log.Fatal(err)
		}
	default:
		log.Fatalf("unknown transport: %q (use stdio or http)", transport)
	}
}

func readBasicAuth() (string, string, error) {
	user := readEnv("BASIC_AUTH_USER", "")
	pass := readEnv("BASIC_AUTH_PASSWORD", "")

	if uf := readEnv("BASIC_AUTH_USER_FILE", ""); uf != "" {
		b, err := os.ReadFile(uf)
		if err != nil {
			return "", "", err
		}
		user = strings.TrimSpace(string(b))
	}
	if pf := readEnv("BASIC_AUTH_PASSWORD_FILE", ""); pf != "" {
		b, err := os.ReadFile(pf)
		if err != nil {
			return "", "", err
		}
		pass = strings.TrimSpace(string(b))
	}
	if (user == "") != (pass == "") {
		return "", "", errors.New("both BASIC_AUTH_USER and BASIC_AUTH_PASSWORD (or *_FILE) must be set together")
	}
	return user, pass, nil
}

func BasicAuthMiddleware(next http.Handler, user, pass string) http.Handler {
	want := base64.StdEncoding.EncodeToString([]byte(user + ":" + pass))
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization")
		if !strings.HasPrefix(auth, "Basic ") {
			w.Header().Set("WWW-Authenticate", "Basic realm=\"mcp\"")
			http.Error(w, "authentication required", http.StatusUnauthorized)
			return
		}
		got := strings.TrimPrefix(auth, "Basic ")
		if got != want {
			w.Header().Set("WWW-Authenticate", "Basic realm=\"mcp\"")
			http.Error(w, "invalid credentials", http.StatusUnauthorized)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func CORSMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		origin := r.Header.Get("Origin")
		if origin != "" {
			w.Header().Set("Access-Control-Allow-Origin", origin)
			w.Header().Set("Vary", "Origin")
			w.Header().Set("Access-Control-Allow-Credentials", "true")
			w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, Mcp-Session-Id")
			w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
			w.Header().Set("Access-Control-Expose-Headers", "Mcp-Session-Id")
		}
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusNoContent)
			return
		}
		next.ServeHTTP(w, r)
	})
}