Engineering

From Graph RAG to Budgeted Graph Walk

An AST and ASG context engine for coding agents, revised for 2026

*Shawn McAllister. This is a revision of the March 2025 explainer that I created for a research repo davila7/ast-asg-graph-rag by Daniel Avila.


In March 2025 I wrote an explainer about using abstract syntax trees and abstract semantic graphs to feed a coding assistant better context. The core of it was a Rust design: parse code with tree-sitter, build a semantic graph of symbols and references, store it in LMDB and LibSQL, and use Monte Carlo Tree Search with a depth-aware UCT to decide which parts of the graph an assistant should see for a given task.

Eighteen months later the graph is still right. Almost everything around it has changed. The consumer of the context is no longer a chat box with a small window; it is an agent with tools, a million-token window behind a prefix cache, and the ability to go and read code itself. The question the 2025 design answered, how do we fit the right code into the prompt, has mostly stopped being asked. The questions that replaced it are harder: how do we keep an agent from reading the wrong code, how do we answer structural questions the agent cannot answer by grepping, and how do we know whether any of it helped.

This revision keeps what held up, replaces what the ecosystem replaced, and adds the two things that did not exist in 2025: a real reward signal, and a GPU that is no longer idle.

What held up

The distinction between an AST and an ASG is the same as it was. An AST is the syntactic shape of one file; an ASG is a directed graph of semantic entities, functions, types, variables, modules, with edges for calls, references, implementations, and data flow. The AST is cheap and local. The ASG is expensive and global, and it is the only one that can answer "what breaks if I change this signature."

That question is the whole reason to build the engine. Agentic search, the pattern where the model greps, reads, and follows its own curiosity, beat embedding retrieval for code during 2025, and it deserved to. But grep cannot find callers through a trait object, cannot follow a generic through three crates, and cannot tell you that a change to a route handler's return type reaches an Inertia page two directories away. Those are graph queries. The ASG answers them in one hop. Everything else in this document is in service of making that hop cheap, current, and trustworthy.

The Rust builder from 2025, tree-sitter parsing into a petgraph DiGraphMap keyed by symbol, also holds up. What changed is where the semantics come from.

What the substrate looks like now

In 2025 the builder derived types and references itself. That was necessary then and is a liability now, because the language servers got there first. rust-analyzer and tsserver expose definitions, references, implementations, and call hierarchy over LSP, and they are maintained by people whose entire job is keeping that analysis correct. Deriving the same facts in a side project means being wrong in all the places they are right.

The modern ASG is therefore three layers. tree-sitter for syntax, because it is fast, incremental, and language-agnostic. LSP for semantics, because the compiler's own view of references and types is the one the code will be judged by. And SCIP, Sourcegraph's code intelligence protocol, as the persisted and interchange form, because an index that can be written once and read by any tool is worth more than a bespoke graph file. The builder's job shrinks to stitching: syntax from tree-sitter, edges from LSP, serialization through SCIP, and the graph in memory for traversal.

pub struct AsgBuilder {
    parser: tree_sitter::Parser,
    lsp: LspClient,              // rust-analyzer, tsserver, or both
    graph: DiGraphMap<SymbolId, EdgeKind>,
    index: ScipWriter,
}

impl AsgBuilder {
    pub fn ingest(&mut self, file: &Path, source: &str, prev: Option<&Tree>) -> Result<Tree> {
        let tree = self.parser.parse(source, prev).ok_or(Error::Parse)?;
        for symbol in symbols_in(&tree, source) {
            let id = self.graph_node(&symbol);
            for r in self.lsp.references(&symbol)? {
                self.graph.add_edge(self.graph_node(&r), id, EdgeKind::References);
            }
            for c in self.lsp.incoming_calls(&symbol)? {
                self.graph.add_edge(self.graph_node(&c), id, EdgeKind::Calls);
            }
            self.index.record(&symbol, &self.graph);
        }
        Ok(tree)
    }
}

Two properties matter more than they did in 2025. The graph must be incremental, because the consumer is now an agent making five hundred edits in an afternoon, and a full re-index per edit is a non-starter; tree-sitter's incremental parse takes the previous tree and a diff, and the LSP servers are already incremental. And the graph must be versioned by commit, because an agent's run has a lineage, and "the graph as of version 12" is a question the run will ask. Git is already the version store; the index just needs to carry the commit it was built from.

Storage, simplified

The 2025 diagram had LMDB for documents, symbols, and references, and LibSQL for a separate vector store and semantic index. LMDB stays; it is the right shape for a symbol table that is read far more than written. The separate vector store goes away, because LibSQL now has native vector search, and a table with a vector column beside the symbol row is simpler than two stores that have to agree. No Postgres, then or now. One file per repository, versioned by commit, rebuildable from the source tree and the index.

From Graph RAG to tools

