Building Cost Controls for AI Agents Querying Azure Cosmos DB

Introduction

AI agents can make database applications much more flexible. Instead of executing only predefined operations, an agent can interpret a request, decide which data it needs, call a database tool, inspect the result, and continue working.

That flexibility introduces a new operational concern: database cost control.

With a traditional application, developers generally know which database operations are executed by each endpoint. An AI agent can make multiple calls depending on the task, generate different queries, retry operations, or retrieve more data than necessary.

For Azure Cosmos DB, this matters because database operations consume Request Units, commonly referred to as RUs.

An AI agent that performs one carefully designed query may be inexpensive. An agent that performs several unnecessary queries or retrieves large amounts of data can consume substantially more resources.

The solution is not to prevent agents from accessing databases. Instead, production systems should place clear cost, query, result-size, and authorization controls around database tools.

Understanding RU Consumption

Azure Cosmos DB measures database operations using Request Units.

The amount of RU consumed by a query depends on factors such as:

Data being processed
Query shape
Indexing
Partition behavior
Number of results
Document size
Number of operations

Therefore, an AI agent should not assume that two logically similar queries have identical costs.

Consider:

User Request
     |
     v
AI Agent
     |
     +--> Query 1
     |
     +--> Query 2
     |
     +--> Query 3
     |
     v
Azure Cosmos DB

The cost of the agent task is the combined cost of the database operations, not just the most expensive individual query.

A useful model is:

Task RU Cost =
Query 1 RU
+ Query 2 RU
+ Query 3 RU
+ ...

This is why cost should be tracked at the agent-task level.

Why AI Agents Need Additional Cost Controls

A normal API endpoint might execute:

Request
  |
  v
One known query
  |
  v
Response

An agent may execute:

Request
  |
  v
Reasoning
  |
  +--> Search customers
  |
  +--> Search orders
  |
  +--> Get product
  |
  +--> Re-query orders
  |
  +--> Retrieve additional details
  |
  v
Final Response

Every unnecessary database operation can increase cost.

The agent may also repeat a query when it could have reused an earlier result.

This means database cost becomes part of agent orchestration.

Start With a Database Cost Budget

The first control should be a task-level budget.

For example:

Agent Task
    |
    v
RU Budget = 100
    |
    +-- Query 1: 8 RU
    +-- Query 2: 15 RU
    +-- Query 3: 12 RU
    |
    v
Remaining = 65 RU

If the agent reaches the defined limit:

Budget Exceeded
      |
      v
Stop Database Calls
      |
      v
Return Controlled Result

The actual budget should be determined from the application's workload and cost requirements.

There is no universal RU threshold that is appropriate for every agent.

Track RU Consumption in the Database Tool

The database tool should record the RU charge returned by each operation.

A C# service can maintain a task-level counter:

public sealed class RuBudget
{
    public double Consumed { get; private set; }

    public double Remaining(double limit)
        => Math.Max(0, limit - Consumed);

    public void Add(double requestCharge)
    {
        Consumed += requestCharge;
    }
}

After executing a query:

var response = await iterator.ReadNextAsync();

budget.Add(response.RequestCharge);

if (budget.Consumed > budgetLimit)
{
    throw new InvalidOperationException(
        "Agent database budget exceeded.");
}

The important design decision is that the database tool, rather than the AI model, should enforce the budget.

Why the Model Should Not Control the Budget

An instruction such as:

Never consume more than 100 RU.

can be useful guidance, but it should not be treated as an enforcement mechanism.

The model can make mistakes.

The application should maintain the authoritative budget:

AI Agent
   |
   v
Database Tool
   |
   +-- Check Budget
   |
   +-- Execute Query
   |
   +-- Record RU
   |
   v
Cosmos DB

This is the same principle used for other security and operational controls: the system enforcing the policy should not depend entirely on the component being controlled.

Limit the Number of Database Calls

RU is not the only useful control.

A task can also have a maximum number of database operations.

For example:

Maximum calls per task: 10

The tool can track the count:

public sealed class DatabaseCallBudget
{
    public int Count { get; private set; }

    public void RegisterCall(int maximum)
    {
        if (++Count > maximum)
        {
            throw new InvalidOperationException(
                "Database call limit exceeded.");
        }
    }
}

This protects against an agent repeatedly calling the database even when individual queries are inexpensive.

Control Result Size

Returning an entire dataset to an AI agent is rarely necessary.

Consider:

SELECT *
FROM c
WHERE c.status = @status

A more focused query may be:

SELECT TOP 20
    c.id,
    c.status,
    c.total
FROM c
WHERE c.status = @status

The application should define sensible result-size limits.

For example:

Maximum records: 50
Maximum response size: application-defined

The appropriate limits depend on the workload.

Result limiting also helps prevent unnecessarily large agent contexts.

Use Projection to Reduce Unnecessary Data

