commit 0ddd0e4598c6b95151042c5298a7c66eacdf65c2 Author: Lukasz Raczylo Date: Sat Mar 7 16:57:12 2026 +0000 Initial implementation of compaction MCP server Ephemeral, per-session context management for LLMs with 8 MCP tools: store, query, status, compact, pin, forget, configure, update. Features: scoring with decay/importance/access, Jaccard dedup, summary promotion, budget-based eviction, auto-detection of client context window via MCP hooks. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a36c95b --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +compaction-mcp diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..62d95c4 --- /dev/null +++ b/go.mod @@ -0,0 +1,17 @@ +module compaction-mcp + +go 1.26.1 + +require github.com/mark3labs/mcp-go v0.45.0 + +require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/invopop/jsonschema v0.13.0 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/spf13/cast v1.7.1 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..72ec30f --- /dev/null +++ b/go.sum @@ -0,0 +1,39 @@ +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= +github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= +github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mark3labs/mcp-go v0.45.0 h1:s0S8qR/9fWaQ3pHxz7pm1uQ0DrswoSnRIxKIjbiQtkc= +github.com/mark3labs/mcp-go v0.45.0/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y= +github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/main.go b/main.go new file mode 100644 index 0000000..9170800 --- /dev/null +++ b/main.go @@ -0,0 +1,70 @@ +package main + +import ( + "context" + "flag" + "fmt" + "os" + "strings" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" +) + +const serverInstructions = `Context compactor — manages working memory within a token budget. + +At session start, call 'configure' with token_budget set to ~40% of your context window. +Example: 200K context window → token_budget = 80000. + +Workflow: +- 'store' important context (always include a summary for efficient compaction) +- 'query' to retrieve stored information instead of re-reading sources +- 'status' to check budget usage +- 'compact' when usage is high — it frees space and identifies items needing summarization +- 'update' to add summaries to items flagged by compaction` + +func main() { + budget := flag.Int("budget", 100000, "Token budget for context storage") + flag.Parse() + + budgetExplicit := false + flag.Visit(func(f *flag.Flag) { + if f.Name == "budget" { + budgetExplicit = true + } + }) + + store := NewStore(*budget) + + hooks := &server.Hooks{} + if !budgetExplicit { + hooks.OnAfterInitialize = append(hooks.OnAfterInitialize, + func(ctx context.Context, id any, req *mcp.InitializeRequest, res *mcp.InitializeResult) { + name := strings.ToLower(req.Params.ClientInfo.Name) + switch { + case strings.Contains(name, "claude"): + // Claude models: 200K context → 40% = 80K budget + store.Configure(80000, nil, 0) + case strings.Contains(name, "cursor"): + store.Configure(60000, nil, 0) + } + }, + ) + } + + s := server.NewMCPServer( + "compactor", + "0.1.0", + server.WithToolCapabilities(false), + server.WithRecovery(), + server.WithInstructions(serverInstructions), + server.WithHooks(hooks), + ) + + registerTools(s, store) + + if err := server.ServeStdio(s); err != nil { + fmt.Fprintf(os.Stderr, "server error: %v\n", err) + os.Exit(1) + } +} diff --git a/store.go b/store.go new file mode 100644 index 0000000..f07f098 --- /dev/null +++ b/store.go @@ -0,0 +1,446 @@ +package main + +import ( + "crypto/rand" + "fmt" + "math" + "sort" + "strings" + "sync" + "time" +) + +type Item struct { + CreatedAt time.Time + AccessedAt time.Time + Content string + Summary string + ID string + Tags []string + Tokens int + AccessCount int + Importance int + Pinned bool +} + +type CompactResult struct { + NeedsSummary []*Item + TokensFreed int + TokensBefore int + TokensAfter int + Evicted int + Summarized int + Deduplicated int +} + +type Store struct { + items map[string]*Item + tokenBudget int + usedTokens int + autoCompactThreshold float64 + mu sync.Mutex + autoCompact bool +} + +func NewStore(tokenBudget int) *Store { + return &Store{ + items: make(map[string]*Item), + tokenBudget: tokenBudget, + autoCompact: true, + autoCompactThreshold: 0.9, + } +} + +func (s *Store) Add(content, summary string, tags []string, importance int) *Item { + s.mu.Lock() + defer s.mu.Unlock() + + if importance < 1 { + importance = 5 + } else if importance > 10 { + importance = 10 + } + + tokens := EstimateTokens(content) + item := &Item{ + ID: newID(), + Content: content, + Summary: summary, + Tags: tags, + Importance: importance, + CreatedAt: time.Now(), + AccessedAt: time.Now(), + Tokens: tokens, + } + + s.items[item.ID] = item + s.usedTokens += tokens + + if s.autoCompact && s.tokenBudget > 0 { + if float64(s.usedTokens)/float64(s.tokenBudget) > s.autoCompactThreshold { + s.compactLocked(0.8) + } + } + + return item +} + +func (s *Store) Get(id string) (*Item, bool) { + s.mu.Lock() + defer s.mu.Unlock() + + item, ok := s.items[id] + if ok { + item.AccessedAt = time.Now() + item.AccessCount++ + } + return item, ok +} + +func (s *Store) Remove(id string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + item, ok := s.items[id] + if !ok { + return false + } + s.usedTokens -= item.Tokens + delete(s.items, id) + return true +} + +func (s *Store) Pin(id string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + item, ok := s.items[id] + if ok { + item.Pinned = true + } + return ok +} + +func (s *Store) UpdateSummary(id, summary string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + item, ok := s.items[id] + if ok { + item.Summary = summary + } + return ok +} + +func (s *Store) Query(query string, tags []string, limit int) []*Item { + s.mu.Lock() + defer s.mu.Unlock() + + if limit <= 0 { + limit = 10 + } + + queryWords := wordSet(query) + var results []*Item + + for _, item := range s.items { + if len(tags) > 0 && !hasAnyTag(item, tags) { + continue + } + item.AccessedAt = time.Now() + item.AccessCount++ + results = append(results, item) + } + + sort.Slice(results, func(i, j int) bool { + return s.queryScore(results[i], queryWords) > s.queryScore(results[j], queryWords) + }) + + if len(results) > limit { + results = results[:limit] + } + return results +} + +func (s *Store) scoreLocked(item *Item) float64 { + if item.Pinned { + return math.MaxFloat64 + } + + age := time.Since(item.AccessedAt).Minutes() + recency := math.Exp(-age / 120.0) // half-life ~2 hours + + importance := float64(item.Importance) / 10.0 + access := math.Log1p(float64(item.AccessCount)) / 5.0 + + var sizePenalty float64 + if s.tokenBudget > 0 { + sizePenalty = float64(item.Tokens) / float64(s.tokenBudget) + } + + return (0.4 * importance) + (0.3 * recency) + (0.2 * access) - (0.1 * sizePenalty) +} + +func (s *Store) queryScore(item *Item, queryWords map[string]struct{}) float64 { + base := s.scoreLocked(item) + if len(queryWords) == 0 { + return base + } + + contentWords := wordSet(item.Content) + if item.Summary != "" { + for w := range wordSet(item.Summary) { + contentWords[w] = struct{}{} + } + } + for _, tag := range item.Tags { + contentWords[strings.ToLower(tag)] = struct{}{} + } + + return base + (0.5 * jaccardSimilarity(queryWords, contentWords)) +} + +func (s *Store) Status() (budget, used, count int, usage float64) { + s.mu.Lock() + defer s.mu.Unlock() + + budget = s.tokenBudget + used = s.usedTokens + count = len(s.items) + if budget > 0 { + usage = float64(used) / float64(budget) + } + return +} + +func (s *Store) BudgetTight() bool { + s.mu.Lock() + defer s.mu.Unlock() + + if s.tokenBudget <= 0 { + return false + } + return float64(s.usedTokens)/float64(s.tokenBudget) > 0.8 +} + +func (s *Store) Configure(budget int, autoCompact *bool, threshold float64) { + s.mu.Lock() + defer s.mu.Unlock() + + if budget > 0 { + s.tokenBudget = budget + } + if autoCompact != nil { + s.autoCompact = *autoCompact + } + if threshold > 0 && threshold <= 1.0 { + s.autoCompactThreshold = threshold + } +} + +func (s *Store) AutoCompact() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.autoCompact +} + +func (s *Store) AutoCompactThreshold() float64 { + s.mu.Lock() + defer s.mu.Unlock() + return s.autoCompactThreshold +} + +func (s *Store) Compact(targetUsage float64) CompactResult { + s.mu.Lock() + defer s.mu.Unlock() + return s.compactLocked(targetUsage) +} + +func (s *Store) compactLocked(targetUsage float64) CompactResult { + result := CompactResult{TokensBefore: s.usedTokens} + + if s.tokenBudget <= 0 { + result.TokensAfter = s.usedTokens + return result + } + + targetTokens := int(float64(s.tokenBudget) * targetUsage) + if s.usedTokens <= targetTokens { + result.TokensAfter = s.usedTokens + return result + } + + // Phase 1: Summary promotion — replace content with summary to save tokens + for _, item := range s.items { + if s.usedTokens <= targetTokens { + break + } + if item.Pinned || item.Summary == "" || item.Content == item.Summary { + continue + } + oldTokens := item.Tokens + item.Content = item.Summary + item.Tokens = EstimateTokens(item.Content) + saved := oldTokens - item.Tokens + if saved > 0 { + s.usedTokens -= saved + result.Summarized++ + result.TokensFreed += saved + } + } + + // Phase 2: Deduplication — merge items with >70% word overlap + ids := make([]string, 0, len(s.items)) + for id := range s.items { + ids = append(ids, id) + } + merged := make(map[string]bool) + + for i := 0; i < len(ids); i++ { + if merged[ids[i]] { + continue + } + a := s.items[ids[i]] + if a == nil || a.Pinned { + continue + } + aWords := wordSet(a.Content) + + for j := i + 1; j < len(ids); j++ { + if merged[ids[j]] { + continue + } + b := s.items[ids[j]] + if b == nil { + continue + } + + if jaccardSimilarity(aWords, wordSet(b.Content)) > 0.7 { + // Keep higher-scoring item, merge tags + if s.scoreLocked(a) >= s.scoreLocked(b) { + a.Tags = mergeTags(a.Tags, b.Tags) + if a.Summary == "" && b.Summary != "" { + a.Summary = b.Summary + } + s.usedTokens -= b.Tokens + result.TokensFreed += b.Tokens + delete(s.items, ids[j]) + merged[ids[j]] = true + } else { + b.Tags = mergeTags(b.Tags, a.Tags) + if b.Summary == "" && a.Summary != "" { + b.Summary = a.Summary + } + s.usedTokens -= a.Tokens + result.TokensFreed += a.Tokens + delete(s.items, ids[i]) + merged[ids[i]] = true + break + } + result.Deduplicated++ + } + } + } + + // Phase 3: Evict lowest-scoring non-pinned items + if s.usedTokens > targetTokens { + evictable := make([]*Item, 0) + for _, item := range s.items { + if !item.Pinned { + evictable = append(evictable, item) + } + } + sort.Slice(evictable, func(i, j int) bool { + return s.scoreLocked(evictable[i]) < s.scoreLocked(evictable[j]) + }) + + for _, item := range evictable { + if s.usedTokens <= targetTokens { + break + } + s.usedTokens -= item.Tokens + result.TokensFreed += item.Tokens + delete(s.items, item.ID) + result.Evicted++ + } + } + + // Collect items that could benefit from LLM summarization + for _, item := range s.items { + if item.Summary == "" && item.Tokens > 100 { + result.NeedsSummary = append(result.NeedsSummary, item) + } + } + + result.TokensAfter = s.usedTokens + return result +} + +// --- helpers --- + +func newID() string { + b := make([]byte, 8) + if _, err := rand.Read(b); err != nil { + panic("crypto/rand: " + err.Error()) + } + return fmt.Sprintf("%x", b) +} + +func wordSet(s string) map[string]struct{} { + words := strings.Fields(strings.ToLower(s)) + set := make(map[string]struct{}, len(words)) + for _, w := range words { + set[w] = struct{}{} + } + return set +} + +func jaccardSimilarity(a, b map[string]struct{}) float64 { + if len(a) == 0 && len(b) == 0 { + return 0 + } + intersection := 0 + for w := range a { + if _, ok := b[w]; ok { + intersection++ + } + } + union := len(a) + len(b) - intersection + if union == 0 { + return 0 + } + return float64(intersection) / float64(union) +} + +func hasAnyTag(item *Item, tags []string) bool { + tagSet := make(map[string]struct{}, len(item.Tags)) + for _, t := range item.Tags { + tagSet[strings.ToLower(t)] = struct{}{} + } + for _, t := range tags { + if _, ok := tagSet[strings.ToLower(t)]; ok { + return true + } + } + return false +} + +func mergeTags(a, b []string) []string { + seen := make(map[string]struct{}) + var result []string + for _, t := range a { + lower := strings.ToLower(t) + if _, ok := seen[lower]; !ok { + seen[lower] = struct{}{} + result = append(result, t) + } + } + for _, t := range b { + lower := strings.ToLower(t) + if _, ok := seen[lower]; !ok { + seen[lower] = struct{}{} + result = append(result, t) + } + } + return result +} diff --git a/store_test.go b/store_test.go new file mode 100644 index 0000000..6ee5ef3 --- /dev/null +++ b/store_test.go @@ -0,0 +1,200 @@ +package main + +import ( + "strings" + "testing" +) + +func TestEstimateTokens(t *testing.T) { + tests := []struct { + input string + want int + }{ + {"", 0}, + {"hi", 1}, + {"hello world", 3}, + {strings.Repeat("a", 100), 25}, + } + for _, tt := range tests { + got := EstimateTokens(tt.input) + if got != tt.want { + t.Errorf("EstimateTokens(%q) = %d, want %d", tt.input, got, tt.want) + } + } +} + +func TestStoreAddAndQuery(t *testing.T) { + s := NewStore(100000) + + item := s.Add("cilium uses netkit on kernel 6.8", "cilium netkit needs 6.8", []string{"cilium", "networking"}, 8) + if item.ID == "" { + t.Fatal("expected non-empty ID") + } + if item.Tokens != EstimateTokens("cilium uses netkit on kernel 6.8") { + t.Errorf("tokens mismatch: got %d", item.Tokens) + } + + // Query by text + results := s.Query("cilium kernel", nil, 10) + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0].ID != item.ID { + t.Error("wrong item returned") + } + + // Query by tag + results = s.Query("", []string{"networking"}, 10) + if len(results) != 1 { + t.Fatalf("expected 1 result by tag, got %d", len(results)) + } + + // Query with non-matching tag + results = s.Query("", []string{"nonexistent"}, 10) + if len(results) != 0 { + t.Fatalf("expected 0 results, got %d", len(results)) + } +} + +func TestStoreRemoveAndPin(t *testing.T) { + s := NewStore(100000) + + item := s.Add("test content", "", nil, 5) + if !s.Pin(item.ID) { + t.Fatal("pin failed") + } + + got, ok := s.Get(item.ID) + if !ok || !got.Pinned { + t.Fatal("item should be pinned") + } + + if !s.Remove(item.ID) { + t.Fatal("remove failed") + } + _, ok = s.Get(item.ID) + if ok { + t.Fatal("item should be gone") + } +} + +func TestCompaction(t *testing.T) { + s := NewStore(200) // small budget: ~200 tokens = ~800 chars + + // Store items that exceed budget + s.Add("alpha bravo charlie delta echo foxtrot golf hotel", "short alpha", []string{"a"}, 3) + s.Add("india juliet kilo lima mike november oscar papa", "", []string{"b"}, 5) + s.Add("quebec romeo sierra tango uniform victor whiskey", "", []string{"c"}, 8) + + _, used, _, _ := s.Status() + if used == 0 { + t.Fatal("expected non-zero usage") + } + + result := s.Compact(0.5) // compact to 50% + if result.TokensAfter > result.TokensBefore { + t.Error("compaction should not increase tokens") + } + if result.TokensFreed == 0 && result.TokensBefore > 100 { + t.Error("expected some tokens freed") + } + + // Item with summary should have been promoted + if result.Summarized == 0 { + t.Log("note: no summary promotions (may depend on budget math)") + } +} + +func TestSummaryPromotion(t *testing.T) { + s := NewStore(100) // very tight budget + + item := s.Add( + "this is a very long content string that takes many tokens to represent in the context window", + "long content, many tokens", + nil, 5, + ) + + result := s.Compact(0.3) // aggressive compaction + if result.Summarized == 0 { + t.Log("note: summary promotion did not trigger") + } + + got, ok := s.Get(item.ID) + if ok && got.Content == "long content, many tokens" { + // Content was promoted to summary + if got.Tokens != EstimateTokens("long content, many tokens") { + t.Errorf("tokens should match promoted summary: got %d", got.Tokens) + } + } +} + +func TestDeduplication(t *testing.T) { + s := NewStore(100) + + s.Add("alpha bravo charlie delta echo foxtrot", "", []string{"a"}, 5) + s.Add("alpha bravo charlie delta echo golf", "", []string{"b"}, 5) + + _, _, countBefore, _ := s.Status() + s.Compact(0.3) + _, _, countAfter, _ := s.Status() + + if countAfter >= countBefore { + t.Log("note: dedup did not reduce count (similarity may be below threshold)") + } +} + +func TestUpdateSummary(t *testing.T) { + s := NewStore(100000) + + item := s.Add("full content here", "", nil, 5) + if item.Summary != "" { + t.Fatal("should have no summary initially") + } + + if !s.UpdateSummary(item.ID, "compact version") { + t.Fatal("update failed") + } + + got, _ := s.Get(item.ID) + if got.Summary != "compact version" { + t.Errorf("summary not updated: got %q", got.Summary) + } +} + +func TestJaccardSimilarity(t *testing.T) { + a := wordSet("hello world foo bar") + b := wordSet("hello world foo baz") + + sim := jaccardSimilarity(a, b) + if sim < 0.5 || sim > 0.7 { + t.Errorf("expected ~0.6 similarity, got %f", sim) + } + + // Identical sets + if jaccardSimilarity(a, a) != 1.0 { + t.Error("identical sets should have similarity 1.0") + } + + // Empty sets + if jaccardSimilarity(wordSet(""), wordSet("")) != 0 { + t.Error("empty sets should have similarity 0") + } +} + +func TestAutoCompact(t *testing.T) { + s := NewStore(50) // very small: 50 tokens = ~200 chars + s.Configure(0, boolPtr(true), 0.5) + + // Add items that should trigger auto-compaction + s.Add("first item with some content padding", "first", nil, 3) + s.Add("second item with more content padding", "second", nil, 3) + s.Add("third item triggering compaction now", "third", nil, 3) + + _, used, _, usage := s.Status() + // Auto-compact should have kept usage reasonable + if usage > 1.0 { + t.Errorf("auto-compact should prevent exceeding budget: used=%d, usage=%.1f%%", used, usage*100) + } +} + +func boolPtr(b bool) *bool { return &b } diff --git a/tokens.go b/tokens.go new file mode 100644 index 0000000..cd7c69a --- /dev/null +++ b/tokens.go @@ -0,0 +1,10 @@ +package main + +// EstimateTokens returns a rough token count for a string. +// Uses ~4 characters per token heuristic. +func EstimateTokens(s string) int { + if len(s) == 0 { + return 0 + } + return (len(s) + 3) / 4 +} diff --git a/tools.go b/tools.go new file mode 100644 index 0000000..c37bf56 --- /dev/null +++ b/tools.go @@ -0,0 +1,267 @@ +package main + +import ( + "context" + "fmt" + "strings" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" +) + +func registerTools(s *server.MCPServer, store *Store) { + s.AddTool(mcp.NewTool("store", + mcp.WithDescription("Store a context item for later retrieval. Offload information from working context. Provide a summary for efficient compaction when budget is tight."), + mcp.WithString("content", mcp.Required(), mcp.Description("The content to store")), + mcp.WithString("summary", mcp.Description("Compact summary of the content (used when budget is tight)")), + mcp.WithString("tags", mcp.Description("Comma-separated tags for categorization and retrieval")), + mcp.WithNumber("importance", mcp.Description("Importance 1-10, default 5. Higher = harder to evict")), + ), handleStore(store)) + + s.AddTool(mcp.NewTool("query", + mcp.WithDescription("Retrieve stored context items by text search and/or tags. Returns most relevant items. Summaries returned when budget is tight."), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithString("query", mcp.Description("Text to search for in stored items")), + mcp.WithString("tags", mcp.Description("Comma-separated tags to filter by")), + mcp.WithNumber("limit", mcp.Description("Max items to return (default 10)")), + ), handleQuery(store)) + + s.AddTool(mcp.NewTool("status", + mcp.WithDescription("Check context budget usage: total budget, tokens used, item count, usage percentage."), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + ), handleStatus(store)) + + s.AddTool(mcp.NewTool("compact", + mcp.WithDescription("Trigger context compaction. Promotes summaries, deduplicates similar items, evicts lowest-scoring items. Returns suggestions for items needing LLM summarization."), + mcp.WithDestructiveHintAnnotation(true), + mcp.WithNumber("target_usage", mcp.Description("Target usage ratio 0.0-1.0 (default 0.7)")), + ), handleCompact(store)) + + s.AddTool(mcp.NewTool("pin", + mcp.WithDescription("Pin a context item to prevent automatic eviction during compaction."), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithString("id", mcp.Required(), mcp.Description("Item ID to pin")), + ), handlePin(store)) + + s.AddTool(mcp.NewTool("forget", + mcp.WithDescription("Remove a context item from storage."), + mcp.WithDestructiveHintAnnotation(true), + mcp.WithString("id", mcp.Required(), mcp.Description("Item ID to remove")), + ), handleForget(store)) + + s.AddTool(mcp.NewTool("configure", + mcp.WithDescription("Configure context budget and auto-compaction settings."), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithNumber("token_budget", mcp.Description("Total token budget for storage")), + mcp.WithBoolean("auto_compact", mcp.Description("Enable/disable auto-compaction")), + mcp.WithNumber("auto_compact_threshold", mcp.Description("Auto-compact triggers at this usage ratio 0.0-1.0")), + ), handleConfigure(store)) + + s.AddTool(mcp.NewTool("update", + mcp.WithDescription("Add or update the summary of an existing item. Use after compaction suggests items that need summarization."), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithString("id", mcp.Required(), mcp.Description("Item ID to update")), + mcp.WithString("summary", mcp.Required(), mcp.Description("New summary for the item")), + ), handleUpdate(store)) +} + +func handleStore(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + content, err := req.RequireString("content") + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + + summary := req.GetString("summary", "") + tagsStr := req.GetString("tags", "") + importance := req.GetInt("importance", 5) + + var tags []string + if tagsStr != "" { + for _, t := range strings.Split(tagsStr, ",") { + if t = strings.TrimSpace(t); t != "" { + tags = append(tags, t) + } + } + } + + item := store.Add(content, summary, tags, importance) + budget, used, count, usage := store.Status() + + return mcp.NewToolResultText(fmt.Sprintf( + "Stored [%s] (%d tokens)\nBudget: %d/%d tokens (%.0f%%), %d items", + item.ID, item.Tokens, used, budget, usage*100, count, + )), nil + } +} + +func handleQuery(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + query := req.GetString("query", "") + tagsStr := req.GetString("tags", "") + limit := req.GetInt("limit", 10) + + var tags []string + if tagsStr != "" { + for _, t := range strings.Split(tagsStr, ",") { + if t = strings.TrimSpace(t); t != "" { + tags = append(tags, t) + } + } + } + + items := store.Query(query, tags, limit) + if len(items) == 0 { + return mcp.NewToolResultText("No matching items found."), nil + } + + tight := store.BudgetTight() + var sb strings.Builder + fmt.Fprintf(&sb, "Found %d items", len(items)) + if tight { + sb.WriteString(" (budget tight, showing summaries where available)") + } + sb.WriteString(":\n\n") + + for _, item := range items { + fmt.Fprintf(&sb, "[%s] importance:%d tokens:%d", item.ID, item.Importance, item.Tokens) + if item.Pinned { + sb.WriteString(" PINNED") + } + sb.WriteString("\n") + if len(item.Tags) > 0 { + fmt.Fprintf(&sb, "Tags: %s\n", strings.Join(item.Tags, ", ")) + } + sb.WriteString("---\n") + if tight && item.Summary != "" { + sb.WriteString(item.Summary) + } else { + sb.WriteString(item.Content) + } + sb.WriteString("\n\n") + } + + return mcp.NewToolResultText(sb.String()), nil + } +} + +func handleStatus(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + budget, used, count, usage := store.Status() + return mcp.NewToolResultText(fmt.Sprintf( + "Budget: %d/%d tokens (%.1f%%)\nItems: %d\nAuto-compact: %v (threshold: %.0f%%)", + used, budget, usage*100, count, + store.AutoCompact(), store.AutoCompactThreshold()*100, + )), nil + } +} + +func handleCompact(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + target := req.GetFloat("target_usage", 0.7) + if target <= 0 || target > 1.0 { + target = 0.7 + } + + result := store.Compact(target) + + var sb strings.Builder + fmt.Fprintf(&sb, "Compaction complete:\n") + fmt.Fprintf(&sb, "- Evicted: %d items\n", result.Evicted) + fmt.Fprintf(&sb, "- Summary promoted: %d items\n", result.Summarized) + fmt.Fprintf(&sb, "- Deduplicated: %d pairs\n", result.Deduplicated) + fmt.Fprintf(&sb, "- Tokens freed: %d\n", result.TokensFreed) + fmt.Fprintf(&sb, "- Budget: %d → %d tokens\n", result.TokensBefore, result.TokensAfter) + + if len(result.NeedsSummary) > 0 { + sb.WriteString("\nItems that would benefit from summarization:\n") + for _, item := range result.NeedsSummary { + preview := item.Content + if len(preview) > 80 { + preview = preview[:80] + "..." + } + fmt.Fprintf(&sb, "- [%s] %d tokens: %q\n", item.ID, item.Tokens, preview) + } + sb.WriteString("\nUse 'update' tool to add summaries to these items.") + } + + return mcp.NewToolResultText(sb.String()), nil + } +} + +func handlePin(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id, err := req.RequireString("id") + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + if store.Pin(id) { + return mcp.NewToolResultText(fmt.Sprintf("Pinned [%s]", id)), nil + } + return mcp.NewToolResultError(fmt.Sprintf("Item [%s] not found", id)), nil + } +} + +func handleForget(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id, err := req.RequireString("id") + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + if store.Remove(id) { + _, used, count, usage := store.Status() + return mcp.NewToolResultText(fmt.Sprintf( + "Removed [%s]. Budget: %d tokens (%.0f%%), %d items", + id, used, usage*100, count, + )), nil + } + return mcp.NewToolResultError(fmt.Sprintf("Item [%s] not found", id)), nil + } +} + +func handleConfigure(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + budget := req.GetInt("token_budget", 0) + threshold := req.GetFloat("auto_compact_threshold", 0) + + // Only set auto_compact if explicitly provided + var autoCompact *bool + if args := req.GetArguments(); args != nil { + if _, ok := args["auto_compact"]; ok { + v := req.GetBool("auto_compact", true) + autoCompact = &v + } + } + + store.Configure(budget, autoCompact, threshold) + + b, u, c, usg := store.Status() + return mcp.NewToolResultText(fmt.Sprintf( + "Configuration updated.\nBudget: %d/%d tokens (%.1f%%), %d items\nAuto-compact: %v (threshold: %.0f%%)", + u, b, usg*100, c, + store.AutoCompact(), store.AutoCompactThreshold()*100, + )), nil + } +} + +func handleUpdate(store *Store) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id, err := req.RequireString("id") + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + summary, err := req.RequireString("summary") + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + if store.UpdateSummary(id, summary) { + return mcp.NewToolResultText(fmt.Sprintf("Updated summary for [%s]", id)), nil + } + return mcp.NewToolResultError(fmt.Sprintf("Item [%s] not found", id)), nil + } +}