The biggest change is in the architecture diagram, and it is a change in who is in charge. In 2025 the engine retrieved chunks and prepended them to the prompt. The model was a passenger. In 2026 the model is the driver, and the engine is a set of tools it can call when grep is not enough.

That means the HTTP and LSP-bridge server from the original diagram becomes an MCP server exposing a small, structural surface:

  • callers(symbol) and callees(symbol): one hop in each direction, with file and line.
  • impact(symbol, change): the transitive set of symbols a signature or type change reaches, ranked by distance.
  • context(query, budget): the budgeted graph walk described below, returning a ranked set of nodes by reference, not by value.
  • graph_at(commit): pin the graph to a lineage version.

Returning by reference matters. In 2025 the engine inlined code into the prompt because the model could not fetch it. Today the model has read; the engine should tell it where to look and why, and let it decide whether to look. That keeps the engine's output small, keeps the model's context under its own control, and makes the engine useful to every agent that speaks MCP, which in August 2026 is all of them.

flowchart TB
    Agent[Coding agent: Claude Code, OMP, Codex, suprnova-coder] -->|MCP| Engine[Context engine]
    Engine --> Tools[callers / callees / impact / context / graph_at]
    Tools --> Walk[Budgeted graph walk]
    Walk --> Graph[(ASG in memory)]
    Graph --> LMDB[(LMDB: symbols, refs)]
    Graph --> LibSQL[(LibSQL: vectors, metadata)]
    Builder[Builder] --> Graph
    TS[tree-sitter] --> Builder
    LSP[rust-analyzer / tsserver] --> Builder
    Builder --> SCIP[(SCIP index, per commit)]
    Watcher[File watcher] --> Builder
    Journal[(Run journal)] -.->|outcomes| Walk

The budgeted graph walk

Here is the part of the 2025 design that survives with a new name. The context tool has a budget, a number of nodes or tokens the agent is willing to spend, and a graph that is far larger than the budget. Something has to choose. The original used MCTS with a depth-aware UCT, and the formula still does the job:

UCT(s, a) = V(s, a) + C · sqrt(ln N(s) / N(s, a)) + α · e^(−β(d − 1)) − γ · sqrt(d)

The first two terms are standard: a value estimate and an exploration bonus. The two depth terms are the part specific to code. The exponential bonus is full strength at depth 1 and decays fast, so the walk strongly prefers the immediate neighbors of the symbol under edit. The square-root penalty grows without bound but slowly, so a deep path is never forbidden, just progressively more expensive. Together they encode something that is true of code and not of game trees: relevance falls off with graph distance, quickly at first, then slowly, and a node five hops out needs a very high value estimate to earn a place in the window.

What changed is the framing. In 2025 this was "MCTS picks what the assistant thinks about," which reads oddly now that reasoning models do their own search at inference time. The honest 2026 description is a budgeted traversal policy: given a seed symbol, a budget, and a value function, walk the graph and return the best set of nodes the budget allows. The formula is the policy. The seed comes from the agent's query, resolved through the same LSP the builder used.

A real reward signal

The 2025 reward was a weighted sum of proxies: context relevance, workspace coverage, change proximity, symbol relevance, semantic similarity. Proxies were all that was available, because nothing downstream said whether the context had helped.

A coding harness with a scorer changes that. If an agent's attempts are evaluated, by tests, by a benchmark, by a verify gate that passes or fails, then every attempt produces a label: the context the agent used, and whether the attempt was accepted. That is a real reward, and V(s, a) can be learned from it. A node that appeared in the context of accepted attempts earns value; a node that appeared in the context of discarded attempts loses it; the depth prior stays as the shape, and the learned value fills in the content.

The "change proximity" and "change history" terms from 2025 become concrete in the same move. A harness keeps a journal of what each attempt touched. The walk's value function reads it. The context for attempt 212 is weighted by what attempts 1 through 211 changed, which is the memory-aware retrieval the original reward was reaching for without a place to store the history.

This is the single largest update in the document, and it only became possible because the harness exists to close the loop. The engine proposes context; the scorer judges the result; the journal carries the verdict back.

Compute, pruned

The 2025 document spent a third of its length on SIMD: std::simd with f32x8, an AVX2 tutorial, a hybrid CPU strategy, and a #[cfg(feature = "cuda")] kernel with a fallback. Most of that can go.

std::simd is still experimental as of 1.98.0. The stable answer is thermite: trait-generic kernels over GenericVector → NumericVector → FloatVector, runtime dispatch with dispatch_dyn!, and a #[thermite::dispatch] attribute that propagates target features across function boundaries, which is the exact mistake that makes hand-rolled target_feature code silently fall back to scalar when the hot loop calls a helper. SSE2 through AVX2 plus FMA on x86, NEON on AArch64, wasm SIMD128, scalar fallback always. No AVX-512 yet, which matters on a server and not on a desk. The cosine kernel is written once:

