memory-mcp/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

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

// MemoryManager handles storage and retrieval of facts with keywords
type MemoryManager struct {
	filePath string
}

// Memory represents a fact with associated keywords
type Memory struct {
	Fact     string
	Keywords []string
}

// NewMemoryManager creates a new memory manager
func NewMemoryManager() (*MemoryManager, error) {
	// Get XDG_STATE_HOME or use default
	stateHome := os.Getenv("XDG_STATE_HOME")
	if stateHome == "" {
		homeDir, err := os.UserHomeDir()
		if err != nil {
			return nil, fmt.Errorf("failed to get home directory: %w", err)
		}
		stateHome = filepath.Join(homeDir, ".local", "state")
	}

	// Create directory if it doesn't exist
	err := os.MkdirAll(stateHome, 0755)
	if err != nil {
		return nil, fmt.Errorf("failed to create state directory: %w", err)
	}

	return &MemoryManager{
		filePath: filepath.Join(stateHome, "memories.txt"),
	}, nil
}

// AddMemory adds a new memory
func (m *MemoryManager) AddMemory(fact string, keywords []string) error {
	if strings.Contains(fact, ";") {
		return errors.New("fact cannot contain semicolon (;)")
	}

	// Format: fact;keyword1,keyword2,...
	line := fact + ";" + strings.Join(keywords, ",") + "\n"

	// Open file in append mode
	file, err := os.OpenFile(m.filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return fmt.Errorf("failed to open memory file: %w", err)
	}
	defer file.Close()

	_, err = file.WriteString(line)
	if err != nil {
		return fmt.Errorf("failed to write memory: %w", err)
	}

	return nil
}

// GetAllMemories returns all stored memories
func (m *MemoryManager) GetAllMemories() ([]Memory, error) {
	return m.loadMemories()
}

// GetUniqueKeywords returns a list of all unique keywords
func (m *MemoryManager) GetUniqueKeywords() ([]string, error) {
	memories, err := m.loadMemories()
	if err != nil {
		return nil, err
	}

	// Use a map to track unique keywords
	keywordMap := make(map[string]bool)
	for _, memory := range memories {
		for _, keyword := range memory.Keywords {
			keywordMap[keyword] = true
		}
	}

	// Convert map keys to slice
	keywords := make([]string, 0, len(keywordMap))
	for keyword := range keywordMap {
		keywords = append(keywords, keyword)
	}

	return keywords, nil
}

// GetMemoriesByKeyword returns all memories that have the given keyword
func (m *MemoryManager) GetMemoriesByKeyword(keyword string) ([]Memory, error) {
	memories, err := m.loadMemories()
	if err != nil {
		return nil, err
	}

	// Filter memories by keyword
	var filtered []Memory
	for _, memory := range memories {
		for _, k := range memory.Keywords {
			if k == keyword {
				filtered = append(filtered, memory)
				break
			}
		}
	}

	return filtered, nil
}

// loadMemories loads all memories from the file
func (m *MemoryManager) loadMemories() ([]Memory, error) {
	// Check if file exists
	if _, err := os.Stat(m.filePath); os.IsNotExist(err) {
		// Return empty slice if file doesn't exist
		return []Memory{}, nil
	}

	// Read file content
	content, err := os.ReadFile(m.filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read memory file: %w", err)
	}

	var memories []Memory
	lines := strings.Split(string(content), "\n")

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		parts := strings.Split(line, ";")
		if len(parts) != 2 {
			continue // Skip invalid lines
		}

		fact := parts[0]
		keywords := strings.Split(parts[1], ",")

		// Filter out empty keywords
		var cleanKeywords []string
		for _, k := range keywords {
			k = strings.TrimSpace(k)
			if k != "" {
				cleanKeywords = append(cleanKeywords, k)
			}
		}

		memories = append(memories, Memory{
			Fact:     fact,
			Keywords: cleanKeywords,
		})
	}

	return memories, nil
}

