.NET: Add a ModeProvider for managing agent modes (#5247)

* Add a ModeProvider for managing agent modes * Fix typo * Fix typo * Fix typo * Address PR comments
2026-06-16 21:04:09 +08:00 · 2026-04-15 09:17:59 +01:00
parent 673f3d9214
commit b4c853ec1b
4 changed files with 513 additions and 1 deletions
@@ -70,7 +70,7 @@ internal static partial class AgentJsonUtilities
    [JsonSerializable(typeof(TextSearchProvider.TextSearchProviderState))]
    [JsonSerializable(typeof(ChatHistoryMemoryProvider.State))]

-    // Harness types
+    // TodoProvider types
    [JsonSerializable(typeof(TodoState))]
    [JsonSerializable(typeof(TodoItem))]
    [JsonSerializable(typeof(TodoItemInput))]
@@ -78,6 +78,9 @@ internal static partial class AgentJsonUtilities
    [JsonSerializable(typeof(List<TodoItem>), TypeInfoPropertyName = "TodoItemList")]
    [JsonSerializable(typeof(List<TodoItemInput>), TypeInfoPropertyName = "TodoItemInputList")]

+    // AgentModeProvider types
+    [JsonSerializable(typeof(AgentModeState))]
+
    [ExcludeFromCodeCoverage]
    internal sealed partial class JsonContext : JsonSerializerContext;
 }
