pi-tree-sitter — Structural Code Access for AI Agents

Vision

A pi extension that gives the pi agent the ability to analyze source code through Tree-sitter, the incremental, fault-tolerant parser. Instead of reading files blindly with grep/text, the agent can query a project's AST (Abstract Syntax Tree) for precise structural information.

Why

Today the pi agent can:

• Read files (read)
• Search text patterns (grep, rg)
• Explore directories (ls, find)

But it is syntax-blind. It cannot answer questions like:

• "Which classes inherit from BaseController?"
• "Where is the process_payment method called?"
• "What public functions does this module expose?"
• "Is there dead code in this directory?"
• "Show me the dependency graph of src/core"

Tree-sitter solves this by producing accurate ASTs even on unparseable code.

Features

Registered Tools (5)

Tool	Description
tree_sitter_languages	List supported languages and available capture groups
tree_sitter_query	Run an S-expression query on a single file
tree_sitter_analyze	Structural directory analysis: classes, methods, functions, imports, docstrings
tree_sitter_capture	Capture specific nodes using pre-built groups (file or directory)
tree_sitter_graph	Generate structured JSON graphs: import dependencies, call graph, inheritance hierarchy

Documentation Comments

Comments preceding a definition (class, method, function) are captured and associated with the corresponding node. This enables:

• Extracting YARD/RDoc docstrings (Ruby), docstrings (Python), JSDoc (JS/TS)
• Understanding inline API documentation of a module
• Checking whether a public function is documented

Supported Languages (9)

Language	Extensions	Capture Groups
Ruby	.rb, .rake, .gemspec, .ru	classes_and_modules, methods, comments, calls, constants, requires
Python	.py, .pyw, .pyi	classes, functions, comments, imports, calls, decorators
JavaScript	.js, .jsx, .mjs, .cjs	classes, functions, variables, comments, imports, exports, calls
TypeScript	.ts, .tsx, .mts, .cts	classes, functions, interfaces, types, comments, imports, exports
Go	.go	functions, methods, structs, interfaces, imports, comments
Rust	.rs	functions, structs, enums, traits, imports (use), macros, comments
C / C++	.c, .h, .cpp, .hpp, .cc, .hh, .cxx, .hxx	functions, classes (C++), comments, includes
HTML	.html, .htm, .shtml, .xhtml	elements, comments
CSS	.css, .scss, .less	rules, comments

More languages can be added by installing their respective npm grammar packages and adding a config in src/languages.ts. Each new language requires ~20 lines.

Graph Generation (tree_sitter_graph)

Generates structured JSON graphs designed for AI agent consumption (not visualization):

• Import graph: which files import/require which modules
• Call graph: which functions call which (with scope detection)
• Inheritance graph: class hierarchy with superclass relationships

Output format:

{
  "type": "imports",
  "nodes": [{"id": "path/file.rb", "type": "file", "name": "file.rb"}],
  "edges": [{"from": "path/file.rb", "to": "module_name", "type": "imports"}],
  "metadata": {"totalFiles": 314, "totalImports": 271}
}

External Configuration

You can customize analysis behavior by creating a .tree-sitter.json file in the project root:

{
  "excludeDirs": ["tmp", "coverage", ".next"],
  "maxFileSize": 2000000,
  "maxDepth": 5
}

Priority (highest to lowest): explicit tool parameters → .tree-sitter.json → built-in defaults.

Built-in default exclude dirs: node_modules, vendor, dist, build, __pycache__, target, .git, .svn, .hg

Verified Performance on Real Projects

pi-tree-sitter has been tested on 8 real open-source projects across 9 languages:

Language	Project	Files	LOC	Grep	Tree-sitter	Value added
Ruby	Rails AR	314	64K	57ms	23.0s	Call graph (1507 edges) + inheritance
TypeScript	Next.js	543	18K	48ms	8.1s	5× structures (class+interface+type)
Rust	Serde	200	42K	200ms	2.8s	Struct/enum/trait separated
Python	Django DB	123	55K	0.4s	2.5s	Top-level vs methods, docstrings
Go	Chi	78	5.5K	100ms	0.36s	Struct/interface distinct
C	fmtlib	15	17K	21ms	0.16s	Precise functions, zero false positives
JavaScript	Express	6	2.7K	9ms	67ms	Require with extracted paths

