A hard dollar limit per job for IChatClient (OpenAI), with a DelegatingHandler and AsyncLocal

[domondi1]

With UseFunctionInvocation(), one GetResponseAsync can turn into several model calls, and a job that fans out into parallel requests has no single place that stops it at a dollar amount. Usage only comes back after each call, so a cap you check in your own code is always a batch behind.

Here's a pattern that needs nothing new from Microsoft.Extensions.AI: an AsyncLocal job scope, a DelegatingHandler that stamps the job's id and budget on every HTTP request the OpenAI client sends, and an OpenAI-compatible gateway that enforces the budget before forwarding. The handler half is useful on its own too, for example to send per-tenant or per-job attribution headers to any proxy.

Run the gateway next to the app (it uses your provider key):

pip install inferrail
export OPENAI_API_KEY=sk-...
inferrail serve --quickstart --app-mode   # http://127.0.0.1:8000

using System.ClientModel;
using System.ClientModel.Primitives;
using Microsoft.Extensions.AI;
using OpenAI;

var http = new HttpClient(new JobBudgetHandler { InnerHandler = new HttpClientHandler() });
var openai = new OpenAIClient(new ApiKeyCredential("unused"), new OpenAIClientOptions
{
    Endpoint = new Uri("http://127.0.0.1:8000/v1"),
    Transport = new HttpClientPipelineTransport(http),
});

IChatClient chat = openai.GetChatClient("gpt-4o-mini").AsIChatClient()
    .AsBuilder().UseFunctionInvocation().Build();
var options = new ChatOptions { MaxOutputTokens = 800, Tools = [/* your AIFunctions */] }; // the reservation uses MaxOutputTokens

static class JobBudget
{
    internal static readonly AsyncLocal<(string Id, string Usd)?> Current = new();
    public static IDisposable Begin(string jobId, string budgetUsd) { Current.Value = (jobId, budgetUsd); return new Scope(); }
    sealed class Scope : IDisposable { public void Dispose() => Current.Value = null; }
}

sealed class JobBudgetHandler : DelegatingHandler
{
    protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken ct)
    {
        if (JobBudget.Current.Value is { } job)
        {
            request.Headers.Add("X-Inferrail-Attribute-Work-Id", job.Id);
            request.Headers.Add("X-Inferrail-Budget-Usd", job.Usd);
        }
        return base.SendAsync(request, ct);
    }
}

Then wrap the job:

using (JobBudget.Begin("job-4812", "0.25"))
{
    var answers = await Task.WhenAll(topics.Select(t => chat.GetResponseAsync($"Research {t}", options)));
}

and check what it cost afterwards:

inferrail work job-4812

Every call with the same job id draws from one budget. Each call reserves its worst-case cost before it's sent, so parallel calls (and parallel tool loops) in one job can't overspend together. A call that doesn't fit is refused with HTTP 402 before it reaches OpenAI, which surfaces as ClientResultException with Status == 402. The default OpenAI retry policy doesn't retry 402, and the two headers aren't forwarded upstream.

The same OpenAIClient works unchanged for Agent Framework (openai.GetChatClient(model).AsIChatClient().AsAIAgent(...)) and Semantic Kernel (builder.AddOpenAIChatCompletion(model, openai), where the refusal shows up as HttpOperationException).

What I tested: .NET 10, Microsoft.Extensions.AI.OpenAI 10.10.1, Microsoft.Agents.AI.OpenAI 1.23.0, Semantic Kernel 1.80.1, inferrail 0.4.12, against a local stub upstream: six parallel tool-calling requests on one job id. With a tight budget, the calls that didn't fit were refused before reaching the upstream and the job's recorded total stayed under the cap. I haven't run it from .NET against the real OpenAI API yet, so reports are welcome.

Caveats:

• Chat completions only. Embeddings should keep their normal client, since the gateway doesn't serve them.
• AsyncLocal flows through await and Task.WhenAll, but not into work you hand to a separate queue or hosted service. Begin the scope where the job actually runs.
• The budget store is a SQLite file on the gateway host, so run one gateway.
• The model needs a price in Inferrail (gpt-4o-mini, gpt-4.1-mini and gpt-4.1 are built in; inferrail models lists the rest and you can add your own prices).

Setup details: guide. I maintain Inferrail (open source, Apache-2.0), so take the suggestion with that in mind. If there's a more idiomatic M.E.AI place for per-request headers than the transport, I'd like to hear it.