Token Budget Allocation and Query Cost Estimation for AI Agents

AI agents querying analytical databases have a spending problem, and it's not the kind that shows up on a monthly invoice right away. A bill that arrives 30 days late reflects it, after an agent has spent a weekend branching, rolling back, and re-querying its way through a dataset nobody asked it to explore that thoroughly. This piece lays out a framework for making agents reason about cost the way a query planner does: estimate first, allocate a budget across the whole plan, then execute.

What "query cost" means inside the database, the primitives agents must reason about

Start with the unit of work. Most columnar analytical engines read data in granules, fixed blocks of rows (commonly around 8,192 rows per granule in systems built this way). Matching a single row inside that block forces the engine to read the entire block. There's no such thing as reading half a granule. An agent filtering for one needle in a haystack still pays for the whole haystack, unless the engine can skip that granule entirely.

That's where primary key design comes in, and it's arguably the single biggest lever on cost. A well-chosen sort order (the ORDER BY clause, essentially) lets the engine skip granules wholesale instead of reading and discarding them. Get this right and query performance can improve by a substantial multiple. Get it wrong, and every query becomes a guessing game.

Columnar storage adds a second wrinkle: cost scales with columns selected, not rows in the table. Cost scales with columns selected, not rows in the table. Querying two columns out of a 100-column table costs about the same as querying two columns out of a 10-column table, because the engine never touches the other 98 or 8 columns in the first place. So SELECT * isn't just lazy, it's disproportionately expensive. It forces the engine to read every column, whether or not the next reasoning step needs them.

Partition pruning is the coarse filter that runs before any of this. A query that filters on the partition key skips whole partitions before the sparse index is even consulted. Think of it as the first bouncer at the door, checking IDs before anyone gets close enough to check granule-level details.

How token budget and query cost are coupled, and why ignoring that coupling is expensive

Every agent query runs two cost meters at once. One meter tracks database cost (bytes scanned, memory pressure, compute time). The other tracks inference cost (tokens spent on schema context, intermediate reasoning, and processing whatever comes back). Agents that only watch one meter are, by definition, only half-aware of what they're spending.

The token side isn't trivial, either. Full-context approaches, where an agent stuffs everything it might need into the prompt, run 25,000 to 100,000+ tokens per interaction. Retrieval-based approaches that pull in only what's relevant can bring that down to roughly 7,000 tokens. That's a wide swing on inference spend alone, before a single row of data gets touched.

Caching changes this math, but only if an agent uses it deliberately. On Claude Sonnet 4.6, a cache read runs $0.30 versus $3.00 for standard input, a sizable difference. The break-even occurs around the third use of the same prefix within the cache's time-to-live window. Schema context, which barely changes query to query, is exactly the kind of prefix that should live in cache rather than get re-sent cold every time.

And then there's the compounding effect: a large result set coming back from a query becomes input tokens for the next reasoning step. An unbounded query doesn't just cost more to run, it costs more to think about afterward. One sloppy query without a LIMIT clause can quietly double or triple the token bill for everything downstream of it.

Put those two meters side by side and the coupling becomes obvious. A cheap database query with an enormous result set can still be an expensive move overall, once the inference side picks up the tab.

A pre-execution estimation framework for agents: reasoning about cost before issuing a query

Query planners inside databases estimate execution cost, rows touched, I/O, memory, before committing to an execution plan. Agents should be doing the same thing before they commit to issuing a query at all. Not after. Before.

Step 1 is schema introspection, treated as a budget-aware first move rather than a formality. The agent inspects table structure, the ORDER BY and PARTITION BY definitions, and whatever materialized views already exist, before it writes a single line of the actual query. But this step costs tokens too, so it needs to stay targeted. Describing the two or three relevant tables costs a fraction of what describing an entire schema costs. Structured, code-like commands for describing schema and pipelines give agents a low-token, repeatable way to do this without re-deriving it from scratch each session.

Step 2 is partition and granule estimation. The agent works out which partitions its filter will actually touch, based on the partition key and whatever time or dimension bounds the query implies. It also checks whether its WHERE clause lines up with the leading columns of the ORDER BY. If it doesn't line up, that's a signal: expect a full-partition scan, and plan accordingly. A cheap metadata query against the system tables that expose part counts and compressed sizes can confirm this before the agent commits to anything expensive.

Step 3 is column selection discipline, which sounds almost too simple to call a step, but it's where a lot of waste hides. Every column added to a SELECT is a cost increment in a columnar system, so agents should ask only for what the immediate reasoning step actually needs. Nullable columns carry a hidden overhead, a separate column tracking which rows are null. An agent querying a wide table full of nullable columns is paying a small tax on each one, and that tax adds up across a big enough query.

Allocating a token budget across a plan of sub-queries rather than a single query

Agents rarely issue one query and call it done. A real task looks more like a sequence: inspect the schema, run an exploratory probe, narrow the search, run the final aggregation, then validate the result. Each of those steps spends both database budget and token budget, and treating them as five isolated events misses the point entirely.