Key finding: tree-sitter typically uses 2-3× fewer tokens than the equivalent grep+read workflow, because the structured JSON replaces dozens of individual file reads.

Architecture

~/.pi/agent/extensions/pi-tree-sitter/
├── index.ts                         # Entry point: registers 5 tools
├── src/
│   ├── parser.ts                    # Parser manager (WASM caching, single-parse)
│   ├── languages.ts                 # 9 language configs + capture groups
│   ├── analyzer.ts                  # Directory analysis with external config
│   ├── enrich.ts                    # Doc comment enrichment
│   ├── queryService.ts              # Centralized file query (single-parse)
│   ├── graph.ts                     # Import/call/inheritance graph generation
│   ├── config.ts                    # .tree-sitter.json loader
│   ├── handlers/
│   │   ├── types.ts                 # Handler interfaces + helpers
│   │   ├── classHandler.ts          # Class/module extraction
│   │   ├── methodHandler.ts         # Method/function extraction
│   │   ├── importHandler.ts         # Import/require/include
│   │   ├── rawHandler.ts            # Generic capture (comments, interfaces, etc.)
│   │   └── registry.ts              # Name → handler lookup
│   ├── tools/
│   │   ├── languages.ts             # tree_sitter_languages tool
│   │   ├── query.ts                 # tree_sitter_query tool
│   │   ├── analyze.ts               # tree_sitter_analyze tool
│   │   ├── capture.ts               # tree_sitter_capture tool
│   │   └── graph.ts                 # tree_sitter_graph tool
│   └── __tests__/
│       ├── setup.ts                 # Vitest setup
│       ├── config.test.ts           # Config tests (6)
│       ├── parser.test.ts           # Parser tests (25)
│       ├── languages.test.ts        # Language tests (30)
│       ├── analyzer.test.ts         # Analyzer tests (17)
│       ├── enrich.test.ts           # Enrichment tests (2)
│       ├── queryService.test.ts     # Query service tests (6)
│       ├── graph.test.ts            # Graph tests (6)
│       ├── index.test.ts            # Schema validation (8)
│       └── fixtures/                # Test fixture files for 7 languages
├── package.json                     # 11 tree-sitter grammars
├── tsconfig.json                    # strict: true
└── vitest.config.ts

Total: 100 passing tests, 8 test files, 0 any types.

Design Decision: Extension vs Skill

Extension is the better choice because:

• Tools must be directly callable by the LLM with typed parameters
• State must be maintained (cached parsers, loaded grammars)
• Native integration with pi.registerTool() and TypeBox schemas
• npm dependencies (tree-sitter, grammars) are managed via package.json

Architecture Highlights

• Handler registry replaces the old switch-case: adding a new capture group requires only a query in languages.ts and optionally a handler in src/handlers/
• queryService centralizes the read → detect → parse → query pattern, avoiding duplication across 3 different code paths
• Single-parse: the same AST rootNode is reused for all queries and doc enrichment. parse() is called exactly once per file, no matter how many capture groups run
• Configurable via .tree-sitter.json or tool parameters (exclude dirs, max file size, recursion depth)

Setup

cd ~/.pi/agent/extensions/pi-tree-sitter
npm install

The extension is auto-discovered by pi on startup from ~/.pi/agent/extensions/.

To install grammars for additional languages:

npm install tree-sitter-<language>

Then add the language config in src/languages.ts (~20 lines).

Development

# Run all 100 tests
npm test

# Watch mode
npm run test:watch

# Run a specific test file
npx vitest run parser.test.ts

# TypeScript strict check
npx tsc --noEmit

References

• Tree-sitter
• web-tree-sitter npm
• Pi Extensions docs