func main() {
	fmt.Println("Memory MCP - Management Control Program")

	// Create memory manager
	memoryManager, err := NewMemoryManager()
	if err != nil {
		fmt.Printf("Error initializing memory manager: %v\n", err)
		os.Exit(1)
	}

	// Create MCP server
	s := server.NewMCPServer(
		"Memory Manager",
		"1.0.0",
		server.WithToolCapabilities(true),
		server.WithRecovery(),
	)

	// Tool to add a new memory
	addMemoryTool := mcp.NewTool("add_memory",
		mcp.WithDescription("Add a new fact with associated keywords"),
		mcp.WithString("fact", mcp.Required(), mcp.Description("The fact to remember (cannot contain semicolons)")),
		mcp.WithString("keywords", mcp.Required(), mcp.Description("Comma-separated list of keywords")),
	)

	s.AddTool(addMemoryTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		fact := request.Params.Arguments["fact"].(string)
		keywordsStr := request.Params.Arguments["keywords"].(string)

		keywords := strings.Split(keywordsStr, ",")
		for i, k := range keywords {
			keywords[i] = strings.TrimSpace(k)
		}

		err := memoryManager.AddMemory(fact, keywords)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("Failed to add memory: %v", err)), nil
		}

		return mcp.NewToolResultText(fmt.Sprintf("Memory added: %s", fact)), nil
	})

	// Tool to list all memories
	listMemoriesTool := mcp.NewTool("list_all_memories",
		mcp.WithDescription("List all stored facts with their keywords"),
	)

	s.AddTool(listMemoriesTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		memories, err := memoryManager.GetAllMemories()
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("Failed to retrieve memories: %v", err)), nil
		}

		if len(memories) == 0 {
			return mcp.NewToolResultText("No memories stored yet."), nil
		}

		var result strings.Builder
		for i, memory := range memories {
			result.WriteString(fmt.Sprintf("%d. Fact: %s\n   Keywords: %s\n",
				i+1, memory.Fact, strings.Join(memory.Keywords, ", ")))
		}

		return mcp.NewToolResultText(result.String()), nil
	})

	// Tool to list unique keywords
	listKeywordsTool := mcp.NewTool("list_unique_keywords",
		mcp.WithDescription("List all unique keywords used across all memories"),
	)

	s.AddTool(listKeywordsTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		keywords, err := memoryManager.GetUniqueKeywords()
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("Failed to retrieve keywords: %v", err)), nil
		}

		if len(keywords) == 0 {
			return mcp.NewToolResultText("No keywords found."), nil
		}

		return mcp.NewToolResultText(fmt.Sprintf("Unique keywords: %s", strings.Join(keywords, ", "))), nil
	})

	// Tool to find memories by keyword
	findByKeywordTool := mcp.NewTool("find_by_keyword",
		mcp.WithDescription("Find all facts associated with a specific keyword"),
		mcp.WithString("keyword", mcp.Required(), mcp.Description("Keyword to search for")),
	)

	s.AddTool(findByKeywordTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		keyword := request.Params.Arguments["keyword"].(string)
		memories, err := memoryManager.GetMemoriesByKeyword(keyword)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("Failed to search memories: %v", err)), nil
		}

		if len(memories) == 0 {
			return mcp.NewToolResultText(fmt.Sprintf("No memories found with keyword '%s'.", keyword)), nil
		}

		var result strings.Builder
		result.WriteString(fmt.Sprintf("Found %d memories with keyword '%s':\n\n", len(memories), keyword))

		for i, memory := range memories {
			result.WriteString(fmt.Sprintf("%d. Fact: %s\n   Keywords: %s\n",
				i+1, memory.Fact, strings.Join(memory.Keywords, ", ")))
		}

		return mcp.NewToolResultText(result.String()), nil
	})

	// Start the server
	fmt.Println("Starting Memory Manager MCP server...")
	if err := server.ServeStdio(s); err != nil {
		fmt.Printf("Server error: %v\n", err)
	}
}