Cap Your LLM Spend With a Hard Kill Switch in 30 Lines

Hey guys, Mr. Technology here.

Pre-counting tokens stops the obvious $11 PDF-paste incidents. It does not stop the agent loop that spins for two hours calling gpt-4o-mini 40,000 times because a JSON schema validator keeps returning 400 and the retry policy is "keep trying until the heat death of the universe." You need a hard kill switch — a process-level budget that aborts mid-stream when the meter passes the cap. Here is the build, 30 lines.

What we are building

A BudgetGuard singleton you wrap around your LLM call. It (1) pre-counts the prompt with tiktoken and rejects early if the estimate alone blows the remaining budget, (2) tracks cumulative spend across the process lifetime, (3) wraps streaming responses and stops the iterator the moment running cost exceeds the cap. There is no queue, no backoff, no "should we maybe ask first."

Step 1 — The guard

import tiktoken
from openai import OpenAI
_ENC = tiktoken.get_encoding("cl100k_base")
PRICING = {
    "gpt-4o":        (0.0025,  0.010),
    "gpt-4o-mini":   (0.00015, 0.0006),
    "claude-sonnet-5": (0.003,  0.015),
}
class BudgetExceeded(RuntimeError):
    pass
class BudgetGuard:
    def __init__(self, cap_usd: float = 5.0):
        self.cap = cap_usd
        self.spent = 0.0
        self.calls = 0
    def _cost(self, model, in_tok, out_tok):
        p_in, p_out = PRICING.get(model, (0, 0))
        return (in_tok / 1e6) * p_in + (out_tok / 1e6) * p_out
    def precheck(self, model, messages, max_out: int = 1024) -> int:
        in_tok = sum(
            len(_ENC.encode(m["content"])) for m in messages
        ) + 4 * len(messages)
        worst = self._cost(model, in_tok, max_out)
        if self.spent + worst > self.cap:
            raise BudgetExceeded(
                f"pre-check fail: ${self.spent:.4f} + worst-case ${worst:.4f}"
                f" > cap ${self.cap:.2f}"
            )
        return in_tok
    def charge(self, model, in_tok, out_tok) -> float:
        c = self._cost(model, in_tok, out_tok)
        self.spent += c
        self.calls += 1
        return c
# Module-level singleton — shared across all calls in this process
BUDGET = BudgetGuard(cap_usd=5.0)

The 4 * len(messages) overhead covers the role/content framing tokens. For Claude Sonnet 5 the real tokenizer overcounts by ~8 percent using cl100k_base; the worst estimate absorbs that buffer.

Step 2 — Non-streaming call

client = OpenAI()
def chat(messages: list[dict]) -> str:
    BUDGET.precheck("gpt-4o-mini", messages)
    r = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
    )
    BUDGET.charge(r.model, r.usage.prompt_tokens, r.usage.completion_tokens)
    return r.choices[0].message.content

That covers the obvious case. The retry loop above dies at precheck on call 51 once cumulative spend plus worst-case crosses $5. No API hit, no charge.

Step 3 — Kill mid-stream

This is the part pre-counting alone cannot do. A runaway agent that gets a 4,000-token completion instead of 50 burns budget on the back-end. Wrap the stream:

def stream_chat(messages: list[dict], model: str = "gpt-4o-mini") -> str:
    in_tok = BUDGET.precheck(model, messages)
    stream = client.chat.completions.create(
        model=model,
        messages=messages,
        stream=True,
        stream_options={"include_usage": True},
    )
    out_tok = 0
    collected = []
    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            collected.append(chunk.choices[0].delta.content)
        if chunk.usage:
            in_tok, out_tok = chunk.usage.prompt_tokens, chunk.usage.completion_tokens
        running_cost = BUDGET._cost(model, in_tok, out_tok)
        if BUDGET.spent + running_cost > BUDGET.cap:
            stream.close()          # stop upstream billing immediately
            BUDGET.charge(model, in_tok, out_tok)
            raise BudgetExceeded(
                f"mid-stream kill at ${BUDGET.spent:.4f} after {BUDGET.calls} calls"
            )
    BUDGET.charge(model, in_tok, out_tok)
    return "".join(collected)

stream.close() is the key line — it terminates the HTTP connection and stops the provider from billing you for the rest of the completion. On OpenAI that is immediate. On Anthropic it closes the SSE connection. Either way, the meter stops.

What you have

• BudgetGuard singleton shared across the process lifetime
• Pre-call gate that kills before any tokens are billed
• Mid-stream abort that catches pathological completions and runaway loops
• BUDGET.spent and BUDGET.calls readable at any time — pipe into your log decorator or Prometheus counter

Ship next

Add BUDGET.reset() on a daily cron so long-running services recycle the budget at midnight. Or read the cap from an env var: BUDGET = BudgetGuard(cap_usd=float(os.getenv("LLM_BUDGET_CAP_USD", "5.0"))) — your platform team dials it from $LLM_BUDGET_CAP_USD=2.0 in production without a redeploy.

Thirty lines. No new infrastructure. One $5 surprise instead of one $800` surprise.

— Mr. Technology

*Requires: tiktoken>=0.7.0, openai>=1.12. The singleton BUDGET lives in the module — pass it as a dependency in FastAPI via Depends(lambda: BUDGET) if you need per-request injection. For Anthropic, map usage.input_tokens / usage.output_tokens to charge() the same way; the pricing dict covers both SDKs.*