diff --git a/.gitignore b/.gitignore index dcd2ddfd..b923c0f9 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,13 @@ cmd/mcp-stdio/mcp-stdio cmd/mcp-http/mcp-http cmd/mcp-http-cli/mcp-http-cli +cmd/mcp-tokens/mcp-tokens + +# Binaries built at the repo root by `go build ./cmd/`. +/mcp-stdio +/mcp-http +/mcp-http-cli +/mcp-tokens # Test binary, built with `go test -c` *.test diff --git a/cmd/mcp-tokens/README.md b/cmd/mcp-tokens/README.md new file mode 100644 index 00000000..e7f0ed53 --- /dev/null +++ b/cmd/mcp-tokens/README.md @@ -0,0 +1,60 @@ +# mcp-tokens + +Reports tiktoken-based token counts for every MCP tool exposed by the codebase +(twprojects + twdesk + twspaces), sorted by cost. Introspects the local Go +source — no HTTP server, no auth. + +Useful for tracking tool-description token budgets across revisions and for +spotting the heaviest tools when trimming. + +## Snapshot the current tree + +```bash +go run ./cmd/mcp-tokens # token table for every tool +go run ./cmd/mcp-tokens -encoding=cl100k_base +go run ./cmd/mcp-tokens -json > tools.json # full export-tools-shaped JSON +``` + +## Diff against a base ref + +`-base=` materialises that ref in a throwaway `git worktree`, runs the +same binary against it, and prints the delta. Your working tree is never +touched, so uncommitted edits in `internal/` are safe — handy for the inner +dev loop while you trim descriptions. + +```bash +go run ./cmd/mcp-tokens -base=main # text summary + per-tool movers +go run ./cmd/mcp-tokens -base=main -format=markdown # GFM table (PR-comment friendly) +go run ./cmd/mcp-tokens -base=main -format=json # structured output for scripts +go run ./cmd/mcp-tokens -base=origin/main # any git ref works +``` + +The diff cleanly handles tools added or removed between the two trees: +missing-on-one-side rows show the absent count as `0`. + +Sample text output: + +``` + base (main) 44,950 + current 33,723 + delta -11,227 (-24.98%) + + per-tool deltas (117): + tool before after delta + -------------------------------------- -------- -------- -------- + twprojects-list_tasks 832 702 -130 + twprojects-update_timelog 608 485 -123 + ... +``` + +## Caveats + +- Counts use OpenAI's `tiktoken` (`o200k_base` by default). Treat them as a + *relative* signal across revisions — Claude's tokenizer differs, but the + percentage delta tracks closely. +- Counts cover the tool name, description, and JSON-marshalled input schema. + They exclude per-message framing the model adds at runtime, so the absolute + total is a lower bound on what the LLM actually sees. +- Diff mode requires `git` and a working `go` toolchain (it shells out to + `go run` inside the temporary worktree). First run of a new base ref pays + one Go-compile cost; subsequent runs reuse Go's build cache. diff --git a/cmd/mcp-tokens/main.go b/cmd/mcp-tokens/main.go new file mode 100644 index 00000000..71faca18 --- /dev/null +++ b/cmd/mcp-tokens/main.go @@ -0,0 +1,423 @@ +// mcp-tokens reports tiktoken-based token counts for every MCP tool exposed +// by the codebase, sorted by cost. It introspects the Go source directly +// (no HTTP server, no auth), so it always reflects the local working tree. +// +// Usage: +// +// go run ./cmd/mcp-tokens # token counts (default) +// go run ./cmd/mcp-tokens -json # full export-tools-shaped JSON +// go run ./cmd/mcp-tokens -base=main # diff vs main, text output +// go run ./cmd/mcp-tokens -base=main -format=markdown +// go run ./cmd/mcp-tokens -encoding=cl100k_base +// +// Diff mode spins up a temporary `git worktree` at the base ref and runs +// the same binary there, so your working tree is never touched — uncommitted +// or staged changes under internal/ are safe. +// +// Token counts use OpenAI's tiktoken; treat the numbers as a *relative* +// signal across revisions, not as absolute Claude figures. +package main + +import ( + "bytes" + "encoding/json" + "flag" + "fmt" + "net/http" + "os" + "os/exec" + "path/filepath" + "sort" + "strings" + + "github.com/localit-io/tiktoken-go" + "github.com/teamwork/mcp/internal/toolsets" + "github.com/teamwork/mcp/internal/twdesk" + "github.com/teamwork/mcp/internal/twprojects" + "github.com/teamwork/mcp/internal/twspaces" +) + +type toolCount struct { + Name string `json:"name"` + Tokens int `json:"tokens"` +} + +type diffRow struct { + Name string `json:"name"` + Before int `json:"before"` + After int `json:"after"` +} + +func (d diffRow) Delta() int { return d.After - d.Before } + +func main() { + asJSON := flag.Bool("json", false, "emit full export-tools JSON instead of token counts") + asCounts := flag.Bool("counts", false, "emit {tool: tokens} map as JSON (used internally by -base)") + encoding := flag.String("encoding", "o200k_base", "tiktoken encoding name (e.g. o200k_base, cl100k_base)") + baseRef := flag.String("base", "", "compare against this git ref (enables diff mode)") + format := flag.String("format", "text", "diff output format: text, markdown, json") + flag.Parse() + + if *format != "text" && *baseRef == "" { + fail("-format=%s requires -base; -format only applies to diff output", *format) + } + + groups := allGroups() + + if *asJSON { + emitJSON(groups) + return + } + + enc, err := tiktoken.GetEncoding(*encoding) + if err != nil { + fail("get encoding %q: %v", *encoding, err) + } + + if *asCounts { + emitCounts(countTools(groups, enc)) + return + } + + if *baseRef != "" { + if err := runDiff(*baseRef, *format, *encoding, enc, groups); err != nil { + fail("%v", err) + } + return + } + + printSnapshot(countTools(groups, enc)) +} + +func emitCounts(rows []toolCount) { + m := make(map[string]int, len(rows)) + for _, r := range rows { + m[r.Name] = r.Tokens + } + if err := json.NewEncoder(os.Stdout).Encode(m); err != nil { + fail("encode counts: %v", err) + } +} + +// allGroups builds every toolset group registered by the servers. The engine +// and HTTP client passed to factories are only dereferenced inside tool +// handlers — schemas are static at registration time, so we can pass +// nil/empty values safely. +func allGroups() []*toolsets.ToolsetGroup { + httpClient := &http.Client{} + return []*toolsets.ToolsetGroup{ + twprojects.DefaultToolsetGroup(false, true, nil), + twdesk.DefaultToolsetGroup(false, httpClient), + twspaces.DefaultToolsetGroup(false, httpClient), + } +} + +func countTools(groups []*toolsets.ToolsetGroup, enc *tiktoken.Tiktoken) []toolCount { + var rows []toolCount + for _, g := range groups { + for _, ts := range g.Toolsets { + for _, tw := range ts.GetAvailableTools() { + t := tw.Tool + payload := map[string]any{ + "name": t.Name, + "description": t.Description, + "input_schema": t.InputSchema, + } + b, err := json.Marshal(payload) + if err != nil { + fail("marshal %s: %v", t.Name, err) + } + rows = append(rows, toolCount{t.Name, len(enc.Encode(string(b), nil, nil))}) + } + } + } + sort.Slice(rows, func(i, j int) bool { + if rows[i].Tokens != rows[j].Tokens { + return rows[i].Tokens > rows[j].Tokens + } + return rows[i].Name < rows[j].Name + }) + return rows +} + +func printSnapshot(rows []toolCount) { + width := 0 + for _, r := range rows { + if len(r.Name) > width { + width = len(r.Name) + } + } + dash := strings.Repeat("-", width) + total := 0 + fmt.Printf("%-*s tokens\n", width, "tool") + fmt.Printf("%s ------\n", dash) + for _, r := range rows { + fmt.Printf("%-*s %6d\n", width, r.Name, r.Tokens) + total += r.Tokens + } + fmt.Printf("%s ------\n", dash) + fmt.Printf("%-*s %6d (%d tools)\n", width, "TOTAL", total, len(rows)) +} + +func emitJSON(groups []*toolsets.ToolsetGroup) { + out := map[string]any{} + for _, g := range groups { + for _, ts := range g.Toolsets { + for _, tw := range ts.GetAvailableTools() { + t := tw.Tool + entry := map[string]any{ + "description": t.Description, + "title": t.Title, + "inputSchema": map[string]any{"jsonSchema": t.InputSchema}, + "type": "dynamic", + } + if t.OutputSchema != nil { + entry["outputSchema"] = map[string]any{"jsonSchema": t.OutputSchema} + } + if t.Annotations != nil { + entry["annotations"] = t.Annotations + } + if t.Meta != nil { + entry["_meta"] = t.Meta + } + out[t.Name] = entry + } + } + } + enc := json.NewEncoder(os.Stdout) + enc.SetIndent("", " ") + if err := enc.Encode(out); err != nil { + fail("encode: %v", err) + } +} + +// runDiff captures token counts for the current tree and the tree at baseRef, +// then emits the diff. The base ref is materialised in a throwaway git worktree +// so the caller's working tree is never modified — uncommitted edits in +// internal/ are safe to leave in place. +func runDiff(baseRef, format, encName string, enc *tiktoken.Tiktoken, groups []*toolsets.ToolsetGroup) error { + if err := exec.Command("git", "rev-parse", "--verify", "--quiet", baseRef+"^{commit}").Run(); err != nil { + return fmt.Errorf("base ref %q not found", baseRef) + } + + current := countTools(groups, enc) + + wt, err := os.MkdirTemp("", "mcp-tokens-base-*") + if err != nil { + return fmt.Errorf("create tmpdir: %w", err) + } + defer func() { + _ = exec.Command("git", "worktree", "remove", "--force", wt).Run() + _ = os.RemoveAll(wt) + }() + + if out, err := exec.Command("git", "worktree", "add", "--detach", wt, baseRef).CombinedOutput(); err != nil { + return fmt.Errorf("git worktree add: %v: %s", err, out) + } + + // Inject our own cmd/mcp-tokens into the base worktree (it may not exist + // there yet) so we can `go run ./cmd/mcp-tokens -json` against base's + // internal/ packages. + cmdDest := filepath.Join(wt, "cmd", "mcp-tokens") + _ = os.RemoveAll(cmdDest) + if err := os.MkdirAll(filepath.Dir(cmdDest), 0o755); err != nil { + return fmt.Errorf("mkdir cmd dir in worktree: %w", err) + } + if out, err := exec.Command("cp", "-R", "cmd/mcp-tokens", cmdDest).CombinedOutput(); err != nil { + return fmt.Errorf("copy cmd/mcp-tokens into worktree: %v: %s", err, out) + } + + // Make sure tiktoken-go is in the base worktree's go.mod (idempotent). + if out, err := runIn(wt, "go", "get", "github.com/localit-io/tiktoken-go").CombinedOutput(); err != nil { + return fmt.Errorf("go get tiktoken-go in base worktree: %v: %s", err, out) + } + + var stdout bytes.Buffer + dump := runIn(wt, "go", "run", "./cmd/mcp-tokens", "-counts", "-encoding", encName) + dump.Stdout = &stdout + dump.Stderr = os.Stderr + if err := dump.Run(); err != nil { + return fmt.Errorf("count tools in base worktree: %w", err) + } + + var baseMap map[string]int + if err := json.Unmarshal(stdout.Bytes(), &baseMap); err != nil { + return fmt.Errorf("parse base counts: %w", err) + } + + base := make([]toolCount, 0, len(baseMap)) + for name, tokens := range baseMap { + base = append(base, toolCount{name, tokens}) + } + + rows := buildDiff(base, current) + switch format { + case "json": + return emitDiffJSON(baseRef, rows) + case "markdown": + emitDiffMarkdown(baseRef, rows) + case "text", "": + emitDiffText(baseRef, rows) + default: + return fmt.Errorf("unknown format %q (want text|markdown|json)", format) + } + return nil +} + +func runIn(dir, name string, args ...string) *exec.Cmd { + c := exec.Command(name, args...) + c.Dir = dir + return c +} + +func buildDiff(before, after []toolCount) []diffRow { + bMap := make(map[string]int, len(before)) + for _, r := range before { + bMap[r.Name] = r.Tokens + } + aMap := make(map[string]int, len(after)) + for _, r := range after { + aMap[r.Name] = r.Tokens + } + seen := make(map[string]bool) + for n := range bMap { + seen[n] = true + } + for n := range aMap { + seen[n] = true + } + rows := make([]diffRow, 0, len(seen)) + for n := range seen { + rows = append(rows, diffRow{n, bMap[n], aMap[n]}) + } + // Most-saved (largest negative delta) first; tiebreak by name. + sort.Slice(rows, func(i, j int) bool { + di, dj := rows[i].Delta(), rows[j].Delta() + if di != dj { + return di < dj + } + return rows[i].Name < rows[j].Name + }) + return rows +} + +func diffTotals(rows []diffRow) (before, after, delta int, pct float64) { + for _, r := range rows { + before += r.Before + after += r.After + } + delta = after - before + if before > 0 { + pct = float64(delta) / float64(before) * 100 + } + return +} + +func emitDiffText(baseRef string, rows []diffRow) { + before, after, delta, pct := diffTotals(rows) + baseLabel := fmt.Sprintf("base (%s)", baseRef) + labelWidth := max(len(baseLabel), len("current"), len("delta")) + fmt.Println() + fmt.Printf(" %-*s %12s\n", labelWidth, baseLabel, commafy(before)) + fmt.Printf(" %-*s %12s\n", labelWidth, "current", commafy(after)) + fmt.Printf(" %-*s %12s (%+.2f%%)\n\n", labelWidth, "delta", signed(delta), pct) + + movers := filterMovers(rows) + if len(movers) == 0 { + fmt.Println(" no per-tool changes") + return + } + + width := 0 + for _, r := range movers { + if len(r.Name) > width { + width = len(r.Name) + } + } + fmt.Printf(" per-tool deltas (%d):\n", len(movers)) + fmt.Printf(" %-*s %8s %8s %8s\n", width, "tool", "before", "after", "delta") + fmt.Printf(" %s %s %s %s\n", + strings.Repeat("-", width), strings.Repeat("-", 8), strings.Repeat("-", 8), strings.Repeat("-", 8)) + for _, r := range movers { + fmt.Printf(" %-*s %8d %8d %+8d\n", width, r.Name, r.Before, r.After, r.Delta()) + } +} + +func emitDiffMarkdown(baseRef string, rows []diffRow) { + before, after, delta, pct := diffTotals(rows) + fmt.Printf("### MCP tool token deltas vs `%s`\n\n", baseRef) + fmt.Println("| metric | tokens |") + fmt.Println("|---|---:|") + fmt.Printf("| base | %s |\n", commafy(before)) + fmt.Printf("| current | %s |\n", commafy(after)) + fmt.Printf("| **delta** | **%s (%+.2f%%)** |\n\n", signed(delta), pct) + + movers := filterMovers(rows) + if len(movers) == 0 { + fmt.Println("_No per-tool changes._") + return + } + fmt.Printf("#### Per-tool deltas (%d)\n\n", len(movers)) + fmt.Println("| tool | before | after | delta |") + fmt.Println("|---|---:|---:|---:|") + for _, r := range movers { + fmt.Printf("| `%s` | %s | %s | **%s** |\n", + r.Name, commafy(r.Before), commafy(r.After), signed(r.Delta())) + } +} + +func emitDiffJSON(baseRef string, rows []diffRow) error { + before, after, delta, pct := diffTotals(rows) + out := map[string]any{ + "base": baseRef, + "before_total": before, + "after_total": after, + "delta": delta, + "delta_pct": pct, + "tools": rows, + } + enc := json.NewEncoder(os.Stdout) + enc.SetIndent("", " ") + return enc.Encode(out) +} + +func filterMovers(rows []diffRow) []diffRow { + out := make([]diffRow, 0, len(rows)) + for _, r := range rows { + if r.Delta() != 0 { + out = append(out, r) + } + } + return out +} + +func commafy(n int) string { + negative := n < 0 + if negative { + n = -n + } + s := fmt.Sprintf("%d", n) + var b strings.Builder + if negative { + b.WriteByte('-') + } + for i, c := range s { + if i > 0 && (len(s)-i)%3 == 0 { + b.WriteByte(',') + } + b.WriteRune(c) + } + return b.String() +} + +func signed(n int) string { + if n > 0 { + return "+" + commafy(n) + } + return commafy(n) +} + +func fail(format string, args ...any) { + fmt.Fprintf(os.Stderr, format+"\n", args...) + os.Exit(1) +} diff --git a/cmd/mcp-tokens/main_test.go b/cmd/mcp-tokens/main_test.go new file mode 100644 index 00000000..68f86539 --- /dev/null +++ b/cmd/mcp-tokens/main_test.go @@ -0,0 +1,79 @@ +package main + +import ( + "strings" + "testing" + + "github.com/localit-io/tiktoken-go" +) + +func TestCountToolsSmoke(t *testing.T) { + enc, err := tiktoken.GetEncoding("o200k_base") + if err != nil { + t.Fatalf("get encoding: %v", err) + } + + rows := countTools(allGroups(), enc) + if len(rows) == 0 { + t.Fatal("expected at least one tool, got 0") + } + + prefixes := map[string]bool{"twprojects-": false, "twdesk-": false, "twspaces-": false} + for _, r := range rows { + if r.Tokens <= 0 { + t.Errorf("tool %q has non-positive token count %d", r.Name, r.Tokens) + } + for p := range prefixes { + if strings.HasPrefix(r.Name, p) { + prefixes[p] = true + } + } + } + for p, seen := range prefixes { + if !seen { + t.Errorf("no tools registered with prefix %q", p) + } + } + + for i := 1; i < len(rows); i++ { + if rows[i-1].Tokens < rows[i].Tokens { + t.Errorf("rows not sorted by tokens desc at %d: %d < %d", i, rows[i-1].Tokens, rows[i].Tokens) + } + } +} + +func TestBuildDiff(t *testing.T) { + before := []toolCount{{"a", 100}, {"b", 50}, {"removed", 200}} + after := []toolCount{{"a", 80}, {"b", 50}, {"added", 30}} + + rows := buildDiff(before, after) + if len(rows) != 4 { + t.Fatalf("want 4 rows (union of names), got %d", len(rows)) + } + + got := map[string]diffRow{} + for _, r := range rows { + got[r.Name] = r + } + if got["a"].Before != 100 || got["a"].After != 80 || got["a"].Delta() != -20 { + t.Errorf("a row wrong: %+v", got["a"]) + } + if got["b"].Delta() != 0 { + t.Errorf("b row should have delta 0: %+v", got["b"]) + } + if got["removed"].Before != 200 || got["removed"].After != 0 { + t.Errorf("removed row should have After=0: %+v", got["removed"]) + } + if got["added"].Before != 0 || got["added"].After != 30 { + t.Errorf("added row should have Before=0: %+v", got["added"]) + } + + // Sort: most-saved first, then by name. "removed" has -200, "a" has -20, + // "b" has 0, "added" has +30. + wantOrder := []string{"removed", "a", "b", "added"} + for i, name := range wantOrder { + if rows[i].Name != name { + t.Errorf("position %d: want %q, got %q", i, name, rows[i].Name) + } + } +} diff --git a/go.mod b/go.mod index 38646968..4a01bc53 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/getsentry/sentry-go/slog v0.46.2 github.com/google/jsonschema-go v0.4.3 github.com/google/uuid v1.6.0 + github.com/localit-io/tiktoken-go v0.2.1 github.com/modelcontextprotocol/go-sdk v1.6.0 github.com/teamwork/desksdkgo v0.0.0-20260420182446-6788c53d670d github.com/teamwork/spacessdkgo v0.0.0-20260422163745-684bf200d31d @@ -40,6 +41,7 @@ require ( github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cihub/seelog v0.0.0-20170130134532-f561c5e57575 // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/ebitengine/purego v0.10.0 // indirect github.com/go-ole/go-ole v1.3.0 // indirect diff --git a/go.sum b/go.sum index c2a9f537..b48b671a 100644 --- a/go.sum +++ b/go.sum @@ -59,6 +59,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8Yc github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38= github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= @@ -109,6 +111,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/linkdata/deadlock v0.5.5 h1:d6O+rzEqasSfamGDA8u7bjtaq7hOX8Ha4Zn36Wxrkvo= github.com/linkdata/deadlock v0.5.5/go.mod h1:tXb28stzAD3trzEEK0UJWC+rZKuobCoPktPYzebb1u0= +github.com/localit-io/tiktoken-go v0.2.1 h1:UQzIa1jaXpsjiWcocjWPYs7hjtE6jxgyt1me1Bkuhlc= +github.com/localit-io/tiktoken-go v0.2.1/go.mod h1:7mPscHTP1Rpgk/vuUQtbUNK3j1iSsKJG7TszsrnRAYk= github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e h1:Q6MvJtQK/iRcRtzAscm/zF23XxJlbECiGPyRicsX+Ak= github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e/go.mod h1:autxFIvghDt3jPTLoqZ9OZ7s9qTGNAWmYCjVFWPX/zg= github.com/minio/simdjson-go v0.4.5 h1:r4IQwjRGmWCQ2VeMc7fGiilu1z5du0gJ/I/FsKwgo5A=