AST-Aware Code Search for Coding Agents

Agentic AI Systems advanced Building with AI ~32 min read

Overview

Coding agents need search tools that match the way developers describe code changes. A user may ask for “async functions that swallow errors”, “React effects that update state”, or “database calls inside request handlers”. Text search can find names and tokens, but it cannot reliably express syntax shape, parent-child relationships, or “this construct inside that construct”. AST-aware search gives the agent a structural query language over the parsed codebase.

The source repository packages an ast-grep skill for a coding assistant. Its reusable mechanism is not the packaging format alone; it is the agent workflow: translate a natural-language code-search request into an example, derive an AST rule, test that rule, then scan the repository. This gives the model a deterministic search substrate instead of asking it to infer structure from raw grep output.

This pattern is useful for large-codebase repair, refactoring, security review, and migration tasks. It complements semantic code retrieval and ordinary text search. Semantic retrieval finds likely relevant files; text search finds exact strings; AST search finds syntactic shapes with metavariables, relational constraints, and language-aware parsing.

Core Concepts

Patterns and Metavariables

An AST search pattern is code-shaped. A pattern such as console.log($ARG) matches call expressions with a captured argument, regardless of the exact argument text. Metavariables let a rule name the parts the agent wants to inspect, report, constrain, or rewrite later.

Relational Rules

Many useful code questions are relational: find a call inside a handler, a return statement inside a conditional, or an await inside a function. Rule systems such as ast-grep support relationships like inside and has, so the query can express containment rather than approximating it with line ranges.

Rule Testing Before Repository Scan

An agent should not run a complex structural rule across a large repository until it has tested the rule against a small example. The skill workflow explicitly makes the model create example code, verify the pattern, inspect the AST when needed, and only then scan the target codebase. That converts an uncertain natural-language request into a checked mechanical query.

Key Concept: AST Search Is a Tool-Use Loop

The agent does not merely call a search binary. It performs a loop: interpret intent, write a candidate structural rule, test it against controlled code, refine the rule, run the scan, and explain the matches with file paths and line numbers.

How It Works

Translate Intent into a Structural Rule

Start with the user’s intent and write one minimal positive example. Then turn that example into an ast-grep pattern or YAML rule. For simple call shapes, a pattern search is enough.

# Find direct logging calls in TypeScript.
ast-grep run --lang ts --pattern 'console.log($$$ARGS)' src

For a coding agent, the important step is to keep the rule inspectable. The model should show the rule it is about to run and state what it expects to match. That makes false positives easier for a human or verifier to catch.

Add Relational Constraints

When the target is a construct inside another construct, use a YAML rule. The example below looks for state updates inside React effect callbacks. The exact rule may need language-specific adjustment, but the shape shows how the agent should move beyond text matching.

id: state-update-inside-effect
language: Tsx
rule:
  pattern: $SETTER($VALUE)
  inside:
    pattern: useEffect($CALLBACK, $$$DEPS)
    stopBy: end
constraints:
  SETTER:
    regex: '^set[A-Z].*'
message: state update inside an effect callback

Debug the Tree Shape

If the rule misses expected examples, inspect the parsed tree rather than guessing. AST tools expose node kinds, fields, and concrete syntax details that tell the agent whether it should match a call_expression, member_expression, function_declaration, or language-specific node kind.

Practical Applications

Refactoring Campaigns

Large migrations often need to find all usages of an old API shape, not just an old function name. AST search can distinguish positional arguments from object options, nested calls from top-level calls, and imports from runtime usages. That lets an agent build a migration inventory before editing.

Security and Reliability Review

Security queries often have structural shape: unsafe SQL construction inside handlers, unchecked deserialization, unbounded user input passed to process execution, or missing error handling after awaited calls. AST-aware rules can make these checks repeatable enough for agent-assisted triage.

Codebase Understanding

Before changing a system, a coding agent can ask structural questions: where are route handlers declared, which functions return promises, which components subscribe to global state, and which tests use a fixture. These answers can seed a plan with concrete file references instead of broad guesses.

Implementation Patterns

Use a Three-Stage Search Ladder

Do not use AST search for every question. A practical agent escalates through a ladder: text search for exact names, semantic search for conceptual locations, and AST search when syntax relationships matter. The ladder keeps searches cheap while reserving structural rules for cases where they add precision.

def choose_search_tool(request):
    if request.has_exact_identifier:
        return "ripgrep"
    if request.describes_conceptual_area:
        return "semantic-code-search"
    if request.describes_syntax_or_containment:
        return "ast-grep"
    return "hybrid"

Require Positive and Negative Examples

For non-trivial rules, the agent should draft at least one example that must match and one that must not match. Running the rule against those examples catches overly broad patterns before a repository-wide scan produces noisy results.

# should match
async function load() {
  const value = await fetchValue()
}

# should not match
function loadSync() {
  return readValue()
}

Return Evidence, Not Just Counts

The scan output should include file path, line number, captured metavariables, and a small code excerpt. A coding agent can then group matches, pick representative examples, and decide whether an edit is mechanical or needs human review.

Performance Considerations

Common Pitfalls

Further Reading