DEV Community

yureki_lab
yureki_lab

Posted on

How I Stopped My AI Agent From Grepping Blindly: Building a Code Graph

TL;DR

My AI coding agent used to re-discover my codebase's structure from scratch every single session — grepping around, opening files "just to check," and occasionally missing a call site during a refactor. I built a small code graph (functions and classes as nodes, calls and imports as edges) and gave the agent a handful of targeted queries instead of raw grep. Here's what worked, what broke, and what I'd do differently.

The Problem

I run Claude Code against a mid-sized codebase (a few hundred files, multiple services) most days. For a long time, every session started the same way: the agent would grep for a symbol, open three or four files "to be safe," grep again with a slightly different pattern, and eventually build a mental model of how things connect — only to throw that model away the moment the session ended.

That rediscovery tax showed up in two ways:

  • Wasted tokens and time. A task like "rename this function and update all call sites" turned into five or six rounds of grep-and-read before the agent even started editing.
  • Missed call sites. Grep is a text search, not a structural one. It misses calls through aliases, re-exports, and indirection. Twice, a "safe" rename shipped with a broken caller that grep simply hadn't turned up, and I only caught it in CI.

The second one is what actually pushed me to fix this. A tool that can't reliably tell you "who calls this function" isn't safe to hand refactors to unsupervised.

The specific incident that made this non-negotiable: I asked the agent to rename a helper function used across a handful of modules. It grepped, found four call sites, updated all four, ran the local tests for the module it was editing, and reported success. CI failed twenty minutes later — a fifth call site lived behind a re-export in a module the grep pattern hadn't matched, because the import used a different local alias. Nothing in the agent's process was wrong given the tools it had. Grep just isn't the right tool for "who calls this," and no amount of cleverer regex fixes that category of miss.

I also want to be upfront that this isn't a novel idea — IDEs have had "find references" for two decades, and language servers do this for humans already. The interesting part wasn't the concept, it was making it cheap and reliable enough that an agent reaches for it by default instead of falling back to grep out of habit.

How I Solved It

The fix wasn't a fancy static-analysis platform — it was a small, boring index that answers a handful of specific questions well.

Step 1: Parse source into a structural index, not just text.

I used tree-sitter to parse each file into an AST and pull out the pieces that matter for navigation: function and class definitions, their locations, and the call/import expressions inside them. Simplified, the extraction looks like this:

from tree_sitter import Parser
from tree_sitter_languages import get_language

def extract_symbols(file_path: str, parser: Parser):
    source = open(file_path, "rb").read()
    tree = parser.parse(source)
    symbols = []

    def walk(node):
        if node.type in ("function_definition", "class_definition"):
            name_node = node.child_by_field_name("name")
            symbols.append({
                "name": name_node.text.decode(),
                "kind": node.type,
                "file": file_path,
                "start_line": node.start_point[0] + 1,
                "end_line": node.end_point[0] + 1,
            })
        for child in node.children:
            walk(child)

    walk(tree.root_node)
    return symbols
Enter fullscreen mode Exit fullscreen mode

Calls and imports get extracted the same way, keyed by which enclosing function they live in. That's the whole graph: nodes are symbols, edges are "calls" and "imports."

Step 2: Store it somewhere queryable, and make it incremental.

The first version rebuilt the entire graph on every run. That was fine at a few hundred files and unbearable once I pointed it at a bigger monorepo — 40+ seconds before the agent could do anything. The fix was to hash each file's contents, store the hash alongside its extracted symbols, and only re-parse files whose hash changed:

import hashlib

def needs_reindex(file_path: str, stored_hashes: dict) -> bool:
    content = open(file_path, "rb").read()
    current_hash = hashlib.sha256(content).hexdigest()
    return stored_hashes.get(file_path) != current_hash
Enter fullscreen mode Exit fullscreen mode

That took a full rebuild down to under a second for a typical single-file edit, since only the changed file (and its direct neighbors) need re-parsing.

One wrinkle: hashing the file is not enough on its own. If file A changes, any edge that points into file A from elsewhere is still valid, but any edge inside file A needs to be dropped and rebuilt. I ended up storing edges keyed by their source file, so invalidation is just "delete all edges where source_file = A, re-extract A, re-insert." Cheap, and it avoids accidentally leaving stale edges around after a file shrinks.

Step 3: Expose queries, not a query language.