The plan needs to be budgeted as a whole, with each step getting a resource envelope sized to how much it's actually likely to teach the agent. That's the core idea behind hierarchical decomposition:

Early steps should be cheap by design, metadata lookups, LIMIT 1 probes, lightweight query-plan checks, run specifically to inform the next step without spending the full budget up front.

• Expensive steps, full aggregations, multi-table joins, only earn their budget after a cheaper step has confirmed the data is worth that spend.
• Every step's result should reduce uncertainty enough to justify what it cost. If a step comes back and doesn't meaningfully narrow anything down, that's the moment to branch the plan or halt it, not push forward on momentum.

Checkpointing ties this together. Before each sub-query fires, the agent should compare what budget remains against the estimated cost of what's left in the plan. If that ratio gets too thin, the agent has two reasonable moves: compress what's left (coarser aggregation, a higher LIMIT, fewer dimensions) or stop and escalate to a human. Neither move is a failure. Continuing to spend past that point, without adjusting, is the actual failure mode.

Schema and infrastructure decisions that make agent cost estimation reliable rather than speculative

A messy schema undermines all of the estimation logic above. Schema design is, by a wide margin, the biggest factor in analytical query performance, capable of swinging results by a large multiple depending on how well it fits the query patterns actually in use. An agent estimating cost against a badly designed schema is estimating noise, and dressing that noise up as a number doesn't make it less noisy. It's estimating noise, and dressing that noise up as a number doesn't make it less noisy.

A few schema decisions turn the outcome from guesswork into estimation.

ORDER BY discipline. The columns agents filter on most, tenant ID, event type, a time range, should lead the sort key. Keeping that key to somewhere around 4 to 6 columns keeps the index tight enough that estimation stays tractable rather than turning into a research project. A mismatched ORDER BY means an agent can't predict which granules get skipped, and every estimate degrades into a worst-case assumption, which defeats the entire purpose of estimating in the first place.

Partition granularity as an anchor. Monthly partitioning gives agents a clean, countable unit. An agent can check the system metadata and know, before running anything expensive, exactly how many partitions a time-bounded query is going to touch.

Low-cardinality encoding for filter columns. Columns with fewer than 10,000 distinct values, encoded this way, filter faster and make cost estimation more tractable. Better statistics mean better estimates, and better estimates are the whole point of this exercise.

Governance, auditability, and enforcing budget limits without breaking agent autonomy

None of this works as policy unless it's enforced somewhere the agent can't talk its way around. That starts with logging, and logging means all of it, not a sample. Agent behavior is stochastic by nature. The expensive query is often the one nobody predicted in advance. Sampling logs is a bet that the interesting failure appears in the sample. It usually doesn't.

Session isolation matters just as much. Agents working with read-only compute over shared data can't mutate production, full stop, which also happens to make per-agent budget accounting possible in the first place. Sandboxed branches extend that same logic to exploration: an agent probing a real schema against a branch that can't touch production gets to run its "check before committing" behavior freely, without anyone worrying about what a runaway query might do to a live table. ObsessionDB, a managed analytical database service built for real-time analytics at terabyte-to-petabyte scale, provides this kind of isolated read-only compute for agents alongside full query logging per session, so every agentic query is auditable after the fact.

None of that replaces enforcement at the infrastructure layer, though, it complements it. A per-query memory ceiling, throwing an exception once a query crosses it, is table stakes. That ceiling can turn soft under memory overcommit settings, or get bypassed through disk spilling, so setting the limit isn't the same as guaranteeing it holds under every configuration. Per-query compute caps, enforced through resource groups or complexity limits, ensure no single agent query gets to monopolize a shared cluster, no matter how confident its plan looked five minutes earlier.

None of this fights against agent autonomy, either. Autonomy inside a budget is still autonomy. It's autonomy without a blank check, which is a different thing entirely from no autonomy at all.

Putting the framework into practice: a worked sequence for an agent querying a large table

Take a concrete scenario: an agent tasked with finding anomalous revenue patterns across tenants, inside an events table running into the hundreds of terabytes. Here's roughly how the framework plays out, step by step.

Step 1, schema introspection, kept cheap and targeted. The agent queries the system tables for just the relevant table, not the full schema catalog. It pulls the ORDER BY, the PARTITION BY, and whatever materialized views already exist. The token cost here is bounded, and because the schema barely changes between queries, that same prefix gets cached and reused for the rest of the session instead of paid for again.

Step 2, partition estimation, before committing to a time range. Before running anything against months of data, the agent checks the partition metadata to see how many partitions a given time window would actually touch. If a "check the last year" request would touch 12 monthly partitions, that's a very different cost profile than a request that touches 3. This step is where the agent decides whether to narrow the time range before spending anything real, and it's a decision that costs almost nothing to make and can save an enormous amount if made well.

That's the shape of the whole framework, really. Cheap steps buy information. Information earns the right to spend more. And the agent, at every point along the way, is holding its remaining budget up against what's left in the plan, deciding whether to keep going, narrow the scope, or hand the decision to a person. The agents built to do this well aren't the ones that run fastest. They're the ones that know, at each step, whether running at all is still worth it.