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:
2026-03-07 16:57:12 +00:00
commit 0ddd0e4598
8 changed files with 1050 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
compaction-mcp
+17
View File
@@ -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
)
+39
View File
@@ -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=
+70
View File
@@ -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)
}
}
+446
View File
@@ -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
View File
@@ -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 }
+10
View File
@@ -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
}
+267
View File
@@ -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
}
}