pi-lens
Version:
Real-time code feedback for pi — LSP, linters, formatters, type-checking, structural analysis & booboo
661 lines (501 loc) • 24.2 kB
Markdown
# Agent tooling improvement plan — `module_report`, `read_symbol`, and `ast_grep_search`
This plan captures practical friction observed while working on #338, where the
bug lived in a detached timer callback and an inline `turn_end` dependency
wrapper rather than in a cleanly named exported function.
## Problem statement
The current tools work well for conventional module structure:
1. `module_report` gives a cheap outline of named symbols.
2. `read_symbol` reads one named symbol body.
3. `ast_grep_search` finds precise structural patterns when the caller knows the
right AST shape.
The weak spot is callback-heavy TypeScript, especially host-adapter files such as
`index.ts`, where important logic is often embedded in:
- event handlers (`pi.on("turn_end", async (...) => { ... })`)
- timer callbacks (`setTimeout(() => { ... })`)
- inline dependency wrappers (`resetLSPService: () => { ... }`)
- fire-and-forget callbacks (`void saveSessionState(...)`)
- anonymous promise callbacks
In #338, `module_report` quickly found `scheduleLSPIdleReset`, but the stale
context capture was in an anonymous wrapper passed from `index.ts` into
`handleTurnEnd`. That required falling back to raw line reads and text search.
## Goals
- Make callback-heavy lifecycle code navigable without reading entire files.
- Let agents jump from a search hit to the smallest useful enclosing code body.
- Let `read_symbol` read important anonymous callback bodies via stable handles.
- Surface risk patterns that frequently cause extension lifecycle bugs.
- Make `ast_grep_search` easier to use for common agent tasks without requiring
every agent to invent AST patterns from scratch.
- Keep search → outline → read workflows token-efficient by returning ready-to-use
read handles wherever possible.
## Non-goals
- Replace `read` entirely. Raw line reads remain the right tool for small files
and arbitrary slices.
- Make `module_report` count as read-guard coverage. It is still shape, not body.
- Build a full control-flow analyzer. The first pass should be structural and
cheap.
## 1. Add `read_enclosing(path, line)` as the search-to-read bridge
The highest-value missing primitive is not another whole-file outline; it is:
> Given a file and line from search output, read the smallest useful semantic
> unit that contains that line.
This is the bridge between `grep` / `ast_grep_search` / diagnostics and a safe,
token-efficient edit.
Proposed tool:
```json
{
"path": "index.ts",
"line": 2147,
"kinds": ["function", "method", "callback", "class", "object_property"]
}
```
Output:
```json
{
"kind": "object_property_callback",
"name": "handleTurnEnd.deps.resetLSPService@2147",
"path": "index.ts",
"startLine": 2146,
"endLine": 2150,
"parentChain": ["pi.on(\"turn_end\")", "handleTurnEnd deps"],
"read": { "path": "index.ts", "offset": 2146, "limit": 5 },
"body": "resetLSPService: () => { ... }"
}
```
Requirements:
- Use tree-sitter ranges, not text heuristics.
- Prefer the smallest useful enclosing node, but allow `kinds` / `maxLines` to
widen to parent functions/classes when needed.
- Record read-guard coverage for the returned range, just like `read_symbol`.
- Return the parent chain so the agent understands local context without reading
the whole file.
- Work for named symbols and anonymous callbacks.
This should be the main target for search results: every search hit should be
able to say "read the enclosing body for this line".
## 2. Add read handles to search-like results
Any tool that reveals a source location should also return ready-to-use read
handles:
```json
{
"file": "index.ts",
"line": 2147,
"context": "resetLSPService: () => { ... }",
"readSlice": { "path": "index.ts", "offset": 2138, "limit": 25 },
"readEnclosing": { "path": "index.ts", "line": 2147 }
}
```
Apply this to:
- `ast_grep_search`
- future `ast_grep_outline` entries
- `lens_diagnostics` findings where line numbers are known
- `lsp_navigation` references/definitions
- module-report callback/risk entries
This keeps the agent from manually deriving offsets and reduces accidental
whole-file reads.
## 3. Teach `module_report` about important anonymous callbacks
Add a `callbacks` or `closures` section to the report. Each entry should look
similar to a symbol entry and include ready-to-use read args.
Example:
```json
{
"name": "pi.on(\"turn_end\")",
"kind": "event_handler",
"startLine": 2108,
"endLine": 2196,
"signature": "async (_event, ctx) => ...",
"flags": ["captures ctx", "async", "lifecycle"],
"read": { "path": "index.ts", "offset": 2108, "limit": 89 }
}
```
Candidate callback kinds:
- `pi.on(<event>, <callback>)`
- `setTimeout(<callback>, ...)`
- `setInterval(<callback>, ...)`
- `setImmediate(<callback>)`
- object-literal function properties passed as dependency objects
- `.then(...)`, `.catch(...)`, `.finally(...)`
- fire-and-forget async expressions (`void somePromise(...)`), when the argument
includes an inline callback or captures lifecycle state
### Implementation notes
- Use tree-sitter extraction, not LSP.
- Keep this bounded to module-level or top-level-function descendants. Avoid
reporting every tiny array callback unless it has risk flags.
- Give each callback a stable synthetic name derived from its enclosing construct
and line number if necessary:
- `pi.on("turn_end")`
- `scheduleLSPIdleReset:setTimeout@78`
- `handleTurnEnd.deps.resetLSPService@2179`
## 4. Add task-focused ranking to `module_report`
A generic outline is useful; a task-aware outline is better. Agents often arrive
with a query like "stale ctx after session replacement" or "timer reset crash".
`module_report` should accept an optional `focus` string and use it to rank
symbols, callbacks, imports, and risk entries.
Example:
```json
{
"path": "index.ts",
"focus": "stale ctx session replacement timer resetLSPService"
}
```
Ranking signals:
- lexical matches in names, signatures, first source line, comments, and import
paths;
- risk-flag matches (`ctx`, `timer`, `async`, `shutdown`, `session`);
- existing importance signals (exported, fanout, complexity, usedBy count);
- proximity to search hits when `module_report` is called after a search.
The output should keep the full outline available but move likely-relevant
entries to `recommendedReads` / `recommendedCallbacks` so agents do not scan a
large JSON response manually.
## 5. Let `read_symbol` read pseudo-symbol handles
Once `module_report` emits callback entries, `read_symbol` should accept their
synthetic names.
Examples:
```json
{ "path": "index.ts", "symbol": "pi.on(\"turn_end\")" }
{ "path": "clients/runtime-turn.ts", "symbol": "scheduleLSPIdleReset:setTimeout@78" }
```
This keeps the existing workflow:
1. `module_report` for shape.
2. Pick a named symbol or pseudo-symbol.
3. `read_symbol` for exact body and read-guard coverage.
### Read-guard behavior
A pseudo-symbol body is still real source text with a concrete line range, so it
should record read coverage just like a named symbol.
## 6. Add closure-capture risk flags
`module_report` should flag callback entries that capture lifecycle-sensitive
objects and can outlive the current event.
High-signal captures:
- `ctx`
- `ctx.ui`
- `ctx.cwd`
- `ctx.signal`
- `pi`
- command/event contexts
- mutable session/runtime identity values
High-signal async boundaries:
- `setTimeout` / `setInterval`
- `setImmediate`
- promise callbacks
- fire-and-forget `void` calls
- callbacks stored in module-level variables or passed into long-lived services
Example flags:
```text
captures ctx
captures ctx.ui
detached timer
may outlive session
uses sessionGeneration guard
no error boundary
```
For #338, the ideal report would have surfaced something like:
```text
⚠ handleTurnEnd.deps.resetLSPService@2179 captures ctx.ui and is passed into a
detached 240s timer via scheduleLSPIdleReset.
```
## 7. Add `ast_grep_search` recipes for common agent tasks
Raw AST patterns are powerful but too easy to get wrong. Add documented recipes
or a thin preset layer for common searches.
Useful presets:
- find event handlers
- find timers
- find callbacks capturing `ctx`
- find callbacks capturing `ctx.ui`
- find empty catch blocks
- find object-literal function properties by name
- find fire-and-forget async calls
- find `resetLSPService` wrappers/callers
Possible interface options:
1. Documentation-only examples in the ast-grep skill.
2. A `preset` parameter on the tool, translated to an AST rule internally.
3. A companion helper command that prints the generated ast-grep rule before
running it.
The safest first step is documentation-only recipes, then graduate the most
used ones into presets.
## 8. Improve local blast-radius reporting for callbacks
`module_report` already has file-level blast radius. For lifecycle bugs, the
more useful relationship is local retention:
```text
callback A is passed to function B
function B stores/schedules it
callback A captures ctx.ui
```
Add a lightweight `retention` or `callbackFlow` hint when a callback is passed
into known scheduler/retainer functions.
Example:
```json
{
"from": "handleTurnEnd.deps.resetLSPService@2179",
"to": "scheduleLSPIdleReset:setTimeout@78",
"reason": "callback passed as resetFn and invoked by detached timer",
"risk": "captures ctx.ui across session replacement"
}
```
This does not need whole-program dataflow initially. A first version can detect
same-file obvious cases and known dependency-object names.
## 9. Codebase convention: prefer named lifecycle handlers
Tooling can help, but the codebase should also make important lifecycle logic
more visible.
Prefer extracting large anonymous handlers into named functions:
```ts
async function handlePiTurnEnd(...) { ... }
function createLspIdleResetDeps(...) { ... }
```
Benefits:
- `module_report` becomes more informative.
- `read_symbol` can target the real body.
- Tests can import smaller units.
- Closure capture becomes explicit in function parameters.
This should be opportunistic, not a giant refactor.
## 10. Generalize beyond TypeScript
The #338 example is TypeScript, but the same navigation problem exists across
both language surfaces pi-lens already depends on:
- tree-sitter WASM grammars used by pi-lens for structural extraction
(`typescript`, `tsx`, `javascript`, `python`, `rust`, `go`, `java`, `kotlin`,
`dart`, `c`, `cpp`, `elixir`, `ruby`, `bash`, `csharp`, `css`, `html`, `json`,
`lua`, `ocaml`, `php`, `swift`, `toml`, `vue`, `yaml`, `zig`)
- ast-grep languages exposed by the tool (`bash`, `c`, `cpp`, `csharp`, `css`,
`elixir`, `go`, `haskell`, `html`, `java`, `javascript`, `json`, `kotlin`,
`lua`, `nix`, `php`, `python`, `ruby`, `rust`, `scala`, `solidity`, `swift`,
`tsx`, `typescript`, `yaml`)
Do not make the callback work TS-only. Build a language-neutral model, then add
per-language extractors incrementally.
### Normalize to semantic roles
Different languages spell the same agent-relevant concepts differently. The
report should normalize them into roles rather than expose only raw AST node
names.
Suggested role vocabulary:
| Role | Examples |
|---|---|
| `function` / `method` | TS function, Python def, Go func, Java method, Rust fn |
| `type` / `class` / `struct` | class, interface, struct, enum, trait, protocol |
| `event_handler` | `pi.on`, DOM listeners, Flask/FastAPI routes, Rails routes/controllers, Java listeners |
| `callback` / `closure` | JS arrow/function expressions, Ruby blocks, Rust closures, Swift trailing closures |
| `task` / `spawn` | `setTimeout`, `asyncio.create_task`, `go f()`, `tokio::spawn`, `Task.Run`, `DispatchQueue.async`, Kotlin `launch` |
| `deferred_cleanup` | Go `defer`, JS `finally`, Python context managers, C# `using`, Java try-with-resources |
| `unsafe_boundary` | Rust `unsafe`, C/C++ raw pointer/free, shell command execution, dynamic eval |
`module_report` should expose the normalized role plus the raw node kind for
debuggability.
### Language-specific high-value patterns
Start with patterns that map to real bugs agents are likely to create or miss:
- **JavaScript/TypeScript/TSX/Vue** — timers, promise callbacks, event handlers,
React hooks/effects, captured stale props/state, detached async `void` calls,
lifecycle `ctx` capture.
- **Python** — decorators/routes (`.get`, `.fixture`), context
managers, `asyncio.create_task`, `threading.Timer`, callbacks passed to
executors/signals, broad/empty `except` blocks.
- **Go** — goroutines, loop-variable capture, `defer` in loops, context
propagation/cancellation, `WaitGroup` closure bodies, HTTP handlers.
- **Rust** — closures passed to iterators and `thread::spawn`/`tokio::spawn`,
captured non-`'static` state, `unsafe` blocks, `Drop`/RAII cleanup.
- **Java/Kotlin/C#** — lambdas, executor/task/coroutine launches, UI/event
listeners, annotations/attributes, try-with-resources/`using` scopes.
- **Ruby/PHP/Elixir** — blocks/anonymous functions, route/controller macros,
supervision/task boundaries, rescue/catch blocks.
- **C/C++/Swift** — callback pointers/lambdas, dispatch queues, RAII/destructor
boundaries, manual memory/resource cleanup, thread entry points.
- **Shell/YAML/JSON/TOML/CSS/HTML** — fewer symbol bodies, but still useful for
structural regions: jobs/steps, scripts, selectors, embedded handlers, and
config blocks with security-sensitive keys.
### Capability matrix
Add a small capability matrix so callers know what a report can honestly provide
for each language:
```text
language symbols imports callbacks captures retention ast_grep_presets
TS yes yes yes yes partial yes
Go yes yes goroutine loop vars partial yes
Python yes yes decorators yes partial yes
YAML blocks n/a jobs n/a n/a yes
```
This avoids implying that every language has the same fidelity on day one.
### Use ast-grep as a fallback/enrichment layer
For languages where tree-sitter symbol extraction is thin or missing, ast-grep
can still provide targeted structural enrichment if the language is supported by
the tool. Examples:
- find `go $CALL` goroutine launches in Go
- find `asyncio.create_task($CALL)` in Python
- find `tokio::spawn($CALL)` in Rust
- find Java/C#/Kotlin executor or coroutine launches
- find YAML CI jobs/steps with shell snippets
This should be best-effort and clearly labeled by source (`tree-sitter`,
`ast-grep`, or both). It should not block `module_report` when ast-grep is
unavailable.
### Embedded language regions
Some important callbacks live inside embedded languages:
- JS/TS inside TSX/Vue/Svelte-like templates
- shell snippets inside YAML CI configs
- SQL strings inside host languages
- HTML event attributes or script tags
Do not attempt full multi-language parsing in phase one, but design the data
model so an entry can say:
```json
{ "language": "yaml", "embeddedLanguage": "bash", "kind": "ci_step_script" }
```
That leaves room for later extraction without changing the schema.
## 11. Expose ast-grep's native `outline` as a low-level tool
`ast-grep 0.44.0` ships `ast-grep outline`, which returns syntax-aware
structure for symbols, imports, exports, and members. It supports JSON output,
file or directory input, item filters (`structure`, `exports`, `imports`,
`all`), presentation levels (`names`, `signatures`, `digest`, `expanded`),
regex matching, public-member filtering, and custom outline rules.
Expose this as an agent tool named `ast_grep_outline`.
Positioning:
- `module_report` remains the preferred pi-lens-aware navigator: it can include
cached imports, used-by refs, complexity/fanout, recommended reads, and blast
radius.
- `ast_grep_outline` is the raw ast-grep view: fast, local, no index, no LSP,
no cross-file semantics, useful for languages or structures where pi-lens's
extractor is weak.
- `module_report` may later use `ast_grep_outline` internally as enrichment, but
the user-facing tools should stay distinct so agents can choose precision vs.
pi-lens context.
Initial schema:
```json
{
"paths": ["src/foo.ts"],
"lang": "typescript",
"items": "structure|exports|imports|all|auto",
"view": "names|signatures|digest|expanded|auto",
"type": ["class", "function"],
"match": "Parser",
"pubMembers": true,
"globs": ["*.ts", "!dist/**"],
"noIgnore": ["hidden"]
}
```
Implementation constraints:
- Use `execFile` / `safeSpawnAsync`, never shell interpolation.
- Constrain paths to the workspace/project root.
- Prefer `--json=compact` for agent consumption.
- Default to `items:auto`, `view:auto`; let agents ask for `expanded` when they
need members.
- Clearly label the result as syntax-only and non-semantic.
Caveat: `ast-grep outline` does not currently solve the #338 anonymous callback
case by itself. On `index.ts`, it reports top-level functions and types well, but
not `pi.on("turn_end")` as a pseudo-symbol. We still need the callback-specific
work above.
## 12. Rename/alias `ast_dump` to `ast_grep_dump`
The current `ast_dump` tool is valuable but under-named: it dumps the
ast-grep/tree-sitter parse tree for a snippet so an agent can discover exact node
kinds and nesting before writing a pattern.
Add `ast_grep_dump` as the preferred name and keep `ast_dump` as a compatibility
alias for at least one release cycle.
Why underscore, not hyphen:
- Existing pi-lens tool names use underscores (`ast_grep_search`,
`ast_grep_replace`, `module_report`).
- Hyphenated names are awkward in schemas and prompt snippets.
Guidance:
- Use `ast_grep_dump` when `ast_grep_search` returns zero matches and the agent
needs to inspect AST node kinds.
- Keep `includeAnonymous=false` by default; anonymous CST nodes are noisy.
- Add `includeAnonymous=true` only when punctuation/operators/field boundaries
matter.
### Should failed `ast_grep_search` automatically run `ast_grep_dump`?
Not by default.
Reasons:
- A failed search over many files does not identify one obvious source snippet to
dump.
- Dumping full-file ASTs can be enormous and token-hostile.
- The useful dump is usually a small representative snippet near the code shape
the agent expects, not the whole candidate file.
- Auto-dump can distract from simpler causes: wrong language, too-strict pattern,
path/glob miss, or project ignore rules.
Better behavior:
1. On zero matches, `ast_grep_search` should return a structured next-step hint:
```text
No matches. Try one simpler pattern, or run ast_grep_dump on a small snippet
from a representative file to inspect node kinds.
```
2. Add an opt-in debug parameter, e.g. `debugOnNoMatch`, with conservative
behavior:
- dump the query AST (`ast-grep run --debug-query`) when possible;
- optionally dump a bounded source snippet only if the caller supplied exactly
one file path and a small `sampleRange` / `sourceSnippet`;
- cap output aggressively.
3. Consider a helper result field:
```json
{
"suggestedDump": {
"tool": "ast_grep_dump",
"lang": "typescript",
"source": "function example() { ... }"
}
}
```
So: failed search should **suggest** or **optionally** produce a dump, not fire an
unbounded dump automatically.
## Related open issues triaged into this plan
Direct fits:
- **#158 — LSP-enriched enclosing symbol names.** Directly supports
`read_enclosing`: use warm `documentSymbol` to name anonymous/structural ranges
more accurately (`arrow_function` → assignment/property/method name).
- **#162 — agent-expose codebase mental model + hybrid ranking.** Feeds
task-focused `module_report` ranking and candidate-file discovery; combine
graph centrality, change frequency, and lexical focus matches.
- **#303 — live word-index search as a pi-lens tool.** First-stage candidate
discovery for the search → outline → read workflow. This should be a `feature`
rather than an `enhancement` because it exposes a new agent tool.
- **#307 — `documentSymbol` fallback when tree-sitter extraction is empty.**
Important for cross-language `module_report` / `read_enclosing` coverage when
tree-sitter grammars or queries are degraded.
- **#325 — middle-man / delegate-only classes.** Belongs in the structural
`module_report` / class-metric layer, not as a naive ast-grep rule; it needs a
whole-class delegation ratio and facade/adapter guards.
- **#236 — LSP-confirmed review-graph edges.** Future enrichment for blast
radius, retention hints, and task ranking; use provenance-stamped LSP edges
where AST-name heuristics are weak.
- **#176 — diagnostics provenance + actionable follow-up chaining.** Aligns with
this plan's `readSlice` / `readEnclosing` handles and fully-parameterized
next-step tool calls on search/diagnostic results.
Reliability prerequisites / adjacent:
- **#255 — lua tree-sitter ERROR after multiple grammars.** Justifies the
capability matrix and fallback/enrichment strategy; some languages need honest
degradation and LSP/ast-grep fallback paths.
- **#282 — ast-grep napi fallback drops utils/TSX/Python rules.** Relevant to
`ast_grep_search`/ast-grep-backed enrichment reliability in spawn-restricted
environments.
- **#177 — tree-sitter WASM provenance sidecars.** Adjacent reliability work for
any plan that leans harder on grammar-specific extraction.
- **#332 — incomplete string escaping ast-grep rule.** Rule-specific, not a
navigation primitive, but useful as a concrete consumer of better
`ast_grep_dump`/recipe ergonomics while authoring and debugging rules.
Not folded into this plan:
- **#268 / #181 / #175 / #179** are primarily diagnostic surfacing,
disposition, and adapter-normalization work. They should adopt `readSlice` /
`readEnclosing` handles once those primitives exist, but they are not blockers
for the navigation-tooling work itself.
## Suggested phases
### Phase 1 — Search-to-read handles
- Add `read_enclosing(path, line)` for named symbols and anonymous callbacks.
- Add `readSlice` / `readEnclosing` handles to `ast_grep_search` results.
- Add the same handles opportunistically to diagnostics and LSP location results.
### Phase 2 — Documentation and recipes
- Add ast-grep recipes for lifecycle/callback searches.
- Add guidance to prefer named lifecycle handlers.
- Add `ast_grep_dump` as the preferred alias for `ast_dump`.
### Phase 3 — Callback extraction in `module_report`
- Emit callback entries for event handlers and timers.
- Include `read` args and basic flags (`async`, `captures ctx`, `detached timer`).
- Keep entries capped and ranked to avoid noise.
- Add a `focus` query string to rank symbols/callbacks by task relevance.
### Phase 4 — Pseudo-symbol support in `read_symbol`
- Allow `read_symbol` to resolve callback handles emitted by `module_report`.
- Record read-guard coverage for callback ranges.
### Phase 5 — Risk and retention hints
- Add lifecycle-sensitive capture flags.
- Add known scheduler/retainer flow hints.
- Add tests with fixtures modeled after #338.
### Phase 6 — Raw ast-grep outline
- Expose `ast_grep_outline` as a syntax-only fallback/enrichment tool.
- Consider using it internally to supplement `module_report` for weak languages.
## Success criteria
Given a file like `index.ts`, an agent should be able to:
1. Search for `resetLSPService` / stale context terms and get `readEnclosing`
handles back with each hit.
2. Call `read_enclosing({ path: "index.ts", line: 2147 })` and read the inline
callback body without manually computing offsets.
3. Run `module_report` with a focus string like `"stale ctx session timer"`.
4. See `pi.on("turn_end")` and its inline `resetLSPService` wrapper as navigable
entries.
5. See a warning that the wrapper captures `ctx.ui` and may be invoked by a
detached timer.
6. Read that callback body with `read_symbol` / pseudo-symbol support if coming
from the outline path.
7. Fix the bug without reading hundreds of unrelated lines.