Skip to content
Datalium StrataDatalium Strata
These docs cover Strata 0.9.0-preview.1.It is a preview, and the API can still change before 1.0.

Add tool middleware

Tool middleware is code that runs around every tool call a ToolRegistry dispatches. In this guide you write four middleware: an audit log that records who called which tool and how long it took, an authorization check that refuses a call based on its caller, a per-call time limit and a concurrency limit. You register them on a registry, run some calls, and see what each kind of refusal or failure looks like to the caller.

Strata ships the hook, not the policies. It has no built-in audit, authorization, timeout or rate-limit middleware, so you write the ones you need.

Before you start

Section titled “Before you start”

Add Strata to your project as Installation shows, and read Define tools. This guide uses its tools, its ToolRegistry and its ToolCallContext.

Middleware runs only for calls that go through a registry. A call you make on a provider directly (provider.DispatchToolAsync in C#, tool.execute_with in Rust) runs no middleware.

How middleware runs

Section titled “How middleware runs”

A middleware receives the call and next, the rest of the chain. It can run code, call next to continue towards the tool, run more code with the result, and return it.

Implement IToolMiddleware from Datalium.Strata.Tools:

public delegate Task<string> ToolCallDelegate(ToolCall call, CancellationToken cancellationToken);


public interface IToolMiddleware
{
    Task<string> InvokeAsync(ToolCall call, ToolCallDelegate next, CancellationToken cancellationToken);


    Task OnRefusedAsync(ToolCall call, ToolNotFoundException refusal, CancellationToken cancellationToken) =>
        Task.CompletedTask;
}

OnRefusedAsync is optional. It is told of a call the registry refused before the chain ran.

A ToolCall record has the tool’s Name, its JSON Arguments, the call’s Context (a ToolCallContext) and the tool’s manifest entry as Descriptor.

A middleware can also:

  • Change the call before passing it on: new arguments with call with { Arguments = ... } or call.with_arguments(...), a new context with call with { Context = ... } or call.with_context(...). It can’t change the tool’s name: next always runs the tool the registry resolved.
  • Answer without calling next, as a cache does. The call counts as succeeded. Demo 26-tool-registry-and-middleware has a cache that does this.
  • Call next more than once, as a retry does. Each call runs the rest of the chain and the tool again.

1. Log each call

Section titled “1. Log each call”

This middleware records who called which tool, the outcome and the time it took. You register it first, so it is outermost and sees what every middleware inside it did.

using System.Diagnostics;
using Datalium.Strata.Tools;


public sealed class AuditMiddleware : IToolMiddleware
{
    public async Task<string> InvokeAsync(ToolCall call, ToolCallDelegate next, CancellationToken cancellationToken)
    {
        var clock = Stopwatch.StartNew();
        try
        {
            var result = await next(call, cancellationToken);
            Write(call, "succeeded", clock.Elapsed);
            return result;
        }
        catch (ToolCallDeniedException denied)
        {
            Write(call, $"denied ({denied.Reason})", clock.Elapsed);
            throw;
        }
        catch (Exception ex)
        {
            Write(call, $"failed ({ex.GetType().Name})", clock.Elapsed);
            throw;
        }
    }


    // Calls the registry refused before the chain ran: not found, or left out of a subset.
    public Task OnRefusedAsync(ToolCall call, ToolNotFoundException refusal, CancellationToken cancellationToken)
    {
        Write(call, refusal is ToolNotPermittedException ? "not permitted" : "not found", TimeSpan.Zero);
        return Task.CompletedTask;
    }


    // call.Arguments holds the JSON arguments. Log them only if they can't hold secrets.
    private static void Write(ToolCall call, string outcome, TimeSpan elapsed) =>
        Console.WriteLine(
            $"[audit] {call.Context.CallerId} called {call.Name}: {outcome} in {elapsed.TotalMilliseconds:0} ms");
}

ToolNotPermittedException derives from ToolNotFoundException, so test for it first.

Strata’s own ToolCallStarted and ToolCallCompleted run events never carry a call’s arguments or result, because they can hold secrets. A middleware has the arguments in hand, so an audit log can record them on purpose. If you do, redact secrets and personal data first, and keep the log somewhere as protected as the data.

2. Refuse a call

Section titled “2. Refuse a call”

To refuse a call, throw ToolCallDeniedException(toolName, reason) in C#, or return StrataError::ToolCallDenied { tool, reason } in Rust. The tool doesn’t run. This middleware lets only named callers use certain tools, based on the caller ID in the call’s context:

public sealed class CallerPolicy(IReadOnlyDictionary<string, string[]> allowedCallers) : IToolMiddleware
{
    public Task<string> InvokeAsync(ToolCall call, ToolCallDelegate next, CancellationToken cancellationToken)
    {
        if (allowedCallers.TryGetValue(call.Name, out var callers) && !callers.Contains(call.Context.CallerId))
        {
            throw new ToolCallDeniedException(call.Name, $"caller '{call.Context.CallerId}' may not call {call.Name}");
        }


        return next(call, cancellationToken);
    }
}

A subset already fixes which tools an agent can call at all. A middleware decides call by call, so it can use anything the call carries: the caller, the run’s metadata such as a tenant, or the arguments. The RefundPolicy in demo 26-tool-registry-and-middleware refuses refunds above a tenant’s limit.

A denial is a different error from not found and not permitted: the tool exists and the caller may use it, but this call was refused. The caller gets ToolCallDeniedException, with ToolName and Reason, or StrataError::ToolCallDenied. The call is reported Refused. When an agent made the call, the model gets {"error":"denied: <reason>"} and the loop goes on. See Build an agent.

3. Set a time limit

Section titled “3. Set a time limit”

A tool call has no time limit of its own, and Strata ships no timeout middleware (Known limitations). This middleware gives every call one:

public sealed class TimeoutMiddleware(TimeSpan limit) : IToolMiddleware
{
    public async Task<string> InvokeAsync(ToolCall call, ToolCallDelegate next, CancellationToken cancellationToken)
    {
        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeout.CancelAfter(limit);
        try
        {
            // The tool gets the linked token. WaitAsync stops waiting even if the tool ignores it.
            return await next(call, timeout.Token).WaitAsync(timeout.Token);
        }
        catch (OperationCanceledException) when (timeout.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
        {
            throw new TimeoutException($"{call.Name} took longer than {limit.TotalSeconds:0} s");
        }
    }
}

The when filter tells the two tokens apart. If the caller’s token fired, the OperationCanceledException passes through and the call counts as cancelled. If only the time limit fired, the call fails with a TimeoutException. A tool that ignores its token keeps running in the background after the middleware stops waiting.

A time-out is reported Failed, not Cancelled, because the run wasn’t cancelled.

To try it, add a slow tool. inventory.recount takes 10 seconds and honors its token, as Define tools recommends:

// In InventoryTools, next to inventory.lookup.
[StrataTool(Name = "inventory.recount", Description = "Counts the units of a SKU on the shelf.")]
public static async Task<int> RecountAsync(string sku, CancellationToken cancellationToken)
{
    // Stands in for a slow call to the warehouse system.
    await Task.Delay(TimeSpan.FromSeconds(10), cancellationToken);
    return Stock.GetValueOrDefault(sku);
}

A tool that ignores its token times out after the same second. In C# it then finishes in the background; in Rust it stops.

4. Limit concurrent calls

Section titled “4. Limit concurrent calls”

This middleware lets at most a fixed number of calls through at once. The others wait for a slot.

public sealed class ConcurrencyLimit(int maxCalls) : IToolMiddleware
{
    private readonly SemaphoreSlim _slots = new(maxCalls, maxCalls);


    public async Task<string> InvokeAsync(ToolCall call, ToolCallDelegate next, CancellationToken cancellationToken)
    {
        await _slots.WaitAsync(cancellationToken);
        try
        {
            return await next(call, cancellationToken);
        }
        finally
        {
            _slots.Release();
        }
    }
}

The limit counts the calls through one middleware instance in one process. A limit shared by several processes, or a rate limit such as calls per minute, needs state outside the process, which your middleware calls.

5. Register the middleware

Section titled “5. Register the middleware”

Add middleware with Use in C# or use_middleware in Rust. Both return the registry, so you can chain them after Register or register.

var tools = new ToolRegistry()
    .Register(new InventoryTools())
    .Register(new BillingTools())
    .Use(new AuditMiddleware())                       // outermost: sees every call first
    .Use(new CallerPolicy(new Dictionary<string, string[]>
    {
        ["billing.refund"] = ["billing-agent"],
    }))
    .Use(new TimeoutMiddleware(TimeSpan.FromSeconds(1)))
    .Use(new ConcurrencyLimit(maxCalls: 4));          // innermost: next to the tool

tools.Middleware lists the middleware in registration order.

Middleware runs in registration order. The first one you add sees the call first and the result last. The order above is deliberate:

  • The audit log is outermost, so it records denials and time-outs from the middleware inside it.
  • The caller check comes before the time limit, so a refused call doesn’t start a timer.
  • The time limit wraps the concurrency limit, so time spent waiting for a slot counts towards the limit. In C#, the slot is released only when the tool really finishes, even if the time limit gave up waiting first.

6. Run some calls

Section titled “6. Run some calls”

Make five calls as two callers. Outside a graph you build the context yourself. Inside a node, start from ToolCallContext.From(ctx) as Define tools shows.

var support = ToolCallContext.None with { CallerId = "support-agent" };
var billing = ToolCallContext.None with { CallerId = "billing-agent" };


Console.WriteLine(await AskAsync(tools, "inventory.lookup", """{"sku":"AB-12"}""", support));
Console.WriteLine(await AskAsync(tools, "billing.refund", """{"orderId":"A-1001","amount":40}""", support));
Console.WriteLine(await AskAsync(tools, "billing.refund", """{"orderId":"A-1001","amount":40}""", billing));
Console.WriteLine(await AskAsync(tools, "inventory.recount", """{"sku":"AB-12"}""", support));
Console.WriteLine(await AskAsync(tools, "inventory.restock", """{"sku":"AB-12"}""", support));


static async Task<string> AskAsync(ToolRegistry tools, string name, string arguments, ToolCallContext context)
{
    try
    {
        var json = await tools.DispatchToolAsync(name, arguments, context);
        return JsonDocument.Parse(json).RootElement.GetProperty("result").ToString();
    }
    catch (ToolCallDeniedException denied)
    {
        return $"denied: {denied.Reason}";
    }
    catch (ToolNotFoundException refused)
    {
        return refused.Message;
    }
    catch (TimeoutException timedOut)
    {
        return $"failed: {timedOut.Message}";
    }
}
[audit] support-agent called inventory.lookup: succeeded in 20 ms
7
[audit] support-agent called billing.refund: denied (caller 'support-agent' may not call billing.refund) in 0 ms
denied: caller 'support-agent' may not call billing.refund
[audit] billing-agent called billing.refund: succeeded in 6 ms
refunded 40.00 on A-1001
[audit] support-agent called inventory.recount: failed (TimeoutException) in 1012 ms
failed: inventory.recount took longer than 1 s
[audit] support-agent called inventory.restock: not found in 0 ms
Tool 'inventory.restock' was not found.

The durations differ from run to run. The call to inventory.restock never entered the chain: the registry refused it, and only the audit log’s on-refused hook saw it.

How middleware fits with the registry

Section titled “How middleware fits with the registry”

Refused calls

Section titled “Refused calls”

A call to a tool the registry doesn’t have, or that a subset left out, runs no middleware. The registry calls each middleware’s OnRefusedAsync or on_refused hook instead, in registration order, and then returns the refusal. A hook can’t turn a refusal into a result. If a hook throws or returns an error, that error replaces the refusal and the hooks after it don’t run, so an audit log that can’t record fails the call.

A denial from a middleware doesn’t reach the on-refused hooks. The middleware outside it see the denial as the error next returned, which is how the audit log above records it.

Subsets

Section titled “Subsets”

A subset copies the registry’s middleware when you make it, as it copies the tools. Add middleware first and make subsets after, and every agent’s calls go through the same chain. Middleware you add to the parent later doesn’t reach an existing subset. Middleware you add to a subset runs inside the parent’s and only for that subset.

In C#, when you register one registry into another, the inner registry’s tools keep its middleware. A call runs the outer chain, then the inner one, then the tool.

Cancellation

Section titled “Cancellation”

If the token has already fired when you make the call, no middleware runs and the call is cancelled. If the run’s token fires while the chain runs, the call counts as cancelled: in C# when the error is an OperationCanceledException and the registry’s token was cancelled, in Rust for any error except ToolCallDenied, which stays a denial.

A middleware can give the tool its own token: pass another token to next in C#, or call with_cancellation on the context in Rust. The time limit above does this in C#.

Outcomes

Section titled “Outcomes”
What the chain did C# Rust Reported
Returned a result, with or without running the tool the result Ok(result) Succeeded
A middleware refused ToolCallDeniedException StrataError::ToolCallDenied Refused, error type denied
A middleware failed (a time-out, a policy service that is down) its own exception StrataError::ToolCallFailed, or any other error Failed
The run’s token fired OperationCanceledException StrataError::ToolCancelled Cancelled

Middleware runs inside the call’s strata.tool span, so the span’s duration, the strata.tool.duration metric and the ToolCallCompleted event include middleware time. A middleware can tag the span: in C#, Activity.Current is the tool span while it runs. A call answered without running the tool gets the span attribute strata.tool.short_circuited = true.

Test a middleware

Section titled “Test a middleware”

In C#, ToolCall has a public constructor and next is a delegate, so you can test a middleware alone: new ToolCall("billing.refund", "{}", ToolCallContext.None with { CallerId = "x" }) and (call, ct) => Task.FromResult("{}"). In Rust you can’t build a Next yourself. Register the middleware on a ToolRegistry with a stub tool and dispatch through it.

Section titled “Related”
  • Define tools: tools, the registry, ToolCallContext, errors and cancellation.
  • Build an agent: what the model is told when a call is refused or fails.
  • Known limitations: tool calls have no built-in time limit.
  • Demo 26-tool-registry-and-middleware (C#, Rust): an audit log, a refund policy, an argument rewrite and a cache.
  • design/23-Tool-Middleware.md: the design of tool middleware in both engines.