Jaliya's Blog: Microsoft Agent Frameworks' Agent Middleware with .NET

In this post, let's have a look at Microsoft Agent Framework's Agent Middleware with .NET. In my previous post, we looked at getting started with Microsoft Agent Framework in .NET by building an agent with tools. In this post, let's see how we can use middleware to intercept and inspect agent runs, function calls, and chat client calls.

Microsoft Agent Framework supports three types of middleware.

Agent Run Middleware: Allows interception of all agent runs, so that input and output can be inspected and/or modified as needed.
Function Calling Middleware: Allows interception of all function calls executed by the agent, so that input and output can be inspected and modified as needed.
Chat Client Middleware: Allows interception of calls to an IChatClient implementation, where an agent is using IChatClient for inference calls, for example, when using ChatClientAgent.

Let's see how we can set up each of these middleware types.

The middleware is registered using the builder pattern. Agent Run Middleware and Function Calling Middleware are applied to the AIAgent via AsBuilder(). Chat Client Middleware is applied to the IChatClient instance via the clientFactory parameter.

AIAgent agent = new AIProjectClient(new Uri(endpoint), credential)
    .AsAIAgent(
        model: deploymentName,
        name: "weekend-planner",
        instructions: """
            You help users plan their weekends and choose the best activities for the given weather.
            If an activity would be unpleasant in weather, don't suggest it.
            Include date of the weekend in response.
            """,
        tools: [
            AIFunctionFactory.Create(GetWeather),
            AIFunctionFactory.Create(GetActivities),
            AIFunctionFactory.Create(GetCurrentDate)
        ],
        clientFactory: (chatClient) => chatClient
            .AsBuilder()
            // Chat Client Middleware
            .Use(getResponseFunc: CustomChatClientMiddleware, getStreamingResponseFunc: CustomChatClientStreamlingMiddleware)
            .Build())
    .AsBuilder()
    // Agent Run Middleware
    .Use(runFunc: CustomAgentRunMiddleware, runStreamingFunc: CustomAgentRunStreamingMiddleware)
    // Function Calling Middleware
    .Use(CustomFunctionCallingMiddleware)
    .Build();

Agent Run Middleware

The Agent Run Middleware intercepts all agent runs. Here we are logging the message count before and after the agent run. The middleware calls agent.RunAsync() to invoke the next middleware in the chain.

async Task<AgentResponse> CustomAgentRunMiddleware(IEnumerable<ChatMessage> messages,
    AgentSession? session,
    AgentRunOptions? options,
    AIAgent agent,
    CancellationToken cancellationToken)
{
    Console.WriteLine($"[AgentRun] Message Count: '{messages.Count()}'.");

    AgentResponse response = await agent.RunAsync(messages, session, options, cancellationToken)
        .ConfigureAwait(false);

    Console.WriteLine($"[AgentRun] Response Message Count: '{response.Messages.Count}'.");

    return response;
}

It also supports a streaming variant using IAsyncEnumerable<AgentResponseUpdate>.

async IAsyncEnumerable<AgentResponseUpdate> CustomAgentRunStreamingMiddleware(IEnumerable<ChatMessage> messages,
    AgentSession? session,
    AgentRunOptions? options,
    AIAgent agent,
    [EnumeratorCancellation] CancellationToken cancellationToken)
{
    Console.WriteLine($"[AgentRunStreaming] Message Count: '{messages.Count()}'.");

    List<AgentResponseUpdate> updates = [];
    await foreach (AgentResponseUpdate update in agent.RunStreamingAsync(messages, session, options, cancellationToken))
    {
        updates.Add(update);
        yield return update;
    }

    Console.WriteLine($"[AgentRunStreaming] Response Message Count: '{updates.ToAgentResponse().Messages.Count}'.");
}

Function Calling Middleware

The Function Calling Middleware intercepts all function calls executed by the agent. Here we are logging the function name being called and the result returned. The middleware uses the next delegate to invoke the next middleware in the chain.

async ValueTask<object?> CustomFunctionCallingMiddleware(
    AIAgent agent,
    FunctionInvocationContext context,
    Func<FunctionInvocationContext, CancellationToken, ValueTask<object?>> next,
    CancellationToken cancellationToken)
{
    Console.WriteLine($"      [FunctionCall] Calling: '{context!.Function.Name}'.");

    object? result = await next(context, cancellationToken);

    Console.WriteLine($"      [FunctionCall] Result: '{context!.Function.Name}' = '<OMITTED>'.");

    return result;
}

Chat Client Middleware

The Chat Client Middleware intercepts calls to IChatClient . This is useful when you want to inspect the raw messages being sent to and received from the underlying LLM. Here we are logging all messages with their content types, distinguishing between TextContent, FunctionCallContent, and FunctionResultContent.

async Task<ChatResponse> CustomChatClientMiddleware(IEnumerable<ChatMessage> messages,
    ChatOptions? options,
    IChatClient client,
    CancellationToken token)
{
    LogChatClientMessages("Messages", messages);

    ChatResponse response = await client.GetResponseAsync(messages, options, token)
        .ConfigureAwait(false);

    LogChatClientMessages("Response", response.Messages);

    return response;
}

And the streaming variant.

async IAsyncEnumerable<ChatResponseUpdate> CustomChatClientStreamlingMiddleware(IEnumerable<ChatMessage> messages,
    ChatOptions? options,
    IChatClient client,
    [EnumeratorCancellation] CancellationToken token)
{
    LogChatClientMessages("Messages", messages);

    List<ChatResponseUpdate> updates = [];
    await foreach (ChatResponseUpdate update in client.GetStreamingResponseAsync(messages, options, token))
    {
        updates.Add(update);
        yield return update;
    }

    LogChatClientMessages("Response", updates.ToChatResponse().Messages);
}

The helper method LogChatClientMessages logs the messages with their content types.

void LogChatClientMessages(string label, IEnumerable<ChatMessage> messages)
{
    Console.WriteLine();
    Console.WriteLine($"   [ChatClient] {label}:");
    foreach (ChatMessage message in messages)
    {
        foreach (AIContent content in message.Contents)
        {
            string detail = content switch
            {
                TextContent text => text.Text,
                FunctionCallContent fc => $"Call: {fc.Name}({string.Join(", ", fc.Arguments?.Select(a => $"{a.Key}={a.Value}") ?? [])})",
                FunctionResultContent fr => $"Result: {fr.CallId} = '<OMITTED>'",
                _ => content.GetType().Name
            };
            Console.WriteLine($"      [{message.Role}] {detail}");
        }
    }
}