My first instinct was to give the agent a general graph query interface — think "write your own Cypher-ish query." That was a mistake I'll get to in a second. What actually worked was a small, fixed set of high-level verbs:

  • find_definition(symbol_name) — where is this defined?
  • find_callers(symbol_name) — who calls this, across the whole repo?
  • trace_call_chain(symbol_name, depth) — walk N levels of callers/callees
  • get_architecture_overview(directory) — what are the major modules and how do they connect?

A trimmed example of find_callers:

def find_callers(graph, symbol_name: str):
    target = graph.get_node(symbol_name)
    if not target:
        return []
    return [
        {"caller": edge.source.name, "file": edge.source.file, "line": edge.line}
        for edge in graph.incoming_edges(target)
        if edge.kind == "calls"
    ]
Enter fullscreen mode Exit fullscreen mode

The agent calls find_callers("process_payment") instead of grepping for process_payment across the repo and manually filtering out comments, string literals, and unrelated matches.

Here's the rough shape of the whole thing:

graph LR
    A[Agent] -->|find_callers / trace_call_chain| B[Query Layer]
    B --> C[Code Graph: nodes + edges]
    C --> D[Incremental Indexer]
    D -->|parses via tree-sitter| E[Source Files]
Enter fullscreen mode Exit fullscreen mode

Lessons Learned

1. Grep-first agents pay a rediscovery tax every session. Without persistent structure, the agent re-derives "what calls what" from scratch every time, even for a codebase it "worked in" yesterday. A cheap index amortizes that cost across sessions instead of paying it repeatedly.

2. A graph doesn't replace reading code — it triages where to read. Early on I let the agent act directly on graph query results without opening the actual file first, for a rename across ~15 call sites. One of those results pointed at a line that had shifted slightly due to a stale index entry, and the edit landed in the wrong spot. Now the graph is explicitly framed (in the tool description) as "here's where to look," not "here's the ground truth" — the agent still opens the file before editing.

3. Staleness kills trust fast, so build incremental re-indexing from day one. The bug above only happened because I hadn't wired up hashing yet. One stale-index incident was enough to make me not trust the tool for a while, which defeats the purpose. Content-hash-based incremental updates fixed it, but I should have built that in from the start instead of bolting it on after getting burned.

4. Query design matters more than graph completeness. My first version exposed a generic query interface so the agent could "ask anything." In practice it almost never used it well — it either wrote overly broad queries that returned noise, or gave up and grepped anyway. Four well-named, narrow verbs (find_definition, find_callers, trace_call_chain, get_architecture_overview) got used correctly on the first try, almost every time. Narrow and well-labeled beats general and powerful, at least for how an LLM actually reaches for tools.

5. The real ROI wasn't speed — it was correctness during refactors. I expected the win to be "fewer tokens spent grepping." The bigger win turned out to be qualitative: refactors that touch many call sites now actually find all of them. That's the difference between shipping a rename and shipping a rename plus a broken build two files away.

6. This works best as a first resort, not the only resort. I still let the agent fall back to grep when a query comes back empty — sometimes the graph genuinely doesn't cover a case (a string-templated call, a symbol defined in a language the parser doesn't support yet). Treating the graph as "try this first" rather than "trust this exclusively" kept it useful even on the days it couldn't answer something.

What's Next

A few things I'm still chewing on:

  • Cross-service tracing. Right now the graph stops at process boundaries. Tracing a call from an HTTP handler in one service to the function it eventually invokes in another means following string-based route names, which tree-sitter alone can't resolve.
  • Dynamic dispatch and duck typing. Anything resolved at runtime (dependency injection, getattr-style dispatch, interface-based polymorphism) is invisible to a static AST walk. I'm looking at a lightweight runtime trace as a fallback for exactly these cases, rather than trying to model every dynamic pattern statically.
  • A fallback for what the graph can't model. For code the graph genuinely can't represent well, I want a secondary search path (semantic/embedding-based) rather than silently falling back to plain grep.

Wrap-up

If your agent is spending half of every session rediscovering how your codebase fits together, a small structural index pays for itself fast — and it doesn't need to be fancy. Four query verbs and a hash-based incremental parser got me most of the value.

If you've built something similar — or found a better way to handle dynamic dispatch in a static graph — I'd genuinely like to hear about it. Follow me here on Dev.to for more of these build logs, and if you're experimenting with Claude Code for agentic workflows, give this pattern a try in your own repo.

Top comments (0)