Suppose the user asks:

Which active orders are above $500?

The agent may only need:

Order ID
Total
Status

There is little reason to retrieve a large document containing:

Customer profile
Shipping address
Payment metadata
Product details
Internal notes
Audit information

A focused projection can keep database and agent processing more efficient.

SELECT
    c.id,
    c.total,
    c.status
FROM c
WHERE c.status = @status
AND c.total > @minimumTotal

The exact cost impact should be measured for the application's data model rather than assumed.

Enforce Partition-Aware Access

Partitioning is particularly important when designing cost controls.

Suppose:

/tenantId

is the partition key.

A query that includes the appropriate partition information can behave differently from one that requires broader data access.

The agent should not be responsible for remembering this requirement on every request.

Instead, the database tool should know the partitioning model.

For example:

Agent
 |
 v
Query Request
 |
 v
Database Tool
 |
 +-- Tenant Context
 +-- Partition Context
 +-- Query Validation
 |
 v
Cosmos DB

This also strengthens tenant isolation.

Do Not Let the Agent Choose Arbitrary Tenants

In a multi-tenant application, a tool such as:

QueryAsync(
    string tenantId,
    string query)

can be dangerous if the agent controls both values.

Instead, derive the tenant from trusted application context.

public sealed record TenantContext(
    string TenantId);

The tool can then apply the tenant boundary independently of the generated request.

This provides both security and cost benefits because an agent cannot intentionally or accidentally scan unrelated tenant data.

Validate Generated Queries

An AI-generated query should not go directly to the database without validation.

A validation layer can inspect:

Query
 |
 +-- Allowed operation?
 +-- Allowed container?
 +-- Required partition information?
 +-- Result limit?
 +-- Allowed fields?
 +-- Allowed operators?
 +-- Parameterized?
 |
 v
Decision

For a read-only agent, the tool might allow only approved query operations.

This reduces the risk of an agent generating an unexpectedly expensive or inappropriate operation.

Use Query Complexity Controls

Some query patterns deserve additional scrutiny.

For example:

Large result sets
Complex joins
Expensive sorting
Broad scans
Unbounded searches
Large cross-partition operations

The application can reject or require additional approval for queries that exceed defined limits.

Conceptually:

Generated Query
      |
      v
Complexity Check
      |
  +---+---+
  |       |
 Pass    Reject
  |       |
  v       v
Execute  Explain

The exact definition of query complexity depends on the database workload.

Add a Cost-Aware Tool Contract

Instead of exposing a generic database query tool, provide a more constrained interface.

For example:

public sealed record QueryRequest(
    string Container,
    string Query,
    IReadOnlyDictionary<string, object> Parameters,
    int MaxResults);

The application can then enforce:

Allowed Container
Allowed Query Pattern
Maximum Results
Maximum Calls
RU Budget
Tenant Scope

The AI selects the tool and supplies the task-specific information, but the application remains responsible for enforcing the constraints.

Cache Repeated Agent Queries

Agents can sometimes ask for the same information more than once.

For example:

Query customer
      |
      v
Agent receives customer
      |
      v
Query customer again

Caching can reduce unnecessary database operations when the data is suitable for caching.

A simple cache key could include:

Tenant
Container
Query
Parameters

However, caching should consider data freshness and authorization.

Never allow cached data from one tenant or authorization context to be returned to another.

Track Cost Per Agent Task

Instead of monitoring only database-level consumption, create an application-level metric:

Agent Task
 |
 +-- Database Calls: 4
 +-- Total RU: 42
 +-- Records Returned: 31
 +-- Duration: Measure

This makes it possible to answer:

Which agent workflows are expensive?
Which tasks generate the most database calls?
Which tools consume the most RU?
Which prompts consistently trigger unnecessary queries?

This is much more actionable than looking only at aggregate database consumption.

Create a Cost Evaluation Dataset

Before deploying an AI database agent, create representative tasks.

For example:

Task	Expected Calls	Cost
Find customer	1	Measure
Find active orders	1	Measure
Customer order summary	2	Measure
Search products	1	Measure
Generate customer report	Multiple	Measure

The expected number of calls should be based on the application's intended design.

Then compare the actual agent behavior.

Test Repeated Runs

AI agents may produce different execution paths for the same request.

Run the same task multiple times where appropriate:

Same Task
   |
   +-- Run 1 -> 18 RU
   +-- Run 2 -> 22 RU
   +-- Run 3 -> 18 RU
   +-- Run 4 -> 41 RU

The exact values above are illustrative only.

Your benchmark should record actual results.

Variation matters because a workflow that is normally inexpensive but occasionally performs many extra queries may still need stronger controls.

Add Alerting

A cost-control system should provide visibility when agent behavior changes.

For example:

Normal
0–50 RU

Warning
50–100 RU

Limit
100+ RU

The thresholds should be configured according to the application's requirements.