#[thermite::dispatch]
fn cosine<V: FloatVector>(a: &[V::Element], b: &[V::Element]) -> V::Element {
    let (mut dot, mut na, mut nb) = (V::zero(), V::zero(), V::zero());
    for (ca, cb) in a.chunks_exact(V::NUM_ELEMENTS).zip(b.chunks_exact(V::NUM_ELEMENTS)) {
        let (x, y) = (V::load_unaligned(ca), V::load_unaligned(cb));
        dot = x.mul_add(y, dot);
        na = x.mul_add(x, na);
        nb = y.mul_add(y, nb);
    }
    dot.sum() / (na.sum() * nb.sum()).sqrt()
}

The CUDA path is replaced by wgpu: WGSL compute shaders over Vulkan, Metal, and DX12, no CUDA toolchain, and it runs on the AMD and Apple machines the 2025 design excluded. The split between the two is by batch size. Per-query scoring, a few thousand candidate nodes against one embedding, stays on thermite; the GPU round-trip costs more than the math. Graph-wide work goes to wgpu: re-embedding a repository after a large diff, periodic clustering, all-pairs similarity for link discovery. That is where there are millions of dot products and the transfer is amortized.

Embeddings themselves come from a named 2026 model, Qwen3-Embedding or a code-specific model of similar vintage, served by whatever inference server already owns the GPU, rather than from an unnamed "embedding store."

The GPU is not idle anymore

This is the second thing that did not exist in 2025. The local inference engines that make a 35B mixture-of-experts model run at 100 tokens per second on a gaming GPU do it by filling VRAM: an expert cache, a KV cache, and elastic reallocation between them. A context engine that grabs a gigabyte for a wgpu job in the middle of a run will stall the model or get evicted.

So GPU maintenance has to be scheduled, not just run. Either it happens between attempts, when the harness knows the worker is parked, or it asks the inference engine's memory API to yield first. A conductor that already tracks worker state has this information; the context engine's background jobs become one more thing the journal tells it when to do.

Prefix caching changes retrieval order

One more 2026 fact with a 2025-shaped consequence. Prompt prefix caching is now standard, and the newer edge engines go further with semantic checkpoints so that tool calls and thinking blocks do not force a full recompute. Retrieval that prepends fresh context every turn invalidates the cache every turn, and the cost shows up as prefill latency on every attempt.

The fix is ordering. Stable material first: system prompt, glossary, lessons, the standing knowledge roots. Volatile material last: the nodes the walk returned for this query. The engine should return context in a form the agent can append rather than prepend, and the agent's harness should treat the walk's output as the tail of the prompt, never the head.

Evaluation, replaced

In 2025 the natural metric was retrieval quality: recall at k, precision at k, whether the right chunks were in the window. None of that survives contact with an agent that can fetch its own chunks. The only metric that matters is downstream: did attempts that used the engine get accepted more often, at lower cost, than attempts that did not, on the same task, with the same model.

That is an ablation, and it needs a lab to run it in: same fixture, same model, same budget, engine on versus engine off, one variable per row. The harness that supplies the reward signal supplies the lab too. The engine's claim, that structural context beats grep for structural questions, is a testable claim now, and it should be tested before it is believed.

Where it sits in a harness

In a harness like the one I am currently building - for Suprnova specifically, the engine is a knowledge root that answers questions instead of listing files. The worker has read, grep, and glob for the cases where they are enough, and callers, impact, and context for the cases where they are not. The conductor pins graph_at to the lineage version before each attempt. The journal feeds the value function. The lab runs the on/off row. None of the engine's internals leak into the conductor; it is a tool behind a protocol, replaceable by a better one without touching the loop.

That is the shape the 2025 design was reaching for, written at a time when the only way to put structure in front of a model was to prepend it and hope. The graph was right. The thing that was missing was a loop that could tell it so.


The March 2025 original, including the full LMDB and LibSQL schema, the petgraph builder, and the SIMD sections this revision removes, remains in the repository under its original filename.

Other projects I have developed over time in this lane

Earlier passes at the same problem, in order:

  • tsg_indexer (May 2025): repository indexing on tree-sitter stack graphs, cross-file definition and reference resolution for 20+ languages, JSON and DOT output, built to sit under an MCP server.
  • code-graph-mcp (July 2025): an MCP server in Python on ast-grep and rustworkx, with definitions, references, callers and callees, complexity and smell detection, and a debounced file watcher, for 25+ languages.
  • fast-context (March 2026): the Rust rebuild, tree-sitter parsing and a petgraph graph with Node and Python bindings, a file watcher, a fast-context-mcp server, Codex skills, and a Claude plugin.
  • cognitive-prompt-architecture and engineered-meta-cognitive-workflow-architecture (2025): the prompt-side and memory-side frameworks the harness mentioned above grew out of.

Comments 0