mirror of
https://github.com/lukaszraczylo/compaction-mcp.git
synced 2026-06-05 23:14:02 +00:00
lukaszraczylo
0ddd0e4598
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.
268 lines
9.5 KiB
Go
268 lines
9.5 KiB
Go
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
|
|
}
|
|
}
|