@@ -0,0 +1,148 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Microsoft.Shared.DiagnosticIds;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// An <see cref="AIContextProvider"/> that tracks the agent's operating mode (e.g., "plan" or "execute")
+/// in the session state and provides tools for querying and switching modes.
+/// </summary>
+/// <remarks>
+/// <para>
+/// The <see cref="AgentModeProvider"/> enables agents to operate in distinct modes during long-running
+/// complex tasks. The current mode is persisted in the session's <see cref="AgentSessionStateBag"/>
+/// and is included in the instructions provided to the agent on each invocation.
+/// </para>
+/// <para>
+/// This provider exposes the following tools to the agent:
+/// <list type="bullet">
+/// <item><description><c>SetMode</c> — Switch the agent's operating mode.</description></item>
+/// <item><description><c>GetMode</c> — Retrieve the agent's current operating mode.</description></item>
+/// </list>
+/// </para>
+/// <para>
+/// Public helper methods <see cref="GetMode"/> and <see cref="SetMode"/> allow external code
+/// to programmatically read and change the mode.
+/// </para>
+/// </remarks>
+[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
+public sealed class AgentModeProvider : AIContextProvider
+{
+    /// <summary>
+    /// The "plan" mode, indicating the agent is planning work.
+    /// </summary>
+    public const string PlanMode = "plan";
+
+    /// <summary>
+    /// The "execute" mode, indicating the agent is executing work.
+    /// </summary>
+    public const string ExecuteMode = "execute";
+
+    private readonly ProviderSessionState<AgentModeState> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AgentModeProvider"/> class.
+    /// </summary>
+    public AgentModeProvider()
+    {
+        this._sessionState = new ProviderSessionState<AgentModeState>(
+            _ => new AgentModeState(),
+            this.GetType().Name,
+            AgentJsonUtilities.DefaultOptions);
+    }
+
+    /// <inheritdoc />
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
+
+    /// <summary>
+    /// Gets the current operating mode from the session state.
+    /// </summary>
+    /// <param name="session">The agent session to read the mode from.</param>
+    /// <returns>The current mode string.</returns>
+    public string GetMode(AgentSession? session)
+    {
+        return this._sessionState.GetOrInitializeState(session).CurrentMode;
+    }
+
+    /// <summary>
+    /// Sets the operating mode in the session state.
+    /// </summary>
+    /// <param name="session">The agent session to update the mode in.</param>
+    /// <param name="mode">The new mode to set.</param>
+    public void SetMode(AgentSession? session, string mode)
+    {
+        if (mode != PlanMode && mode != ExecuteMode)
+        {
+            throw new ArgumentException($"Invalid mode: {mode}. Supported modes are \"{PlanMode}\" and \"{ExecuteMode}\".", nameof(mode));
+        }
+
+        AgentModeState state = this._sessionState.GetOrInitializeState(session);
+        state.CurrentMode = mode;
+        this._sessionState.SaveState(session, state);
+    }
+
+    /// <inheritdoc />
+    protected override ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    {
+        AgentModeState state = this._sessionState.GetOrInitializeState(context.Session);
+
+        string instructions = $"""
+            You are currently operating in "{state.CurrentMode}" mode.
+            Available modes:
+            - "plan": Use this mode when analyzing requirements, breaking down tasks, and creating plans.
+            - "execute": Use this mode when implementing changes, writing code, and carrying out planned work.
+            Use the SetMode tool to switch between modes as your work progresses. Only use SetMode if the user explicitly instructs you to change modes.
+            Use the GetMode tool to check your current operating mode.
+            """;
+
+        return new ValueTask<AIContext>(new AIContext
+        {
+            Instructions = instructions,
+            Tools = this.CreateTools(state, context.Session),
+        });
+    }
+
+    private AITool[] CreateTools(AgentModeState state, AgentSession? session)
+    {
+        var serializerOptions = AgentJsonUtilities.DefaultOptions;
+
+        return
+        [
+            AIFunctionFactory.Create(
+                (string mode) =>
+                {
+                    if (mode != PlanMode && mode != ExecuteMode)
+                    {
+                        throw new ArgumentException($"Invalid mode: {mode}. Supported modes are \"{PlanMode}\" and \"{ExecuteMode}\".", nameof(mode));
+                    }
+
+                    state.CurrentMode = mode;
+                    this._sessionState.SaveState(session, state);
+                    return $"Mode changed to \"{mode}\".";
+                },
+                new AIFunctionFactoryOptions
+                {
+                    Name = "SetMode",
+                    Description = "Switch the agent's operating mode. Supported modes: \"plan\" and \"execute\".",
+                    SerializerOptions = serializerOptions,
+                }),
+
+            AIFunctionFactory.Create(
+                () => state.CurrentMode,
+                new AIFunctionFactoryOptions
+                {
+                    Name = "GetMode",
+                    Description = "Get the agent's current operating mode.",
+                    SerializerOptions = serializerOptions,
+                }),
+        ];
+    }
+}
@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Diagnostics.CodeAnalysis;
+using System.Text.Json.Serialization;
+using Microsoft.Shared.DiagnosticIds;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Represents the state of the agent's operating mode, stored in the session's <see cref="AgentSessionStateBag"/>.
+/// </summary>
+[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
+internal sealed class AgentModeState
+{
+    /// <summary>
+    /// Gets or sets the current operating mode of the agent.
+    /// </summary>
+    [JsonPropertyName("currentMode")]
+    public string CurrentMode { get; set; } = AgentModeProvider.PlanMode;
+}
@@ -0,0 +1,341 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text.Json;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Moq;
+
+namespace Microsoft.Agents.AI.UnitTests;
+
+/// <summary>
+/// Unit tests for the <see cref="AgentModeProvider"/> class.
+/// </summary>
+public class AgentModeProviderTests
+{
+    #region ProvideAIContextAsync Tests
+
+    /// <summary>
+    /// Verify that the provider returns tools and instructions.
+    /// </summary>
+    [Fact]
+    public async Task ProvideAIContextAsync_ReturnsToolsAndInstructionsAsync()
+    {
+        // Arrange
+        var provider = new AgentModeProvider();
+        var agent = new Mock<AIAgent>().Object;
+        var session = new ChatClientAgentSession();
+#pragma warning disable MAAI001
+        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
+#pragma warning restore MAAI001
+
+        // Act
+        AIContext result = await provider.InvokingAsync(context);
+
+        // Assert
+        Assert.NotNull(result.Instructions);
+        Assert.NotNull(result.Tools);
+        Assert.Equal(2, result.Tools!.Count());
+    }
+
+    /// <summary>
+    /// Verify that the instructions include the current mode.
+    /// </summary>
+    [Fact]
+    public async Task ProvideAIContextAsync_InstructionsIncludeCurrentModeAsync()
+    {
+        // Arrange
+        var provider = new AgentModeProvider();
+        var agent = new Mock<AIAgent>().Object;
+        var session = new ChatClientAgentSession();
+#pragma warning disable MAAI001
+        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
+#pragma warning restore MAAI001
+
+        // Act
+        AIContext result = await provider.InvokingAsync(context);
+
+        // Assert
+        Assert.Contains("plan", result.Instructions);
+    }
+
+    #endregion
+
+    #region SetMode Tool Tests
+
+    /// <summary>
+    /// Verify that SetMode changes the mode.
+    /// </summary>
+    [Fact]
+    public async Task SetMode_ChangesModeAsync()
+    {
+        // Arrange
+        var (tools, state) = await CreateToolsWithStateAsync();
+        AIFunction setMode = GetTool(tools, "SetMode");
+
+        // Act
+        await setMode.InvokeAsync(new AIFunctionArguments() { ["mode"] = "execute" });
+
+        // Assert
+        Assert.Equal("execute", state.CurrentMode);
+    }
+
+    /// <summary>
+    /// Verify that SetMode returns a confirmation message.
+    /// </summary>
+    [Fact]
+    public async Task SetMode_ReturnsConfirmationAsync()
+    {
+        // Arrange
+        var (tools, _) = await CreateToolsWithStateAsync();
+        AIFunction setMode = GetTool(tools, "SetMode");
+
+        // Act
+        object? result = await setMode.InvokeAsync(new AIFunctionArguments() { ["mode"] = "execute" });
+
+        // Assert
+        Assert.Equal("Mode changed to \"execute\".", GetStringResult(result));
+    }
+
+    /// <summary>
+    /// Verify that SetMode with an unsupported value throws and does not persist the mode.
+    /// </summary>
+    [Fact]
+    public async Task SetMode_InvalidMode_ThrowsAsync()
+    {
+        // Arrange
+        var (tools, provider, session) = await CreateToolsWithProviderAndSessionAsync();
+        AIFunction setMode = GetTool(tools, "SetMode");
+        AIFunction getMode = GetTool(tools, "GetMode");
+
+        // Act & Assert
+        await Assert.ThrowsAsync<ArgumentException>(async () =>
+            await setMode.InvokeAsync(new AIFunctionArguments() { ["mode"] = "foo" }));
+
+        // Verify mode was not changed from default
+        object? currentMode = await getMode.InvokeAsync(new AIFunctionArguments());
+        Assert.Equal(AgentModeProvider.PlanMode, GetStringResult(currentMode));
+    }
+
+    #endregion
+
+    #region GetMode Tool Tests
+
+    /// <summary>
+    /// Verify that GetMode returns the default mode.
+    /// </summary>
+    [Fact]
+    public async Task GetMode_ReturnsDefaultModeAsync()
+    {
+        // Arrange
+        var (tools, _) = await CreateToolsWithStateAsync();
+        AIFunction getMode = GetTool(tools, "GetMode");
+
+        // Act
+        object? result = await getMode.InvokeAsync(new AIFunctionArguments());
+
+        // Assert
+        Assert.Equal("plan", GetStringResult(result));
+    }
+
+    /// <summary>
+    /// Verify that GetMode returns the mode after SetMode.
+    /// </summary>
+    [Fact]
+    public async Task GetMode_ReturnsUpdatedModeAfterSetAsync()
+    {
+        // Arrange
+        var (tools, _) = await CreateToolsWithStateAsync();
+        AIFunction setMode = GetTool(tools, "SetMode");
+        AIFunction getMode = GetTool(tools, "GetMode");
+
+        // Act
+        await setMode.InvokeAsync(new AIFunctionArguments() { ["mode"] = "execute" });
+        object? result = await getMode.InvokeAsync(new AIFunctionArguments());
+
+        // Assert
+        Assert.Equal("execute", GetStringResult(result));
+    }
+
+    #endregion
+
+    #region Public Helper Method Tests
+
+    /// <summary>
+    /// Verify that the public GetMode helper returns the default mode.
+    /// </summary>
+    [Fact]
+    public void PublicGetMode_ReturnsDefaultMode()
+    {
+        // Arrange
+        var provider = new AgentModeProvider();
+        var session = new ChatClientAgentSession();
+
+        // Act
+        string mode = provider.GetMode(session);
+
+        // Assert
+        Assert.Equal(AgentModeProvider.PlanMode, mode);
+    }
+
+    /// <summary>
+    /// Verify that the public SetMode helper changes the mode.
+    /// </summary>
+    [Fact]
+    public void PublicSetMode_ChangesMode()
+    {
+        // Arrange
+        var provider = new AgentModeProvider();
+        var session = new ChatClientAgentSession();
+
+        // Act
+        provider.SetMode(session, AgentModeProvider.ExecuteMode);
+        string mode = provider.GetMode(session);
+
+        // Assert
+        Assert.Equal(AgentModeProvider.ExecuteMode, mode);
+    }
+
+    /// <summary>
+    /// Verify that the public SetMode helper throws for an unsupported value and does not persist the mode.
+    /// </summary>
+    [Fact]
+    public void PublicSetMode_InvalidMode_Throws()
+    {
+        // Arrange
+        var provider = new AgentModeProvider();
+        var session = new ChatClientAgentSession();
+
+        // Act & Assert
+        Assert.Throws<ArgumentException>(() => provider.SetMode(session, "foo"));
+
+        // Verify mode was not changed from default
+        string mode = provider.GetMode(session);
+        Assert.Equal(AgentModeProvider.PlanMode, mode);
+    }
+
+    /// <summary>
+    /// Verify that public helper changes are reflected in tool results.
+    /// </summary>
+    [Fact]
+    public async Task PublicSetMode_ReflectedInToolResultsAsync()
+    {
+        // Arrange
+        var provider = new AgentModeProvider();
+        var agent = new Mock<AIAgent>().Object;
+        var session = new ChatClientAgentSession();
+
+        // Set mode via public helper
+        provider.SetMode(session, AgentModeProvider.ExecuteMode);
+
+#pragma warning disable MAAI001
+        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
+#pragma warning restore MAAI001
+
+        // Act
+        AIContext result = await provider.InvokingAsync(context);
+        AIFunction getMode = GetTool(result.Tools!, "GetMode");
+        object? modeResult = await getMode.InvokeAsync(new AIFunctionArguments());
+
+        // Assert
+        Assert.Equal("execute", GetStringResult(modeResult));
+        Assert.Contains("execute", result.Instructions);
+    }
+
+    #endregion
+
+    #region State Persistence Tests
+
+    /// <summary>
+    /// Verify that state persists across invocations.
+    /// </summary>
+    [Fact]
+    public async Task State_PersistsAcrossInvocationsAsync()
+    {
+        // Arrange
+        var provider = new AgentModeProvider();
+        var agent = new Mock<AIAgent>().Object;
+        var session = new ChatClientAgentSession();
+#pragma warning disable MAAI001
+        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
+#pragma warning restore MAAI001
+
+        // Act — first invocation changes mode
+        AIContext result1 = await provider.InvokingAsync(context);
+        AIFunction setMode = GetTool(result1.Tools!, "SetMode");
+        await setMode.InvokeAsync(new AIFunctionArguments() { ["mode"] = "execute" });
+
+        // Second invocation should see the updated mode
+        AIContext result2 = await provider.InvokingAsync(context);
+        AIFunction getMode = GetTool(result2.Tools!, "GetMode");
+        object? modeResult = await getMode.InvokeAsync(new AIFunctionArguments());
+
+        // Assert
+        Assert.Equal("execute", GetStringResult(modeResult));
+        Assert.Contains("execute", result2.Instructions);
+    }
+
+    #endregion
+
+    #region Constants Tests
+
+    /// <summary>
+    /// Verify that mode constants have expected values.
+    /// </summary>
+    [Fact]
+    public void ModeConstants_HaveExpectedValues()
+    {
+        // Assert
+        Assert.Equal("plan", AgentModeProvider.PlanMode);
+        Assert.Equal("execute", AgentModeProvider.ExecuteMode);
+    }
+
+    #endregion
+
+    #region Helper Methods
+
+    private static async Task<(IEnumerable<AITool> Tools, AgentModeState State)> CreateToolsWithStateAsync()
+    {
+        var provider = new AgentModeProvider();
+        var agent = new Mock<AIAgent>().Object;
+        var session = new ChatClientAgentSession();
+#pragma warning disable MAAI001
+        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
+#pragma warning restore MAAI001
+
+        AIContext result = await provider.InvokingAsync(context);
+
+        // Retrieve the state from the session to verify mutations
+        session.StateBag.TryGetValue<AgentModeState>("AgentModeProvider", out var state, AgentJsonUtilities.DefaultOptions);
+
+        return (result.Tools!, state!);
+    }
+
+    private static async Task<(IEnumerable<AITool> Tools, AgentModeProvider Provider, AgentSession Session)> CreateToolsWithProviderAndSessionAsync()
+    {
+        var provider = new AgentModeProvider();
+        var agent = new Mock<AIAgent>().Object;
+        var session = new ChatClientAgentSession();
+#pragma warning disable MAAI001
+        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
+#pragma warning restore MAAI001
+
+        AIContext result = await provider.InvokingAsync(context);
+        return (result.Tools!, provider, session);
+    }
+
+    private static AIFunction GetTool(IEnumerable<AITool> tools, string name)
+    {
+        return (AIFunction)tools.First(t => t is AIFunction f && f.Name == name);
+    }
+
+    private static string GetStringResult(object? result)
+    {
+        var element = Assert.IsType<JsonElement>(result);
+        return element.GetString()!;
+    }
+
+    #endregion
+}