← back to Cli Printing Press
fix(cli): include cobratree tools in MCP token scoring (#552)
74eb87692cb040a7b782684442a7557d0530e509 · 2026-05-03 19:27:05 -0700 · Trevin Chow
Files touched
M internal/pipeline/mcp_size.goM internal/pipeline/mcp_size_test.go
Diff
commit 74eb87692cb040a7b782684442a7557d0530e509
Author: Trevin Chow <trevin@trevinchow.com>
Date: Sun May 3 19:27:05 2026 -0700
fix(cli): include cobratree tools in MCP token scoring (#552)
---
internal/pipeline/mcp_size.go | 270 +++++++++++++++++++++++++++++++++++++
internal/pipeline/mcp_size_test.go | 80 +++++++++++
2 files changed, 350 insertions(+)
diff --git a/internal/pipeline/mcp_size.go b/internal/pipeline/mcp_size.go
index 14db3806..5498b17e 100644
--- a/internal/pipeline/mcp_size.go
+++ b/internal/pipeline/mcp_size.go
@@ -1,11 +1,16 @@
package pipeline
import (
+ "go/ast"
+ "go/parser"
+ "go/token"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
+
+ "github.com/mvanhorn/cli-printing-press/v3/internal/naming"
)
// MCPTokenEstimate reports the approximate token weight of a generated MCP
@@ -153,6 +158,14 @@ func estimateMCPTokens(dir string) MCPTokenEstimate {
})
}
+ if strings.Contains(src, "cobratree.RegisterAll(") {
+ runtimeTools := estimateCobratreeRuntimeTokens(dir)
+ for _, tool := range runtimeTools {
+ totalChars += tool.Chars
+ perTool = append(perTool, tool)
+ }
+ }
+
est := MCPTokenEstimate{
TotalChars: totalChars,
TotalTokens: totalChars / 4,
@@ -174,6 +187,263 @@ func estimateMCPTokens(dir string) MCPTokenEstimate {
return est
}
+// cobratreeFrameworkCommands mirrors the generated cobratree classify
+// template's framework skip set. The scorer cannot import generated code
+// from a printed CLI, so it keeps the same names here for static estimates.
+var cobratreeFrameworkCommands = map[string]bool{
+ "about": true,
+ "agent-context": true,
+ "api": true,
+ "auth": true,
+ "completion": true,
+ "doctor": true,
+ "feedback": true,
+ "help": true,
+ "profile": true,
+ "search": true,
+ "sql": true,
+ "version": true,
+ "which": true,
+}
+
+type cobraCommandLiteral struct {
+ use string
+ short string
+ long string
+ hidden bool
+ annotations map[string]string
+ runnable bool
+}
+
+type cobraSourceCommand struct {
+ literal cobraCommandLiteral
+ hasLiteral bool
+ children []string
+}
+
+type mcpCobraCommandKind int
+
+const (
+ mcpCobraNovel mcpCobraCommandKind = iota
+ mcpCobraEndpoint
+ mcpCobraFramework
+ mcpCobraHidden
+)
+
+func estimateCobratreeRuntimeTokens(dir string) []MCPToolSize {
+ cliDir := filepath.Join(dir, "internal", "cli")
+ files := listGoFiles(cliDir)
+ if len(files) == 0 {
+ return nil
+ }
+
+ commands := map[string]*cobraSourceCommand{}
+ for _, path := range files {
+ file, err := parser.ParseFile(token.NewFileSet(), path, nil, 0)
+ if err != nil {
+ continue
+ }
+ collectCobraSourceCommands(file, commands)
+ }
+ return reachableCobratreeRuntimeTools(commands)
+}
+
+// collectCobraSourceCommands builds a constructor graph for generated Cobra
+// commands. The runtime walker starts from RootCmd/newRootCmd and follows
+// AddCommand edges, so the scorer does the same instead of counting every
+// runnable command literal that happens to exist in internal/cli.
+func collectCobraSourceCommands(file *ast.File, commands map[string]*cobraSourceCommand) {
+ for _, decl := range file.Decls {
+ fn, ok := decl.(*ast.FuncDecl)
+ if !ok || fn.Body == nil || fn.Name == nil {
+ continue
+ }
+ fnName := fn.Name.Name
+ if fnName != "RootCmd" && fnName != "newRootCmd" && (!strings.HasPrefix(fnName, "new") || !strings.HasSuffix(fnName, "Cmd")) {
+ continue
+ }
+ rec := sourceCommandRecord(commands, fnName)
+ ast.Inspect(fn.Body, func(n ast.Node) bool {
+ switch node := n.(type) {
+ case *ast.FuncLit:
+ return false
+ case *ast.CompositeLit:
+ if !rec.hasLiteral && isCobraCommandLiteral(node) {
+ rec.literal = parseCobraCommandLiteral(node)
+ rec.hasLiteral = true
+ }
+ case *ast.CallExpr:
+ if !isAddCommandCall(node) {
+ return true
+ }
+ for _, arg := range node.Args {
+ if child := cobraConstructorCallName(arg); child != "" {
+ rec.children = append(rec.children, child)
+ }
+ }
+ }
+ return true
+ })
+ }
+}
+
+func sourceCommandRecord(commands map[string]*cobraSourceCommand, name string) *cobraSourceCommand {
+ rec := commands[name]
+ if rec == nil {
+ rec = &cobraSourceCommand{}
+ commands[name] = rec
+ }
+ return rec
+}
+
+func reachableCobratreeRuntimeTools(commands map[string]*cobraSourceCommand) []MCPToolSize {
+ var tools []MCPToolSize
+ visited := map[string]bool{}
+ var walk func(string)
+ walk = func(fnName string) {
+ rec := commands[fnName]
+ if rec == nil {
+ return
+ }
+ for _, childName := range rec.children {
+ if visited[childName] {
+ continue
+ }
+ visited[childName] = true
+ child := commands[childName]
+ if child == nil || !child.hasLiteral {
+ continue
+ }
+ kind := cobratreeCommandKind(child.literal)
+ if kind == mcpCobraHidden || kind == mcpCobraFramework {
+ continue
+ }
+ if kind == mcpCobraNovel && child.literal.runnable {
+ if tool, ok := estimateCobratreeCommandTool(child.literal); ok {
+ tools = append(tools, tool)
+ }
+ }
+ walk(childName)
+ }
+ }
+ walk("RootCmd")
+ walk("newRootCmd")
+ return tools
+}
+
+func isAddCommandCall(call *ast.CallExpr) bool {
+ sel, ok := call.Fun.(*ast.SelectorExpr)
+ return ok && sel.Sel != nil && sel.Sel.Name == "AddCommand"
+}
+
+func cobraConstructorCallName(expr ast.Expr) string {
+ call, ok := expr.(*ast.CallExpr)
+ if !ok {
+ return ""
+ }
+ ident, ok := call.Fun.(*ast.Ident)
+ if !ok || !strings.HasPrefix(ident.Name, "new") || !strings.HasSuffix(ident.Name, "Cmd") {
+ return ""
+ }
+ return ident.Name
+}
+
+func parseCobraCommandLiteral(lit *ast.CompositeLit) cobraCommandLiteral {
+ cmd := cobraCommandLiteral{annotations: map[string]string{}}
+ for _, elt := range lit.Elts {
+ kv, ok := elt.(*ast.KeyValueExpr)
+ if !ok {
+ continue
+ }
+ key, ok := kv.Key.(*ast.Ident)
+ if !ok {
+ continue
+ }
+ switch key.Name {
+ case "Use":
+ cmd.use = stringLiteralValue(kv.Value)
+ case "Short":
+ cmd.short = stringLiteralValue(kv.Value)
+ case "Long":
+ cmd.long = stringLiteralValue(kv.Value)
+ case "Hidden":
+ cmd.hidden = identBoolValue(kv.Value)
+ case "Annotations":
+ cmd.annotations = normalizedStringMapLiteral(kv.Value)
+ case "Run", "RunE":
+ cmd.runnable = true
+ }
+ }
+ return cmd
+}
+
+func estimateCobratreeCommandTool(cmd cobraCommandLiteral) (MCPToolSize, bool) {
+ name := mcpCobraUseName(cmd.use)
+ if name == "" || cmd.hidden || !cmd.runnable {
+ return MCPToolSize{}, false
+ }
+ if cobratreeCommandKind(cmd) != mcpCobraNovel {
+ return MCPToolSize{}, false
+ }
+ toolName := naming.SnakeIdentifier(name)
+ if toolName == "" {
+ return MCPToolSize{}, false
+ }
+ description := cmd.long
+ if description == "" {
+ description = cmd.short
+ }
+ if description == "" {
+ description = "Run `" + name + "` through the companion CLI binary."
+ }
+ chars := len(toolName) + len(description)
+ return MCPToolSize{
+ Name: "cobratree:" + toolName,
+ Chars: chars,
+ Tokens: chars / 4,
+ }, true
+}
+
+func cobratreeCommandKind(cmd cobraCommandLiteral) mcpCobraCommandKind {
+ name := mcpCobraUseName(cmd.use)
+ if name == "" || cmd.hidden || annotationIsTrueValue(cmd.annotations["mcp:hidden"]) {
+ return mcpCobraHidden
+ }
+ if strings.TrimSpace(cmd.annotations["pp:endpoint"]) != "" {
+ return mcpCobraEndpoint
+ }
+ if cobratreeFrameworkCommands[name] {
+ return mcpCobraFramework
+ }
+ return mcpCobraNovel
+}
+
+func annotationIsTrueValue(v string) bool {
+ return v == "true" || v == "1" || v == "yes"
+}
+
+func mcpCobraUseName(use string) string {
+ fields := strings.Fields(use)
+ if len(fields) == 0 {
+ return ""
+ }
+ return fields[0]
+}
+
+func identBoolValue(expr ast.Expr) bool {
+ ident, ok := expr.(*ast.Ident)
+ return ok && ident.Name == "true"
+}
+
+func normalizedStringMapLiteral(expr ast.Expr) map[string]string {
+ raw := stringMapLiteral(expr)
+ out := map[string]string{}
+ for key, value := range raw {
+ out[key] = strings.ToLower(strings.TrimSpace(value))
+ }
+ return out
+}
+
// scoreMCPTokenEfficiency scores 0-10 based on the token weight of the
// generated MCP surface. Returns (score, scored) where scored is false
// for CLIs without an MCP surface so the dimension can be excluded from
diff --git a/internal/pipeline/mcp_size_test.go b/internal/pipeline/mcp_size_test.go
index 2a3cb72e..a4d039a8 100644
--- a/internal/pipeline/mcp_size_test.go
+++ b/internal/pipeline/mcp_size_test.go
@@ -30,6 +30,13 @@ func RegisterTools(s *server.MCPServer) {
return dir
}
+func writeMCPCLISource(t *testing.T, dir, name, content string) {
+ t.Helper()
+ cliDir := filepath.Join(dir, "internal", "cli")
+ require.NoError(t, os.MkdirAll(cliDir, 0o755))
+ writeFile(t, filepath.Join(cliDir, name), content)
+}
+
func TestEstimateMCPTokens_NoMCPSurface(t *testing.T) {
dir := t.TempDir()
est := estimateMCPTokens(dir)
@@ -81,6 +88,79 @@ func TestEstimateMCPTokens_MultipleToolsTopHeaviest(t *testing.T) {
assert.Greater(t, est.TopHeaviest[1].Chars, est.TopHeaviest[2].Chars)
}
+func TestEstimateMCPTokens_IncludesCobratreeRuntimeTools(t *testing.T) {
+ dir := writeMCPTools(t, `
+ s.AddTool(mcplib.NewTool("typed_get", mcplib.WithDescription("Typed endpoint.")), nil)
+ cobratree.RegisterAll(s, cli.RootCmd(), cobratree.SiblingCLIPath)
+`)
+ writeMCPCLISource(t, dir, "root.go", `package cli
+
+import "github.com/spf13/cobra"
+
+func RootCmd() *cobra.Command {
+ rootCmd := &cobra.Command{Use: "demo-pp-cli"}
+ rootCmd.AddCommand(newDigestCmd())
+ rootCmd.AddCommand(newTrendsCmd())
+ rootCmd.AddCommand(newEndpointCmd())
+ rootCmd.AddCommand(newAuthCmd())
+ rootCmd.AddCommand(newHiddenCmd())
+ return rootCmd
+}
+
+func newDigestCmd() *cobra.Command {
+ return &cobra.Command{
+ Use: "digest",
+ Short: "Build a compact daily digest.",
+ RunE: func(cmd *cobra.Command, args []string) error { return nil },
+ }
+}
+
+func newTrendsCmd() *cobra.Command {
+ return &cobra.Command{
+ Use: "trends",
+ Long: "Compare trending entities across synced data and explain what changed.",
+ RunE: func(cmd *cobra.Command, args []string) error { return nil },
+ }
+}
+
+func newEndpointCmd() *cobra.Command {
+ return &cobra.Command{
+ Use: "typed-endpoint",
+ Short: "Already covered by a typed MCP tool.",
+ Annotations: map[string]string{"pp:endpoint": "items.get"},
+ RunE: func(cmd *cobra.Command, args []string) error { return nil },
+ }
+}
+
+func newAuthCmd() *cobra.Command {
+ return &cobra.Command{
+ Use: "auth",
+ Short: "Framework command skipped by cobratree.",
+ RunE: func(cmd *cobra.Command, args []string) error { return nil },
+ }
+}
+
+func newHiddenCmd() *cobra.Command {
+ return &cobra.Command{
+ Use: "debug-hidden",
+ Short: "Hidden from MCP.",
+ Annotations: map[string]string{"mcp:hidden": "true"},
+ RunE: func(cmd *cobra.Command, args []string) error { return nil },
+ }
+}
+`)
+
+ est := estimateMCPTokens(dir)
+ require.Equal(t, 3, est.ToolCount, "typed tool plus two cobratree runtime tools should all count")
+ names := make([]string, 0, len(est.PerTool))
+ for _, tool := range est.PerTool {
+ names = append(names, tool.Name)
+ }
+ assert.Contains(t, names, "typed_get")
+ assert.Contains(t, names, "cobratree:digest")
+ assert.Contains(t, names, "cobratree:trends")
+}
+
func TestScoreMCPTokenEfficiency_FullMarksForLeanSurface(t *testing.T) {
dir := writeMCPTools(t, `
s.AddTool(mcplib.NewTool("get", mcplib.WithDescription("get item")), nil)
← e4e66bd3 fix(cli): dry-run narrative examples in strict validation (#
·
back to Cli Printing Press
·
feat(skills): /printing-press-reprint orchestrator (#553) f362dd44 →