An alert might contain:

Agent Task: CustomerSummary
RU Consumed: Measured
Database Calls: Measured
Result Count: Measured
Status: Budget Warning

Do not include sensitive query parameters or customer data in logs unless necessary.

Common Mistakes

Relying Only on Prompts

A prompt saying "use as few database calls as possible" is not a cost-control mechanism.

Setting Only a Database-Level Budget

Aggregate database monitoring does not tell you which agent workflow caused the cost.

Track consumption per task.

Allowing Unlimited Result Sets

Large responses can increase both database and agent processing costs.

Ignoring Duplicate Queries

Agents may repeat operations unnecessarily.

Allowing Cross-Partition Searches Without Review

Broad searches can behave very differently from targeted queries.

Giving the Agent Full Database Access

Expose only the operations and containers required for the workflow.

Using Fixed RU Benchmarks

RU consumption depends on workload characteristics. A number measured in one dataset should not automatically be treated as a universal expectation.

Troubleshooting

Agent Consumes Too Many RUs

Start with the execution trace.

Task
 |
 +-- Query 1
 +-- Query 2
 +-- Query 3
 +-- Query 4

Identify whether the problem is:

Too many calls
Large result sets
Poor query filtering
Unnecessary repeated queries
Incorrect partition usage

Agent Repeats the Same Query

Check whether the agent receives the previous result in a form it can reuse.

Application-level caching can also help where appropriate.

Budget Is Exceeded Too Quickly

Determine whether the budget is being applied per query or across the complete task.

The most useful control is generally a cumulative task budget.

Query Is Correct but Too Expensive

Compare it with a known efficient query.

Review:

Filters
Projection
Partition usage
Result limits
Indexing
Query shape

Do not automatically assume that the AI query is wrong simply because its RU charge is higher.

Best Practices

Treat RU consumption as a first-class agent metric.
Enforce budgets in application code.
Track cumulative RU per task.
Limit database calls per task.
Limit result size.
Use focused projections.
Enforce tenant and partition context outside the model.
Validate generated queries before execution.
Restrict accessible containers and operations.
Cache repeated reads when appropriate.
Track database cost alongside latency and result size.
Test repeated executions.
Add warnings and hard limits.
Keep sensitive data out of diagnostic logs.
Review expensive agent workflows regularly.

Advantages

Prevents uncontrolled AI-driven database consumption.
Makes agent database usage measurable.
Helps identify inefficient agent workflows.
Adds an additional safety boundary around database tools.
Encourages efficient query design.
Supports predictable operational behavior.
Can reduce unnecessary database calls.

Disadvantages

Cost controls add implementation complexity.
Strict budgets can interrupt legitimate long-running tasks.
RU consumption varies by workload and data distribution.
Query validation can become complicated for flexible agent requirements.
Caching introduces freshness and authorization considerations.
Too many restrictions can reduce the usefulness of an AI agent.

A Practical Cost-Control Architecture

A production-oriented design can look like this:

                         User
                           |
                           v
                       AI Agent
                           |
                           v
                    Database Tool
                           |
             +-------------+-------------+
             |             |             |
             v             v             v
        Authorization   Validation   Cost Budget
             |             |             |
             +-------------+-------------+
                           |
                           v
                   Query Execution
                           |
                           v
                    Azure Cosmos DB
                           |
                 +---------+---------+
                 |                   |
                 v                   v
              Results            RU Charge
                 |                   |
                 +---------+---------+
                           |
                           v
                    Task Telemetry

This architecture keeps the model responsible for reasoning while the application remains responsible for enforcement.

Example Agent Execution Flow

Suppose the user asks:

Show the five most recent active orders
for customer 1001.

A well-designed workflow might be:

User Request
     |
     v
AI Agent
     |
     v
search_orders
     |
     v
Validation
     |
     +-- Tenant Context
     +-- Customer Context
     +-- Maximum Results = 5
     |
     v
Cosmos DB
     |
     v
Results + RU Charge
     |
     v
Budget Tracker
     |
     v
AI Agent
     |
     v
Final Response

The agent does not need unrestricted database access.

It needs a well-defined capability that performs the required operation.

Conclusion

AI agents can make database applications more flexible, but autonomous database access introduces a cost-management problem that traditional APIs do not always have. An agent may perform several queries, repeat operations, retrieve unnecessarily large results, or generate a query that is more expensive than the equivalent application-written operation.

Azure Cosmos DB's RU model makes this behavior measurable, which is useful for building explicit controls.

A production-ready AI database architecture should enforce task-level RU budgets, database-call limits, result-size limits, query validation, tenant restrictions, and appropriate authorization. These controls should exist in the application and database tool layer rather than relying solely on instructions given to the AI model.

The most important principle is simple: let the AI agent decide what information it needs, but let the application decide how much database access it is allowed to consume.