mirror of
https://github.com/lukaszraczylo/compaction-mcp.git
synced 2026-06-05 23:14:02 +00:00
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.
This commit is contained in:
@@ -0,0 +1 @@
|
||||
compaction-mcp
|
||||
@@ -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
|
||||
)
|
||||
@@ -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=
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
+200
@@ -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 }
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user