Running LLM integration tests in CI without a hard limit is how teams blow their API budget on a stuck retry loop. Capping token spend in CI test runs ensures a single bad commit can’t silently rack up thousands of completions. This guide walks through a concrete pattern to enforce a per-run token ceiling using environment-injected budgets, a lightweight client wrapper, and CI configuration.
Step 1: Route every LLM call through one guarded client
The first mistake is letting individual test modules import the OpenAI SDK directly. You lose the ability to intercept usage and enforce a ceiling. Create a single llm_client.py module that wraps an OpenAI-compatible client and tracks cumulative tokens.
Point LLM_BASE_URL at any OpenAI-compatible gateway. For example, n4n.ai exposes one endpoint that fronts 240+ models and returns accurate per-token usage, which makes the accounting below exact rather than estimated.
import os
import openai
class BudgetExceeded(Exception):
pass
class GuardedClient:
def __init__(self, budget_tokens=None):
self.client = openai.OpenAI(
base_url=os.getenv("LLM_BASE_URL", "https://api.openai.com/v1"),
api_key=os.getenv("LLM_API_KEY")
)
self.budget = budget_tokens or int(os.getenv("CI_TOKEN_BUDGET", "0"))
self.used = 0
def chat(self, **kwargs):
if self.budget and self.used >= self.budget:
raise BudgetExceeded(f"Token budget {self.budget} exhausted before call")
resp = self.client.chat.completions.create(**kwargs)
self.used += resp.usage.total_tokens
if self.budget and self.used >= self.budget:
raise BudgetExceeded(f"Token budget {self.budget} exceeded after call")
return respThe resp.usage.total_tokens field is returned by every compliant endpoint. If your gateway supports provider cache-control hints, forward them in extra_headers so repeated prompt prefixes cost less, but the meter still counts correctly.
Choose a budget that reflects your test matrix. A suite with 200 test functions each making one 2k-token call needs roughly 400k tokens. Set CI_TOKEN_BUDGET to 1.5x that to allow retries but cap runaway. Capping token spend in CI test runs is about bounding tail risk, not micro-optimizing per call.
Step 2: Inject the budget from CI environment
Hard-coding a budget in source control invites drift. Read it from an environment variable so CI owns the policy. Locally, you can simulate a tight run:
CI_TOKEN_BUDGET=50000 python -m pytest tests/integrationIn GitHub Actions, declare it at the job level:
jobs:
llm-tests:
runs-on: ubuntu-latest
env:
CI_TOKEN_BUDGET: 50000
LLM_BASE_URL: ${{ secrets.LLM_BASE_URL }}
LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
steps:
- uses: actions/checkout@v4
- run: pip install openai pytest
- run: pytest tests/integrationKeep the value in CI secrets or vars, not in the YAML, if you want per-branch overrides. The key point for capping token spend in CI test runs is that the number is external to the test logic.
Step 3: Accumulate usage across parallel workers
pytest-xdist spins up multiple processes. Each GuardedClient instance would only see its own tokens. Use a session-scoped file or a Redis key. A temp file is enough for a single runner:
# conftest.py
import pytest, json, tempfile, os
@pytest.fixture(scope="session")
def token_ledger():
path = os.path.join(tempfile.gettempdir(), "ci_token_ledger.json")
if not os.path.exists(path):
with open(path, "w") as f:
json.dump({"used": 0}, f)
class Ledger:
def add(self, n):
with open(path) as f: data = json.load(f)
data["used"] += n
with open(path, "w") as f: json.dump(data, f)
def total(self):
with open(path) as f: return json.load(f)["used"]
return Ledger()File locking for correctness
Parallel workers can race on the ledger file. Use fcntl to serialize writes:
import fcntl
def add(self, n):
with open(path, "r+") as f:
fcntl.flock(f, fcntl.LOCK_EX)
data = json.load(f)
data["used"] += n
f.seek(0); f.truncate(); json.dump(data, f)
fcntl.flock(f, fcntl.LOCK_UN)Wire the ledger into the client:
def chat(self, ledger, **kwargs):
if self.budget and ledger.total() >= self.budget:
raise BudgetExceeded("budget exhausted")
resp = self.client.chat.completions.create(**kwargs)
ledger.add(resp.usage.total_tokens)
if self.budget and ledger.total() >= self.budget:
raise BudgetExceeded("budget exceeded")
return respNow the cap is shared. This is essential when capping token spend in CI test runs that exercise many modules concurrently.
Step 4: Abort the entire suite on breach
A single test raising BudgetExceeded fails that test, but the rest of the suite keeps burning tokens. You want an immediate halt. Use a session flag and pytest.exit.
# conftest.py
import pytest
def pytest_sessionfinish(session, exitstatus):
if getattr(session.config, "budget_breached", False):
pytest.exit("Token budget breached", returncode=2)In your test wrapper, catch the exception and set the flag:
def test_some_flow(token_ledger, request):
client = GuardedClient()
try:
client.chat(ledger=token_ledger, model="gpt-4o-mini",
messages=[{"role":"user","content":"Summarize this"}])
except BudgetExceeded:
request.config.budget_breached = True
raiseThe exact plumbing matters less than the behavior: one breach stops the job.
Step 5: Add a deliberate overrun test to verify the cap
You cannot trust a guard you have never seen trigger. Write a test that sets a tiny budget and expects the exception:
def test_budget_cap_enforced(token_ledger):
client = GuardedClient(budget_tokens=10)
with pytest.raises(BudgetExceeded):
client.chat(ledger=token_ledger, model="gpt-4o-mini",
messages=[{"role":"user","content":"hi"}])Run it locally with CI_TOKEN_BUDGET=10. The test should pass (the raise is caught) and the session should exit with code 2 because the flag is set. That confirms capping token spend in CI test runs actually works before you rely on it in production pipelines.
Step 6: Handle retries and streaming without bypassing the cap
LLM calls fail transiently. If you wrap calls in tenacity, ensure the retry re-checks the budget before each attempt. Otherwise a retried call after a budget breach will still fire.
from tenacity import retry, stop_after_attempt
@retry(stop=stop_after_attempt(3))
def safe_chat(client, ledger, **kwargs):
if ledger.total() >= client.budget:
raise BudgetExceeded("no retries after breach")
return client.chat(ledger=ledger, **kwargs)If you use streaming responses, token usage is often only available in the final chunk. Accumulate there, not on the first byte. Also set a timeout-minutes on the CI job so a hung connection can’t sit idle while you assume the cap is working:
jobs:
llm-tests:
timeout-minutes: 10Step 7: Make the budget a required check
In GitHub, mark the llm-tests job as a required status check for the main branch. A breached budget fails the job, which blocks merge. This closes the loop: capping token spend in CI test runs is now a policy, not a suggestion.
Verify success
After implementing, validate with two runs:
- Normal run with
CI_TOKEN_BUDGET=50000. All integration tests pass, ledger file shows total under cap, exit code 0. - Attack run with
CI_TOKEN_BUDGET=10. The overrun test triggers, session exits with code 2, and CI log contains “Token budget breached”. No further API calls are made after the breach.
Check your gateway’s usage dashboard after the run to confirm billed tokens match the ledger. If you routed through a metered endpoint, the numbers should align to the token. If they don’t, your client is missing a code path (streaming, function calls, or embeddings) that also consumes tokens.
Capping token spend in CI test runs is not exotic. It is a few dozen lines of guard code plus CI config. Ship it before your next flaky test teaches you the lesson expensively.