Merge branch 'main' into feature-xunit3-mtp-upgrade

dotnet/Directory.Packages.props

@@ -19,8 +19,8 @@
     <PackageVersion Include="Aspire.Microsoft.Azure.Cosmos" Version="$(AspireAppHostSdkVersion)" />
     <PackageVersion Include="CommunityToolkit.Aspire.OllamaSharp" Version="13.0.0" />
     <!-- Azure.* -->
-    <PackageVersion Include="Azure.AI.Projects" Version="1.2.0-beta.5" />
-    <PackageVersion Include="Azure.AI.Projects.OpenAI" Version="1.0.0-beta.5" />
+    <PackageVersion Include="Azure.AI.Projects" Version="2.0.0-beta.1" />
+    <PackageVersion Include="Azure.AI.Projects.OpenAI" Version="2.0.0-beta.1" />
     <PackageVersion Include="Azure.AI.Agents.Persistent" Version="1.2.0-beta.8" />
     <PackageVersion Include="Azure.AI.OpenAI" Version="2.8.0-beta.1" />
     <PackageVersion Include="Azure.Identity" Version="1.17.1" />
@@ -35,7 +35,7 @@
     <!-- System.* -->
     <PackageVersion Include="Microsoft.Bcl.AsyncInterfaces" Version="10.0.3" />
     <PackageVersion Include="Microsoft.Bcl.HashCode" Version="6.0.0" />
-    <PackageVersion Include="System.ClientModel" Version="1.8.1" />
+    <PackageVersion Include="System.ClientModel" Version="1.9.0" />
     <PackageVersion Include="System.CodeDom" Version="10.0.0" />
     <PackageVersion Include="System.Collections.Immutable" Version="10.0.1" />
     <PackageVersion Include="System.CommandLine" Version="2.0.0-rc.2.25502.107" />
@@ -94,7 +94,7 @@
     <PackageVersion Include="Microsoft.SemanticKernel.Agents.AzureAI" Version="1.67.0-preview" />
     <PackageVersion Include="Microsoft.SemanticKernel.Plugins.OpenApi" Version="1.67.0" />
     <!-- Agent SDKs -->
-    <PackageVersion Include="GitHub.Copilot.SDK" Version="0.1.23" />
+    <PackageVersion Include="GitHub.Copilot.SDK" Version="0.1.29" />
     <PackageVersion Include="Microsoft.Agents.CopilotStudio.Client" Version="1.3.171-beta" />
     <!-- M365 Agents SDK -->
     <PackageVersion Include="AdaptiveCards" Version="3.1.0" />
@@ -185,4 +185,4 @@
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
     </PackageReference>
   </ItemGroup>
-</Project>
+</Project>

dotnet/samples/01-get-started/04_memory/Program.cs

@@ -89,6 +89,7 @@ namespace SampleApp
     internal sealed class UserInfoMemory : AIContextProvider
     {
         private readonly ProviderSessionState<UserInfo> _sessionState;
+        private IReadOnlyList<string>? _stateKeys;
         private readonly IChatClient _chatClient;
 
         public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? stateInitializer = null)
@@ -99,7 +100,7 @@ namespace SampleApp
             this._chatClient = chatClient;
         }
 
-        public override string StateKey => this._sessionState.StateKey;
+        public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
         public UserInfo GetUserInfo(AgentSession session)
             => this._sessionState.GetOrInitializeState(session);

dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs

@@ -79,6 +79,7 @@ namespace SampleApp
     internal sealed class VectorChatHistoryProvider : ChatHistoryProvider
     {
         private readonly ProviderSessionState<State> _sessionState;
+        private IReadOnlyList<string>? _stateKeys;
         private readonly VectorStore _vectorStore;
 
         public VectorChatHistoryProvider(
@@ -92,7 +93,7 @@ namespace SampleApp
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
         }
 
-        public override string StateKey => this._sessionState.StateKey;
+        public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
         public string GetSessionDbKey(AgentSession session)
             => this._sessionState.GetOrInitializeState(session).SessionDbKey;

dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Evaluations_Step01_RedTeaming/Program.cs

@@ -60,7 +60,7 @@ Console.WriteLine();
 
 // Submit the red team run to the service
 Console.WriteLine("Submitting red team run...");
-RedTeam redTeamRun = await aiProjectClient.RedTeams.CreateAsync(redTeamConfig);
+RedTeam redTeamRun = await aiProjectClient.RedTeams.CreateAsync(redTeamConfig, options: null);
 
 Console.WriteLine($"Red team run created: {redTeamRun.Name}");
 Console.WriteLine($"Status: {redTeamRun.Status}");

dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Step22_MemorySearch/Program.cs

@@ -35,7 +35,7 @@ string userScope = $"user_{Environment.MachineName}";
 AIProjectClient aiProjectClient = new(new Uri(endpoint), new AzureCliCredential());
 
 // Create the Memory Search tool configuration
-MemorySearchTool memorySearchTool = new(memoryStoreName, userScope)
+MemorySearchPreviewTool memorySearchTool = new(memoryStoreName, userScope)
 {
     // Optional: Configure how quickly new memories are indexed (in seconds)
     UpdateDelay = 1,

dotnet/samples/03-workflows/Declarative/HostedWorkflow/Program.cs

@@ -88,7 +88,9 @@ internal sealed class Program
     {
         string workflowYaml = File.ReadAllText("MathChat.yaml");
 
+#pragma warning disable AAIP001 // WorkflowAgentDefinition is experimental
         WorkflowAgentDefinition workflowAgentDefinition = WorkflowAgentDefinition.FromYaml(workflowYaml);
+#pragma warning restore AAIP001
 
         return
             await agentClient.CreateAgentAsync(

dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs

@@ -36,6 +36,8 @@ public abstract class AIContextProvider
     private static IEnumerable<ChatMessage> DefaultNoopFilter(IEnumerable<ChatMessage> messages)
         => messages;
 
+    private IReadOnlyList<string>? _stateKeys;
+
     /// <summary>
     /// Initializes a new instance of the <see cref="AIContextProvider"/> class.
     /// </summary>
@@ -68,14 +70,15 @@ public abstract class AIContextProvider
     protected Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> StoreInputResponseMessageFilter { get; }
 
     /// <summary>
-    /// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
+    /// Gets the set of keys used to store the provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
     /// <remarks>
-    /// The default value is the name of the concrete type (e.g. <c>"TextSearchProvider"</c>).
-    /// Implementations may override this to provide a custom key, for example when multiple
-    /// instances of the same provider type are used in the same session.
+    /// The default value is a single-element set containing the name of the concrete type (e.g. <c>"TextSearchProvider"</c>).
+    /// Implementations may override this to provide custom keys, for example when multiple
+    /// instances of the same provider type are used in the same session, or when a provider
+    /// stores state under more than one key.
     /// </remarks>
-    public virtual string StateKey => this.GetType().Name;
+    public virtual IReadOnlyList<string> StateKeys => this._stateKeys ??= [this.GetType().Name];
 
     /// <summary>
     /// Called at the start of agent invocation to provide additional context.

dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs

@@ -45,6 +45,7 @@ public abstract class ChatHistoryProvider
     private static IEnumerable<ChatMessage> DefaultNoopFilter(IEnumerable<ChatMessage> messages)
         => messages;
 
+    private IReadOnlyList<string>? _stateKeys;
     private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? _provideOutputMessageFilter;
     private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storeInputRequestMessageFilter;
     private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storeInputResponseMessageFilter;
@@ -66,14 +67,15 @@ public abstract class ChatHistoryProvider
     }
 
     /// <summary>
-    /// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
+    /// Gets the set of keys used to store the provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
     /// <remarks>
-    /// The default value is the name of the concrete type (e.g. <c>"InMemoryChatHistoryProvider"</c>).
-    /// Implementations may override this to provide a custom key, for example when multiple
-    /// instances of the same provider type are used in the same session.
+    /// The default value is a single-element set containing the name of the concrete type (e.g. <c>"InMemoryChatHistoryProvider"</c>).
+    /// Implementations may override this to provide custom keys, for example when multiple
+    /// instances of the same provider type are used in the same session, or when a provider
+    /// stores state under more than one key.
     /// </remarks>
-    public virtual string StateKey => this.GetType().Name;
+    public virtual IReadOnlyList<string> StateKeys => this._stateKeys ??= [this.GetType().Name];
 
     /// <summary>
     /// Called at the start of agent invocation to provide messages for the next agent invocation.

dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs

@@ -27,6 +27,7 @@ namespace Microsoft.Agents.AI;
 public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
 {
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="InMemoryChatHistoryProvider"/> class.
@@ -50,7 +51,7 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <summary>
     /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied.

dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs

@@ -39,7 +39,7 @@ public static partial class AzureAIProjectChatClientExtensions
     /// <exception cref="InvalidOperationException">The agent with the specified name was not found.</exception>
     /// <remarks>
     /// When instantiating a <see cref="ChatClientAgent"/> by using an <see cref="AgentReference"/>, minimal information will be available about the agent in the instance level, and any logic that relies
-    /// on <see cref="AIAgent.GetService(Type, object?)"/> to retrieve information about the agent like <see cref="AgentVersion" /> will receive <see langword="null"/> as the result.
+    /// on <see cref="AIAgent.GetService{TService}(object?)"/> to retrieve information about the agent like <see cref="AgentVersion" /> will receive <see langword="null"/> as the result.
     /// </remarks>
     public static ChatClientAgent AsAIAgent(
         this AIProjectClient aiProjectClient,
@@ -355,28 +355,27 @@ public static partial class AzureAIProjectChatClientExtensions
     private static readonly ModelReaderWriterOptions s_modelWriterOptionsWire = new("W");
 
     /// <summary>
-    /// Asynchronously retrieves an agent record by name using the Protocol method with user-agent header.
+    /// Asynchronously retrieves an agent record by name using the protocol method to inject user-agent headers.
     /// </summary>
     private static async Task<AgentRecord> GetAgentRecordByNameAsync(AIProjectClient aiProjectClient, string agentName, CancellationToken cancellationToken)
     {
         ClientResult protocolResponse = await aiProjectClient.Agents.GetAgentAsync(agentName, cancellationToken.ToRequestOptions(false)).ConfigureAwait(false);
         var rawResponse = protocolResponse.GetRawResponse();
         AgentRecord? result = ModelReaderWriter.Read<AgentRecord>(rawResponse.Content, s_modelWriterOptionsWire, AzureAIProjectsOpenAIContext.Default);
-        return ClientResult.FromOptionalValue(result, rawResponse).Value!
-            ?? throw new InvalidOperationException($"Agent with name '{agentName}' not found.");
+        return result ?? throw new InvalidOperationException($"Agent with name '{agentName}' not found.");
     }
 
     /// <summary>
-    /// Asynchronously creates an agent version using the Protocol method with user-agent header.
+    /// Asynchronously creates an agent version using the protocol method to inject user-agent headers.
     /// </summary>
     private static async Task<AgentVersion> CreateAgentVersionWithProtocolAsync(AIProjectClient aiProjectClient, string agentName, AgentVersionCreationOptions creationOptions, CancellationToken cancellationToken)
     {
-        using BinaryContent protocolRequest = BinaryContent.Create(ModelReaderWriter.Write(creationOptions, ModelReaderWriterOptions.Json, AzureAIProjectsContext.Default));
-        ClientResult protocolResponse = await aiProjectClient.Agents.CreateAgentVersionAsync(agentName, protocolRequest, cancellationToken.ToRequestOptions(false)).ConfigureAwait(false);
-
+        BinaryData serializedOptions = ModelReaderWriter.Write(creationOptions, s_modelWriterOptionsWire, AzureAIProjectsContext.Default);
+        BinaryContent content = BinaryContent.Create(serializedOptions);
+        ClientResult protocolResponse = await aiProjectClient.Agents.CreateAgentVersionAsync(agentName, content, foundryFeatures: null, cancellationToken.ToRequestOptions(false)).ConfigureAwait(false);
         var rawResponse = protocolResponse.GetRawResponse();
         AgentVersion? result = ModelReaderWriter.Read<AgentVersion>(rawResponse.Content, s_modelWriterOptionsWire, AzureAIProjectsOpenAIContext.Default);
-        return ClientResult.FromValue(result, rawResponse).Value!;
+        return result ?? throw new InvalidOperationException($"Failed to create agent version for agent '{agentName}'.");
     }
 
     private static async Task<ChatClientAgent> CreateAIAgentAsync(

dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs

@@ -22,6 +22,7 @@ namespace Microsoft.Agents.AI;
 public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
 {
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly CosmosClient _cosmosClient;
     private readonly Container _container;
     private readonly bool _ownsClient;
@@ -114,7 +115,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <summary>
     /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using a connection string.

dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs

@@ -32,6 +32,7 @@ public sealed class FoundryMemoryProvider : AIContextProvider
     private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
 
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly string _contextPrompt;
     private readonly string _memoryStoreName;
     private readonly int _maxMemories;
@@ -82,7 +83,7 @@ public sealed class FoundryMemoryProvider : AIContextProvider
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     private static Func<AgentSession?, State> ValidateStateInitializer(Func<AgentSession?, State> stateInitializer) =>
         session =>

dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs

@@ -27,6 +27,7 @@ public sealed class Mem0Provider : MessageAIContextProvider
     private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
 
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly string _contextPrompt;
     private readonly bool _enableSensitiveTelemetryData;
 
@@ -72,7 +73,7 @@ public sealed class Mem0Provider : MessageAIContextProvider
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     private static Func<AgentSession?, State> ValidateStateInitializer(Func<AgentSession?, State> stateInitializer) =>
         session =>

dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs

@@ -12,6 +12,7 @@ namespace Microsoft.Agents.AI.Workflows;
 internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
 {
     private readonly ProviderSessionState<StoreState> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="WorkflowChatHistoryProvider"/> class.
@@ -30,7 +31,7 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     internal sealed class StoreState
     {

dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs

@@ -112,7 +112,7 @@ public sealed partial class ChatClientAgent : AIAgent
         this.ChatHistoryProvider = options?.ChatHistoryProvider ?? new InMemoryChatHistoryProvider();
         this.AIContextProviders = this._agentOptions?.AIContextProviders as IReadOnlyList<AIContextProvider> ?? this._agentOptions?.AIContextProviders?.ToList();
 
-        // Validate that no two providers share the same StateKey, since they would overwrite each other's state in the session.
+        // Validate that no two providers share any StateKeys, since they would overwrite each other's state in the session.
         this._aiContextProviderStateKeys = ValidateAndCollectStateKeys(this._agentOptions?.AIContextProviders, this.ChatHistoryProvider);
 
         this._logger = (loggerFactory ?? chatClient.GetService<ILoggerFactory>() ?? NullLoggerFactory.Instance).CreateLogger<ChatClientAgent>();
@@ -824,11 +824,17 @@ public sealed partial class ChatClientAgent : AIAgent
                     $"Only {nameof(ChatClientAgentSession.ConversationId)} or {nameof(this.ChatHistoryProvider)} may be used, but not both. The current {nameof(ChatClientAgentSession)} has a {nameof(ChatClientAgentSession.ConversationId)} indicating server-side chat history management, but an override {nameof(this.ChatHistoryProvider)} was provided via {nameof(AgentRunOptions.AdditionalProperties)}.");
             }
 
-            // Validate that the override provider's StateKey does not clash with any AIContextProvider's StateKey.
-            if (overrideProvider is not null && this._aiContextProviderStateKeys.Contains(overrideProvider.StateKey))
+            // Validate that the override provider's StateKeys do not clash with any AIContextProvider's StateKeys.
+            if (overrideProvider is not null)
             {
-                throw new InvalidOperationException(
-                    $"The ChatHistoryProvider '{overrideProvider.GetType().Name}' uses the state key '{overrideProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state.");
+                foreach (var key in overrideProvider.StateKeys)
+                {
+                    if (this._aiContextProviderStateKeys.Contains(key))
+                    {
+                        throw new InvalidOperationException(
+                            $"The ChatHistoryProvider '{overrideProvider.GetType().Name}' uses state key '{key}' which is already used by one of the configured AIContextProviders. Each provider must use unique state keys to avoid overwriting each other's state.");
+                    }
+                }
             }
 
             provider = overrideProvider;
@@ -879,7 +885,7 @@ public sealed partial class ChatClientAgent : AIAgent
     private string GetLoggingAgentName() => this.Name ?? "UnnamedAgent";
 
     /// <summary>
-    /// Validates that all configured providers have unique <see cref="AIContextProvider.StateKey"/> values
+    /// Validates that all configured providers have unique <see cref="AIContextProvider.StateKeys"/> values
     /// and returns a <see cref="HashSet{T}"/> of the AIContextProvider state keys.
     /// </summary>
     private static HashSet<string> ValidateAndCollectStateKeys(IEnumerable<AIContextProvider>? aiContextProviders, ChatHistoryProvider? chatHistoryProvider)
@@ -890,10 +896,13 @@ public sealed partial class ChatClientAgent : AIAgent
         {
             foreach (var provider in aiContextProviders)
             {
-                if (!stateKeys.Add(provider.StateKey))
+                foreach (var key in provider.StateKeys)
                 {
-                    throw new InvalidOperationException(
-                        $"Multiple providers use the same state key '{provider.StateKey}'. Each provider must use a unique state key to avoid overwriting each other's state.");
+                    if (!stateKeys.Add(key))
+                    {
+                        throw new InvalidOperationException(
+                            $"Multiple providers use the same state key '{key}'. Each provider must use a unique state key to avoid overwriting each other's state.");
+                    }
                 }
             }
         }
@@ -905,11 +914,16 @@ public sealed partial class ChatClientAgent : AIAgent
                 $"The default {nameof(InMemoryChatHistoryProvider)} uses the state key '{nameof(InMemoryChatHistoryProvider)}', which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{nameof(InMemoryChatHistoryProvider)}' as its state key, or provide a custom ChatHistoryProvider with a unique state key.");
         }
 
-        if (chatHistoryProvider is not null
-            && stateKeys.Contains(chatHistoryProvider.StateKey))
+        if (chatHistoryProvider is not null)
         {
-            throw new InvalidOperationException(
-                $"The ChatHistoryProvider '{chatHistoryProvider.GetType().Name}' uses the state key '{chatHistoryProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{chatHistoryProvider.StateKey}' as its state key, or reconfigure the custom ChatHistoryProvider with a unique state key.");
+            foreach (var key in chatHistoryProvider.StateKeys)
+            {
+                if (stateKeys.Contains(key))
+                {
+                    throw new InvalidOperationException(
+                        $"The ChatHistoryProvider '{chatHistoryProvider.GetType().Name}' uses state key '{key}' which is already used by one of the configured AIContextProviders. Each provider must use unique state keys to avoid overwriting each other's state. To resolve this, either configure different state keys for the AIContextProvider that shares keys with the ChatHistoryProvider, or reconfigure the custom ChatHistoryProvider with unique state keys.");
+                }
+            }
         }
 
         return stateKeys;

dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs

@@ -54,6 +54,7 @@ public sealed class ChatHistoryMemoryProvider : MessageAIContextProvider, IDispo
     private const string ContentEmbeddingField = "ContentEmbedding";
 
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
 
 #pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal
     private readonly VectorStore _vectorStore;
@@ -128,7 +129,7 @@ public sealed class ChatHistoryMemoryProvider : MessageAIContextProvider, IDispo
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <inheritdoc />
     protected override async ValueTask<AIContext> ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)

dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs

@@ -40,6 +40,7 @@ public sealed class TextSearchProvider : MessageAIContextProvider
     private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available.";
 
     private readonly ProviderSessionState<TextSearchProviderState> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly Func<string, CancellationToken, Task<IEnumerable<TextSearchResult>>> _searchAsync;
     private readonly ILogger<TextSearchProvider>? _logger;
     private readonly AITool[] _tools;
@@ -88,7 +89,7 @@ public sealed class TextSearchProvider : MessageAIContextProvider
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <inheritdoc />
     protected override async ValueTask<AIContext> ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)

dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentCreateTests.cs

@@ -132,10 +132,15 @@ public class AzureAIAgentsPersistentCreateTests
         }
     }
 
-    [Theory]
-    [InlineData("CreateWithChatClientAgentOptionsAsync")]
-    [InlineData("CreateWithFoundryOptionsAsync")]
-    public async Task CreateAgent_CreatesAgentWithCodeInterpreterAsync(string createMechanism)
+    [Fact]
+    public Task CreateAgent_CreatesAgentWithCodeInterpreter_ChatClientAgentOptionsAsync()
+        => this.CreateAgent_CreatesAgentWithCodeInterpreterAsync("CreateWithChatClientAgentOptionsAsync");
+
+    [RetryFact(Constants.RetryCount, Constants.RetryDelay)]
+    public Task CreateAgent_CreatesAgentWithCodeInterpreter_FoundryOptionsAsync()
+        => this.CreateAgent_CreatesAgentWithCodeInterpreterAsync("CreateWithFoundryOptionsAsync");
+
+    private async Task CreateAgent_CreatesAgentWithCodeInterpreterAsync(string createMechanism)
     {
         // Arrange.
         const string AgentInstructions = """

dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs

@@ -43,23 +43,25 @@ public class InMemoryChatHistoryProviderTests
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new InMemoryChatHistoryProvider();
 
         // Assert
-        Assert.Equal("InMemoryChatHistoryProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("InMemoryChatHistoryProvider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new InMemoryChatHistoryProvider(new() { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Fact]

dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs

@@ -467,7 +467,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithModelAndOptions_CreatesValidAgentAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", instructions: "Test instructions");
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", instructions: "Test instructions");
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -475,7 +475,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-model", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -490,7 +490,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithModelAndOptions_WithClientFactory_AppliesFactoryCorrectlyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", instructions: "Test instructions");
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", instructions: "Test instructions");
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -499,7 +499,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         TestChatClient? testChatClient = null;
 
         // Act
-        var agent = await client.CreateAIAgentAsync(
+        var agent = await testClient.Client.CreateAIAgentAsync(
             "test-model",
             options,
             clientFactory: (innerClient) => testChatClient = new TestChatClient(innerClient));
@@ -560,12 +560,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithDefinition_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -582,12 +582,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
 
         var definitionResponse = GeneratePromptDefinitionResponse(definition, null);
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -602,12 +602,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     {
         // Arrange
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definition);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definition);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -628,12 +628,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
 
         // Create a response definition with the same tool
         var definitionResponse = GeneratePromptDefinitionResponse(definition, definition.Tools.Select(t => t.AsAITool()).ToList());
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -667,12 +667,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
             definitionResponse.Tools.Add(tool);
         }
 
-        AIProjectClient client = this.CreateTestAgentClient(agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentDefinitionResponse: definitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -803,10 +803,10 @@ public sealed class AzureAIProjectChatClientExtensionsTests
 
         var definitionResponse = GeneratePromptDefinitionResponse(new PromptAgentDefinition("test-model") { Instructions = "Test instructions" }, tools);
 
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
 
         // Act
-        var agent = await client.CreateAIAgentAsync(
+        var agent = await testClient.Client.CreateAIAgentAsync(
             "test-agent",
             "test-model",
             "Test instructions",
@@ -831,14 +831,14 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithDefinitionTools_CreatesAgentAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test instructions" };
         definition.Tools.Add(ResponseTool.CreateFunctionTool("async_tool", BinaryData.FromString("{}"), strictModeEnabled: false));
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -885,7 +885,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         var sharepointOptions = new SharePointGroundingToolOptions();
         sharepointOptions.ProjectConnections.Add(new ToolProjectConnection("connection-id"));
 
-        var structuredOutputs = new StructuredOutputDefinition("name", "description", BinaryData.FromString(AIJsonUtilities.CreateJsonSchema(new { id = "test" }.GetType()).ToString()), false);
+        var structuredOutputs = new StructuredOutputDefinition("name", "description", new Dictionary<string, BinaryData> { ["schema"] = BinaryData.FromString(AIJsonUtilities.CreateJsonSchema(new { id = "test" }.GetType()).ToString()) }, false);
 
         // Add tools to the definition
         definition.Tools.Add(ResponseTool.CreateFunctionTool("create_tool", BinaryData.FromString("{}"), strictModeEnabled: false));
@@ -902,12 +902,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         // Generate agent definition response with the tools
         var definitionResponse = GeneratePromptDefinitionResponse(definition, definition.Tools.Select(t => t.AsAITool()).ToList());
 
-        AIProjectClient client = this.CreateTestAgentClient(agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentDefinitionResponse: definitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -942,12 +942,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         var definitionResponse = new PromptAgentDefinition("test-model") { Instructions = "Test" };
         definitionResponse.Tools.Add(functionTool);
 
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -961,7 +961,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithDeclarativeFunctionFromDefinition_AcceptsDeclarativeFunctionAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
 
         // Create a declarative function (not invocable) using AIFunctionFactory.CreateDeclaration
@@ -974,7 +974,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -1001,12 +1001,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         var definitionResponse = new PromptAgentDefinition("test-model") { Instructions = "Test" };
         definitionResponse.Tools.Add(declarativeFunction.AsOpenAIResponseTool() ?? throw new InvalidOperationException());
 
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -1027,12 +1027,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test instructions" };
 
         var definitionResponse = GeneratePromptDefinitionResponse(definition, null);
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-agent", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -1083,7 +1083,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
             new PromptAgentDefinition("test-model") { Instructions = "Test" },
             tools);
 
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
 
         var options = new ChatClientAgentOptions
         {
@@ -1092,7 +1092,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        var agent = await client.CreateAIAgentAsync("test-model", options);
+        var agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -1278,14 +1278,14 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithClientFactory_ReceivesCorrectUnderlyingClientAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
         IChatClient? receivedClient = null;
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync(
+        var agent = await testClient.Client.CreateAIAgentAsync(
             "test-agent",
             options,
             clientFactory: (innerClient) =>
@@ -1340,10 +1340,10 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         const string AgentName = "test-agent";
         const string Model = "test-model";
         const string Instructions = "Test instructions";
-        AIProjectClient client = this.CreateTestAgentClient(AgentName, Instructions);
+        using var testClient = CreateTestAgentClientWithHandler(AgentName, Instructions);
 
         // Act
-        var agent = await client.CreateAIAgentAsync(
+        var agent = await testClient.Client.CreateAIAgentAsync(
             AgentName,
             Model,
             Instructions,
@@ -1367,12 +1367,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
 
         var agentDefinitionResponse = GeneratePromptDefinitionResponse(definition, null);
-        AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: agentDefinitionResponse);
+        using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: agentDefinitionResponse);
 
         var options = new AgentVersionCreationOptions(definition);
 
         // Act
-        var agent = await client.CreateAIAgentAsync(
+        var agent = await testClient.Client.CreateAIAgentAsync(
             "test-agent",
             options,
             clientFactory: (innerClient) => new TestChatClient(innerClient));
@@ -1390,7 +1390,8 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     #region User-Agent Header Tests
 
     /// <summary>
-    /// Verifies that the user-agent header is added to both synchronous and asynchronous requests made by agent creation methods.
+    /// Verifies that the MEAI user-agent header is added to CreateAIAgentAsync POST requests
+    /// via the protocol method's RequestOptions pipeline policy.
     /// </summary>
     [Fact]
     public async Task CreateAIAgentAsync_UserAgentHeaderAddedToRequestsAsync()
@@ -1398,9 +1399,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         using var httpHandler = new HttpHandlerAssert(request =>
         {
             Assert.Equal("POST", request.Method.Method);
-            Assert.Contains("MEAI", request.Headers.UserAgent.ToString());
 
-            return new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(TestDataUtil.GetAgentResponseJson(), Encoding.UTF8, "application/json") };
+            // Verify MEAI user-agent header is present on CreateAgentVersion POST request
+            Assert.True(request.Headers.TryGetValues("User-Agent", out var userAgentValues));
+            Assert.Contains(userAgentValues, v => v.Contains("MEAI"));
+
+            return new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(TestDataUtil.GetAgentVersionResponseJson(), Encoding.UTF8, "application/json") };
         });
 
 #pragma warning disable CA5399
@@ -1940,7 +1944,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithTextResponseFormat_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -1952,7 +1956,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -1966,7 +1970,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithJsonResponseFormatWithoutSchema_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -1978,7 +1982,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -1992,7 +1996,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithJsonResponseFormatWithSchema_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         JsonElement schemaElement = AIJsonUtilities.CreateJsonSchema(typeof(TestSchema));
         var jsonFormat = ChatResponseFormat.ForJsonSchema(schemaElement, "test_schema", "A test schema");
         var options = new ChatClientAgentOptions
@@ -2006,7 +2010,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2020,7 +2024,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithJsonResponseFormatWithSchemaAndStrictMode_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         JsonElement schemaElement = AIJsonUtilities.CreateJsonSchema(typeof(TestSchema));
         var jsonFormat = ChatResponseFormat.ForJsonSchema(schemaElement, "test_schema", "A test schema");
         var additionalProps = new AdditionalPropertiesDictionary
@@ -2039,7 +2043,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2053,7 +2057,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithJsonResponseFormatWithSchemaAndStrictModeFalse_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         JsonElement schemaElement = AIJsonUtilities.CreateJsonSchema(typeof(TestSchema));
         var jsonFormat = ChatResponseFormat.ForJsonSchema(schemaElement, "test_schema", "A test schema");
         var additionalProps = new AdditionalPropertiesDictionary
@@ -2072,7 +2076,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2090,7 +2094,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithRawRepresentationFactory_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -2102,7 +2106,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2116,7 +2120,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithRawRepresentationFactoryReturningNull_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -2128,7 +2132,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2142,7 +2146,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithRawRepresentationFactoryReturningNonCreateResponseOptions_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -2154,7 +2158,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2172,7 +2176,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithDescription_SetsDescriptionAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient(description: "Test description");
+        using var testClient = CreateTestAgentClientWithHandler(description: "Test description");
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -2181,7 +2185,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2195,7 +2199,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithoutDescription_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var options = new ChatClientAgentOptions
         {
             Name = "test-agent",
@@ -2203,7 +2207,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2688,7 +2692,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
     public async Task CreateAIAgentAsync_WithHostedToolTypes_CreatesAgentSuccessfullyAsync()
     {
         // Arrange
-        AIProjectClient client = this.CreateTestAgentClient();
+        using var testClient = CreateTestAgentClientWithHandler();
         var webSearchTool = new HostedWebSearchTool();
 
         var options = new ChatClientAgentOptions
@@ -2702,7 +2706,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         };
 
         // Act
-        ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+        ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
 
         // Assert
         Assert.NotNull(agent);
@@ -2855,6 +2859,54 @@ public sealed class AzureAIProjectChatClientExtensionsTests
         return new FakeAgentClient(agentName, instructions, description, agentDefinitionResponse);
     }
 
+    /// <summary>
+    /// Creates a test AIProjectClient backed by an HTTP handler that returns canned responses.
+    /// Used for tests that exercise the protocol-method code path (CreateAgentVersion).
+    /// The returned client must be disposed to clean up the underlying HttpClient/handler.
+    /// </summary>
+    private static DisposableTestClient CreateTestAgentClientWithHandler(string? agentName = null, string? instructions = null, string? description = null, AgentDefinition? agentDefinitionResponse = null)
+    {
+        var responseJson = TestDataUtil.GetAgentVersionResponseJson(agentName, agentDefinitionResponse, instructions, description);
+
+        var httpHandler = new HttpHandlerAssert(_ =>
+            new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(responseJson, Encoding.UTF8, "application/json") });
+
+#pragma warning disable CA5399
+        var httpClient = new HttpClient(httpHandler);
+#pragma warning restore CA5399
+
+        var client = new AIProjectClient(
+            new Uri("https://test.openai.azure.com/"),
+            new FakeAuthenticationTokenProvider(),
+            new() { Transport = new HttpClientPipelineTransport(httpClient) });
+
+        return new DisposableTestClient(client, httpClient, httpHandler);
+    }
+
+    /// <summary>
+    /// Wraps an AIProjectClient and its disposable dependencies for deterministic cleanup.
+    /// </summary>
+    private sealed class DisposableTestClient : IDisposable
+    {
+        private readonly HttpClient _httpClient;
+        private readonly HttpHandlerAssert _httpHandler;
+
+        public DisposableTestClient(AIProjectClient client, HttpClient httpClient, HttpHandlerAssert httpHandler)
+        {
+            this.Client = client;
+            this._httpClient = httpClient;
+            this._httpHandler = httpHandler;
+        }
+
+        public AIProjectClient Client { get; }
+
+        public void Dispose()
+        {
+            this._httpClient.Dispose();
+            this._httpHandler.Dispose();
+        }
+    }
+
     /// <summary>
     /// Creates a test AgentRecord for testing.
     /// </summary>
@@ -3039,25 +3091,13 @@ public sealed class AzureAIProjectChatClientExtensionsTests
                 return Task.FromResult(ClientResult.FromValue(ModelReaderWriter.Read<AgentRecord>(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200)));
             }
 
-            public override ClientResult CreateAgentVersion(string agentName, BinaryContent content, RequestOptions? options = null)
-            {
-                var responseJson = this.GetAgentVersionResponseJson();
-                return ClientResult.FromValue(ModelReaderWriter.Read<AgentVersion>(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200, BinaryData.FromString(responseJson)));
-            }
-
-            public override ClientResult<AgentVersion> CreateAgentVersion(string agentName, AgentVersionCreationOptions? options = null, CancellationToken cancellationToken = default)
+            public override ClientResult<AgentVersion> CreateAgentVersion(string agentName, AgentVersionCreationOptions? options = null, string? foundryFeatures = null, CancellationToken cancellationToken = default)
             {
                 var responseJson = this.GetAgentVersionResponseJson();
                 return ClientResult.FromValue(ModelReaderWriter.Read<AgentVersion>(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200));
             }
 
-            public override Task<ClientResult> CreateAgentVersionAsync(string agentName, BinaryContent content, RequestOptions? options = null)
-            {
-                var responseJson = this.GetAgentVersionResponseJson();
-                return Task.FromResult<ClientResult>(ClientResult.FromValue(ModelReaderWriter.Read<AgentVersion>(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200, BinaryData.FromString(responseJson))));
-            }
-
-            public override Task<ClientResult<AgentVersion>> CreateAgentVersionAsync(string agentName, AgentVersionCreationOptions? options = null, CancellationToken cancellationToken = default)
+            public override Task<ClientResult<AgentVersion>> CreateAgentVersionAsync(string agentName, AgentVersionCreationOptions? options = null, string? foundryFeatures = null, CancellationToken cancellationToken = default)
             {
                 var responseJson = this.GetAgentVersionResponseJson();
                 return Task.FromResult(ClientResult.FromValue(ModelReaderWriter.Read<AgentVersion>(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200)));

dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientTests.cs

@@ -22,7 +22,7 @@ public class AzureAIProjectChatClientTests
         var requestTriggered = false;
         using var httpHandler = new HttpHandlerAssert(async (request) =>
         {
-            if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+            if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
             {
                 requestTriggered = true;
 
@@ -71,7 +71,7 @@ public class AzureAIProjectChatClientTests
         var requestTriggered = false;
         using var httpHandler = new HttpHandlerAssert(async (request) =>
         {
-            if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+            if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
             {
                 requestTriggered = true;
 
@@ -120,7 +120,7 @@ public class AzureAIProjectChatClientTests
         var requestTriggered = false;
         using var httpHandler = new HttpHandlerAssert(async (request) =>
         {
-            if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+            if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
             {
                 requestTriggered = true;
 
@@ -169,7 +169,7 @@ public class AzureAIProjectChatClientTests
         var requestTriggered = false;
         using var httpHandler = new HttpHandlerAssert(async (request) =>
         {
-            if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+            if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
             {
                 requestTriggered = true;
 

dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs

@@ -152,7 +152,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
 
     [Fact]
     [Trait("Category", "CosmosDB")]
-    public void StateKey_ReturnsDefaultKey_WhenNoStateKeyProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoStateKeyProvided()
     {
         // Arrange & Act
         this.SkipIfEmulatorNotAvailable();
@@ -161,12 +161,13 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
             _ => new CosmosChatHistoryProvider.State("test-conversation"));
 
         // Assert
-        Assert.Equal("CosmosChatHistoryProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("CosmosChatHistoryProvider", provider.StateKeys);
     }
 
     [Fact]
     [Trait("Category", "CosmosDB")]
-    public void StateKey_ReturnsCustomKey_WhenSetViaConstructor()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaConstructor()
     {
         // Arrange & Act
         this.SkipIfEmulatorNotAvailable();
@@ -176,7 +177,8 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
             stateKey: "custom-key");
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Fact]

dotnet/tests/Microsoft.Agents.AI.GitHub.Copilot.UnitTests/GitHubCopilotAgentTests.cs

@@ -109,7 +109,7 @@ public sealed class GitHubCopilotAgentTests
         var hooks = new SessionHooks();
         var infiniteSessions = new InfiniteSessionConfig();
         var systemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = "Be helpful" };
-        PermissionHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
+        PermissionRequestHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
         UserInputHandler userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
         var mcpServers = new Dictionary<string, object> { ["server1"] = new McpLocalServerConfig() };
 
@@ -160,7 +160,7 @@ public sealed class GitHubCopilotAgentTests
         var hooks = new SessionHooks();
         var infiniteSessions = new InfiniteSessionConfig();
         var systemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = "Be helpful" };
-        PermissionHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
+        PermissionRequestHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
         UserInputHandler userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
         var mcpServers = new Dictionary<string, object> { ["server1"] = new McpLocalServerConfig() };
 

dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs

@@ -67,17 +67,18 @@ public sealed class Mem0ProviderTests : IDisposable
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" }));
 
         // Assert
-        Assert.Equal("Mem0Provider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("Mem0Provider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new Mem0Provider(
@@ -86,7 +87,8 @@ public sealed class Mem0ProviderTests : IDisposable
             new Mem0ProviderOptions { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Fact]
@@ -419,7 +421,7 @@ public sealed class Mem0ProviderTests : IDisposable
     }
 
     [Fact]
-    public async Task StateKey_CanBeConfiguredViaOptionsAsync()
+    public async Task StateKeys_CanBeConfiguredViaOptionsAsync()
     {
         // Arrange
         this._handler.EnqueueJsonResponse("[]");

dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs

@@ -380,7 +380,7 @@ public class AIContextProviderChatClientTests
     /// </summary>
     private sealed class TestAIContextProvider : AIContextProvider
     {
-        private readonly string _stateKey;
+        private readonly IReadOnlyList<string> _stateKeys;
         private readonly IEnumerable<ChatMessage> _provideMessages;
         private readonly string? _provideInstructions;
         private readonly IEnumerable<AITool>? _provideTools;
@@ -389,7 +389,7 @@ public class AIContextProviderChatClientTests
 
         public InvokedContext? LastInvokedContext { get; private set; }
 
-        public override string StateKey => this._stateKey;
+        public override IReadOnlyList<string> StateKeys => this._stateKeys;
 
         public TestAIContextProvider(
             string stateKey,
@@ -397,7 +397,7 @@ public class AIContextProviderChatClientTests
             string? provideInstructions = null,
             IEnumerable<AITool>? provideTools = null)
         {
-            this._stateKey = stateKey;
+            this._stateKeys = [stateKey];
             this._provideMessages = provideMessages ?? [];
             this._provideInstructions = provideInstructions;
             this._provideTools = provideTools;

dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs

@@ -105,8 +105,8 @@ public partial class ChatClientAgentTests
                 ChatHistoryProvider = historyProvider
             }));
 
-        Assert.Contains("SharedKey", ex.Message);
-        Assert.Contains(nameof(ChatHistoryProvider), ex.Message);
+        Assert.Contains("ChatHistoryProvider", ex.Message);
+        Assert.Contains("state key 'SharedKey'", ex.Message);
     }
 
     /// <summary>
@@ -159,11 +159,11 @@ public partial class ChatClientAgentTests
         var ex = await Assert.ThrowsAsync<InvalidOperationException>(() =>
             agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }));
 
-        Assert.Contains("SharedKey", ex.Message);
+        Assert.Contains("state key 'SharedKey'", ex.Message);
     }
 
     /// <summary>
-    /// Verify that RunAsync succeeds when an override ChatHistoryProvider uses the same StateKey as the default ChatHistoryProvider.
+    /// Verify that RunAsync succeeds when an override ChatHistoryProvider uses the same StateKeys as the default ChatHistoryProvider.
     /// </summary>
     [Fact]
     public async Task RunAsync_SucceedsWhenOverrideChatHistoryProviderSharesKeyWithDefaultAsync()
@@ -192,6 +192,102 @@ public partial class ChatClientAgentTests
         await agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties });
     }
 
+    /// <summary>
+    /// Verify that the constructor throws when two multi-key AIContextProviders have an overlapping key.
+    /// </summary>
+    [Fact]
+    public void Constructor_ThrowsWhenMultiKeyAIContextProvidersOverlap()
+    {
+        // Arrange
+        var chatClient = new Mock<IChatClient>().Object;
+        var provider1 = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+        var provider2 = new MultiKeyTestAIContextProvider("Key2", "SharedKey");
+
+        // Act & Assert
+        var ex = Assert.Throws<InvalidOperationException>(() =>
+            new ChatClientAgent(chatClient, options: new()
+            {
+                AIContextProviders = [provider1, provider2]
+            }));
+
+        Assert.Contains("state key 'SharedKey'", ex.Message);
+    }
+
+    /// <summary>
+    /// Verify that the constructor throws when a multi-key ChatHistoryProvider has an overlapping key with an AIContextProvider.
+    /// </summary>
+    [Fact]
+    public void Constructor_ThrowsWhenMultiKeyChatHistoryProviderOverlapsWithAIContextProvider()
+    {
+        // Arrange
+        var chatClient = new Mock<IChatClient>().Object;
+        var contextProvider = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+        var historyProvider = new MultiKeyTestChatHistoryProvider("Key2", "SharedKey");
+
+        // Act & Assert
+        var ex = Assert.Throws<InvalidOperationException>(() =>
+            new ChatClientAgent(chatClient, options: new()
+            {
+                AIContextProviders = [contextProvider],
+                ChatHistoryProvider = historyProvider
+            }));
+
+        Assert.Contains("state key 'SharedKey'", ex.Message);
+    }
+
+    /// <summary>
+    /// Verify that the constructor succeeds when multi-key providers have no overlapping keys.
+    /// </summary>
+    [Fact]
+    public void Constructor_SucceedsWithMultiKeyProvidersWithUniqueKeys()
+    {
+        // Arrange
+        var chatClient = new Mock<IChatClient>().Object;
+        var contextProvider1 = new MultiKeyTestAIContextProvider("Key1", "Key2");
+        var contextProvider2 = new MultiKeyTestAIContextProvider("Key3", "Key4");
+        var historyProvider = new MultiKeyTestChatHistoryProvider("Key5", "Key6");
+
+        // Act & Assert - should not throw
+        _ = new ChatClientAgent(chatClient, options: new()
+        {
+            AIContextProviders = [contextProvider1, contextProvider2],
+            ChatHistoryProvider = historyProvider
+        });
+    }
+
+    /// <summary>
+    /// Verify that RunAsync throws when a multi-key override ChatHistoryProvider has an overlapping key with an AIContextProvider.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_ThrowsWhenMultiKeyOverrideChatHistoryProviderClashesWithAIContextProviderAsync()
+    {
+        // Arrange
+        Mock<IChatClient> mockService = new();
+        mockService.Setup(
+            s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
+
+        var contextProvider = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+        var overrideHistoryProvider = new MultiKeyTestChatHistoryProvider("Key2", "SharedKey");
+
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            AIContextProviders = [contextProvider]
+        });
+
+        // Act & Assert
+        ChatClientAgentSession? session = await agent.CreateSessionAsync() as ChatClientAgentSession;
+        AdditionalPropertiesDictionary additionalProperties = new();
+        additionalProperties.Add<ChatHistoryProvider>(overrideHistoryProvider);
+
+        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }));
+
+        Assert.Contains("state key 'SharedKey'", ex.Message);
+    }
+
     #endregion
 
     #region RunAsync Tests
@@ -489,6 +585,7 @@ public partial class ChatClientAgentTests
             .ReturnsAsync(new ChatResponse(responseMessages));
 
         var mockProvider = new Mock<AIContextProvider>(null, null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -560,6 +657,7 @@ public partial class ChatClientAgentTests
             .Throws(new InvalidOperationException("downstream failure"));
 
         var mockProvider = new Mock<AIContextProvider>(null, null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -618,6 +716,7 @@ public partial class ChatClientAgentTests
             .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
 
         var mockProvider = new Mock<AIContextProvider>(null, null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -678,7 +777,7 @@ public partial class ChatClientAgentTests
 
         // Provider 1: adds a system message and a tool
         var mockProvider1 = new Mock<AIContextProvider>(null, null, null);
-        mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockProvider1
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -697,7 +796,7 @@ public partial class ChatClientAgentTests
         // Provider 2: adds another system message and verifies it receives accumulated context from provider 1
         AIContext? provider2ReceivedContext = null;
         var mockProvider2 = new Mock<AIContextProvider>(null, null, null);
-        mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+        mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
         mockProvider2
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -785,7 +884,7 @@ public partial class ChatClientAgentTests
             .ThrowsAsync(new InvalidOperationException("downstream failure"));
 
         var mockProvider1 = new Mock<AIContextProvider>(null, null, null);
-        mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockProvider1
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -802,7 +901,7 @@ public partial class ChatClientAgentTests
             .Returns(new ValueTask());
 
         var mockProvider2 = new Mock<AIContextProvider>(null, null, null);
-        mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+        mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
         mockProvider2
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -870,7 +969,7 @@ public partial class ChatClientAgentTests
             .Returns(ToAsyncEnumerableAsync(responseUpdates));
 
         var mockProvider1 = new Mock<AIContextProvider>(null, null, null);
-        mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockProvider1
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -887,7 +986,7 @@ public partial class ChatClientAgentTests
             .Returns(new ValueTask());
 
         var mockProvider2 = new Mock<AIContextProvider>(null, null, null);
-        mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+        mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
         mockProvider2
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1829,6 +1928,7 @@ public partial class ChatClientAgentTests
             .Returns(ToAsyncEnumerableAsync(responseUpdates));
 
         var mockProvider = new Mock<AIContextProvider>(null, null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1908,6 +2008,7 @@ public partial class ChatClientAgentTests
             .Throws(new InvalidOperationException("downstream failure"));
 
         var mockProvider = new Mock<AIContextProvider>(null, null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1965,7 +2066,17 @@ public partial class ChatClientAgentTests
 
     private sealed class TestAIContextProvider(string stateKey) : AIContextProvider
     {
-        public override string StateKey => stateKey;
+        private readonly IReadOnlyList<string> _stateKeys = [stateKey];
+
+        public override IReadOnlyList<string> StateKeys => this._stateKeys;
+
+        protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+            => new(context.AIContext);
+    }
+
+    private sealed class MultiKeyTestAIContextProvider(params string[] stateKeys) : AIContextProvider
+    {
+        public override IReadOnlyList<string> StateKeys => stateKeys;
 
         protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
             => new(context.AIContext);
@@ -1973,7 +2084,20 @@ public partial class ChatClientAgentTests
 
     private sealed class TestChatHistoryProvider(string stateKey) : ChatHistoryProvider
     {
-        public override string StateKey => stateKey;
+        private readonly IReadOnlyList<string> _stateKeys = [stateKey];
+
+        public override IReadOnlyList<string> StateKeys => this._stateKeys;
+
+        protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+            => new(context.RequestMessages);
+
+        protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+            => default;
+    }
+
+    private sealed class MultiKeyTestChatHistoryProvider(params string[] stateKeys) : ChatHistoryProvider
+    {
+        public override IReadOnlyList<string> StateKeys => stateKeys;
 
         protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
             => new(context.RequestMessages);

dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs

@@ -339,7 +339,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         // Create a mock chat history provider that would normally provide messages
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -347,7 +347,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         // Create a mock AI context provider that would normally provide context
         var mockContextProvider = new Mock<AIContextProvider>(null, null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -408,7 +408,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         // Create a mock chat history provider that would normally provide messages
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -416,7 +416,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         // Create a mock AI context provider that would normally provide context
         var mockContextProvider = new Mock<AIContextProvider>(null, null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -639,7 +639,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         List<ChatMessage> capturedMessagesAddedToProvider = [];
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -648,7 +648,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         AIContextProvider.InvokedContext? capturedInvokedContext = null;
         var mockContextProvider = new Mock<AIContextProvider>(null, null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -703,7 +703,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         List<ChatMessage> capturedMessagesAddedToProvider = [];
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -712,7 +712,7 @@ public class ChatClientAgent_BackgroundResponsesTests
 
         AIContextProvider.InvokedContext? capturedInvokedContext = null;
         var mockContextProvider = new Mock<AIContextProvider>(null, null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())

dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs

@@ -186,6 +186,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
                 It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
 
         Mock<ChatHistoryProvider> mockChatHistoryProvider = new(null, null, null);
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -241,6 +242,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
                 It.IsAny<CancellationToken>())).Throws(new InvalidOperationException("Test Error"));
 
         Mock<ChatHistoryProvider> mockChatHistoryProvider = new(null, null, null);
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -430,6 +432,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
 
         // Arrange a chat history provider to override the factory provided one.
         Mock<ChatHistoryProvider> mockOverrideChatHistoryProvider = new(null, null, null);
+        mockOverrideChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockOverrideChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -443,6 +446,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
         // Arrange a chat history provider to provide to the agent at construction time.
         // This one shouldn't be used since it is being overridden.
         Mock<ChatHistoryProvider> mockAgentOptionsChatHistoryProvider = new(null, null, null);
+        mockAgentOptionsChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockAgentOptionsChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())

dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs

@@ -39,17 +39,18 @@ public sealed class TextSearchProviderTests
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new TextSearchProvider((_, _) => Task.FromResult<IEnumerable<TextSearchProvider.TextSearchResult>>([]));
 
         // Assert
-        Assert.Equal("TextSearchProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("TextSearchProvider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new TextSearchProvider(
@@ -57,7 +58,8 @@ public sealed class TextSearchProviderTests
             new TextSearchProviderOptions { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Theory]

dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs

@@ -56,7 +56,7 @@ public class ChatHistoryMemoryProviderTests
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new ChatHistoryMemoryProvider(
@@ -66,11 +66,12 @@ public class ChatHistoryMemoryProviderTests
             _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }));
 
         // Assert
-        Assert.Equal("ChatHistoryMemoryProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("ChatHistoryMemoryProvider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new ChatHistoryMemoryProvider(
@@ -81,7 +82,8 @@ public class ChatHistoryMemoryProviderTests
             new ChatHistoryMemoryProviderOptions { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Fact]

dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeCodeGenTest.cs

@@ -14,7 +14,7 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests;
 public sealed class DeclarativeCodeGenTest(ITestOutputHelper output) : WorkflowTest(output)
 {
     [Theory]
-    [InlineData("CheckSystem.yaml", "CheckSystem.json")]
+    [InlineData("CheckSystem.yaml", "CheckSystem.json", Skip = "Temporarily skipped")]
     [InlineData("SendActivity.yaml", "SendActivity.json")]
     [InlineData("InvokeAgent.yaml", "InvokeAgent.json")]
     [InlineData("InvokeAgent.yaml", "InvokeAgent.json", true)]

dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeWorkflowTest.cs

@@ -15,7 +15,7 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests;
 public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : WorkflowTest(output)
 {
     [Theory]
-    [InlineData("CheckSystem.yaml", "CheckSystem.json")]
+    [InlineData("CheckSystem.yaml", "CheckSystem.json", Skip = "Temporarily skipped")]
     [InlineData("ConversationMessages.yaml", "ConversationMessages.json")]
     [InlineData("ConversationMessages.yaml", "ConversationMessages.json", true)]
     [InlineData("InputArguments.yaml", "InputArguments.json")]
@@ -33,7 +33,7 @@ public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : Workflow
     public Task ValidateScenarioAsync(string workflowFileName, string testcaseFileName, bool externalConveration = false) =>
         this.RunWorkflowAsync(GetWorkflowPath(workflowFileName, isSample: true), testcaseFileName, externalConveration);
 
-    [Theory]
+    [Theory(Skip = "Multi-turn tests hang in CI - needs investigation")]
     [InlineData("ConfirmInput.yaml", "ConfirmInput.json", false)]
     [InlineData("RequestExternalInput.yaml", "RequestExternalInput.json", false)]
     public Task ValidateMultiTurnAsync(string workflowFileName, string testcaseFileName, bool isSample) =>

dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs

@@ -1,9 +1,23 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System.Threading.Tasks;
 using AgentConformance.IntegrationTests;
 
 namespace OpenAIAssistant.IntegrationTests;
 
 public class OpenAIAssistantStructuredOutputRunTests() : StructuredOutputRunTests<OpenAIAssistantFixture>(() => new())
 {
+    private const string SkipReason = "Fails intermittently on the build agent/CI";
+
+    [Fact(Skip = SkipReason)]
+    public override Task RunWithResponseFormatReturnsExpectedResultAsync() =>
+        base.RunWithResponseFormatReturnsExpectedResultAsync();
+
+    [Fact(Skip = SkipReason)]
+    public override Task RunWithGenericTypeReturnsExpectedResultAsync() =>
+        base.RunWithGenericTypeReturnsExpectedResultAsync();
+
+    [Fact(Skip = SkipReason)]
+    public override Task RunWithPrimitiveTypeReturnsExpectedResultAsync() =>
+        base.RunWithPrimitiveTypeReturnsExpectedResultAsync();
 }

python/packages/ag-ui/agent_framework_ag_ui/_run_common.py

@@ -372,6 +372,15 @@ def _emit_usage(content: Content) -> list[BaseEvent]:
     return [CustomEvent(name="usage", value=usage_details)]
 
 
+def _emit_oauth_consent(content: Content) -> list[BaseEvent]:
+    """Emit an OAuth consent request as a custom event so frontends can render a consent link."""
+    return (
+        [CustomEvent(name="oauth_consent_request", value={"consent_link": content.consent_link})]
+        if content.consent_link
+        else []
+    )
+
+
 def _emit_content(
     content: Any,
     flow: FlowState,
@@ -391,5 +400,7 @@ def _emit_content(
         return _emit_approval_request(content, flow, predictive_handler, require_confirmation)
     if content_type == "usage":
         return _emit_usage(content)
+    if content_type == "oauth_consent_request":
+        return _emit_oauth_consent(content)
     logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type)
     return []

python/packages/ag-ui/tests/ag_ui/test_run.py

@@ -4,6 +4,7 @@
 
 import pytest
 from ag_ui.core import (
+    CustomEvent,
     TextMessageEndEvent,
     TextMessageStartEvent,
     ToolCallArgsEvent,
@@ -871,3 +872,26 @@ class TestTextMessageEventBalancing:
 
         assert len(start_events) == 2
         assert len(end_events) == 2
+
+
+def test_emit_oauth_consent_request():
+    """Test that oauth_consent_request content emits a CustomEvent."""
+    content = Content.from_oauth_consent_request(
+        consent_link="https://login.microsoftonline.com/consent",
+    )
+    flow = FlowState()
+    events = _emit_content(content, flow)
+
+    assert len(events) == 1
+    assert isinstance(events[0], CustomEvent)
+    assert events[0].name == "oauth_consent_request"
+    assert events[0].value == {"consent_link": "https://login.microsoftonline.com/consent"}
+
+
+def test_emit_oauth_consent_request_no_link():
+    """Test that oauth_consent_request without a consent_link emits no events."""
+    content = Content("oauth_consent_request")
+    flow = FlowState()
+    events = _emit_content(content, flow)
+
+    assert len(events) == 0

python/packages/anthropic/agent_framework_anthropic/_chat_client.py

@@ -894,6 +894,7 @@ class AnthropicClient(
                     usage_details.append(Content.from_usage(usage_details=details))
 
                 return ChatResponseUpdate(
+                    role="assistant",
                     response_id=event.message.id,
                     contents=[
                         *self._parse_contents_from_anthropic(event.message.content),

python/packages/anthropic/tests/test_anthropic_client.py

@@ -1044,6 +1044,128 @@ async def test_inner_get_response_ignores_options_stream_streaming(mock_anthropi
     assert mock_anthropic_client.beta.messages.create.call_args.kwargs["stream"] is True
 
 
+def test_process_stream_event_message_start_sets_assistant_role(mock_anthropic_client: MagicMock) -> None:
+    """Test that message_start streaming event sets role='assistant'.
+
+    This is critical: without role='assistant', _process_update cannot detect
+    a role boundary between a prior tool message and the new assistant turn,
+    causing tool_use blocks to collapse into a user-role message and triggering
+    Anthropic's '`tool_use` blocks can only be in `assistant` messages' error.
+    """
+    client = create_test_anthropic_client(mock_anthropic_client)
+
+    mock_event = MagicMock()
+    mock_event.type = "message_start"
+    mock_event.message.id = "msg_abc"
+    mock_event.message.role = "assistant"
+    mock_event.message.model = "claude-3-5-sonnet-20241022"
+    mock_event.message.content = []
+    mock_event.message.stop_reason = None
+    mock_event.message.usage = None
+
+    result = client._process_stream_event(mock_event)
+
+    assert result is not None
+    assert result.role == "assistant"
+
+
+def test_process_stream_event_message_start_role_prevents_tool_use_collapse() -> None:
+    """Regression test: tool_use blocks must not end up in a user-role message.
+
+    Simulates two consecutive streaming tool-call iterations:
+      Iteration 1: assistant emits tool_use → framework appends tool result (role=tool)
+      Iteration 2: assistant starts a new message_start → must create a NEW message
+
+    Without role='assistant' on the message_start update, _process_update sees
+    update.role=None (falsy) and appends to the last message (role='tool'),
+    producing {"role": "user", "content": [tool_result, tool_use]} which
+    Anthropic rejects with HTTP 400.
+    """
+    from agent_framework import ChatResponse, ChatResponseUpdate, Content, Message
+
+    # Simulate what the streaming tool loop produces after iteration 1:
+    # an existing 'tool' message is the last in the response
+    existing_tool_message = Message(
+        role="tool",
+        contents=[Content.from_function_result(call_id="call_1", result="some result")],
+    )
+
+    response = ChatResponse(messages=[existing_tool_message])
+
+    # Now simulate the message_start update from iteration 2 — WITH role set
+    message_start_update = ChatResponseUpdate(
+        role="assistant",
+        response_id="msg_iter2",
+    )
+
+    # Simulate a content_block_start carrying a tool_use — no role on this one (correct)
+    tool_use_update = ChatResponseUpdate(
+        contents=[
+            Content.from_function_call(
+                call_id="call_2",
+                name="get_weather",
+                arguments={"location": "NYC"},
+            )
+        ],
+    )
+
+    # Apply updates exactly as from_updates / _process_update would
+    from agent_framework._types import _process_update
+
+    _process_update(response, message_start_update)
+    _process_update(response, tool_use_update)
+
+    # Must have TWO messages: the original tool message + a new assistant message
+    assert len(response.messages) == 2, "tool_use from iteration 2 collapsed into the tool message from iteration 1"
+    assert response.messages[0].role == "tool"
+    assert response.messages[1].role == "assistant"
+
+    # The assistant message must contain the tool_use, not the tool result
+    assert response.messages[1].contents[0].type == "function_call"
+    assert response.messages[1].contents[0].call_id == "call_2"
+
+
+def test_process_stream_event_message_start_without_role_reproduces_bug() -> None:
+    """Documents the original bug: missing role causes tool_use to collapse into tool message.
+
+    This test demonstrates WHY the fix (adding role='assistant') was necessary.
+    It intentionally reproduces the broken behavior when role is absent.
+    """
+    from agent_framework import ChatResponse, ChatResponseUpdate, Content, Message
+    from agent_framework._types import _process_update
+
+    existing_tool_message = Message(
+        role="tool",
+        contents=[Content.from_function_result(call_id="call_1", result="some result")],
+    )
+    response = ChatResponse(messages=[existing_tool_message])
+
+    # message_start WITHOUT role (the original broken state)
+    message_start_update = ChatResponseUpdate(
+        role=None,
+        response_id="msg_iter2",
+    )
+    tool_use_update = ChatResponseUpdate(
+        contents=[
+            Content.from_function_call(
+                call_id="call_2",
+                name="get_weather",
+                arguments={"location": "NYC"},
+            )
+        ],
+    )
+
+    _process_update(response, message_start_update)
+    _process_update(response, tool_use_update)
+
+    # BUG: only 1 message — tool_use collapsed into the tool message
+    assert len(response.messages) == 1, "Expected bug: should still be 1 message without the fix"
+    # The single message has role='tool' but contains a function_call — invalid for Anthropic API
+    assert response.messages[0].role == "tool"
+    has_function_call = any(c.type == "function_call" for c in response.messages[0].contents)
+    assert has_function_call, "Expected bug: function_call leaked into tool message"
+
+
 # Integration Tests
 
 

python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py

@@ -87,10 +87,11 @@ from azure.ai.agents.models import (
     ToolApproval,
     ToolDefinition,
     ToolOutput,
+    VectorStoreDataSource,
 )
 from pydantic import BaseModel
 
-from ._shared import AzureAISettings, to_azure_ai_agent_tools
+from ._shared import AzureAISettings, resolve_file_ids, to_azure_ai_agent_tools
 
 if sys.version_info >= (3, 13):
     from typing import TypeVar  # type: ignore # pragma: no cover
@@ -219,9 +220,21 @@ class AzureAIAgentClient(
     # region Hosted Tool Factory Methods
 
     @staticmethod
-    def get_code_interpreter_tool() -> CodeInterpreterTool:
+    def get_code_interpreter_tool(
+        *,
+        file_ids: list[str | Content] | None = None,
+        data_sources: list[VectorStoreDataSource] | None = None,
+    ) -> CodeInterpreterTool:
         """Create a code interpreter tool configuration for Azure AI Agents.
 
+        Keyword Args:
+            file_ids: List of uploaded file IDs or Content objects to make available to
+                the code interpreter. Accepts plain strings or Content.from_hosted_file()
+                instances. The underlying SDK raises ValueError if both file_ids and
+                data_sources are provided.
+            data_sources: List of vector store data sources for enterprise file search.
+                Mutually exclusive with file_ids.
+
         Returns:
             A CodeInterpreterTool instance ready to pass to ChatAgent.
 
@@ -230,10 +243,21 @@ class AzureAIAgentClient(
 
                 from agent_framework.azure import AzureAIAgentClient
 
+                # Basic code interpreter
                 tool = AzureAIAgentClient.get_code_interpreter_tool()
+
+                # With uploaded file IDs
+                tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc123"])
+
+                # With Content objects
+                from agent_framework import Content
+
+                tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=[Content.from_hosted_file("file-abc123")])
+
                 agent = ChatAgent(client, tools=[tool])
         """
-        return CodeInterpreterTool()
+        resolved = resolve_file_ids(file_ids)
+        return CodeInterpreterTool(file_ids=resolved, data_sources=data_sources)
 
     @staticmethod
     def get_file_search_tool(

python/packages/azure-ai/agent_framework_azure_ai/_client.py

@@ -37,12 +37,13 @@ from agent_framework.openai._responses_client import RawOpenAIResponsesClient
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
     ApproximateLocation,
+    CodeInterpreterContainerAuto,
     CodeInterpreterTool,
-    CodeInterpreterToolAuto,
+    FoundryFeaturesOptInKeys,
     ImageGenTool,
     MCPTool,
     PromptAgentDefinition,
-    PromptAgentDefinitionText,
+    PromptAgentDefinitionTextOptions,
     RaiConfig,
     Reasoning,
     WebSearchPreviewTool,
@@ -50,7 +51,7 @@ from azure.ai.projects.models import (
 from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool
 from azure.core.exceptions import ResourceNotFoundError
 
-from ._shared import AzureAISettings, create_text_format_config
+from ._shared import AzureAISettings, create_text_format_config, resolve_file_ids
 
 if sys.version_info >= (3, 13):
     from typing import TypeVar  # type: ignore # pragma: no cover
@@ -78,6 +79,9 @@ class AzureAIProjectAgentOptions(OpenAIResponsesOptions, total=False):
     reasoning: Reasoning  # type: ignore[misc]
     """Configuration for enabling reasoning capabilities (requires azure.ai.projects.models.Reasoning)."""
 
+    foundry_features: FoundryFeaturesOptInKeys | str
+    """Optional Foundry preview feature opt-in for agent version creation."""
+
 
 AzureAIClientOptionsT = TypeVar(
     "AzureAIClientOptionsT",
@@ -392,7 +396,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
             # response_format is accessed from chat_options or additional_properties
             # since the base class excludes it from run_options
             if chat_options and (response_format := chat_options.get("response_format")):
-                args["text"] = PromptAgentDefinitionText(format=create_text_format_config(response_format))
+                args["text"] = PromptAgentDefinitionTextOptions(format=create_text_format_config(response_format))
 
             # Combine instructions from messages and options
             # instructions is accessed from chat_options since the base class excludes it from run_options
@@ -404,11 +408,15 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
             if combined_instructions:
                 args["instructions"] = "".join(combined_instructions)
 
-            created_agent = await self.project_client.agents.create_version(
-                agent_name=self.agent_name,
-                definition=PromptAgentDefinition(**args),
-                description=self.agent_description,
-            )
+            create_version_kwargs: dict[str, Any] = {
+                "agent_name": self.agent_name,
+                "definition": PromptAgentDefinition(**args),
+                "description": self.agent_description,
+            }
+            if foundry_features := run_options.get("foundry_features"):
+                create_version_kwargs["foundry_features"] = foundry_features
+
+            created_agent = await self.project_client.agents.create_version(**create_version_kwargs)
 
             self.agent_version = created_agent.version
             self.warn_runtime_tools_and_structure_changed = True
@@ -500,6 +508,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
             "temperature": ("temperature",),
             "top_p": ("top_p",),
             "reasoning": ("reasoning",),
+            "foundry_features": ("foundry_features",),
         }
 
         for run_keys in agent_level_option_to_run_keys.values():
@@ -526,9 +535,9 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
             run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"]))
 
         if not self._is_application_endpoint:
-            # Application-scoped response APIs do not support "agent" property.
+            # Application-scoped response APIs do not support "agent_reference" property.
             agent_reference = await self._get_agent_reference_or_create(run_options, instructions, options)
-            run_options["extra_body"] = {"agent": agent_reference}
+            run_options["extra_body"] = {"agent_reference": agent_reference}
 
         # Remove only keys that map to this client's declared options TypedDict.
         self._remove_agent_level_run_options(run_options, options)
@@ -588,6 +597,68 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
         """Get the current conversation ID from chat options or kwargs."""
         return options.get("conversation_id") or kwargs.get("conversation_id") or self.conversation_id
 
+    @override
+    def _parse_response_from_openai(
+        self,
+        response: Any,
+        options: dict[str, Any],
+    ) -> ChatResponse:
+        """Parse an Azure AI Responses API response, handling Azure-specific output item types."""
+        result = super()._parse_response_from_openai(response, options)
+
+        if result.messages:
+            for item in response.output:
+                if item.type == "oauth_consent_request":
+                    consent_link = item.consent_link
+                    if consent_link and not consent_link.startswith("https://"):
+                        logger.warning("Skipping oauth_consent_request with non-HTTPS consent_link: %s", item)
+                        consent_link = ""
+                    if consent_link:
+                        result.messages[0].contents.append(
+                            Content.from_oauth_consent_request(
+                                consent_link=consent_link,
+                                raw_representation=item,
+                            )
+                        )
+                    else:
+                        logger.warning("Received oauth_consent_request output without consent_link: %s", item)
+
+        return result
+
+    @override
+    def _parse_chunk_from_openai(
+        self,
+        event: Any,
+        options: dict[str, Any],
+        function_call_ids: dict[int, tuple[str, str]],
+    ) -> ChatResponseUpdate:
+        """Parse an Azure AI streaming event, handling Azure-specific event types."""
+        # Intercept output_item.added events for Azure-specific item types
+        if event.type == "response.output_item.added" and event.item.type == "oauth_consent_request":
+            event_item = event.item
+            consent_link = event_item.consent_link
+            if consent_link and not consent_link.startswith("https://"):
+                logger.warning("Skipping oauth_consent_request with non-HTTPS consent_link: %s", event_item)
+                consent_link = ""
+            contents: list[Content] = []
+            if consent_link:
+                contents.append(
+                    Content.from_oauth_consent_request(
+                        consent_link=consent_link,
+                        raw_representation=event_item,
+                    )
+                )
+            else:
+                logger.warning("Received oauth_consent_request output without consent_link: %s", event_item)
+            return ChatResponseUpdate(
+                contents=contents,
+                role="assistant",
+                model_id=self.model_id,
+                raw_representation=event,
+            )
+
+        return super()._parse_chunk_from_openai(event, options, function_call_ids)
+
     def _prepare_messages_for_azure_ai(self, messages: Sequence[Message]) -> tuple[list[Message], str | None]:
         """Prepare input from messages and convert system/developer messages to instructions."""
         result: list[Message] = []
@@ -830,14 +901,16 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
     @staticmethod
     def get_code_interpreter_tool(  # type: ignore[override]
         *,
-        file_ids: list[str] | None = None,
+        file_ids: list[str | Content] | None = None,
         container: Literal["auto"] | dict[str, Any] = "auto",
         **kwargs: Any,
     ) -> CodeInterpreterTool:
         """Create a code interpreter tool configuration for Azure AI Projects.
 
         Keyword Args:
-            file_ids: Optional list of file IDs to make available to the code interpreter.
+            file_ids: Optional list of file IDs or Content objects to make available to
+                the code interpreter. Accepts plain strings or Content.from_hosted_file()
+                instances.
             container: Container configuration. Use "auto" for automatic container management.
                 Note: Custom container settings from this parameter are not used by Azure AI Projects;
                 use file_ids instead.
@@ -857,7 +930,8 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
         # Extract file_ids from container if provided as dict and file_ids not explicitly set
         if file_ids is None and isinstance(container, dict):
             file_ids = container.get("file_ids")
-        tool_container = CodeInterpreterToolAuto(file_ids=file_ids if file_ids else None)
+        resolved = resolve_file_ids(file_ids)
+        tool_container = CodeInterpreterContainerAuto(file_ids=resolved)
         return CodeInterpreterTool(container=tool_container, **kwargs)
 
     @staticmethod

python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py

@@ -18,7 +18,6 @@ from agent_framework._sessions import AgentSession, BaseContextProvider, Session
 from agent_framework._settings import load_settings
 from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
 from azure.ai.projects.aio import AIProjectClient
-from azure.ai.projects.models import ItemParam, ResponsesAssistantMessageItemParam, ResponsesUserMessageItemParam
 
 from ._shared import AzureAISettings
 
@@ -149,7 +148,7 @@ class FoundryMemoryProvider(BaseContextProvider):
         # On first run, retrieve static memories (user profile memories)
         if not state.get("initialized"):
             try:
-                static_search_result = await self.project_client.memory_stores.search_memories(
+                static_search_result = await self.project_client.beta.memory_stores.search_memories(
                     name=self.memory_store_name,
                     scope=self.scope or context.session_id,  # type: ignore[arg-type]
                 )
@@ -169,15 +168,15 @@ class FoundryMemoryProvider(BaseContextProvider):
         if not has_input:
             return
 
-        # Convert input messages to ItemParam format for search
+        # Convert input messages to memory search item format
         items = [
-            ItemParam({"type": "text", "text": msg.text})
+            {"type": "text", "text": msg.text}
             for msg in context.input_messages
             if msg and msg.text and msg.text.strip()
         ]
 
         try:
-            search_result = await self.project_client.memory_stores.search_memories(
+            search_result = await self.project_client.beta.memory_stores.search_memories(
                 name=self.memory_store_name,
                 scope=self.scope or context.session_id,  # type: ignore[arg-type]
                 items=items,
@@ -224,24 +223,24 @@ class FoundryMemoryProvider(BaseContextProvider):
         if context.response and context.response.messages:
             messages_to_store.extend(context.response.messages)
 
-        # Filter and convert messages to ItemParam format
-        items: list[ResponsesUserMessageItemParam | ResponsesAssistantMessageItemParam] = []
+        # Filter and convert messages to memory update item format
+        items: list[dict[str, str]] = []
         for message in messages_to_store:
             if message.role in {"user", "assistant", "system"} and message.text and message.text.strip():
                 if message.role == "user":
-                    items.append(ResponsesUserMessageItemParam(content=message.text))
+                    items.append({"role": "user", "type": "message", "content": message.text})
                 elif message.role == "assistant":
-                    items.append(ResponsesAssistantMessageItemParam(content=message.text))
+                    items.append({"role": "assistant", "type": "message", "content": message.text})
 
         if not items:
             return
 
         try:
             # Fire and forget - don't wait for the update to complete
-            update_poller = await self.project_client.memory_stores.begin_update_memories(
+            update_poller = await self.project_client.beta.memory_stores.begin_update_memories(
                 name=self.memory_store_name,
                 scope=self.scope or context.session_id,  # type: ignore[arg-type]
-                items=items,  # type: ignore[arg-type]
+                items=items,
                 previous_update_id=state.get("previous_update_id"),
                 update_delay=self.update_delay,
             )

python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import logging
 import sys
-from collections.abc import Callable, MutableMapping, Sequence
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
 from typing import Any, Generic
 
 from agent_framework import (
@@ -21,10 +21,9 @@ from agent_framework._tools import ToolTypes
 from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
-    AgentReference,
     AgentVersionDetails,
     PromptAgentDefinition,
-    PromptAgentDefinitionText,
+    PromptAgentDefinitionTextOptions,
 )
 from azure.ai.projects.models import (
     FunctionTool as AzureFunctionTool,
@@ -200,13 +199,14 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
         response_format = opts.get("response_format")
         rai_config = opts.get("rai_config")
         reasoning = opts.get("reasoning")
+        foundry_features = opts.get("foundry_features")
 
         args: dict[str, Any] = {"model": resolved_model}
 
         if instructions:
             args["instructions"] = instructions
         if response_format and isinstance(response_format, (type, dict)):
-            args["text"] = PromptAgentDefinitionText(
+            args["text"] = PromptAgentDefinitionTextOptions(
                 format=create_text_format_config(response_format)  # type: ignore[arg-type]
             )
         if rai_config:
@@ -241,11 +241,15 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
         if all_tools_for_azure:
             args["tools"] = to_azure_ai_tools(all_tools_for_azure)
 
-        created_agent = await self._project_client.agents.create_version(
-            agent_name=name,
-            definition=PromptAgentDefinition(**args),
-            description=description,
-        )
+        create_version_kwargs: dict[str, Any] = {
+            "agent_name": name,
+            "definition": PromptAgentDefinition(**args),
+            "description": description,
+        }
+        if foundry_features:
+            create_version_kwargs["foundry_features"] = foundry_features
+
+        created_agent = await self._project_client.agents.create_version(**create_version_kwargs)
 
         return self._to_chat_agent_from_details(
             created_agent,
@@ -259,7 +263,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
         self,
         *,
         name: str | None = None,
-        reference: AgentReference | None = None,
+        reference: Mapping[str, str | None] | None = None,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
         default_options: OptionsCoT | None = None,
         middleware: Sequence[MiddlewareTypes] | None = None,
@@ -272,7 +276,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
 
         Args:
             name: The name of the agent to retrieve (fetches latest version).
-            reference: Reference containing the agent's name and optionally a specific version.
+            reference: Mapping containing the agent's ``name`` and optionally a specific ``version``.
             tools: Tools to make available to the agent. Required if the agent has function tools.
             default_options: A TypedDict containing default chat options for the agent.
                 These options are applied to every run unless overridden.
@@ -287,12 +291,15 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
         """
         existing_agent: AgentVersionDetails
 
-        if reference and reference.version:
+        reference_name = str(reference.get("name")) if reference and reference.get("name") else None
+        reference_version = str(reference.get("version")) if reference and reference.get("version") else None
+
+        if reference_name and reference_version:
             # Fetch specific version
             existing_agent = await self._project_client.agents.get_version(
-                agent_name=reference.name, agent_version=reference.version
+                agent_name=reference_name, agent_version=reference_version
             )
-        elif agent_name := (reference.name if reference else name):
+        elif agent_name := (reference_name if reference_name else name):
             # Fetch latest version
             details = await self._project_client.agents.get(agent_name=agent_name)
             existing_agent = details.versions.latest

python/packages/azure-ai/agent_framework_azure_ai/_shared.py

@@ -8,6 +8,7 @@ from collections.abc import Mapping, MutableMapping, Sequence
 from typing import Any, cast
 
 from agent_framework import (
+    Content,
     FunctionTool,
 )
 from agent_framework.exceptions import IntegrationInvalidRequestException
@@ -18,9 +19,9 @@ from azure.ai.agents.models import (
 from azure.ai.projects.models import (
     CodeInterpreterTool,
     MCPTool,
-    ResponseTextFormatConfigurationJsonObject,
-    ResponseTextFormatConfigurationJsonSchema,
-    ResponseTextFormatConfigurationText,
+    TextResponseFormatConfigurationResponseFormatJsonObject,
+    TextResponseFormatConfigurationResponseFormatText,
+    TextResponseFormatJsonSchema,
     Tool,
     WebSearchPreviewTool,
 )
@@ -109,6 +110,47 @@ def _extract_project_connection_id(additional_properties: dict[str, Any] | None)
     return None
 
 
+def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None:
+    """Resolve a list of file ID values that may include Content objects.
+
+    Accepts plain strings and Content objects with type "hosted_file", extracting
+    the file_id from each. This enables users to pass Content.from_hosted_file()
+    alongside plain file ID strings.
+
+    Args:
+        file_ids: Sequence of file ID strings or Content objects, or None.
+
+    Returns:
+        A list of resolved file ID strings, or None if input is None or empty.
+
+    Raises:
+        ValueError: If a Content object has an unsupported type (not "hosted_file").
+    """
+    if not file_ids:
+        return None
+
+    resolved: list[str] = []
+    for item in file_ids:
+        if isinstance(item, str):
+            if not item:
+                raise ValueError("file_ids must not contain empty strings.")
+            resolved.append(item)
+        elif isinstance(item, Content):
+            if item.type != "hosted_file":
+                raise ValueError(
+                    f"Unsupported Content type '{item.type}' for code interpreter file_ids. "
+                    "Only Content.from_hosted_file() is supported."
+                )
+            if item.file_id is None:
+                raise ValueError(
+                    "Content.from_hosted_file() item is missing a file_id. "
+                    "Ensure the Content object has a valid file_id before using it in file_ids."
+                )
+            resolved.append(item.file_id)
+
+    return resolved if resolved else None
+
+
 def to_azure_ai_agent_tools(
     tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None,
     run_options: dict[str, Any] | None = None,
@@ -421,9 +463,9 @@ def _prepare_mcp_tool_dict_for_azure_ai(tool_dict: dict[str, Any]) -> MCPTool:
 def create_text_format_config(
     response_format: type[BaseModel] | Mapping[str, Any],
 ) -> (
-    ResponseTextFormatConfigurationJsonSchema
-    | ResponseTextFormatConfigurationJsonObject
-    | ResponseTextFormatConfigurationText
+    TextResponseFormatJsonSchema
+    | TextResponseFormatConfigurationResponseFormatJsonObject
+    | TextResponseFormatConfigurationResponseFormatText
 ):
     """Convert response_format into Azure text format configuration."""
     if isinstance(response_format, type) and issubclass(response_format, BaseModel):
@@ -431,7 +473,7 @@ def create_text_format_config(
         # Ensure additionalProperties is explicitly false to satisfy Azure validation
         if isinstance(schema, dict):
             schema.setdefault("additionalProperties", False)
-        return ResponseTextFormatConfigurationJsonSchema(
+        return TextResponseFormatJsonSchema(
             name=response_format.__name__,
             schema=schema,
             strict=True,
@@ -452,11 +494,11 @@ def create_text_format_config(
                 config_kwargs["strict"] = format_config["strict"]
             if "description" in format_config:
                 config_kwargs["description"] = format_config["description"]
-            return ResponseTextFormatConfigurationJsonSchema(**config_kwargs)
+            return TextResponseFormatJsonSchema(**config_kwargs)
         if format_type == "json_object":
-            return ResponseTextFormatConfigurationJsonObject()
+            return TextResponseFormatConfigurationResponseFormatJsonObject()
         if format_type == "text":
-            return ResponseTextFormatConfigurationText()
+            return TextResponseFormatConfigurationResponseFormatText()
 
     raise IntegrationInvalidRequestException("response_format must be a Pydantic model or mapping.")
 

python/packages/azure-ai/tests/test_azure_ai_agent_client.py

@@ -855,6 +855,110 @@ async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_file_search_with_
     assert run_options["tool_resources"] == {"file_search": {"vector_store_ids": ["vs-123"]}}
 
 
+async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_code_interpreter_with_file_ids(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_tools_for_azure_ai with CodeInterpreterTool with file_ids from get_code_interpreter_tool()."""
+
+    client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+    code_interpreter_tool = client.get_code_interpreter_tool(file_ids=["file-123", "file-456"])
+
+    run_options: dict[str, Any] = {}
+    result = await client._prepare_tools_for_azure_ai([code_interpreter_tool], run_options)  # type: ignore
+
+    assert len(result) == 1
+    assert result[0] == {"type": "code_interpreter"}
+    assert "tool_resources" in run_options
+    assert "code_interpreter" in run_options["tool_resources"]
+    assert sorted(run_options["tool_resources"]["code_interpreter"]["file_ids"]) == ["file-123", "file-456"]
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_basic() -> None:
+    """Test get_code_interpreter_tool returns CodeInterpreterTool without files."""
+    from azure.ai.agents.models import CodeInterpreterTool
+
+    tool = AzureAIAgentClient.get_code_interpreter_tool()
+    assert isinstance(tool, CodeInterpreterTool)
+    assert len(tool.file_ids) == 0
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_file_ids() -> None:
+    """Test get_code_interpreter_tool forwards file_ids to the SDK."""
+    from azure.ai.agents.models import CodeInterpreterTool
+
+    tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc", "file-def"])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert "file-abc" in tool.file_ids
+    assert "file-def" in tool.file_ids
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_data_sources() -> None:
+    """Test get_code_interpreter_tool forwards data_sources to the SDK."""
+    from azure.ai.agents.models import CodeInterpreterTool, VectorStoreDataSource
+
+    ds = VectorStoreDataSource(asset_identifier="test-asset-id", asset_type="id_asset")
+    tool = AzureAIAgentClient.get_code_interpreter_tool(data_sources=[ds])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert "test-asset-id" in tool.data_sources
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_mutually_exclusive() -> None:
+    """Test get_code_interpreter_tool raises ValueError when both file_ids and data_sources are provided."""
+    from azure.ai.agents.models import VectorStoreDataSource
+
+    ds = VectorStoreDataSource(asset_identifier="test-asset-id", asset_type="id_asset")
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc"], data_sources=[ds])
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_content() -> None:
+    """Test get_code_interpreter_tool accepts Content.from_hosted_file in file_ids."""
+    from agent_framework import Content
+    from azure.ai.agents.models import CodeInterpreterTool
+
+    content = Content.from_hosted_file("file-content-123")
+    tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert "file-content-123" in tool.file_ids
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_mixed_file_ids() -> None:
+    """Test get_code_interpreter_tool accepts a mix of strings and Content objects."""
+    from agent_framework import Content
+    from azure.ai.agents.models import CodeInterpreterTool
+
+    content = Content.from_hosted_file("file-from-content")
+    tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-plain", content])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert "file-plain" in tool.file_ids
+    assert "file-from-content" in tool.file_ids
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_content_unsupported_type() -> None:
+    """Test get_code_interpreter_tool raises ValueError for unsupported Content types."""
+    from agent_framework import Content
+
+    content = Content.from_hosted_vector_store("vs-123")
+    with pytest.raises(ValueError, match="Unsupported Content type"):
+        AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_content_missing_file_id() -> None:
+    """Test get_code_interpreter_tool raises ValueError when Content.file_id is None."""
+    from agent_framework import Content
+
+    content = Content(type="hosted_file")
+    with pytest.raises(ValueError, match="missing a file_id"):
+        AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_empty_string_file_id() -> None:
+    """Test get_code_interpreter_tool raises ValueError for empty string file_ids."""
+    with pytest.raises(ValueError, match="must not contain empty strings"):
+        AzureAIAgentClient.get_code_interpreter_tool(file_ids=[""])
+
+
 async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals(
     mock_agents_client: MagicMock,
 ) -> None:

python/packages/azure-ai/tests/test_azure_ai_client.py

@@ -28,12 +28,12 @@ from agent_framework.openai._responses_client import RawOpenAIResponsesClient
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
     ApproximateLocation,
+    CodeInterpreterContainerAuto,
     CodeInterpreterTool,
-    CodeInterpreterToolAuto,
     FileSearchTool,
     ImageGenTool,
     MCPTool,
-    ResponseTextFormatConfigurationJsonSchema,
+    TextResponseFormatJsonSchema,
     WebSearchPreviewTool,
 )
 from azure.core.exceptions import ResourceNotFoundError
@@ -427,7 +427,7 @@ async def test_prepare_options_basic(mock_project_client: MagicMock) -> None:
         run_options = await client._prepare_options(messages, {})
 
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
 
 
 @pytest.mark.parametrize(
@@ -465,7 +465,7 @@ async def test_prepare_options_with_application_endpoint(
 
     if expects_agent:
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
     else:
         assert "extra_body" not in run_options
 
@@ -507,7 +507,7 @@ async def test_prepare_options_with_application_project_client(
 
     if expects_agent:
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
     else:
         assert "extra_body" not in run_options
 
@@ -979,10 +979,10 @@ async def test_agent_creation_with_response_format(
     assert hasattr(created_definition, "text")
     assert created_definition.text is not None
 
-    # Check that the format is a ResponseTextFormatConfigurationJsonSchema
+    # Check that the format is a TextResponseFormatJsonSchema
     assert hasattr(created_definition.text, "format")
     format_config = created_definition.text.format
-    assert isinstance(format_config, ResponseTextFormatConfigurationJsonSchema)
+    assert isinstance(format_config, TextResponseFormatJsonSchema)
 
     # Check the schema name matches the model class name
     assert format_config.name == "ResponseFormatModel"
@@ -1040,7 +1040,7 @@ async def test_agent_creation_with_mapping_response_format(
     assert hasattr(created_definition, "text")
     assert created_definition.text is not None
     format_config = created_definition.text.format
-    assert isinstance(format_config, ResponseTextFormatConfigurationJsonSchema)
+    assert isinstance(format_config, TextResponseFormatJsonSchema)
     assert format_config.name == runtime_schema["title"]
     assert format_config.schema == runtime_schema
     assert format_config.strict is True
@@ -1110,7 +1110,7 @@ async def test_prepare_options_excludes_response_format(
         assert "text_format" not in run_options
         # But extra_body should contain agent reference
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
 
 
 async def test_prepare_options_keeps_values_for_unsupported_option_keys(
@@ -1254,7 +1254,7 @@ def test_from_azure_ai_tools_mcp() -> None:
 
 def test_from_azure_ai_tools_code_interpreter() -> None:
     """Test from_azure_ai_tools with Code Interpreter tool."""
-    ci_tool = CodeInterpreterTool(container=CodeInterpreterToolAuto(file_ids=["file-1"]))
+    ci_tool = CodeInterpreterTool(container=CodeInterpreterContainerAuto(file_ids=["file-1"]))
     parsed_tools = from_azure_ai_tools([ci_tool])
     assert len(parsed_tools) == 1
     assert parsed_tools[0]["type"] == "code_interpreter"
@@ -1685,6 +1685,35 @@ def test_get_code_interpreter_tool_with_file_ids() -> None:
     assert tool["container"]["file_ids"] == ["file-123", "file-456"]
 
 
+def test_get_code_interpreter_tool_with_content() -> None:
+    """Test get_code_interpreter_tool accepts Content.from_hosted_file in file_ids."""
+    from agent_framework import Content
+
+    content = Content.from_hosted_file("file-content-123")
+    tool = AzureAIClient.get_code_interpreter_tool(file_ids=[content])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert tool["container"]["file_ids"] == ["file-content-123"]
+
+
+def test_get_code_interpreter_tool_with_mixed_file_ids() -> None:
+    """Test get_code_interpreter_tool accepts a mix of strings and Content objects."""
+    from agent_framework import Content
+
+    content = Content.from_hosted_file("file-from-content")
+    tool = AzureAIClient.get_code_interpreter_tool(file_ids=["file-plain", content])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert sorted(tool["container"]["file_ids"]) == ["file-from-content", "file-plain"]
+
+
+def test_get_code_interpreter_tool_content_unsupported_type() -> None:
+    """Test get_code_interpreter_tool raises ValueError for unsupported Content types."""
+    from agent_framework import Content
+
+    content = Content.from_hosted_vector_store("vs-123")
+    with pytest.raises(ValueError, match="Unsupported Content type"):
+        AzureAIClient.get_code_interpreter_tool(file_ids=[content])
+
+
 def test_get_file_search_tool_basic() -> None:
     """Test get_file_search_tool returns FileSearchTool."""
     tool = AzureAIClient.get_file_search_tool(vector_store_ids=["vs-123"])
@@ -2145,4 +2174,103 @@ def test_build_url_citation_content_with_dict(mock_project_client: MagicMock) ->
     assert "get_url" not in ann.get("additional_properties", {})
 
 
+# region OAuth Consent
+
+
+def test_parse_chunk_with_oauth_consent_request(mock_project_client: MagicMock) -> None:
+    """Test that a streaming oauth_consent_request output item is parsed into oauth_consent_request content.
+
+    This reproduces the bug from issue #3950 where the event was logged as "Unparsed event"
+    and silently discarded, causing the agent run to complete with zero content.
+    """
+    client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+    chat_options: dict[str, Any] = {}
+    function_call_ids: dict[int, tuple[str, str]] = {}
+
+    mock_item = MagicMock()
+    mock_item.type = "oauth_consent_request"
+    mock_item.consent_link = "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc123"
+
+    mock_event = MagicMock()
+    mock_event.type = "response.output_item.added"
+    mock_event.item = mock_item
+    mock_event.output_index = 0
+
+    update = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
+
+    assert len(update.contents) == 1
+    consent_content = update.contents[0]
+    assert consent_content.type == "oauth_consent_request"
+    assert consent_content.consent_link == "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc123"
+    assert consent_content.user_input_request is True
+
+
+def test_parse_response_with_oauth_consent_output_item(mock_project_client: MagicMock) -> None:
+    """Test that a non-streaming oauth_consent_request output item is parsed correctly."""
+    client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+
+    mock_item = MagicMock()
+    mock_item.type = "oauth_consent_request"
+    mock_item.consent_link = "https://login.microsoftonline.com/consent?code=abc"
+
+    mock_response = MagicMock()
+    mock_response.output = [mock_item]
+    mock_response.output_parsed = None
+    mock_response.metadata = {}
+    mock_response.id = "resp-oauth-1"
+    mock_response.model = "test-model"
+    mock_response.created_at = 1000000000
+    mock_response.usage = None
+    mock_response.status = "completed"
+
+    response = client._parse_response_from_openai(mock_response, {})
+
+    assert len(response.messages) > 0
+    consent_contents = [c for c in response.messages[0].contents if c.type == "oauth_consent_request"]
+    assert len(consent_contents) == 1
+    assert consent_contents[0].consent_link == "https://login.microsoftonline.com/consent?code=abc"
+
+
+def test_parse_chunk_oauth_consent_no_link(mock_project_client: MagicMock) -> None:
+    """Test that a streaming oauth_consent_request with no consent_link produces empty contents."""
+    client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+
+    mock_item = MagicMock()
+    mock_item.type = "oauth_consent_request"
+    mock_item.consent_link = ""
+
+    mock_event = MagicMock()
+    mock_event.type = "response.output_item.added"
+    mock_event.item = mock_item
+    mock_event.output_index = 0
+
+    update = client._parse_chunk_from_openai(mock_event, {}, {})
+
+    assert not any(c.type == "oauth_consent_request" for c in update.contents)
+
+
+def test_parse_response_oauth_consent_no_link(mock_project_client: MagicMock) -> None:
+    """Test that a non-streaming oauth_consent_request with no consent_link appends no content."""
+    client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+
+    mock_item = MagicMock()
+    mock_item.type = "oauth_consent_request"
+    mock_item.consent_link = None
+
+    mock_response = MagicMock()
+    mock_response.output = [mock_item]
+    mock_response.output_parsed = None
+    mock_response.metadata = {}
+    mock_response.id = "resp-oauth-2"
+    mock_response.model = "test-model"
+    mock_response.created_at = 1000000000
+    mock_response.usage = None
+    mock_response.status = "completed"
+
+    response = client._parse_response_from_openai(mock_response, {})
+
+    consent_contents = [c for c in response.messages[0].contents if c.type == "oauth_consent_request"]
+    assert len(consent_contents) == 0
+
+
 # endregion

python/packages/azure-ai/tests/test_foundry_memory_provider.py

@@ -17,9 +17,10 @@ from agent_framework_azure_ai._foundry_memory_provider import FoundryMemoryProvi
 def mock_project_client() -> AsyncMock:
     """Create a mock AIProjectClient."""
     mock_client = AsyncMock()
-    mock_client.memory_stores = AsyncMock()
-    mock_client.memory_stores.search_memories = AsyncMock()
-    mock_client.memory_stores.begin_update_memories = AsyncMock()
+    mock_client.beta = AsyncMock()
+    mock_client.beta.memory_stores = AsyncMock()
+    mock_client.beta.memory_stores.search_memories = AsyncMock()
+    mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
     mock_client.__aenter__ = AsyncMock(return_value=mock_client)
     mock_client.__aexit__ = AsyncMock()
     return mock_client
@@ -146,7 +147,7 @@ class TestBeforeRun:
         mem2.memory_item.content = "User is based in Seattle"
         mock_search_result = Mock()
         mock_search_result.memories = [mem1, mem2]
-        mock_project_client.memory_stores.search_memories.return_value = mock_search_result
+        mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -161,7 +162,7 @@ class TestBeforeRun:
         )
 
         # Should call search_memories twice: once for static, once for contextual
-        assert mock_project_client.memory_stores.search_memories.call_count == 2
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
         # Static memories should be cached
         assert len(session.state[provider.source_id]["static_memories"]) == 2
         assert session.state[provider.source_id]["initialized"] is True
@@ -181,7 +182,7 @@ class TestBeforeRun:
         contextual_result.memories = [contextual_mem]
         contextual_result.search_id = "search-123"
 
-        mock_project_client.memory_stores.search_memories.side_effect = [static_result, contextual_result]
+        mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -208,7 +209,7 @@ class TestBeforeRun:
         """Empty input messages → only static search performed, no contextual search."""
         static_result = Mock()
         static_result.memories = []
-        mock_project_client.memory_stores.search_memories.return_value = static_result
+        mock_project_client.beta.memory_stores.search_memories.return_value = static_result
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -223,14 +224,14 @@ class TestBeforeRun:
         )
 
         # Should only call search_memories once for static memories
-        assert mock_project_client.memory_stores.search_memories.call_count == 1
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
         assert provider.source_id not in ctx.context_messages
 
     async def test_empty_search_results_no_messages(self, mock_project_client: AsyncMock) -> None:
         """Empty search results → no messages added."""
         mock_search_result = Mock()
         mock_search_result.memories = []
-        mock_project_client.memory_stores.search_memories.return_value = mock_search_result
+        mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -255,7 +256,7 @@ class TestBeforeRun:
         contextual_result = Mock()
         contextual_result.memories = []
 
-        mock_project_client.memory_stores.search_memories.side_effect = [static_result, contextual_result]
+        mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -269,24 +270,24 @@ class TestBeforeRun:
         await provider.before_run(  # type: ignore[arg-type]
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
-        assert mock_project_client.memory_stores.search_memories.call_count == 2
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
 
         # Reset mock for second call
-        mock_project_client.memory_stores.search_memories.reset_mock()
+        mock_project_client.beta.memory_stores.search_memories.reset_mock()
         contextual_result2 = Mock()
         contextual_result2.memories = []
-        mock_project_client.memory_stores.search_memories.return_value = contextual_result2
+        mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
 
         # Second call - should only search contextual, not static
         ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
         await provider.before_run(  # type: ignore[arg-type]
             agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
         )
-        assert mock_project_client.memory_stores.search_memories.call_count == 1
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
 
     async def test_handles_search_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
         """Search exception is logged but doesn't fail the operation."""
-        mock_project_client.memory_stores.search_memories.side_effect = Exception("API error")
+        mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -315,7 +316,7 @@ class TestAfterRun:
         """Stores input+response messages via begin_update_memories."""
         mock_poller = Mock()
         mock_poller.update_id = "update-456"
-        mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+        mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -330,8 +331,8 @@ class TestAfterRun:
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        mock_project_client.memory_stores.begin_update_memories.assert_awaited_once()
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         assert call_kwargs["name"] == "test_store"
         assert call_kwargs["scope"] == "user_123"
         assert len(call_kwargs["items"]) == 2
@@ -342,7 +343,7 @@ class TestAfterRun:
     async def test_only_stores_user_assistant_system(self, mock_project_client: AsyncMock) -> None:
         """Only stores user/assistant/system messages with text."""
         mock_poller = Mock()
-        mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+        mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -363,7 +364,7 @@ class TestAfterRun:
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         items = call_kwargs["items"]
         assert len(items) == 2
         assert items[0]["content"] == "hello"
@@ -390,12 +391,12 @@ class TestAfterRun:
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        mock_project_client.memory_stores.begin_update_memories.assert_not_awaited()
+        mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
 
     async def test_uses_configured_update_delay(self, mock_project_client: AsyncMock) -> None:
         """Uses the configured update_delay parameter."""
         mock_poller = Mock()
-        mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+        mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -411,7 +412,7 @@ class TestAfterRun:
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         assert call_kwargs["update_delay"] == 60
 
     async def test_uses_previous_update_id_for_incremental_updates(self, mock_project_client: AsyncMock) -> None:
@@ -421,7 +422,7 @@ class TestAfterRun:
         mock_poller2 = Mock()
         mock_poller2.update_id = "update-2"
 
-        mock_project_client.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
+        mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -446,13 +447,13 @@ class TestAfterRun:
             agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
         )
 
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         assert call_kwargs["previous_update_id"] == "update-1"
         assert session.state[provider.source_id]["previous_update_id"] == "update-2"
 
     async def test_handles_update_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
         """Update exception is logged but doesn't fail the operation."""
-        mock_project_client.memory_stores.begin_update_memories.side_effect = Exception("API error")
+        mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,

python/packages/azure-ai/tests/test_provider.py

@@ -8,7 +8,6 @@ from agent_framework import Agent, FunctionTool
 from agent_framework._mcp import MCPTool
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
-    AgentReference,
     AgentVersionDetails,
     PromptAgentDefinition,
 )
@@ -345,7 +344,7 @@ async def test_provider_get_agent_with_reference(mock_project_client: MagicMock)
     mock_project_client.agents = AsyncMock()
     mock_project_client.agents.get_version.return_value = mock_agent_version
 
-    agent_reference = AgentReference(name="test-agent", version="1.0")
+    agent_reference = {"name": "test-agent", "version": "1.0"}
     agent = await provider.get_agent(reference=agent_reference)
 
     assert isinstance(agent, Agent)

python/packages/claude/agent_framework_claude/__init__.py

@@ -2,7 +2,7 @@
 
 import importlib.metadata
 
-from ._agent import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings
+from ._agent import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings, RawClaudeAgent
 
 try:
     __version__ = importlib.metadata.version(__name__)
@@ -13,5 +13,6 @@ __all__ = [
     "ClaudeAgent",
     "ClaudeAgentOptions",
     "ClaudeAgentSettings",
+    "RawClaudeAgent",
     "__version__",
 ]

python/packages/claude/agent_framework_claude/_agent.py

@@ -27,6 +27,7 @@ from agent_framework import (
     normalize_tools,
 )
 from agent_framework.exceptions import AgentException
+from agent_framework.observability import AgentTelemetryLayer
 from claude_agent_sdk import (
     AssistantMessage,
     ClaudeSDKClient,
@@ -171,8 +172,11 @@ OptionsT = TypeVar(
 )
 
 
-class ClaudeAgent(BaseAgent, Generic[OptionsT]):
-    """Claude Agent using Claude Code CLI.
+class RawClaudeAgent(BaseAgent, Generic[OptionsT]):
+    """Claude Agent using Claude Code CLI without telemetry layers.
+
+    This is the core Claude agent implementation without OpenTelemetry instrumentation.
+    For most use cases, prefer :class:`ClaudeAgent` which includes telemetry support.
 
     Wraps the Claude Agent SDK to provide agentic capabilities including
     tool use, session management, and streaming responses.
@@ -188,45 +192,13 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
 
         .. code-block:: python
 
-            from agent_framework_claude import ClaudeAgent
+            from agent_framework.anthropic import RawClaudeAgent
 
-            async with ClaudeAgent(
+            async with RawClaudeAgent(
                 instructions="You are a helpful assistant.",
             ) as agent:
                 response = await agent.run("Hello!")
                 print(response.text)
-
-        With streaming:
-
-        .. code-block:: python
-
-            async with ClaudeAgent() as agent:
-                async for update in agent.run("Write a poem"):
-                    print(update.text, end="", flush=True)
-
-        With session management:
-
-        .. code-block:: python
-
-            async with ClaudeAgent() as agent:
-                session = agent.create_session()
-                await agent.run("Remember my name is Alice", session=session)
-                response = await agent.run("What's my name?", session=session)
-                # Claude will remember "Alice" from the same session
-
-        With Agent Framework tools:
-
-        .. code-block:: python
-
-            from agent_framework import tool
-
-            @tool
-            def greet(name: str) -> str:
-                \"\"\"Greet someone by name.\"\"\"
-                return f"Hello, {name}!"
-
-            async with ClaudeAgent(tools=[greet]) as agent:
-                response = await agent.run("Greet Alice")
     """
 
     AGENT_PROVIDER_NAME: ClassVar[str] = "anthropic.claude"
@@ -246,7 +218,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
-        """Initialize a ClaudeAgent instance.
+        """Initialize a RawClaudeAgent instance.
 
         Args:
             instructions: System prompt for the agent.
@@ -343,7 +315,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
                 normalized = normalize_tools(tool)
                 self._custom_tools.extend(normalized)
 
-    async def __aenter__(self) -> ClaudeAgent[OptionsT]:
+    async def __aenter__(self) -> RawClaudeAgent[OptionsT]:
         """Start the agent when entering async context."""
         await self.start()
         return self
@@ -568,61 +540,19 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
             return ""
         return "\n".join([msg.text or "" for msg in messages])
 
-    @overload
-    def run(
-        self,
-        messages: AgentRunInputs | None = None,
-        *,
-        stream: Literal[True],
-        session: AgentSession | None = None,
-        options: OptionsT | MutableMapping[str, Any] | None = None,
-        **kwargs: Any,
-    ) -> AsyncIterable[AgentResponseUpdate]: ...
+    @property
+    def default_options(self) -> dict[str, Any]:
+        """Expose options with ``instructions`` key.
 
-    @overload
-    async def run(
-        self,
-        messages: AgentRunInputs | None = None,
-        *,
-        stream: Literal[False] = ...,
-        session: AgentSession | None = None,
-        options: OptionsT | MutableMapping[str, Any] | None = None,
-        **kwargs: Any,
-    ) -> AgentResponse[Any]: ...
-
-    def run(
-        self,
-        messages: AgentRunInputs | None = None,
-        *,
-        stream: bool = False,
-        session: AgentSession | None = None,
-        options: OptionsT | MutableMapping[str, Any] | None = None,
-        **kwargs: Any,
-    ) -> AsyncIterable[AgentResponseUpdate] | Awaitable[AgentResponse[Any]]:
-        """Run the agent with the given messages.
-
-        Args:
-            messages: The messages to process.
-
-        Keyword Args:
-            stream: If True, returns an async iterable of updates. If False (default),
-                returns an awaitable AgentResponse.
-            session: The conversation session. If session has service_session_id set,
-                the agent will resume that session.
-            options: Runtime options (model, permission_mode can be changed per-request).
-            kwargs: Additional keyword arguments.
-
-        Returns:
-            When stream=True: An ResponseStream for streaming updates.
-            When stream=False: An Awaitable[AgentResponse] with the complete response.
+        Maps ``system_prompt`` to ``instructions`` for compatibility with
+        :class:`AgentTelemetryLayer`, which reads the system prompt from
+        the ``instructions`` key.
         """
-        response = ResponseStream(
-            self._get_stream(messages, session=session, options=options, **kwargs),
-            finalizer=self._finalize_response,
-        )
-        if stream:
-            return response
-        return response.get_final_response()
+        opts = dict(self._default_options)
+        system_prompt = opts.pop("system_prompt", None)
+        if system_prompt is not None:
+            opts["instructions"] = system_prompt
+        return opts
 
     def _finalize_response(self, updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]:
         """Build AgentResponse and propagate structured_output as value.
@@ -636,6 +566,61 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
         structured_output = getattr(self, "_structured_output", None)
         return AgentResponse.from_updates(updates, value=structured_output)
 
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        """Run the agent with the given messages.
+
+        Args:
+            messages: The messages to process.
+
+        Keyword Args:
+            stream: If True, returns an async iterable of updates. If False (default),
+                returns an awaitable AgentResponse.
+            session: The conversation session. If session has service_session_id set,
+                the agent will resume that session.
+            kwargs: Additional keyword arguments including 'options' for runtime options
+                (model, permission_mode can be changed per-request).
+
+        Returns:
+            When stream=True: An ResponseStream for streaming updates.
+            When stream=False: An Awaitable[AgentResponse] with the complete response.
+        """
+        options = kwargs.pop("options", None)
+        response = ResponseStream(
+            self._get_stream(messages, session=session, options=options, **kwargs),
+            finalizer=self._finalize_response,
+        )
+
+        if stream:
+            return response
+        return response.get_final_response()
+
     async def _get_stream(
         self,
         messages: AgentRunInputs | None = None,
@@ -721,3 +706,25 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
 
         # Store structured output for the finalizer
         self._structured_output = structured_output
+
+
+class ClaudeAgent(AgentTelemetryLayer, RawClaudeAgent[OptionsT], Generic[OptionsT]):
+    """Claude Agent with OpenTelemetry instrumentation.
+
+    This is the recommended agent class for most use cases. It includes
+    OpenTelemetry-based telemetry for observability. For a minimal
+    implementation without telemetry, use :class:`RawClaudeAgent`.
+
+    Examples:
+        Basic usage with context manager:
+
+        .. code-block:: python
+
+            from agent_framework.anthropic import ClaudeAgent
+
+            async with ClaudeAgent(
+                instructions="You are a helpful assistant.",
+            ) as agent:
+                response = await agent.run("Hello!")
+                print(response.text)
+    """

python/packages/claude/tests/test_claude_agent.py

@@ -945,3 +945,191 @@ class TestClaudeAgentStructuredOutput:
             with pytest.raises(AgentException) as exc_info:
                 await agent.run("Hello")
             assert "Something went wrong" in str(exc_info.value)
+
+
+# region Test ClaudeAgent Telemetry
+
+
+class TestClaudeAgentTelemetry:
+    """Tests for ClaudeAgent OpenTelemetry instrumentation."""
+
+    @staticmethod
+    async def _create_async_generator(items: list[Any]) -> Any:
+        """Helper to create async generator from list."""
+        for item in items:
+            yield item
+
+    def _create_mock_client(self, messages: list[Any]) -> MagicMock:
+        """Create a mock ClaudeSDKClient that yields given messages."""
+        mock_client = MagicMock()
+        mock_client.connect = AsyncMock()
+        mock_client.disconnect = AsyncMock()
+        mock_client.query = AsyncMock()
+        mock_client.set_model = AsyncMock()
+        mock_client.set_permission_mode = AsyncMock()
+        mock_client.receive_response = MagicMock(return_value=self._create_async_generator(messages))
+        return mock_client
+
+    def _create_standard_messages(self) -> list[Any]:
+        """Create a standard set of mock messages for testing."""
+        from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
+        from claude_agent_sdk.types import StreamEvent
+
+        return [
+            StreamEvent(
+                event={
+                    "type": "content_block_delta",
+                    "delta": {"type": "text_delta", "text": "Hello!"},
+                },
+                uuid="event-1",
+                session_id="session-123",
+            ),
+            AssistantMessage(
+                content=[TextBlock(text="Hello!")],
+                model="claude-sonnet",
+            ),
+            ResultMessage(
+                subtype="success",
+                duration_ms=100,
+                duration_api_ms=50,
+                is_error=False,
+                num_turns=1,
+                session_id="session-123",
+            ),
+        ]
+
+    async def test_run_emits_span_when_instrumentation_enabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Test that run() creates an OpenTelemetry span when instrumentation is enabled."""
+        from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+        messages = self._create_standard_messages()
+        mock_client = self._create_mock_client(messages)
+
+        monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+        with (
+            patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+            patch("agent_framework.observability._get_span") as mock_get_span,
+        ):
+            mock_span = MagicMock()
+            mock_get_span.return_value.__enter__ = MagicMock(return_value=mock_span)
+            mock_get_span.return_value.__exit__ = MagicMock(return_value=False)
+
+            agent = ClaudeAgent(name="test-agent")
+            response = await agent.run("Hello")
+
+            assert response.text == "Hello!"
+            mock_get_span.assert_called_once()
+            call_kwargs = mock_get_span.call_args[1]
+            assert call_kwargs["attributes"]["gen_ai.agent.name"] == "test-agent"
+            assert call_kwargs["attributes"]["gen_ai.operation.name"] == "invoke_agent"
+
+    async def test_run_skips_telemetry_when_instrumentation_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Test that run() skips telemetry when instrumentation is disabled."""
+        from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+        messages = self._create_standard_messages()
+        mock_client = self._create_mock_client(messages)
+
+        monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", False)
+
+        with (
+            patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+            patch("agent_framework.observability._get_span") as mock_get_span,
+        ):
+            agent = ClaudeAgent(name="test-agent")
+            response = await agent.run("Hello")
+
+            assert response.text == "Hello!"
+            mock_get_span.assert_not_called()
+
+    async def test_run_stream_emits_span_when_instrumentation_enabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Test that run(stream=True) creates a span when instrumentation is enabled."""
+        from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+        messages = self._create_standard_messages()
+        mock_client = self._create_mock_client(messages)
+
+        monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+        with (
+            patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+            patch("agent_framework.observability.get_tracer") as mock_get_tracer,
+        ):
+            mock_span = MagicMock()
+            mock_tracer = MagicMock()
+            mock_tracer.start_span.return_value = mock_span
+            mock_get_tracer.return_value = mock_tracer
+
+            agent = ClaudeAgent(name="stream-agent")
+            updates: list[AgentResponseUpdate] = []
+            async for update in agent.run("Hello", stream=True):
+                updates.append(update)
+
+            assert len(updates) == 1
+            mock_tracer.start_span.assert_called_once()
+            span_name = mock_tracer.start_span.call_args[0][0]
+            assert "stream-agent" in span_name
+            assert "invoke_agent" in span_name
+
+    async def test_run_captures_exception_in_span(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Test that exceptions during run() are captured in the telemetry span."""
+        from agent_framework.exceptions import AgentException
+        from agent_framework.observability import OBSERVABILITY_SETTINGS
+        from claude_agent_sdk import ResultMessage
+
+        error_messages = [
+            ResultMessage(
+                subtype="error",
+                duration_ms=100,
+                duration_api_ms=50,
+                is_error=True,
+                num_turns=0,
+                session_id="error-session",
+                result="Model not found",
+            ),
+        ]
+        mock_client = self._create_mock_client(error_messages)
+
+        monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+        with (
+            patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+            patch("agent_framework.observability._get_span") as mock_get_span,
+            patch("agent_framework.observability.capture_exception") as mock_capture_exc,
+        ):
+            mock_span = MagicMock()
+            mock_get_span.return_value.__enter__ = MagicMock(return_value=mock_span)
+            mock_get_span.return_value.__exit__ = MagicMock(return_value=False)
+
+            agent = ClaudeAgent(name="error-agent")
+            with pytest.raises(AgentException):
+                await agent.run("Hello")
+
+            mock_capture_exc.assert_called_once()
+            exc_kwargs = mock_capture_exc.call_args[1]
+            assert exc_kwargs["span"] is mock_span
+            assert isinstance(exc_kwargs["exception"], AgentException)
+
+    async def test_telemetry_uses_correct_provider_name(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Test that telemetry uses AGENT_PROVIDER_NAME as provider."""
+        from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+        messages = self._create_standard_messages()
+        mock_client = self._create_mock_client(messages)
+
+        monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+        with (
+            patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+            patch("agent_framework.observability._get_span") as mock_get_span,
+        ):
+            mock_span = MagicMock()
+            mock_get_span.return_value.__enter__ = MagicMock(return_value=mock_span)
+            mock_get_span.return_value.__exit__ = MagicMock(return_value=False)
+
+            agent = ClaudeAgent(name="test-agent")
+            await agent.run("Hello")
+
+            call_kwargs = mock_get_span.call_args[1]
+            assert call_kwargs["attributes"]["gen_ai.provider.name"] == "anthropic.claude"

python/packages/core/agent_framework/_agents.py

@@ -1051,10 +1051,11 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
             else:
                 final_tools.append(tool)  # type: ignore
 
+        existing_names = {name for t in final_tools if (name := _get_tool_name(t)) is not None}
         for mcp_server in self.mcp_tools:
             if not mcp_server.is_connected:
                 await self._async_exit_stack.enter_async_context(mcp_server)
-            final_tools.extend(mcp_server.functions)
+            final_tools.extend(f for f in mcp_server.functions if f.name not in existing_names)
 
         # Merge runtime kwargs into additional_function_arguments so they're available
         # in function middleware context and tool invocation.

python/packages/core/agent_framework/_types.py

@@ -345,6 +345,7 @@ ContentType = Literal[
     "shell_command_output",
     "function_approval_request",
     "function_approval_response",
+    "oauth_consent_request",
 ]
 
 
@@ -498,6 +499,8 @@ class Content:
         function_call: Content | None = None,
         user_input_request: bool | None = None,
         approved: bool | None = None,
+        # OAuth consent fields
+        consent_link: str | None = None,
         # Common fields
         annotations: Sequence[Annotation] | None = None,
         additional_properties: MutableMapping[str, Any] | None = None,
@@ -546,6 +549,7 @@ class Content:
         self.function_call = function_call
         self.user_input_request = user_input_request
         self.approved = approved
+        self.consent_link = consent_link
 
     @classmethod
     def from_text(
@@ -1122,6 +1126,37 @@ class Content:
             raw_representation=raw_representation,
         )
 
+    @classmethod
+    def from_oauth_consent_request(
+        cls: type[ContentT],
+        consent_link: str,
+        *,
+        annotations: Sequence[Annotation] | None = None,
+        additional_properties: MutableMapping[str, Any] | None = None,
+        raw_representation: Any = None,
+    ) -> ContentT:
+        """Create OAuth consent request content.
+
+        Args:
+            consent_link: The URL the user must visit to complete OAuth consent.
+
+        Keyword Args:
+            annotations: Optional annotations.
+            additional_properties: Optional additional properties.
+            raw_representation: Optional raw representation from the provider.
+
+        Returns:
+            A new Content instance with type ``oauth_consent_request``.
+        """
+        return cls(
+            "oauth_consent_request",
+            consent_link=consent_link,
+            user_input_request=True,
+            annotations=annotations,
+            additional_properties=additional_properties,
+            raw_representation=raw_representation,
+        )
+
     def to_function_approval_response(
         self,
         approved: bool,
@@ -1176,6 +1211,7 @@ class Content:
             "user_input_request",
             "approved",
             "id",
+            "consent_link",
             "additional_properties",
         )
 

python/packages/core/agent_framework/anthropic/__init__.py

@@ -11,6 +11,7 @@ Supported classes:
 - AnthropicChatOptions
 - ClaudeAgent
 - ClaudeAgentOptions
+- RawClaudeAgent
 """
 
 import importlib
@@ -21,6 +22,7 @@ _IMPORTS: dict[str, tuple[str, str]] = {
     "AnthropicChatOptions": ("agent_framework_anthropic", "agent-framework-anthropic"),
     "ClaudeAgent": ("agent_framework_claude", "agent-framework-claude"),
     "ClaudeAgentOptions": ("agent_framework_claude", "agent-framework-claude"),
+    "RawClaudeAgent": ("agent_framework_claude", "agent-framework-claude"),
 }
 
 

python/packages/core/pyproject.toml

@@ -34,8 +34,7 @@ dependencies = [
     # connectors and functions
     "openai>=1.99.0",
     "azure-identity>=1,<2",
-    # Pinned to 2.0.0b3 - breaking changes in 2.0.0b4, unpin once upgrades complete
-    "azure-ai-projects == 2.0.0b3",
+    "azure-ai-projects == 2.0.0b4",
     "mcp[ws]>=1.24.0,<2",
     "packaging>=24.1",
 ]
@@ -105,6 +104,7 @@ extend = "../../pyproject.toml"
 
 [tool.pyright]
 extends = "../../pyproject.toml"
+include = ["tests/workflow"]
 
 [tool.mypy]
 plugins = ['pydantic.mypy']

python/packages/core/tests/core/test_agents.py

@@ -755,6 +755,49 @@ async def test_chat_agent_with_local_mcp_tools(client: SupportsChatGetResponse)
             pass
 
 
+async def test_mcp_tools_not_duplicated_when_passed_as_runtime_tools(chat_client_base: Any) -> None:
+    """Test that MCP tool functions from self.mcp_tools are not duplicated when already present in runtime tools."""
+    captured_options: list[dict[str, Any]] = []
+
+    original_inner = chat_client_base._inner_get_response
+
+    async def capturing_inner(
+        *, messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
+    ) -> ChatResponse:
+        captured_options.append(dict(options))
+        return await original_inner(messages=messages, options=options, **kwargs)
+
+    chat_client_base._inner_get_response = capturing_inner
+
+    # Create FunctionTool instances that simulate expanded MCP functions
+    mcp_func_a = FunctionTool(func=lambda: "a", name="tool_a", description="Tool A")
+    mcp_func_b = FunctionTool(func=lambda: "b", name="tool_b", description="Tool B")
+
+    # Create a mock MCP tool that is already connected (simulates turn 2)
+    mock_mcp_tool = MagicMock(spec=MCPTool)
+    mock_mcp_tool.is_connected = True
+    mock_mcp_tool.functions = [mcp_func_a, mcp_func_b]
+    mock_mcp_tool.__aenter__ = AsyncMock(return_value=mock_mcp_tool)
+    mock_mcp_tool.__aexit__ = AsyncMock(return_value=None)
+
+    # Agent has the MCP tool in its constructor (stored in self.mcp_tools)
+    agent = Agent(client=chat_client_base, name="TestAgent", tools=[mock_mcp_tool])
+
+    # Simulate AG-UI turn 2: pass already-expanded MCP functions + a client tool as runtime tools
+    client_tool = FunctionTool(func=lambda: "client", name="client_tool", description="Client tool")
+    runtime_tools = [mcp_func_a, mcp_func_b, client_tool]
+
+    await agent.run("hello", tools=runtime_tools)
+
+    # Verify the chat client received each tool exactly once
+    assert len(captured_options) >= 1
+    tool_names = [t.name for t in captured_options[0]["tools"]]
+    assert tool_names.count("tool_a") == 1, f"tool_a duplicated: {tool_names}"
+    assert tool_names.count("tool_b") == 1, f"tool_b duplicated: {tool_names}"
+    assert "client_tool" in tool_names
+    assert len(tool_names) == 3
+
+
 async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> None:
     """Verify tool execution receives 'session' inside **kwargs when function is called by client."""
 

python/packages/core/tests/core/test_types.py

@@ -3424,3 +3424,30 @@ class TestResponseStreamEdgeCases:
 
 
 # endregion
+
+
+# region OAuth Consent Content
+
+
+def test_oauth_consent_request_creation():
+    """Test Content.from_oauth_consent_request creates the correct content."""
+    content = Content.from_oauth_consent_request(
+        consent_link="https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc",
+    )
+    assert content.type == "oauth_consent_request"
+    assert content.consent_link == "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc"
+    assert content.user_input_request is True
+
+
+def test_oauth_consent_request_serialization_roundtrip():
+    """Test that oauth_consent_request content serializes and includes consent_link."""
+    content = Content.from_oauth_consent_request(
+        consent_link="https://login.microsoftonline.com/consent",
+    )
+    d = content.to_dict()
+    assert d["type"] == "oauth_consent_request"
+    assert d["consent_link"] == "https://login.microsoftonline.com/consent"
+    assert d["user_input_request"] is True
+
+
+# endregion

python/packages/core/tests/workflow/test_agent_executor.py

@@ -2,19 +2,20 @@
 
 import logging
 from collections.abc import AsyncIterable, Awaitable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal, overload
 
 import pytest
-
 from agent_framework import (
     AgentExecutor,
     AgentResponse,
     AgentResponseUpdate,
+    AgentRunInputs,
     AgentSession,
     BaseAgent,
     Content,
     Message,
     ResponseStream,
+    WorkflowEvent,
     WorkflowRunState,
 )
 from agent_framework._workflows._agent_executor import AgentExecutorResponse
@@ -32,26 +33,56 @@ class _CountingAgent(BaseAgent):
         super().__init__(**kwargs)
         self.call_count = 0
 
+    @overload
     def run(
         self,
-        messages: str | Message | list[str] | list[Message] | None = None,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> (
+        Awaitable[AgentResponse[Any]]
+        | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
+    ):
         self.call_count += 1
         if stream:
 
             async def _stream() -> AsyncIterable[AgentResponseUpdate]:
                 yield AgentResponseUpdate(
-                    contents=[Content.from_text(text=f"Response #{self.call_count}: {self.name}")]
+                    contents=[
+                        Content.from_text(
+                            text=f"Response #{self.call_count}: {self.name}"
+                        )
+                    ]
                 )
 
             return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
 
         async def _run() -> AgentResponse:
-            return AgentResponse(messages=[Message("assistant", [f"Response #{self.call_count}: {self.name}"])])
+            return AgentResponse(
+                messages=[
+                    Message("assistant", [f"Response #{self.call_count}: {self.name}"])
+                ]
+            )
 
         return _run()
 
@@ -63,13 +94,36 @@ class _StreamingHookAgent(BaseAgent):
         super().__init__(**kwargs)
         self.result_hook_called = False
 
+    @overload
     def run(
         self,
-        messages: str | Message | list[str] | list[Message] | None = None,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
+        session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> (
+        Awaitable[AgentResponse[Any]]
+        | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
+    ):
         if stream:
 
             async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -78,13 +132,15 @@ class _StreamingHookAgent(BaseAgent):
                     role="assistant",
                 )
 
-            async def _mark_result_hook_called(response: AgentResponse) -> AgentResponse:
+            async def _mark_result_hook_called(
+                response: AgentResponse,
+            ) -> AgentResponse:
                 self.result_hook_called = True
                 return response
 
-            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates).with_result_hook(
-                _mark_result_hook_called
-            )
+            return ResponseStream(
+                _stream(), finalizer=AgentResponse.from_updates
+            ).with_result_hook(_mark_result_hook_called)
 
         async def _run() -> AgentResponse:
             return AgentResponse(messages=[Message("assistant", ["hook test"])])
@@ -92,7 +148,9 @@ class _StreamingHookAgent(BaseAgent):
         return _run()
 
 
-async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> None:
+async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> (
+    None
+):
     """AgentExecutor should call get_final_response() so stream result hooks execute."""
     agent = _StreamingHookAgent(id="hook_agent", name="HookAgent")
     executor = AgentExecutor(agent, id="hook_exec")
@@ -159,7 +217,9 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
 
     executor_state = executor_states[executor.id]  # type: ignore[index]
     assert "cache" in executor_state, "Checkpoint should store executor cache state"
-    assert "agent_session" in executor_state, "Checkpoint should store executor session state"
+    assert "agent_session" in executor_state, (
+        "Checkpoint should store executor session state"
+    )
 
     # Verify session state structure
     session_state = executor_state["agent_session"]  # type: ignore[index]
@@ -180,11 +240,15 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     assert restored_agent.call_count == 0
 
     # Build new workflow with the restored executor
-    wf_resume = SequentialBuilder(participants=[restored_executor], checkpoint_storage=storage).build()
+    wf_resume = SequentialBuilder(
+        participants=[restored_executor], checkpoint_storage=storage
+    ).build()
 
     # Resume from checkpoint
     resumed_output: AgentExecutorResponse | None = None
-    async for ev in wf_resume.run(checkpoint_id=restore_checkpoint.checkpoint_id, stream=True):
+    async for ev in wf_resume.run(
+        checkpoint_id=restore_checkpoint.checkpoint_id, stream=True
+    ):
         if ev.type == "output":
             resumed_output = ev.data  # type: ignore[assignment]
         if ev.type == "status" and ev.state in (
@@ -278,7 +342,7 @@ async def test_agent_executor_run_streaming_with_stream_kwarg_does_not_raise() -
     workflow = SequentialBuilder(participants=[executor]).build()
 
     # stream=True at workflow level triggers streaming mode (returns async iterable)
-    events = []
+    events: list[WorkflowEvent] = []
     async for event in workflow.run("hello", stream=True):
         events.append(event)
     assert len(events) > 0
@@ -288,10 +352,13 @@ async def test_agent_executor_run_streaming_with_stream_kwarg_does_not_raise() -
 @pytest.mark.parametrize("reserved_kwarg", ["session", "stream", "messages"])
 async def test_prepare_agent_run_args_strips_reserved_kwargs(reserved_kwarg: str, caplog: "LogCaptureFixture") -> None:
     """_prepare_agent_run_args must remove reserved kwargs and log a warning."""
-    raw = {reserved_kwarg: "should-be-stripped", "custom_key": "keep-me"}
+    raw: dict[str, Any] = {
+        reserved_kwarg: "should-be-stripped",
+        "custom_key": "keep-me",
+    }
 
     with caplog.at_level(logging.WARNING):
-        run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)
+        run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
 
     assert reserved_kwarg not in run_kwargs
     assert "custom_key" in run_kwargs
@@ -302,8 +369,8 @@ async def test_prepare_agent_run_args_strips_reserved_kwargs(reserved_kwarg: str
 
 async def test_prepare_agent_run_args_preserves_non_reserved_kwargs() -> None:
     """Non-reserved workflow kwargs should pass through unchanged."""
-    raw = {"custom_param": "value", "another": 42}
-    run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)
+    raw: dict[str, Any] = {"custom_param": "value", "another": 42}
+    run_kwargs, _options = AgentExecutor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
     assert run_kwargs["custom_param"] == "value"
     assert run_kwargs["another"] == 42
 
@@ -312,10 +379,10 @@ async def test_prepare_agent_run_args_strips_all_reserved_kwargs_at_once(
     caplog: "LogCaptureFixture",
 ) -> None:
     """All reserved kwargs should be stripped when supplied together, each emitting a warning."""
-    raw = {"session": "x", "stream": True, "messages": [], "custom": 1}
+    raw: dict[str, Any] = {"session": "x", "stream": True, "messages": [], "custom": 1}
 
     with caplog.at_level(logging.WARNING):
-        run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)
+        run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
 
     assert "session" not in run_kwargs
     assert "stream" not in run_kwargs
@@ -324,7 +391,11 @@ async def test_prepare_agent_run_args_strips_all_reserved_kwargs_at_once(
     assert options is not None
     assert options["additional_function_arguments"]["custom"] == 1
 
-    warned_keys = {r.message.split("'")[1] for r in caplog.records if "reserved" in r.message.lower()}
+    warned_keys = {
+        r.message.split("'")[1]
+        for r in caplog.records
+        if "reserved" in r.message.lower()
+    }
     assert warned_keys == {"session", "stream", "messages"}
 
 

python/packages/core/tests/workflow/test_agent_executor_tool_calls.py

@@ -3,7 +3,7 @@
 """Tests for AgentExecutor handling of tool calls and results in streaming mode."""
 
 from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
-from typing import Any
+from typing import Any, Literal, overload
 
 from typing_extensions import Never
 
@@ -13,6 +13,7 @@ from agent_framework import (
     AgentExecutorResponse,
     AgentResponse,
     AgentResponseUpdate,
+    AgentRunInputs,
     AgentSession,
     BaseAgent,
     ChatResponse,
@@ -37,18 +38,38 @@ class _ToolCallingAgent(BaseAgent):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
 
+    @overload
     def run(
         self,
-        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         if stream:
             return ResponseStream(self._run_stream_impl(), finalizer=AgentResponse.from_updates)
 
-        async def _run() -> AgentResponse:
+        async def _run() -> AgentResponse[Any]:
             return AgentResponse(messages=[Message("assistant", ["done"])])
 
         return _run()
@@ -111,6 +132,7 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
     # First event: text update
     assert events[0].data is not None
     assert events[0].data.contents[0].type == "text"
+    assert events[0].data.contents[0].text is not None
     assert "Let me search" in events[0].data.contents[0].text
 
     # Second event: function call
@@ -129,6 +151,7 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
     # Fourth event: final text
     assert events[3].data is not None
     assert events[3].data.contents[0].type == "text"
+    assert events[3].data.contents[0].text is not None
     assert "sunny" in events[3].data.contents[0].text
 
 

python/packages/core/tests/workflow/test_agent_utils.py

@@ -1,9 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from collections.abc import AsyncIterable
-from typing import Any
+from collections.abc import Awaitable
+from typing import Any, Literal, overload
 
-from agent_framework import AgentResponse, AgentResponseUpdate, AgentSession, Message
+from agent_framework import AgentResponse, AgentResponseUpdate, AgentRunInputs, AgentSession, ResponseStream
 from agent_framework._workflows._agent_utils import resolve_agent_id
 
 
@@ -11,40 +11,23 @@ class MockAgent:
     """Mock agent for testing agent utilities."""
 
     def __init__(self, agent_id: str, name: str | None = None) -> None:
-        self._id = agent_id
-        self._name = name
+        self.id: str = agent_id
+        self.name: str | None = name
+        self.description: str | None = None
 
-    @property
-    def id(self) -> str:
-        return self._id
-
-    @property
-    def name(self) -> str | None:
-        return self._name
-
-    @property
-    def display_name(self) -> str:
-        """Returns the display name of the agent."""
-        ...
-
-    @property
-    def description(self) -> str | None:
-        """Returns the description of the agent."""
-        ...
-
-    def run(
-        self,
-        messages: str | Message | list[str] | list[Message] | None = None,
-        *,
-        stream: bool = False,
-        session: AgentSession | None = None,
-        **kwargs: Any,
-    ) -> AgentResponse | AsyncIterable[AgentResponseUpdate]: ...
+    @overload
+    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(self, messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
     def create_session(self, **kwargs: Any) -> AgentSession:
         """Creates a new conversation session for the agent."""
         ...
 
+    def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
+        return AgentSession()
+
 
 def test_resolve_agent_id_with_name() -> None:
     """Test that resolve_agent_id returns name when agent has a name."""

python/packages/core/tests/workflow/test_checkpoint.py

@@ -5,6 +5,7 @@ import tempfile
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
+from typing import Any
 
 import pytest
 
@@ -24,7 +25,7 @@ class _TestToolApprovalRequest:
     """Request data for tool approval in tests."""
 
     tool_name: str
-    arguments: dict
+    arguments: dict[str, Any]
     timestamp: datetime
 
 
@@ -41,7 +42,7 @@ class _TestApprovalRequest:
     """Approval request data for tests."""
 
     action: str
-    params: tuple
+    params: tuple[Any, ...]
 
 
 @dataclass
@@ -78,8 +79,8 @@ def test_workflow_checkpoint_custom_values():
         workflow_name="test-workflow-456",
         graph_signature_hash="test-hash-456",
         timestamp=custom_timestamp,
-        messages={"executor1": [{"data": "test"}]},
-        pending_request_info_events={"req123": {"data": "test"}},
+        messages={"executor1": [{"data": "test"}]},  # type: ignore[arg-type]  # raw dict for serialization test
+        pending_request_info_events={"req123": {"data": "test"}},  # type: ignore[arg-type]  # raw dict for serialization test
         state={"key": "value"},
         iteration_count=5,
         metadata={"test": True},
@@ -103,7 +104,7 @@ def test_workflow_checkpoint_to_dict():
         checkpoint_id="test-id",
         workflow_name="test-workflow",
         graph_signature_hash="test-hash",
-        messages={"executor1": [{"data": "test"}]},
+        messages={"executor1": [{"data": "test"}]},  # type: ignore[arg-type]  # raw dict for serialization test
         state={"key": "value"},
         iteration_count=5,
     )
@@ -161,8 +162,8 @@ async def test_memory_checkpoint_storage_save_and_load():
     checkpoint = WorkflowCheckpoint(
         workflow_name="test-workflow",
         graph_signature_hash="test-hash",
-        messages={"executor1": [{"data": "hello"}]},
-        pending_request_info_events={"req123": {"data": "test"}},
+        messages={"executor1": [{"data": "hello"}]},  # type: ignore[arg-type]  # raw dict for serialization test
+        pending_request_info_events={"req123": {"data": "test"}},  # type: ignore[arg-type]  # raw dict for serialization test
     )
 
     # Save checkpoint
@@ -776,9 +777,9 @@ async def test_file_checkpoint_storage_save_and_load():
         checkpoint = WorkflowCheckpoint(
             workflow_name="test-workflow",
             graph_signature_hash="test-hash",
-            messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]},
+            messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]},  # type: ignore[arg-type]  # raw dict for serialization test
             state={"key": "value"},
-            pending_request_info_events={"req123": {"data": "test"}},
+            pending_request_info_events={"req123": {"data": "test"}},  # type: ignore[arg-type]  # raw dict for serialization test
         )
 
         # Save checkpoint
@@ -904,9 +905,9 @@ async def test_file_checkpoint_storage_json_serialization():
         checkpoint = WorkflowCheckpoint(
             workflow_name="test-workflow",
             graph_signature_hash="test-hash",
-            messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},
+            messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},  # type: ignore[arg-type]  # raw dict for serialization test
             state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None},
-            pending_request_info_events={"req123": {"data": "test"}},
+            pending_request_info_events={"req123": {"data": "test"}},  # type: ignore[arg-type]  # raw dict for serialization test
         )
 
         # Save and load

python/packages/core/tests/workflow/test_checkpoint_encode.py

@@ -3,11 +3,11 @@
 import json
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import Any
+from typing import Any, cast
 
 from agent_framework._workflows._checkpoint_encoding import (
-    _PICKLE_MARKER,
-    _TYPE_MARKER,
+    _PICKLE_MARKER,  # pyright: ignore[reportPrivateUsage]
+    _TYPE_MARKER,  # pyright: ignore[reportPrivateUsage]
     encode_checkpoint_value,
 )
 
@@ -185,8 +185,9 @@ def test_encode_list_of_dataclasses() -> None:
     result = encode_checkpoint_value(data)
 
     assert isinstance(result, list)
-    assert len(result) == 2
-    for item in result:
+    result_list = cast(list[Any], result)
+    assert len(result_list) == 2
+    for item in result_list:
         assert _PICKLE_MARKER in item
 
 

python/packages/core/tests/workflow/test_edge.py

@@ -4,6 +4,8 @@ from dataclasses import dataclass
 from typing import Any
 from unittest.mock import patch
 
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
 import pytest
 
 from agent_framework import (
@@ -275,6 +277,7 @@ async def test_single_edge_group_send_message_with_condition_pass() -> None:
     success = await edge_runner.send_message(message, state, ctx)
     assert success is True
     assert target.call_count == 1
+    assert target.last_message is not None
     assert target.last_message.data == "test"
 
 
@@ -301,7 +304,7 @@ async def test_single_edge_group_send_message_with_condition_fail() -> None:
     assert target.call_count == 0
 
 
-async def test_single_edge_group_tracing_success(span_exporter) -> None:
+async def test_single_edge_group_tracing_success(span_exporter: InMemorySpanExporter) -> None:
     """Test that single edge group processing creates proper success spans."""
     source = MockExecutor(id="source_executor")
     target = MockExecutor(id="target_executor")
@@ -352,7 +355,7 @@ async def test_single_edge_group_tracing_success(span_exporter) -> None:
     assert link.context.span_id == int("00f067aa0ba902b7", 16)
 
 
-async def test_single_edge_group_tracing_condition_failure(span_exporter) -> None:
+async def test_single_edge_group_tracing_condition_failure(span_exporter: InMemorySpanExporter) -> None:
     """Test that single edge group processing creates proper spans for condition failures."""
     source = MockExecutor(id="source_executor")
     target = MockExecutor(id="target_executor")
@@ -386,7 +389,7 @@ async def test_single_edge_group_tracing_condition_failure(span_exporter) -> Non
     assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value
 
 
-async def test_single_edge_group_tracing_type_mismatch(span_exporter) -> None:
+async def test_single_edge_group_tracing_type_mismatch(span_exporter: InMemorySpanExporter) -> None:
     """Test that single edge group processing creates proper spans for type mismatches."""
     source = MockExecutor(id="source_executor")
     target = MockExecutor(id="target_executor")
@@ -421,7 +424,7 @@ async def test_single_edge_group_tracing_type_mismatch(span_exporter) -> None:
     assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value
 
 
-async def test_single_edge_group_tracing_target_mismatch(span_exporter) -> None:
+async def test_single_edge_group_tracing_target_mismatch(span_exporter: InMemorySpanExporter) -> None:
     """Test that single edge group processing creates proper spans for target mismatches."""
     source = MockExecutor(id="source_executor")
     target = MockExecutor(id="target_executor")
@@ -775,7 +778,7 @@ async def test_source_edge_group_with_selection_func_send_message_with_target_in
     assert success is False
 
 
-async def test_fan_out_edge_group_tracing_success(span_exporter) -> None:
+async def test_fan_out_edge_group_tracing_success(span_exporter: InMemorySpanExporter) -> None:
     """Test that fan-out edge group processing creates proper success spans."""
     source = MockExecutor(id="source_executor")
     target1 = MockExecutor(id="target_executor_1")
@@ -827,7 +830,7 @@ async def test_fan_out_edge_group_tracing_success(span_exporter) -> None:
     assert link.context.span_id == int("00f067aa0ba902b7", 16)
 
 
-async def test_fan_out_edge_group_tracing_with_target(span_exporter) -> None:
+async def test_fan_out_edge_group_tracing_with_target(span_exporter: InMemorySpanExporter) -> None:
     """Test that fan-out edge group processing creates proper spans for targeted messages."""
     source = MockExecutor(id="source_executor")
     target1 = MockExecutor(id="target_executor_1")
@@ -994,7 +997,7 @@ async def test_target_edge_group_send_message_with_invalid_data() -> None:
     assert success is False
 
 
-async def test_fan_in_edge_group_tracing_buffered(span_exporter) -> None:
+async def test_fan_in_edge_group_tracing_buffered(span_exporter: InMemorySpanExporter) -> None:
     """Test that fan-in edge group processing creates proper spans for buffered messages."""
     source1 = MockExecutor(id="source_executor_1")
     source2 = MockExecutor(id="source_executor_2")
@@ -1086,7 +1089,7 @@ async def test_fan_in_edge_group_tracing_buffered(span_exporter) -> None:
     assert link.context.span_id == int("00f067aa0ba902b8", 16)
 
 
-async def test_fan_in_edge_group_tracing_type_mismatch(span_exporter) -> None:
+async def test_fan_in_edge_group_tracing_type_mismatch(span_exporter: InMemorySpanExporter) -> None:
     """Test that fan-in edge group processing creates proper spans for type mismatches."""
     source1 = MockExecutor(id="source_executor_1")
     source2 = MockExecutor(id="source_executor_2")

python/packages/core/tests/workflow/test_executor.py

@@ -3,8 +3,6 @@
 from dataclasses import dataclass
 
 import pytest
-from typing_extensions import Never
-
 from agent_framework import (
     Executor,
     Message,
@@ -16,6 +14,7 @@ from agent_framework import (
     handler,
     response_handler,
 )
+from typing_extensions import Never
 
 
 # Module-level types for string forward reference tests
@@ -59,7 +58,7 @@ def test_executor_handler_without_annotations():
         class MockExecutorWithOneHandlerWithoutAnnotations(Executor):  # type: ignore
             """A mock executor with one handler that does not implement any annotations."""
 
-            @handler
+            @handler  # pyright: ignore[reportUnknownArgumentType]
             async def handle(self, message, ctx) -> None:  # type: ignore
                 """A mock handler that does not implement any annotations."""
                 pass
@@ -156,7 +155,11 @@ async def test_executor_invoked_event_contains_input_data():
     workflow = WorkflowBuilder(start_executor=upper).add_edge(upper, collector).build()
 
     events = await workflow.run("hello world")
-    invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
+    invoked_events = [
+        e
+        for e in events
+        if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
+    ]
 
     assert len(invoked_events) == 2
 
@@ -190,10 +193,16 @@ async def test_executor_completed_event_contains_sent_messages():
     sender = MultiSenderExecutor(id="sender")
     collector = CollectorExecutor(id="collector")
 
-    workflow = WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
+    workflow = (
+        WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
+    )
 
     events = await workflow.run("hello")
-    completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
+    completed_events = [
+        e
+        for e in events
+        if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
+    ]
 
     # Sender should have completed with the sent messages
     sender_completed = next(e for e in completed_events if e.executor_id == "sender")
@@ -201,7 +210,9 @@ async def test_executor_completed_event_contains_sent_messages():
     assert sender_completed.data == ["hello-first", "hello-second"]
 
     # Collector should have completed with no sent messages (None)
-    collector_completed_events = [e for e in completed_events if e.executor_id == "collector"]
+    collector_completed_events = [
+        e for e in completed_events if e.executor_id == "collector"
+    ]
     # Collector is called twice (once per message from sender)
     assert len(collector_completed_events) == 2
     for collector_completed in collector_completed_events:
@@ -220,7 +231,11 @@ async def test_executor_completed_event_includes_yielded_outputs():
     workflow = WorkflowBuilder(start_executor=executor).build()
 
     events = await workflow.run("test")
-    completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
+    completed_events = [
+        e
+        for e in events
+        if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
+    ]
 
     assert len(completed_events) == 1
     assert completed_events[0].executor_id == "yielder"
@@ -248,7 +263,9 @@ async def test_executor_events_with_complex_message_types():
 
     class ProcessorExecutor(Executor):
         @handler
-        async def handle(self, request: Request, ctx: WorkflowContext[Response]) -> None:
+        async def handle(
+            self, request: Request, ctx: WorkflowContext[Response]
+        ) -> None:
             response = Response(results=[request.query.upper()] * request.limit)
             await ctx.send_message(response)
 
@@ -260,13 +277,23 @@ async def test_executor_events_with_complex_message_types():
     processor = ProcessorExecutor(id="processor")
     collector = CollectorExecutor(id="collector")
 
-    workflow = WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
+    workflow = (
+        WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
+    )
 
     input_request = Request(query="hello", limit=3)
     events = await workflow.run(input_request)
 
-    invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
-    completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
+    invoked_events = [
+        e
+        for e in events
+        if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
+    ]
+    completed_events = [
+        e
+        for e in events
+        if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
+    ]
 
     # Check processor invoked event has the Request object
     processor_invoked = next(e for e in invoked_events if e.executor_id == "processor")
@@ -275,7 +302,9 @@ async def test_executor_events_with_complex_message_types():
     assert processor_invoked.data.limit == 3
 
     # Check processor completed event has the Response object
-    processor_completed = next(e for e in completed_events if e.executor_id == "processor")
+    processor_completed = next(
+        e for e in completed_events if e.executor_id == "processor"
+    )
     assert processor_completed.data is not None
     assert len(processor_completed.data) == 1
     assert isinstance(processor_completed.data[0], Response)
@@ -361,7 +390,9 @@ def test_executor_workflow_output_types_property():
     # Test executor with union workflow output types
     class UnionWorkflowOutputExecutor(Executor):
         @handler
-        async def handle(self, text: str, ctx: WorkflowContext[int, str | bool]) -> None:
+        async def handle(
+            self, text: str, ctx: WorkflowContext[int, str | bool]
+        ) -> None:
             pass
 
     executor = UnionWorkflowOutputExecutor(id="union_workflow_output")
@@ -372,11 +403,15 @@ def test_executor_workflow_output_types_property():
     # Test executor with multiple handlers having different workflow output types
     class MultiHandlerWorkflowExecutor(Executor):
         @handler
-        async def handle_string(self, text: str, ctx: WorkflowContext[int, str]) -> None:
+        async def handle_string(
+            self, text: str, ctx: WorkflowContext[int, str]
+        ) -> None:
             pass
 
         @handler
-        async def handle_number(self, num: int, ctx: WorkflowContext[bool, float]) -> None:
+        async def handle_number(
+            self, num: int, ctx: WorkflowContext[bool, float]
+        ) -> None:
             pass
 
     executor = MultiHandlerWorkflowExecutor(id="multi_workflow")
@@ -430,7 +465,9 @@ def test_executor_output_types_includes_response_handlers():
             pass
 
         @response_handler
-        async def handle_response(self, original_request: str, response: bool, ctx: WorkflowContext[float]) -> None:
+        async def handle_response(
+            self, original_request: str, response: bool, ctx: WorkflowContext[float]
+        ) -> None:
             pass
 
     executor = RequestResponseExecutor(id="request_response")
@@ -452,7 +489,10 @@ def test_executor_workflow_output_types_includes_response_handlers():
 
         @response_handler
         async def handle_response(
-            self, original_request: str, response: bool, ctx: WorkflowContext[float, bool]
+            self,
+            original_request: str,
+            response: bool,
+            ctx: WorkflowContext[float, bool],
         ) -> None:
             pass
 
@@ -509,7 +549,10 @@ def test_executor_response_handler_union_output_types():
 
         @response_handler
         async def handle_response(
-            self, original_request: str, response: bool, ctx: WorkflowContext[int | str | float, bool | int]
+            self,
+            original_request: str,
+            response: bool,
+            ctx: WorkflowContext[int | str | float, bool | int],
         ) -> None:
             pass
 
@@ -531,7 +574,9 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
     """Test that executor_invoked event (type='executor_invoked').data captures original input, not mutated input."""
 
     @executor(id="Mutator")
-    async def mutator(messages: list[Message], ctx: WorkflowContext[list[Message]]) -> None:
+    async def mutator(
+        messages: list[Message], ctx: WorkflowContext[list[Message]]
+    ) -> None:
         # The handler mutates the input list by appending new messages
         original_len = len(messages)
         messages.append(Message(role="assistant", text="Added by executor"))
@@ -546,7 +591,11 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
     events = await workflow.run(input_messages)
 
     # Find the invoked event for the Mutator executor
-    invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
+    invoked_events = [
+        e
+        for e in events
+        if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
+    ]
     assert len(invoked_events) == 1
     mutator_invoked = invoked_events[0]
 
@@ -577,8 +626,8 @@ class TestHandlerExplicitTypes:
         exec_instance = ExplicitInputExecutor(id="explicit_input")
 
         # Handler should be registered for str (explicit), not Any (introspected)
-        assert str in exec_instance._handlers
-        assert len(exec_instance._handlers) == 1
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        assert len(exec_instance._handlers) == 1  # pyright: ignore[reportPrivateUsage]
 
         # Can handle str messages
         assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
@@ -596,8 +645,8 @@ class TestHandlerExplicitTypes:
         exec_instance = ExplicitOutputExecutor(id="explicit_output")
 
         # Handler spec should have int as output type (explicit)
-        handler_func = exec_instance._handlers[str]
-        assert handler_func._handler_spec["output_types"] == [int]
+        handler_func = exec_instance._handlers[str]  # pyright: ignore[reportPrivateUsage]
+        assert handler_func._handler_spec["output_types"] == [int]  # pyright: ignore[reportFunctionMemberAccess]
 
         # Executor output_types property should reflect explicit type
         assert int in exec_instance.output_types
@@ -615,16 +664,20 @@ class TestHandlerExplicitTypes:
         exec_instance = ExplicitBothExecutor(id="explicit_both")
 
         # Handler should be registered for dict (explicit input type)
-        assert dict in exec_instance._handlers
-        assert len(exec_instance._handlers) == 1
+        assert dict in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        assert len(exec_instance._handlers) == 1  # pyright: ignore[reportPrivateUsage]
 
         # Output type should be list (explicit)
-        handler_func = exec_instance._handlers[dict]
-        assert handler_func._handler_spec["output_types"] == [list]
+        handler_func = exec_instance._handlers[dict]  # pyright: ignore[reportPrivateUsage]
+        assert handler_func._handler_spec["output_types"] == [list]  # pyright: ignore[reportFunctionMemberAccess]
 
         # Verify can_handle
-        assert exec_instance.can_handle(WorkflowMessage(data={"key": "value"}, source_id="mock"))
-        assert not exec_instance.can_handle(WorkflowMessage(data="string", source_id="mock"))
+        assert exec_instance.can_handle(
+            WorkflowMessage(data={"key": "value"}, source_id="mock")
+        )
+        assert not exec_instance.can_handle(
+            WorkflowMessage(data="string", source_id="mock")
+        )
 
     def test_handler_with_explicit_union_input_type(self):
         """Test that explicit union input_type is handled correctly."""
@@ -639,13 +692,15 @@ class TestHandlerExplicitTypes:
 
         # Handler should be registered for the union type
         # The union type itself is stored as the key
-        assert len(exec_instance._handlers) == 1
+        assert len(exec_instance._handlers) == 1  # pyright: ignore[reportPrivateUsage]
 
         # Can handle both str and int messages
         assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
         assert exec_instance.can_handle(WorkflowMessage(data=42, source_id="mock"))
         # Cannot handle float
-        assert not exec_instance.can_handle(WorkflowMessage(data=3.14, source_id="mock"))
+        assert not exec_instance.can_handle(
+            WorkflowMessage(data=3.14, source_id="mock")
+        )
 
     def test_handler_with_explicit_union_output_type(self):
         """Test that explicit union output is normalized to a list."""
@@ -674,8 +729,8 @@ class TestHandlerExplicitTypes:
         exec_instance = PrecedenceExecutor(id="precedence")
 
         # Should use explicit input type (bytes), not introspected (str)
-        assert bytes in exec_instance._handlers
-        assert str not in exec_instance._handlers
+        assert bytes in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        assert str not in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Should use explicit output type (float), not introspected (int)
         assert float in exec_instance.output_types
@@ -692,7 +747,7 @@ class TestHandlerExplicitTypes:
         exec_instance = IntrospectedExecutor(id="introspected")
 
         # Should use introspected types
-        assert str in exec_instance._handlers
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
         assert int in exec_instance.output_types
 
     def test_handler_explicit_mode_requires_input(self):
@@ -705,13 +760,13 @@ class TestHandlerExplicitTypes:
                 pass
 
         exec_input = OnlyInputExecutor(id="only_input")
-        assert bytes in exec_input._handlers  # Explicit
+        assert bytes in exec_input._handlers  # pyright: ignore[reportPrivateUsage]  # Explicit
         assert exec_input.output_types == []  # No output types (not introspected)
 
         # Only explicit output without input should raise error
         with pytest.raises(ValueError, match="must specify 'input' type"):
 
-            class OnlyOutputExecutor(Executor):
+            class OnlyOutputExecutor(Executor):  # pyright: ignore[reportUnusedClass]
                 @handler(output=float)
                 async def handle(self, message: str, ctx: WorkflowContext[int]) -> None:
                     pass
@@ -719,9 +774,11 @@ class TestHandlerExplicitTypes:
         # Only explicit workflow_output without input should raise error
         with pytest.raises(ValueError, match="must specify 'input' type"):
 
-            class OnlyWorkflowOutputExecutor(Executor):
+            class OnlyWorkflowOutputExecutor(Executor):  # pyright: ignore[reportUnusedClass]
                 @handler(workflow_output=bool)
-                async def handle(self, message: str, ctx: WorkflowContext[int, str]) -> None:
+                async def handle(
+                    self, message: str, ctx: WorkflowContext[int, str]
+                ) -> None:
                     pass
 
     def test_handler_explicit_input_type_allows_no_message_annotation(self):
@@ -734,8 +791,7 @@ class TestHandlerExplicitTypes:
 
         exec_instance = NoAnnotationExecutor(id="no_annotation")
 
-        # Should work with explicit input_type
-        assert str in exec_instance._handlers
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
         assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
 
     def test_handler_multiple_handlers_mixed_explicit_and_introspected(self):
@@ -747,15 +803,17 @@ class TestHandlerExplicitTypes:
                 pass
 
             @handler
-            async def handle_introspected(self, message: float, ctx: WorkflowContext[bool]) -> None:
+            async def handle_introspected(
+                self, message: float, ctx: WorkflowContext[bool]
+            ) -> None:
                 pass
 
         exec_instance = MixedExecutor(id="mixed")
 
         # Should have both handlers
-        assert len(exec_instance._handlers) == 2
-        assert str in exec_instance._handlers  # Explicit
-        assert float in exec_instance._handlers  # Introspected
+        assert len(exec_instance._handlers) == 2  # pyright: ignore[reportPrivateUsage]
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]  # Explicit
+        assert float in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]  # Introspected
 
         # Should have both output types
         assert int in exec_instance.output_types  # Explicit
@@ -772,8 +830,10 @@ class TestHandlerExplicitTypes:
         exec_instance = StringRefExecutor(id="string_ref")
 
         # Should resolve the string to the actual type
-        assert ForwardRefMessage in exec_instance._handlers
-        assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock"))
+        assert ForwardRefMessage in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        assert exec_instance.can_handle(
+            WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock")
+        )
 
     def test_handler_with_string_forward_reference_union(self):
         """Test that string forward references work with union types."""
@@ -786,8 +846,12 @@ class TestHandlerExplicitTypes:
         exec_instance = StringUnionExecutor(id="string_union")
 
         # Should handle both types
-        assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock"))
-        assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock"))
+        assert exec_instance.can_handle(
+            WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock")
+        )
+        assert exec_instance.can_handle(
+            WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock")
+        )
 
     def test_handler_with_string_forward_reference_output_type(self):
         """Test that string forward references work for output_type."""
@@ -813,8 +877,8 @@ class TestHandlerExplicitTypes:
         exec_instance = ExplicitWorkflowOutputExecutor(id="explicit_workflow_output")
 
         # Handler spec should have bool as workflow_output_type (explicit)
-        handler_func = exec_instance._handlers[str]
-        assert handler_func._handler_spec["workflow_output_types"] == [bool]
+        handler_func = exec_instance._handlers[str]  # pyright: ignore[reportPrivateUsage]
+        assert handler_func._handler_spec["workflow_output_types"] == [bool]  # pyright: ignore[reportFunctionMemberAccess]
 
         # Executor workflow_output_types property should reflect explicit type
         assert bool in exec_instance.workflow_output_types
@@ -826,13 +890,14 @@ class TestHandlerExplicitTypes:
 
         class PrecedenceExecutor(Executor):
             @handler(input=int, output=float, workflow_output=str)
-            async def handle(self, message: int, ctx: WorkflowContext[int, bool]) -> None:
+            async def handle(
+                self, message: int, ctx: WorkflowContext[int, bool]
+            ) -> None:
                 pass
 
         exec_instance = PrecedenceExecutor(id="precedence")
 
-        # All types should come from explicit params
-        assert int in exec_instance._handlers
+        assert int in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
         assert float in exec_instance.output_types
         assert str in exec_instance.workflow_output_types
         # Introspected types should NOT be present
@@ -849,8 +914,7 @@ class TestHandlerExplicitTypes:
 
         exec_instance = AllExplicitExecutor(id="all_explicit")
 
-        # Check input type
-        assert str in exec_instance._handlers
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
         assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
 
         # Check output_type
@@ -894,7 +958,9 @@ class TestHandlerExplicitTypes:
             async def handle(self, message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
                 pass
 
-        exec_instance = StringUnionWorkflowOutputExecutor(id="string_union_workflow_output")
+        exec_instance = StringUnionWorkflowOutputExecutor(
+            id="string_union_workflow_output"
+        )
 
         # Should resolve both types from string union
         assert ForwardRefTypeA in exec_instance.workflow_output_types
@@ -905,10 +971,14 @@ class TestHandlerExplicitTypes:
 
         class IntrospectedWorkflowOutputExecutor(Executor):
             @handler
-            async def handle(self, message: str, ctx: WorkflowContext[int, bool]) -> None:
+            async def handle(
+                self, message: str, ctx: WorkflowContext[int, bool]
+            ) -> None:
                 pass
 
-        exec_instance = IntrospectedWorkflowOutputExecutor(id="introspected_workflow_output")
+        exec_instance = IntrospectedWorkflowOutputExecutor(
+            id="introspected_workflow_output"
+        )
 
         # Should use introspected types from WorkflowContext[int, bool]
         assert int in exec_instance.output_types

python/packages/core/tests/workflow/test_executor_future.py

@@ -34,8 +34,8 @@ class TestExecutorFutureAnnotations:
                 pass
 
         exec_instance = MyExecutor(id="test")
-        assert str in exec_instance._handlers
-        spec = exec_instance._handler_specs[0]
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is str
         assert spec["output_types"] == [MyTypeA]
         assert spec["workflow_output_types"] == [MyTypeB]
@@ -49,8 +49,8 @@ class TestExecutorFutureAnnotations:
                 pass
 
         exec_instance = MyExecutor(id="test")
-        assert int in exec_instance._handlers
-        spec = exec_instance._handler_specs[0]
+        assert int in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is int
         assert spec["output_types"] == [MyTypeA]
 
@@ -63,7 +63,7 @@ class TestExecutorFutureAnnotations:
                 pass
 
         exec_instance = MyExecutor(id="test")
-        spec = exec_instance._handler_specs[0]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] == dict[str, Any]
         assert spec["output_types"] == [list[str]]
 
@@ -76,8 +76,8 @@ class TestExecutorFutureAnnotations:
                 pass
 
         exec_instance = MyExecutor(id="test")
-        assert str in exec_instance._handlers
-        spec = exec_instance._handler_specs[0]
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["output_types"] == []
         assert spec["workflow_output_types"] == []
 
@@ -86,12 +86,12 @@ class TestExecutorFutureAnnotations:
 
         class MyExecutor(Executor):
             @handler(input=str, output=MyTypeA)
-            async def example(self, input, ctx) -> None:
+            async def example(self, input, ctx) -> None:  # type: ignore[no-untyped-def]
                 pass
 
         exec_instance = MyExecutor(id="test")
-        assert str in exec_instance._handlers
-        spec = exec_instance._handler_specs[0]
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is str
         assert spec["output_types"] == [MyTypeA]
 
@@ -104,8 +104,8 @@ class TestExecutorFutureAnnotations:
                 pass
 
         exec_instance = MyExecutor(id="test")
-        assert str in exec_instance._handlers
-        spec = exec_instance._handler_specs[0]
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["output_types"] == [MyTypeA, MyTypeB]
         assert spec["workflow_output_types"] == [MyTypeC]
 
@@ -118,7 +118,7 @@ class TestExecutorFutureAnnotations:
         """
         with pytest.raises(ValueError):
 
-            class Bad(Executor):
-                @handler
-                async def example(self, input: NonExistentType, ctx: WorkflowContext[MyTypeA, MyTypeB]) -> None:  # noqa: F821
+            class Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler  # pyright: ignore[reportUnknownArgumentType]
+                async def example(self, input: NonExistentType, ctx: WorkflowContext[MyTypeA, MyTypeB]) -> None:  # noqa: F821  # type: ignore[name-defined]
                     pass

python/packages/core/tests/workflow/test_full_conversation.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from collections.abc import AsyncIterable, Awaitable, Sequence
-from typing import Any
+from collections.abc import AsyncIterable, Awaitable
+from typing import Any, Literal, overload
 
 import pytest
 from pydantic import PrivateAttr
@@ -13,6 +13,7 @@ from agent_framework import (
     AgentExecutorResponse,
     AgentResponse,
     AgentResponseUpdate,
+    AgentRunInputs,
     AgentSession,
     BaseAgent,
     Content,
@@ -34,14 +35,32 @@ class _SimpleAgent(BaseAgent):
         super().__init__(**kwargs)
         self._reply_text = reply_text
 
+    @overload
     def run(
         self,
-        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         if stream:
 
             async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -81,14 +100,32 @@ class _ToolHistoryAgent(BaseAgent):
             Message(role="assistant", contents=[Content.from_text(text=self._summary_text)]),
         ]
 
+    @overload
     def run(
         self,
-        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         if stream:
 
             async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -165,14 +202,32 @@ class _CaptureAgent(BaseAgent):
         super().__init__(**kwargs)
         self._reply_text = reply_text
 
+    @overload
     def run(
         self,
-        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         # Normalize and record messages for verification
         norm: list[Message] = []
         if messages:
@@ -260,7 +315,7 @@ class _RoundTripCoordinator(Executor):
     async def handle_response(
         self,
         response: AgentExecutorResponse,
-        ctx: WorkflowContext[Never, dict[str, Any]],
+        ctx: WorkflowContext[AgentExecutorRequest, dict[str, Any]],
     ) -> None:
         self._seen += 1
         if self._seen == 1:
@@ -314,14 +369,32 @@ class _SessionIdCapturingAgent(BaseAgent):
 
     _captured_service_session_id: str | None = PrivateAttr(default="NOT_CAPTURED")
 
+    @overload
     def run(
         self,
-        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         self._captured_service_session_id = session.service_session_id if session else None
 
         async def _run() -> AgentResponse:
@@ -342,7 +415,7 @@ class _FullHistoryReplayCoordinator(Executor):
     async def handle(
         self,
         response: AgentExecutorResponse,
-        ctx: WorkflowContext[Never, Any],
+        ctx: WorkflowContext[AgentExecutorRequest, Any],
     ) -> None:
         full_conv = list(response.full_conversation or response.agent_response.messages)
         full_conv.append(Message(role="user", text="follow-up"))

python/packages/core/tests/workflow/test_function_executor.py

@@ -48,12 +48,12 @@ class TestFunctionExecutor:
         func_exec = FunctionExecutor(process_string)
 
         # Check that handler was registered
-        assert len(func_exec._handlers) == 1
-        assert str in func_exec._handlers
+        assert len(func_exec._handlers) == 1  # pyright: ignore[reportPrivateUsage]
+        assert str in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Check handler spec was created
-        assert len(func_exec._handler_specs) == 1
-        spec = func_exec._handler_specs[0]
+        assert len(func_exec._handler_specs) == 1  # pyright: ignore[reportPrivateUsage]
+        spec = func_exec._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["name"] == "process_string"
         assert spec["message_type"] is str
         assert spec["output_types"] == [str]
@@ -67,10 +67,10 @@ class TestFunctionExecutor:
 
         assert isinstance(process_int, FunctionExecutor)
         assert process_int.id == "test_executor"
-        assert int in process_int._handlers
+        assert int in process_int._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Check spec
-        spec = process_int._handler_specs[0]
+        spec = process_int._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is int
         assert spec["output_types"] == [int]
 
@@ -78,7 +78,7 @@ class TestFunctionExecutor:
         """Test @executor decorator uses function name as default ID."""
 
         @executor
-        async def my_function(data: dict, ctx: WorkflowContext[Any]) -> None:
+        async def my_function(data: dict[str, Any], ctx: WorkflowContext[Any]) -> None:
             await ctx.send_message(data)
 
         assert my_function.id == "my_function"
@@ -92,7 +92,7 @@ class TestFunctionExecutor:
 
         assert isinstance(no_parens_function, FunctionExecutor)
         assert no_parens_function.id == "no_parens_function"
-        assert str in no_parens_function._handlers
+        assert str in no_parens_function._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Also test with single parameter function
         @executor
@@ -101,7 +101,7 @@ class TestFunctionExecutor:
 
         assert isinstance(simple_no_parens, FunctionExecutor)
         assert simple_no_parens.id == "simple_no_parens"
-        assert int in simple_no_parens._handlers
+        assert int in simple_no_parens._handlers  # pyright: ignore[reportPrivateUsage]
 
     def test_union_output_types(self):
         """Test that union output types are properly inferred for both messages and workflow outputs."""
@@ -113,7 +113,7 @@ class TestFunctionExecutor:
             else:
                 await ctx.send_message(text.upper())
 
-        spec = multi_output._handler_specs[0]
+        spec = multi_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert set(spec["output_types"]) == {str, int}
         assert spec["workflow_output_types"] == []  # No workflow outputs defined
 
@@ -127,7 +127,7 @@ class TestFunctionExecutor:
             else:
                 await ctx.yield_output(data.upper())
 
-        workflow_spec = multi_workflow_output._handler_specs[0]
+        workflow_spec = multi_workflow_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert workflow_spec["output_types"] == []  # None means no message outputs
         assert set(workflow_spec["workflow_output_types"]) == {str, int, bool}
 
@@ -139,7 +139,7 @@ class TestFunctionExecutor:
             # This executor doesn't send any messages
             pass
 
-        spec = no_output._handler_specs[0]
+        spec = no_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["output_types"] == []
         assert spec["workflow_output_types"] == []  # No workflow outputs defined
 
@@ -150,7 +150,7 @@ class TestFunctionExecutor:
         async def any_output(data: str, ctx: WorkflowContext[Any]) -> None:
             await ctx.send_message("result")
 
-        spec = any_output._handler_specs[0]
+        spec = any_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["output_types"] == [Any]
         assert spec["workflow_output_types"] == []  # No workflow outputs defined
 
@@ -160,7 +160,7 @@ class TestFunctionExecutor:
             await ctx.send_message("message")
             await ctx.yield_output("workflow_output")
 
-        both_spec = any_both_output._handler_specs[0]
+        both_spec = any_both_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert both_spec["output_types"] == [Any]
         assert both_spec["workflow_output_types"] == [Any]
 
@@ -228,11 +228,11 @@ class TestFunctionExecutor:
             await ctx.yield_output(result)
 
         # Verify type inference for both executors
-        upper_spec = to_upper._handler_specs[0]
+        upper_spec = to_upper._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert upper_spec["output_types"] == [str]
         assert upper_spec["workflow_output_types"] == []  # No workflow outputs
 
-        reverse_spec = reverse_text._handler_specs[0]
+        reverse_spec = reverse_text._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert reverse_spec["output_types"] == [Any]  # First parameter is Any
         assert reverse_spec["workflow_output_types"] == [str]  # Second parameter is str
 
@@ -270,7 +270,7 @@ class TestFunctionExecutor:
             await ctx.send_message(message)
 
         with pytest.raises(ValueError, match="Handler for type .* already registered"):
-            func_exec._register_instance_handler(
+            func_exec._register_instance_handler(  # pyright: ignore[reportPrivateUsage]
                 name="second",
                 func=second_handler,
                 message_type=str,
@@ -287,7 +287,7 @@ class TestFunctionExecutor:
             result = {item: len(item) for item in items}
             await ctx.send_message(result)
 
-        spec = process_list._handler_specs[0]
+        spec = process_list._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] == list[str]
         assert spec["output_types"] == [dict[str, int]]
 
@@ -300,10 +300,10 @@ class TestFunctionExecutor:
 
         assert isinstance(process_simple, FunctionExecutor)
         assert process_simple.id == "simple_processor"
-        assert str in process_simple._handlers
+        assert str in process_simple._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Check spec - single parameter functions have no output types since they can't send messages
-        spec = process_simple._handler_specs[0]
+        spec = process_simple._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is str
         assert spec["output_types"] == []
         assert spec["ctx_annotation"] is None
@@ -316,7 +316,7 @@ class TestFunctionExecutor:
             return data * 2
 
         func_exec = FunctionExecutor(valid_single)
-        assert int in func_exec._handlers
+        assert int in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Single parameter with missing type annotation should still fail
         async def no_annotation(data):  # type: ignore
@@ -349,7 +349,7 @@ class TestFunctionExecutor:
 
         # For testing purposes, we can check that the handler is registered correctly
         assert double_value.can_handle(WorkflowMessage(data=5, source_id="mock"))
-        assert int in double_value._handlers
+        assert int in double_value._handlers  # pyright: ignore[reportPrivateUsage]
 
     def test_sync_function_basic(self):
         """Test basic synchronous function support."""
@@ -360,10 +360,10 @@ class TestFunctionExecutor:
 
         assert isinstance(process_sync, FunctionExecutor)
         assert process_sync.id == "sync_processor"
-        assert str in process_sync._handlers
+        assert str in process_sync._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Check spec - sync single parameter functions have no output types
-        spec = process_sync._handler_specs[0]
+        spec = process_sync._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is str
         assert spec["output_types"] == []
         assert spec["ctx_annotation"] is None
@@ -378,10 +378,10 @@ class TestFunctionExecutor:
 
         assert isinstance(sync_with_ctx, FunctionExecutor)
         assert sync_with_ctx.id == "sync_with_ctx"
-        assert int in sync_with_ctx._handlers
+        assert int in sync_with_ctx._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Check spec - sync functions with context can infer output types
-        spec = sync_with_ctx._handler_specs[0]
+        spec = sync_with_ctx._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is int
         assert spec["output_types"] == [int]
 
@@ -404,18 +404,18 @@ class TestFunctionExecutor:
             return data.upper()
 
         func_exec = FunctionExecutor(valid_sync)
-        assert str in func_exec._handlers
+        assert str in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Valid sync function with two parameters
         def valid_sync_with_ctx(data: int, ctx: WorkflowContext[str]):
             return str(data)
 
         func_exec2 = FunctionExecutor(valid_sync_with_ctx)
-        assert int in func_exec2._handlers
+        assert int in func_exec2._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Sync function with missing type annotation should still fail
-        def no_annotation(data):  # type: ignore
-            return data
+        def no_annotation(data):  # type: ignore  # pyright: ignore[reportUnknownVariableType]
+            return data  # pyright: ignore[reportUnknownVariableType]
 
         with pytest.raises(ValueError, match="type annotation for the message"):
             FunctionExecutor(no_annotation)  # type: ignore
@@ -457,11 +457,11 @@ class TestFunctionExecutor:
             await ctx.yield_output(result)
 
         # Verify type inference for sync and async functions
-        sync_spec = to_upper_sync._handler_specs[0]
+        sync_spec = to_upper_sync._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert sync_spec["output_types"] == [str]
         assert sync_spec["workflow_output_types"] == []  # No workflow outputs
 
-        async_spec = reverse_async._handler_specs[0]
+        async_spec = reverse_async._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert async_spec["output_types"] == [Any]  # First parameter is Any
         assert async_spec["workflow_output_types"] == [str]  # Second parameter is str
 
@@ -471,8 +471,8 @@ class TestFunctionExecutor:
 
         # For integration testing, we mainly verify that the handlers are properly registered
         # and the functions are wrapped correctly
-        assert str in to_upper_sync._handlers
-        assert str in reverse_async._handlers
+        assert str in to_upper_sync._handlers  # pyright: ignore[reportPrivateUsage]
+        assert str in reverse_async._handlers  # pyright: ignore[reportPrivateUsage]
 
     async def test_sync_function_thread_execution(self):
         """Test that sync functions run in thread pool and don't block the event loop."""
@@ -491,13 +491,13 @@ class TestFunctionExecutor:
             return data.upper()
 
         # Verify the function is wrapped and registered
-        assert str in blocking_function._handlers
+        assert str in blocking_function._handlers  # pyright: ignore[reportPrivateUsage]
 
         # For a more complete test, we'd need to create a full workflow context,
         # but for now we can verify that the function was properly wrapped
         # and that sync functions store the correct metadata
-        assert not blocking_function._is_async
-        assert not blocking_function._has_context
+        assert not blocking_function._is_async  # pyright: ignore[reportPrivateUsage]
+        assert not blocking_function._has_context  # pyright: ignore[reportPrivateUsage]
 
         # The actual thread execution test would require a full workflow setup,
         # but the important thing is that asyncio.to_thread is used in the wrapper
@@ -506,7 +506,7 @@ class TestFunctionExecutor:
         """Test that @executor decorator properly rejects @staticmethod with clear error."""
         with pytest.raises(ValueError) as exc_info:
 
-            class Example:
+            class Example:  # pyright: ignore[reportUnusedClass]
                 @executor
                 @staticmethod
                 async def bad_handler(data: str) -> str:
@@ -519,7 +519,7 @@ class TestFunctionExecutor:
         """Test that @executor decorator properly rejects @classmethod with clear error."""
         with pytest.raises(ValueError) as exc_info:
 
-            class Example:
+            class Example:  # pyright: ignore[reportUnusedClass]
                 @executor
                 @classmethod
                 async def bad_handler(cls, data: str) -> str:
@@ -570,8 +570,8 @@ class TestExecutorExplicitTypes:
             pass
 
         # Handler should be registered for str (explicit)
-        assert str in process._handlers
-        assert len(process._handlers) == 1
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert len(process._handlers) == 1  # pyright: ignore[reportPrivateUsage]
 
         # Can handle str messages
         assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
@@ -586,7 +586,7 @@ class TestExecutorExplicitTypes:
             pass
 
         # Handler spec should have int as output type (explicit), not str (introspected)
-        spec = process._handler_specs[0]
+        spec = process._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["output_types"] == [int]
 
         # Executor output_types property should reflect explicit type
@@ -601,11 +601,11 @@ class TestExecutorExplicitTypes:
             pass
 
         # Handler should be registered for dict (explicit input type)
-        assert dict in process._handlers
-        assert len(process._handlers) == 1
+        assert dict in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert len(process._handlers) == 1  # pyright: ignore[reportPrivateUsage]
 
         # Output type should be list (explicit)
-        spec = process._handler_specs[0]
+        spec = process._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["output_types"] == [list]
 
         # Verify can_handle
@@ -620,7 +620,7 @@ class TestExecutorExplicitTypes:
             pass
 
         # Handler should be registered for the union type
-        assert len(process._handlers) == 1
+        assert len(process._handlers) == 1  # pyright: ignore[reportPrivateUsage]
 
         # Can handle both str and int messages
         assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
@@ -648,8 +648,8 @@ class TestExecutorExplicitTypes:
             pass
 
         # Should use explicit input type (bytes), not introspected (str)
-        assert bytes in process._handlers
-        assert str not in process._handlers
+        assert bytes in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert str not in process._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Should use explicit output type (float), not introspected (int)
         assert float in process.output_types
@@ -663,7 +663,7 @@ class TestExecutorExplicitTypes:
             pass
 
         # Should use introspected types
-        assert str in process._handlers
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
         assert int in process.output_types
 
     def test_executor_partial_explicit_types(self):
@@ -674,7 +674,7 @@ class TestExecutorExplicitTypes:
         async def process_input(message: str, ctx: WorkflowContext[int]) -> None:
             pass
 
-        assert bytes in process_input._handlers  # Explicit
+        assert bytes in process_input._handlers  # Explicit  # pyright: ignore[reportPrivateUsage]
         assert int in process_input.output_types  # Introspected
 
         # Only explicit output_type, introspect input_type
@@ -682,7 +682,7 @@ class TestExecutorExplicitTypes:
         async def process_output(message: str, ctx: WorkflowContext[int]) -> None:
             pass
 
-        assert str in process_output._handlers  # Introspected
+        assert str in process_output._handlers  # Introspected  # pyright: ignore[reportPrivateUsage]
         assert float in process_output.output_types  # Explicit
         assert int not in process_output.output_types  # Not introspected when explicit provided
 
@@ -694,7 +694,7 @@ class TestExecutorExplicitTypes:
             pass
 
         # Should work with explicit input_type
-        assert str in process._handlers
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
         assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
 
     def test_executor_explicit_types_with_id(self):
@@ -705,7 +705,7 @@ class TestExecutorExplicitTypes:
             pass
 
         assert process.id == "custom_id"
-        assert bytes in process._handlers
+        assert bytes in process._handlers  # pyright: ignore[reportPrivateUsage]
         assert int in process.output_types
 
     def test_executor_explicit_types_with_single_param_function(self):
@@ -713,10 +713,10 @@ class TestExecutorExplicitTypes:
 
         @executor(input=str)
         async def process(message):  # type: ignore[no-untyped-def]
-            return message.upper()
+            return message.upper()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
 
         # Should work with explicit input_type
-        assert str in process._handlers
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
         assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
         assert not process.can_handle(WorkflowMessage(data=42, source_id="mock"))
 
@@ -727,7 +727,7 @@ class TestExecutorExplicitTypes:
         def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
             pass
 
-        assert int in process._handlers
+        assert int in process._handlers  # pyright: ignore[reportPrivateUsage]
         assert str in process.output_types
 
     def test_function_executor_constructor_with_explicit_types(self):
@@ -736,10 +736,10 @@ class TestExecutorExplicitTypes:
         async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
             pass
 
-        func_exec = FunctionExecutor(process, id="test", input=dict, output=list)
+        func_exec = FunctionExecutor(process, id="test", input=dict, output=list)  # pyright: ignore[reportUnknownArgumentType]
 
-        assert dict in func_exec._handlers
-        spec = func_exec._handler_specs[0]
+        assert dict in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = func_exec._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is dict
         assert spec["output_types"] == [list]
 
@@ -766,7 +766,7 @@ class TestExecutorExplicitTypes:
             pass
 
         # Should resolve the string to the actual type
-        assert FuncExecForwardRefMessage in process._handlers
+        assert FuncExecForwardRefMessage in process._handlers  # pyright: ignore[reportPrivateUsage]
         assert process.can_handle(WorkflowMessage(data=FuncExecForwardRefMessage("hello"), source_id="mock"))
 
     def test_executor_with_string_forward_reference_union(self):
@@ -798,7 +798,7 @@ class TestExecutorExplicitTypes:
             pass
 
         # Handler spec should have bool as workflow_output_type (explicit)
-        spec = process._handler_specs[0]
+        spec = process._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["workflow_output_types"] == [bool]
 
         # Executor workflow_output_types property should reflect explicit type
@@ -826,7 +826,7 @@ class TestExecutorExplicitTypes:
             pass
 
         # Check input type
-        assert str in process._handlers
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
         assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
 
         # Check output_type
@@ -892,6 +892,6 @@ class TestExecutorExplicitTypes:
             workflow_output=bool,
         )
 
-        assert str in exec_instance._handlers
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
         assert int in exec_instance.output_types
         assert bool in exec_instance.workflow_output_types

python/packages/core/tests/workflow/test_function_executor_future.py

@@ -19,10 +19,10 @@ class TestFunctionExecutorFutureAnnotations:
 
         assert isinstance(process_future, FunctionExecutor)
         assert process_future.id == "future_test"
-        assert int in process_future._handlers
+        assert int in process_future._handlers  # pyright: ignore[reportPrivateUsage]
 
         # Check spec
-        spec = process_future._handler_specs[0]
+        spec = process_future._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] is int
         assert spec["output_types"] == [int]
 
@@ -34,6 +34,6 @@ class TestFunctionExecutorFutureAnnotations:
             await ctx.send_message(["done"])
 
         assert isinstance(process_complex, FunctionExecutor)
-        spec = process_complex._handler_specs[0]
+        spec = process_complex._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
         assert spec["message_type"] == dict[str, Any]
         assert spec["output_types"] == [list[str]]

python/packages/core/tests/workflow/test_request_info_mixin.py

@@ -794,7 +794,7 @@ class TestResponseHandlerExplicitTypes:
         """Test response_handler with explicit request and response types."""
 
         @response_handler(request=str, response=int)
-        async def test_handler(self, original_request, response, ctx) -> None:
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
             pass
 
         spec = test_handler._response_handler_spec  # type: ignore[reportAttributeAccessIssue]
@@ -806,7 +806,7 @@ class TestResponseHandlerExplicitTypes:
         """Test response_handler with explicit output and workflow_output types."""
 
         @response_handler(request=str, response=int, output=bool, workflow_output=float)
-        async def test_handler(self, original_request, response, ctx) -> None:
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
             pass
 
         spec = test_handler._response_handler_spec  # type: ignore[reportAttributeAccessIssue]
@@ -818,8 +818,8 @@ class TestResponseHandlerExplicitTypes:
     def test_response_handler_with_union_types(self):
         """Test response_handler with union types."""
 
-        @response_handler(request=str | int, response=bool | float)
-        async def test_handler(self, original_request, response, ctx) -> None:
+        @response_handler(request=str | int, response=bool | float)  # pyright: ignore[reportArgumentType]
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
             pass
 
         spec = test_handler._response_handler_spec  # type: ignore[reportAttributeAccessIssue]
@@ -830,7 +830,7 @@ class TestResponseHandlerExplicitTypes:
         """Test response_handler with string forward references."""
 
         @response_handler(request="str", response="int")
-        async def test_handler(self, original_request, response, ctx) -> None:
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
             pass
 
         spec = test_handler._response_handler_spec  # type: ignore[reportAttributeAccessIssue]
@@ -842,7 +842,7 @@ class TestResponseHandlerExplicitTypes:
         with pytest.raises(ValueError, match="must specify 'request' type"):
 
             @response_handler(response=int)
-            async def test_handler(self, original_request, response, ctx) -> None:
+            async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:  # pyright: ignore[reportUnusedFunction]
                 pass
 
     def test_response_handler_explicit_missing_response_raises_error(self):
@@ -850,7 +850,7 @@ class TestResponseHandlerExplicitTypes:
         with pytest.raises(ValueError, match="must specify 'response' type"):
 
             @response_handler(request=str)
-            async def test_handler(self, original_request, response, ctx) -> None:
+            async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:  # pyright: ignore[reportUnusedFunction]
                 pass
 
     def test_response_handler_explicit_only_output_raises_error(self):
@@ -858,7 +858,7 @@ class TestResponseHandlerExplicitTypes:
         with pytest.raises(ValueError, match="must specify 'request' type"):
 
             @response_handler(output=bool)
-            async def test_handler(self, original_request, response, ctx) -> None:
+            async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:  # pyright: ignore[reportUnusedFunction]
                 pass
 
     def test_executor_with_explicit_response_handlers(self):
@@ -873,7 +873,7 @@ class TestResponseHandlerExplicitTypes:
                 pass
 
             @response_handler(request=str, response=int, output=bool)
-            async def handle_explicit(self, original_request, response, ctx) -> None:
+            async def handle_explicit(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
                 pass
 
         executor = TestExecutor()
@@ -907,7 +907,7 @@ class TestResponseHandlerExplicitTypes:
                 pass
 
             @response_handler(request=str, response=int)
-            async def handle_response(self, original_request, response, ctx) -> None:
+            async def handle_response(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
                 self.handled_request = original_request
                 self.handled_response = response
 
@@ -942,7 +942,7 @@ class TestResponseHandlerExplicitTypes:
 
             # Explicit type handler
             @response_handler(request=dict, response=bool)
-            async def handle_explicit(self, original_request, response, ctx) -> None:
+            async def handle_explicit(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
                 pass
 
         executor = TestExecutor()

python/packages/core/tests/workflow/test_runner.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from dataclasses import dataclass
+from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -113,7 +114,7 @@ async def test_runner_run_until_convergence():
     assert result is not None and result == 10
 
     # iteration count shouldn't be reset after convergence
-    assert runner._iteration == 10  # type: ignore
+    assert runner._iteration == 10  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_run_until_convergence_not_completed():
@@ -173,7 +174,7 @@ async def test_runner_run_iteration_preserves_message_order_per_edge_runner() ->
     for index in range(5):
         await ctx.send_message(WorkflowMessage(data=MockMessage(data=index), source_id="source"))
 
-    await runner._run_iteration()
+    await runner._run_iteration()  # pyright: ignore[reportPrivateUsage]
 
     assert edge_runner.received == [0, 1, 2, 3, 4]
 
@@ -213,7 +214,7 @@ async def test_runner_run_iteration_delivers_different_edge_runners_concurrently
 
     await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="source"))
 
-    iteration_task = asyncio.create_task(runner._run_iteration())
+    iteration_task = asyncio.create_task(runner._run_iteration())  # pyright: ignore[reportPrivateUsage]
 
     await blocking_edge_runner.started.wait()
     await asyncio.wait_for(probe_edge_runner.probe_completed.wait(), timeout=2.0)
@@ -280,7 +281,7 @@ async def test_fanout_edge_runner_delivers_to_multiple_targets_concurrently() ->
     # Queue a message from source (will be delivered to both targets via FanOut)
     await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id=source.id))
 
-    iteration_task = asyncio.create_task(runner._run_iteration())
+    iteration_task = asyncio.create_task(runner._run_iteration())  # pyright: ignore[reportPrivateUsage]
 
     # Wait for the blocking executor to start
     await blocking_target.started.wait()
@@ -477,11 +478,11 @@ async def test_runner_reset_iteration_count():
     ctx = InProcRunnerContext()
 
     runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
-    runner._iteration = 10
+    runner._iteration = 10  # pyright: ignore[reportPrivateUsage]
 
     runner.reset_iteration_count()
 
-    assert runner._iteration == 0
+    assert runner._iteration == 0  # pyright: ignore[reportPrivateUsage]
 
 
 class CheckpointingContext(InProcRunnerContext):
@@ -501,18 +502,19 @@ class CheckpointingContext(InProcRunnerContext):
         graph_signature_hash: str,
         state: State,
         previous_checkpoint_id: str | None,
-        iteration: int,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
     ) -> str:
         checkpoint = WorkflowCheckpoint(
             workflow_name=workflow_name,
             graph_signature_hash=graph_signature_hash,
-            state=state.export(),
+            state=state.export_state(),
             previous_checkpoint_id=previous_checkpoint_id,
-            iteration_count=iteration,
+            iteration_count=iteration_count,
         )
         return await self._storage.save(checkpoint)
 
-    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:  # pyright: ignore[reportIncompatibleMethodOverride]
         try:
             return await self._storage.load(checkpoint_id)
         except WorkflowCheckpointException:
@@ -537,7 +539,8 @@ class FailingCheckpointContext(InProcRunnerContext):
         graph_signature_hash: str,
         state: State,
         previous_checkpoint_id: str | None,
-        iteration: int,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
     ) -> str:
         raise RuntimeError("Simulated checkpoint failure")
 
@@ -609,8 +612,8 @@ async def test_runner_restore_from_checkpoint_with_external_storage():
     # Restore using external storage
     await runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage=storage)
 
-    assert runner._resumed_from_checkpoint is True
-    assert runner._iteration == 5
+    assert runner._resumed_from_checkpoint is True  # pyright: ignore[reportPrivateUsage]
+    assert runner._iteration == 5  # pyright: ignore[reportPrivateUsage]
     assert state.get("test_key") == "test_value"
 
 
@@ -684,7 +687,7 @@ async def test_runner_restore_executor_states_invalid_states_type():
     runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     with pytest.raises(WorkflowCheckpointException, match="not a dictionary"):
-        await runner._restore_executor_states()
+        await runner._restore_executor_states()  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_restore_executor_states_invalid_executor_id_type():
@@ -698,7 +701,7 @@ async def test_runner_restore_executor_states_invalid_executor_id_type():
     runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     with pytest.raises(WorkflowCheckpointException, match="not a string"):
-        await runner._restore_executor_states()
+        await runner._restore_executor_states()  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_restore_executor_states_invalid_state_type():
@@ -712,7 +715,7 @@ async def test_runner_restore_executor_states_invalid_state_type():
     runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     with pytest.raises(WorkflowCheckpointException, match="not a dict"):
-        await runner._restore_executor_states()
+        await runner._restore_executor_states()  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_restore_executor_states_invalid_state_keys():
@@ -726,7 +729,7 @@ async def test_runner_restore_executor_states_invalid_state_keys():
     runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     with pytest.raises(WorkflowCheckpointException, match="not a dict"):
-        await runner._restore_executor_states()
+        await runner._restore_executor_states()  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_restore_executor_states_missing_executor():
@@ -739,7 +742,7 @@ async def test_runner_restore_executor_states_missing_executor():
     runner = Runner([], {}, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     with pytest.raises(WorkflowCheckpointException, match="not found during state restoration"):
-        await runner._restore_executor_states()
+        await runner._restore_executor_states()  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_set_executor_state_invalid_existing_states():
@@ -752,7 +755,7 @@ async def test_runner_set_executor_state_invalid_existing_states():
     runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     with pytest.raises(WorkflowCheckpointException, match="not a dictionary"):
-        await runner._set_executor_state("executor_a", {"key": "value"})
+        await runner._set_executor_state("executor_a", {"key": "value"})  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_with_pre_loop_events():
@@ -779,7 +782,7 @@ class EventEmittingExecutor(Executor):
     """An executor that emits events during execution."""
 
     @handler
-    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, str]) -> None:
         # Emit event during processing
         await ctx.yield_output(f"processed-{message.data}")
         if message.data < 3:
@@ -831,7 +834,7 @@ async def test_runner_restore_executor_states_no_states():
     runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     # Should complete without error when no executor states exist
-    await runner._restore_executor_states()
+    await runner._restore_executor_states()  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_runner_checkpoint_with_resumed_flag():
@@ -853,7 +856,7 @@ async def test_runner_checkpoint_with_resumed_flag():
     state = State()
 
     runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
-    runner._mark_resumed(5)
+    runner._mark_resumed(5)  # pyright: ignore[reportPrivateUsage]
 
     # Add a message to trigger the checkpoint creation path
     await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id="START"))
@@ -870,7 +873,7 @@ async def test_runner_checkpoint_with_resumed_flag():
         pass
 
     # After completing, resumed flag should be reset
-    assert runner._resumed_from_checkpoint is False
+    assert runner._resumed_from_checkpoint is False  # pyright: ignore[reportPrivateUsage]
 
 
 class ExecutorThatFailsWithEvents(Executor):
@@ -883,7 +886,7 @@ class ExecutorThatFailsWithEvents(Executor):
         self._iteration_count = 0
 
     @handler
-    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, str]) -> None:
         self._iteration_count += 1
         # First emit an output event to the workflow context
         await ctx.yield_output(f"output-before-failure-{message.data}")
@@ -951,7 +954,7 @@ class SlowEventEmittingExecutor(Executor):
         self.current_iteration = 0
 
     @handler
-    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, str]) -> None:
         self.current_iteration += 1
         # Emit output event
         await ctx.yield_output(f"iteration-{self.current_iteration}")

python/packages/core/tests/workflow/test_state.py

@@ -61,9 +61,9 @@ class TestSuperstepCaching:
         state.set("key", "value")
 
         # Value is in pending
-        assert "key" in state._pending
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
         # Value is NOT in committed
-        assert "key" not in state._committed
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
         # But get() still returns it
         assert state.get("key") == "value"
 
@@ -72,14 +72,14 @@ class TestSuperstepCaching:
         state.set("key", "value")
 
         # Before commit: in pending, not committed
-        assert "key" in state._pending
-        assert "key" not in state._committed
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
 
         state.commit()
 
         # After commit: in committed, pending cleared
-        assert "key" not in state._pending
-        assert "key" in state._committed
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
+        assert "key" in state._committed  # pyright: ignore[reportPrivateUsage]
         assert state.get("key") == "value"
 
     def test_discard_clears_pending_without_committing(self) -> None:
@@ -108,7 +108,7 @@ class TestSuperstepCaching:
         # get() returns pending value, not committed
         assert state.get("key") == "pending_value"
         # But committed still has old value
-        assert state._committed["key"] == "committed_value"
+        assert state._committed["key"] == "committed_value"  # pyright: ignore[reportPrivateUsage]
 
     def test_multiple_sets_before_commit(self) -> None:
         state = State()
@@ -130,13 +130,13 @@ class TestDeleteWithSuperstepCaching:
         state = State()
         state.set("key", "value")
         # Key only in pending, not committed
-        assert "key" in state._pending
-        assert "key" not in state._committed
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
 
         state.delete("key")
 
         # Should be removed from pending
-        assert "key" not in state._pending
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
         assert state.get("key") is None
         assert state.has("key") is False
 
@@ -148,14 +148,14 @@ class TestDeleteWithSuperstepCaching:
         state.delete("key")
 
         # Key should be marked for deletion in pending (sentinel)
-        assert "key" in state._pending
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
         # get() should return default (not the sentinel!)
         assert state.get("key") is None
         assert state.get("key", "default") == "default"
         # has() should return False
         assert state.has("key") is False
         # But committed still has it until commit()
-        assert "key" in state._committed
+        assert "key" in state._committed  # pyright: ignore[reportPrivateUsage]
 
     def test_delete_committed_key_removed_on_commit(self) -> None:
         state = State()
@@ -166,8 +166,8 @@ class TestDeleteWithSuperstepCaching:
         state.commit()
 
         # Now it should be gone from committed too
-        assert "key" not in state._committed
-        assert "key" not in state._pending
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
 
     def test_delete_key_in_both_pending_and_committed(self) -> None:
         """Test delete when key exists in both pending (modified) and committed."""
@@ -177,8 +177,8 @@ class TestDeleteWithSuperstepCaching:
 
         # Modify the key (now in both pending and committed)
         state.set("key", "modified")
-        assert state._pending["key"] == "modified"
-        assert state._committed["key"] == "original"
+        assert state._pending["key"] == "modified"  # pyright: ignore[reportPrivateUsage]
+        assert state._committed["key"] == "original"  # pyright: ignore[reportPrivateUsage]
 
         # Delete should mark for deletion from committed
         state.delete("key")
@@ -189,8 +189,8 @@ class TestDeleteWithSuperstepCaching:
 
         # After commit, key should be fully removed
         state.commit()
-        assert "key" not in state._committed
-        assert "key" not in state._pending
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
 
     def test_discard_after_delete_restores_committed_value(self) -> None:
         state = State()
@@ -238,12 +238,12 @@ class TestFailureScenarios:
         state.set("key3", "value3")
 
         # Before commit - nothing in committed
-        assert len(state._committed) == 0
+        assert len(state._committed) == 0  # pyright: ignore[reportPrivateUsage]
 
         state.commit()
 
         # After commit - all three values committed together
-        assert state._committed == {"key1": "value1", "key2": "value2", "key3": "value3"}
+        assert state._committed == {"key1": "value1", "key2": "value2", "key3": "value3"}  # pyright: ignore[reportPrivateUsage]
 
     def test_repeated_supersteps_are_isolated(self) -> None:
         """Test that each superstep's changes are isolated until committed."""
@@ -300,4 +300,4 @@ class TestExportImport:
 
         # Pending is still there
         assert state.get("pending_key") == "pending_value"
-        assert "pending_key" in state._pending
+        assert "pending_key" in state._pending  # pyright: ignore[reportPrivateUsage]

python/packages/core/tests/workflow/test_typing_utils.py

@@ -36,32 +36,32 @@ def test_normalize_type_to_list_none() -> None:
 
 def test_normalize_type_to_list_union_pipe_syntax() -> None:
     """Test normalize_type_to_list with union types using | syntax."""
-    result = normalize_type_to_list(str | int)
+    result = normalize_type_to_list(str | int)  # pyright: ignore[reportArgumentType]
     assert set(result) == {str, int}
 
-    result = normalize_type_to_list(str | int | bool)
+    result = normalize_type_to_list(str | int | bool)  # pyright: ignore[reportArgumentType]
     assert set(result) == {str, int, bool}
 
 
 def test_normalize_type_to_list_union_typing_syntax() -> None:
     """Test normalize_type_to_list with Union[] from typing module."""
-    result = normalize_type_to_list(Union[str, int])
+    result = normalize_type_to_list(Union[str, int])  # pyright: ignore[reportArgumentType]
     assert set(result) == {str, int}
 
-    result = normalize_type_to_list(Union[str, int, bool])
+    result = normalize_type_to_list(Union[str, int, bool])  # pyright: ignore[reportArgumentType]
     assert set(result) == {str, int, bool}
 
 
 def test_normalize_type_to_list_optional() -> None:
     """Test normalize_type_to_list with Optional types (Union[T, None])."""
     # Optional[str] is Union[str, None]
-    result = normalize_type_to_list(Optional[str])
+    result = normalize_type_to_list(Optional[str])  # pyright: ignore[reportArgumentType]
     assert str in result
     assert type(None) in result
     assert len(result) == 2
 
     # str | None is equivalent
-    result = normalize_type_to_list(str | None)
+    result = normalize_type_to_list(str | None)  # pyright: ignore[reportArgumentType]
     assert str in result
     assert type(None) in result
     assert len(result) == 2
@@ -77,7 +77,7 @@ def test_normalize_type_to_list_custom_types() -> None:
     result = normalize_type_to_list(CustomMessage)
     assert result == [CustomMessage]
 
-    result = normalize_type_to_list(CustomMessage | str)
+    result = normalize_type_to_list(CustomMessage | str)  # pyright: ignore[reportArgumentType]
     assert set(result) == {CustomMessage, str}
 
 
@@ -96,7 +96,7 @@ def test_resolve_type_annotation_actual_types() -> None:
     """Test resolve_type_annotation passes through actual types unchanged."""
     assert resolve_type_annotation(str) is str
     assert resolve_type_annotation(int) is int
-    assert resolve_type_annotation(str | int) == str | int
+    assert resolve_type_annotation(str | int) == str | int  # pyright: ignore[reportArgumentType]
 
 
 def test_resolve_type_annotation_string_builtin() -> None:

python/packages/core/tests/workflow/test_validation.py

@@ -484,8 +484,8 @@ def test_handler_ctx_missing_annotation_raises() -> None:
     # Validation now happens at handler registration time, not workflow build time
     with pytest.raises(ValueError) as exc:
 
-        class BadExecutor(Executor):
-            @handler
+        class BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
             async def handle(self, message: str, ctx) -> None:  # type: ignore[no-untyped-def]
                 pass
 
@@ -496,8 +496,8 @@ def test_handler_ctx_invalid_t_out_entries_raises() -> None:
     # Validation now happens at handler registration time, not workflow build time
     with pytest.raises(ValueError) as exc:
 
-        class BadExecutor(Executor):
-            @handler
+        class BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
             async def handle(self, message: str, ctx: WorkflowContext[123]) -> None:  # type: ignore[valid-type]
                 pass
 
@@ -555,7 +555,7 @@ def test_output_validation_with_valid_output_executors():
     )
 
     assert workflow is not None
-    assert workflow._output_executors == ["executor2"]
+    assert workflow._output_executors == ["executor2"]  # pyright: ignore[reportPrivateUsage]
 
 
 def test_output_validation_with_multiple_valid_output_executors():
@@ -572,7 +572,7 @@ def test_output_validation_with_multiple_valid_output_executors():
     )
 
     assert workflow is not None
-    assert set(workflow._output_executors) == {"executor1", "executor3"}
+    assert set(workflow._output_executors) == {"executor1", "executor3"}  # pyright: ignore[reportPrivateUsage]
 
 
 def test_output_validation_fails_for_nonexistent_executor():

python/packages/core/tests/workflow/test_viz.py

@@ -2,6 +2,9 @@
 
 """Tests for the workflow visualization module."""
 
+from pathlib import Path
+from typing import Any
+
 import pytest
 
 from agent_framework import Executor, WorkflowBuilder, WorkflowContext, WorkflowExecutor, WorkflowViz, handler
@@ -25,7 +28,7 @@ class ListStrTargetExecutor(Executor):
 
 
 @pytest.fixture
-def basic_sub_workflow():
+def basic_sub_workflow() -> dict[str, Any]:
     """Fixture that creates a basic sub-workflow setup for testing."""
     # Create a sub-workflow
     sub_exec1 = MockExecutor(id="sub_exec1")
@@ -98,7 +101,7 @@ def test_workflow_viz_export_dot():
     assert '"executor1" -> "executor2"' in content
 
 
-def test_workflow_viz_export_dot_with_filename(tmp_path):
+def test_workflow_viz_export_dot_with_filename(tmp_path: Path):
     """Test exporting workflow as DOT format with specified filename."""
     executor1 = MockExecutor(id="executor1")
     executor2 = MockExecutor(id="executor2")
@@ -203,7 +206,7 @@ def test_workflow_viz_graphviz_binary_not_found():
         mock_source_class.return_value = mock_source
 
         # Import the ExecutableNotFound exception for the test
-        from graphviz.backend.execute import ExecutableNotFound
+        from graphviz.backend.execute import ExecutableNotFound  # type: ignore[import-not-found]
 
         mock_source.render.side_effect = ExecutableNotFound("failed to execute PosixPath('dot')")
 
@@ -329,7 +332,7 @@ def test_workflow_viz_mermaid_fan_in_edge_group():
     assert "s2 --> t" not in mermaid
 
 
-def test_workflow_viz_sub_workflow_digraph(basic_sub_workflow):
+def test_workflow_viz_sub_workflow_digraph(basic_sub_workflow: dict[str, Any]):
     """Test that WorkflowViz can visualize sub-workflows in DOT format."""
     main_workflow = basic_sub_workflow["main_workflow"]
 
@@ -353,7 +356,7 @@ def test_workflow_viz_sub_workflow_digraph(basic_sub_workflow):
     assert '"workflow_executor_1/sub_exec1" -> "workflow_executor_1/sub_exec2"' in dot_content
 
 
-def test_workflow_viz_sub_workflow_mermaid(basic_sub_workflow):
+def test_workflow_viz_sub_workflow_mermaid(basic_sub_workflow: dict[str, Any]):
     """Test that WorkflowViz can visualize sub-workflows in Mermaid format."""
     main_workflow = basic_sub_workflow["main_workflow"]
 

python/packages/core/tests/workflow/test_workflow.py

@@ -4,7 +4,7 @@ import asyncio
 import tempfile
 from collections.abc import AsyncIterable, Awaitable, Sequence
 from dataclasses import dataclass, field
-from typing import Any, cast
+from typing import Any, Literal, cast, overload
 from uuid import uuid4
 
 import pytest
@@ -13,6 +13,7 @@ from agent_framework import (
     AgentExecutor,
     AgentResponse,
     AgentResponseUpdate,
+    AgentRunInputs,
     AgentSession,
     BaseAgent,
     Content,
@@ -474,7 +475,7 @@ class StateTrackingExecutor(Executor):
     ) -> None:
         """Handle the message and track it in workflow state."""
         # Get existing messages from workflow state
-        existing_messages = ctx.get_state("processed_messages") or []
+        existing_messages: list[str] = ctx.get_state("processed_messages") or []
 
         # Record this message
         message_record = f"{message.run_id}:{message.data}"
@@ -833,6 +834,26 @@ class _StreamingTestAgent(BaseAgent):
         super().__init__(**kwargs)
         self._reply_text = reply_text
 
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
     def run(
         self,
         messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
@@ -883,8 +904,10 @@ async def test_agent_streaming_vs_non_streaming() -> None:
         stream_events.append(event)
 
     # Filter for agent events
-    agent_response = [
-        cast(AgentResponse, e.data) for e in stream_events if e.type == "output" and isinstance(e.data, AgentResponse)
+    agent_response: list[AgentResponse[Any]] = [
+        cast(AgentResponse[Any], e.data)  # pyright: ignore[reportUnknownMemberType]
+        for e in stream_events
+        if e.type == "output" and isinstance(e.data, AgentResponse)
     ]
     agent_response_updates = [
         e.data for e in stream_events if e.type == "output" and isinstance(e.data, AgentResponseUpdate)

python/packages/core/tests/workflow/test_workflow_agent.py

@@ -2,7 +2,7 @@
 
 import uuid
 from collections.abc import Awaitable, Sequence
-from typing import Any
+from typing import Any, Literal, overload
 
 import pytest
 from typing_extensions import Never
@@ -713,6 +713,14 @@ class TestWorkflowAgent:
             def create_session(self, **kwargs: Any) -> AgentSession:
                 return AgentSession()
 
+            def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
+                return AgentSession()
+
+            @overload
+            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+            @overload
+            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
             def run(
                 self,
                 messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
@@ -801,6 +809,14 @@ class TestWorkflowAgent:
             def create_session(self, **kwargs: Any) -> AgentSession:
                 return AgentSession()
 
+            def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
+                return AgentSession()
+
+            @overload
+            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+            @overload
+            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
             def run(
                 self,
                 messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
@@ -1207,7 +1223,7 @@ class TestWorkflowAgentMergeUpdates:
         ]
 
         # Compare using role.value for Role enum
-        actual_sequence_normalized = [(t, r.value if hasattr(r, "value") else r) for t, r in content_sequence]
+        actual_sequence_normalized = [(t, r.value if hasattr(r, "value") else r) for t, r in content_sequence]  # type: ignore[union-attr]
 
         assert actual_sequence_normalized == expected_sequence, (
             f"FunctionResultContent should come immediately after FunctionCallContent. "

python/packages/core/tests/workflow/test_workflow_builder.py

@@ -1,7 +1,8 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from collections.abc import AsyncIterator, Awaitable
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal, overload
 
 import pytest
 
@@ -9,10 +10,12 @@ from agent_framework import (
     AgentExecutor,
     AgentResponse,
     AgentResponseUpdate,
+    AgentRunInputs,
     AgentSession,
     BaseAgent,
     Executor,
     Message,
+    ResponseStream,
     WorkflowBuilder,
     WorkflowContext,
     WorkflowValidationError,
@@ -21,22 +24,49 @@ from agent_framework import (
 
 
 class DummyAgent(BaseAgent):
-    def run(self, messages=None, *, stream: bool = False, session: AgentSession | None = None, **kwargs):  # type: ignore[override]
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         if stream:
-            return self._run_stream_impl()
+            return ResponseStream[AgentResponseUpdate, AgentResponse[Any]](self._run_stream_impl())
         return self._run_impl(messages)
 
-    async def _run_impl(self, messages=None) -> AgentResponse:
+    async def _run_impl(self, messages: AgentRunInputs | None = None) -> AgentResponse:
         norm: list[Message] = []
         if messages:
-            for m in messages:  # type: ignore[iteration-over-optional]
+            for m in messages:  # type: ignore[union-attr]
                 if isinstance(m, Message):
                     norm.append(m)
                 elif isinstance(m, str):
                     norm.append(Message(role="user", text=m))
         return AgentResponse(messages=norm)
 
-    async def _run_stream_impl(self):  # type: ignore[override]
+    async def _run_stream_impl(self) -> AsyncIterator[AgentResponseUpdate]:
         # Minimal async generator
         yield AgentResponseUpdate()
 
@@ -202,7 +232,7 @@ def test_with_output_from_returns_builder():
     builder = WorkflowBuilder(output_executors=[executor_a], start_executor=executor_a)
 
     # Verify builder was created with output_executors
-    assert builder._output_executors == [executor_a]
+    assert builder._output_executors == [executor_a]  # pyright: ignore[reportPrivateUsage]
 
 
 def test_with_output_from_with_executor_instances():

python/packages/core/tests/workflow/test_workflow_context.py

@@ -84,7 +84,7 @@ async def test_executor_emits_normal_event() -> None:
 
 class _TestEvent(WorkflowEvent):
     def __init__(self, data: Any = None) -> None:
-        super().__init__("test_event", data=data)
+        super().__init__("test_event", data=data)  # type: ignore[arg-type]
 
 
 async def test_workflow_context_type_annotations_no_parameter() -> None:
@@ -244,8 +244,8 @@ async def test_workflow_context_missing_annotation_error() -> None:
     # Test class-based executor with missing ctx annotation
     with pytest.raises(ValueError, match="must have a WorkflowContext"):
 
-        class _BadExecutor(Executor):
-            @handler
+        class _BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
             async def bad_handler(self, text: str, ctx) -> None:  # type: ignore[no-untyped-def]
                 pass
 
@@ -264,8 +264,8 @@ async def test_workflow_context_invalid_type_parameter_error() -> None:
     # Test class-based executor with invalid type parameter
     with pytest.raises(ValueError, match="invalid type entry"):
 
-        class _BadExecutor(Executor):
-            @handler
+        class _BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
             async def bad_handler(self, text: str, ctx: WorkflowContext[456]) -> None:  # type: ignore[valid-type]
                 pass
 

python/packages/core/tests/workflow/test_workflow_kwargs.py

@@ -1,13 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from collections.abc import AsyncIterable, Awaitable, Sequence
-from typing import Annotated, Any
+from collections.abc import AsyncIterable, Awaitable
+from typing import Annotated, Any, Literal, overload
 
 import pytest
 
 from agent_framework import (
     AgentResponse,
     AgentResponseUpdate,
+    AgentRunInputs,
     AgentSession,
     BaseAgent,
     Content,
@@ -50,14 +51,19 @@ class _KwargsCapturingAgent(BaseAgent):
         super().__init__(name=name, description="Test agent for kwargs capture")
         self.captured_kwargs = []
 
+    @overload
+    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
     def run(
         self,
-        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         self.captured_kwargs.append(dict(kwargs))
         if stream:
 
@@ -83,15 +89,20 @@ class _OptionsAwareAgent(BaseAgent):
         self.captured_options = []
         self.captured_kwargs = []
 
+    @overload
+    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
     def run(
         self,
-        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        messages: AgentRunInputs | None = None,
         *,
         stream: bool = False,
         session: AgentSession | None = None,
         options: dict[str, Any] | None = None,
         **kwargs: Any,
-    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         self.captured_options.append(dict(options) if options is not None else None)
         self.captured_kwargs.append(dict(kwargs))
         if stream:
@@ -189,15 +200,15 @@ async def test_sequential_run_options_does_not_conflict_with_agent_options() ->
             break
 
     assert len(agent.captured_options) >= 1
-    captured_options = agent.captured_options[0]
+    captured_options: dict[str, Any] | None = agent.captured_options[0]
     assert captured_options is not None
     assert captured_options.get("store") is False
 
-    additional_args = captured_options.get("additional_function_arguments")
+    additional_args: Any = captured_options.get("additional_function_arguments")
     assert isinstance(additional_args, dict)
-    assert additional_args.get("source") == "workflow-options"
-    assert additional_args.get("custom_data") == custom_data
-    assert additional_args.get("user_token") == user_token
+    assert additional_args.get("source") == "workflow-options"  # pyright: ignore[reportUnknownMemberType]
+    assert additional_args.get("custom_data") == custom_data  # pyright: ignore[reportUnknownMemberType]
+    assert additional_args.get("user_token") == user_token  # pyright: ignore[reportUnknownMemberType]
 
     # "options" should be passed once via the dedicated options parameter,
     # not duplicated in **kwargs.
@@ -225,13 +236,13 @@ async def test_sequential_run_additional_function_arguments_flattened() -> None:
             break
 
     assert len(agent.captured_options) >= 1
-    captured_options = agent.captured_options[0]
+    captured_options: dict[str, Any] | None = agent.captured_options[0]
     assert captured_options is not None
 
-    additional_args = captured_options.get("additional_function_arguments")
+    additional_args: Any = captured_options.get("additional_function_arguments")
     assert isinstance(additional_args, dict)
-    assert additional_args.get("custom_data") == custom_data
-    assert additional_args.get("user_token") == user_token
+    assert additional_args.get("custom_data") == custom_data  # pyright: ignore[reportUnknownMemberType]
+    assert additional_args.get("user_token") == user_token  # pyright: ignore[reportUnknownMemberType]
     assert "additional_function_arguments" not in additional_args
 
     assert len(agent.captured_kwargs) >= 1
@@ -255,14 +266,14 @@ async def test_sequential_run_additional_function_arguments_merges_with_options(
             break
 
     assert len(agent.captured_options) >= 1
-    captured_options = agent.captured_options[0]
+    captured_options: dict[str, Any] | None = agent.captured_options[0]
     assert captured_options is not None
 
-    additional_args = captured_options.get("additional_function_arguments")
+    additional_args: Any = captured_options.get("additional_function_arguments")
     assert isinstance(additional_args, dict)
-    assert additional_args.get("source") == "workflow-options"
-    assert additional_args.get("custom_data") == {"session_id": "abc123"}
-    assert additional_args.get("user_token") == {"user_name": "alice"}
+    assert additional_args.get("source") == "workflow-options"  # pyright: ignore[reportUnknownMemberType]
+    assert additional_args.get("custom_data") == {"session_id": "abc123"}  # pyright: ignore[reportUnknownMemberType]
+    assert additional_args.get("user_token") == {"user_name": "alice"}  # pyright: ignore[reportUnknownMemberType]
     assert "additional_function_arguments" not in additional_args
 
 
@@ -463,14 +474,19 @@ async def test_kwargs_preserved_on_response_continuation() -> None:
             self.captured_kwargs = []
             self._asked = False
 
+        @overload
+        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+        @overload
+        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
         def run(
             self,
-            messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+            messages: AgentRunInputs | None = None,
             *,
             stream: bool = False,
             session: AgentSession | None = None,
             **kwargs: Any,
-        ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+        ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
             self.captured_kwargs.append(dict(kwargs))
             if not self._asked:
                 self._asked = True
@@ -521,14 +537,19 @@ async def test_kwargs_overridden_on_response_continuation() -> None:
             self.captured_kwargs = []
             self._asked = False
 
+        @overload
+        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+        @overload
+        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
         def run(
             self,
-            messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+            messages: AgentRunInputs | None = None,
             *,
             stream: bool = False,
             session: AgentSession | None = None,
             **kwargs: Any,
-        ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+        ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
             self.captured_kwargs.append(dict(kwargs))
             if not self._asked:
                 self._asked = True
@@ -583,14 +604,19 @@ async def test_kwargs_empty_value_passed_on_continuation() -> None:
             self.captured_kwargs = []
             self._asked = False
 
+        @overload
+        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+        @overload
+        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
         def run(
             self,
-            messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+            messages: AgentRunInputs | None = None,
             *,
             stream: bool = False,
             session: AgentSession | None = None,
             **kwargs: Any,
-        ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+        ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
             self.captured_kwargs.append(dict(kwargs))
             if not self._asked:
                 self._asked = True
@@ -690,8 +716,8 @@ async def test_handoff_kwargs_flow_to_agents() -> None:
 
     workflow = (
         HandoffBuilder(termination_condition=lambda conv: len(conv) >= 4)
-        .participants([agent1, agent2])
-        .with_start_agent(agent1)
+        .participants([agent1, agent2])  # type: ignore[list-item]
+        .with_start_agent(agent1)  # type: ignore[arg-type]
         .with_autonomous_mode()
         .build()
     )

python/packages/core/tests/workflow/test_workflow_observability.py

@@ -109,7 +109,7 @@ async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter)
             {
                 "id": "test-workflow-123",
                 "max_iterations": 100,
-                "model_dump_json": lambda self: '{"id": "test-workflow-123", "type": "mock"}',
+                "model_dump_json": lambda self: '{"id": "test-workflow-123", "type": "mock"}',  # pyright: ignore[reportUnknownLambdaType]
             },
         )(),
     )
@@ -122,7 +122,7 @@ async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter)
         },
     ) as workflow_span:
         workflow_span.add_event(OtelAttr.WORKFLOW_STARTED)
-        sending_attributes = {
+        sending_attributes: dict[str, str | int] = {
             OtelAttr.MESSAGE_TYPE: "ResponseMessage",
             OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID: "target-789",
         }
@@ -231,7 +231,7 @@ async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> No
 
 @pytest.mark.parametrize("enable_instrumentation", [False], indirect=True)
 async def test_trace_context_disabled_when_tracing_disabled(
-    enable_instrumentation, span_exporter: InMemorySpanExporter
+    enable_instrumentation: bool, span_exporter: InMemorySpanExporter
 ) -> None:
     """Test that no trace context is added when tracing is disabled."""
     # Tracing should be disabled by default
@@ -313,7 +313,7 @@ async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter)
     span_exporter.clear()
 
     # Run workflow (this should create run spans)
-    events = []
+    events: list[Any] = []
     async for event in workflow.run("test input", stream=True):
         events.append(event)
 

python/packages/core/tests/workflow/test_workflow_states.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from typing import Any
+
 import pytest
 from typing_extensions import Never
 
@@ -36,16 +38,16 @@ async def test_executor_failed_and_workflow_failed_events_streaming():
             events.append(ev)
 
     # executor_failed event (type='executor_failed') should be emitted before workflow failed event
-    executor_failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
+    executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
     assert executor_failed_events, "executor_failed event should be emitted when start executor fails"
     assert executor_failed_events[0].executor_id == "f"
     assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
 
     # Workflow-level failure and FAILED status should be surfaced
-    failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
+    failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
     assert failed_events
     assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)
-    status = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
+    status: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
     assert status and status[-1].state == WorkflowRunState.FAILED
     assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status)
 
@@ -94,13 +96,13 @@ async def test_executor_failed_event_from_second_executor_in_chain():
             events.append(ev)
 
     # executor_failed event should be emitted for the failing executor
-    executor_failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
+    executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
     assert executor_failed_events, "executor_failed event should be emitted when second executor fails"
     assert executor_failed_events[0].executor_id == "failing"
     assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
 
     # Workflow-level failure should also be surfaced
-    failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
+    failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
     assert failed_events
     assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)
 

python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py

@@ -388,11 +388,15 @@ class DeclarativeWorkflowState:
             from System.Globalization import CultureInfo
 
             original_culture = CultureInfo.CurrentCulture
-            CultureInfo.CurrentCulture = CultureInfo("en-US")
+            original_ui_culture = CultureInfo.CurrentUICulture
+            en_us_culture = CultureInfo("en-US")
+            CultureInfo.CurrentCulture = en_us_culture
+            CultureInfo.CurrentUICulture = en_us_culture
             try:
                 return engine.eval(formula, symbols=symbols)
             finally:
                 CultureInfo.CurrentCulture = original_culture
+                CultureInfo.CurrentUICulture = original_ui_culture
         except ValueError as e:
             error_msg = str(e)
             # Handle undefined variable errors gracefully by returning None

python/packages/declarative/tests/test_powerfx_yaml_compatibility.py

@@ -493,6 +493,31 @@ class TestPowerFxUndefinedVariables:
         result = state.eval("=Local.Something.Nested.Deep")
         assert result is None
 
+    async def test_undefined_variable_returns_none_with_non_english_ui_culture(self, mock_state):
+        """Test that undefined variables return None even when CurrentUICulture is non-English.
+
+        Regression test for #4321: on non-English systems, CurrentUICulture causes
+        PowerFx to emit localized error messages that don't match the English
+        string guards ("isn't recognized", "Name isn't valid"), crashing the workflow.
+        The fix sets CurrentUICulture to en-US alongside CurrentCulture before eval.
+        """
+        from System.Globalization import CultureInfo
+
+        state = DeclarativeWorkflowState(mock_state)
+        state.initialize()
+
+        # Simulate a non-English UI culture (e.g. Italian)
+        original_ui_culture = CultureInfo.CurrentUICulture
+        CultureInfo.CurrentUICulture = CultureInfo("it-IT")
+        try:
+            # Should return None, not raise ValueError with Italian error text
+            result = state.eval("=Local.StatusConversationId")
+            assert result is None
+            # Verify the production code restored CurrentUICulture after eval
+            assert str(CultureInfo.CurrentUICulture) == str(CultureInfo("it-IT"))
+        finally:
+            CultureInfo.CurrentUICulture = original_ui_culture
+
 
 class TestStringInterpolation:
     """Test string interpolation patterns."""

python/packages/orchestrations/agent_framework_orchestrations/_magentic.py

@@ -14,6 +14,7 @@ from typing import Any, ClassVar, TypeVar, cast
 
 from agent_framework import (
     AgentResponse,
+    AgentSession,
     Message,
     SupportsAgentRun,
 )
@@ -559,6 +560,7 @@ class StandardMagenticManager(MagenticManagerBase):
         )
 
         self._agent: SupportsAgentRun = agent
+        self._session: AgentSession = self._agent.create_session()
         self.task_ledger: _MagenticTaskLedger | None = task_ledger
 
         # Prompts may be overridden if needed
@@ -587,7 +589,7 @@ class StandardMagenticManager(MagenticManagerBase):
         The agent's run method is called which applies the agent's configured options
         (temperature, seed, instructions, etc.).
         """
-        response: AgentResponse = await self._agent.run(messages)
+        response: AgentResponse = await self._agent.run(messages, session=self._session)
         if not response.messages:
             raise RuntimeError("Agent returned no messages in response.")
         if len(response.messages) > 1:
@@ -730,6 +732,7 @@ class StandardMagenticManager(MagenticManagerBase):
         state: dict[str, Any] = {}
         if self.task_ledger is not None:
             state["task_ledger"] = self.task_ledger.to_dict()
+        state["agent_session"] = self._session.to_dict()
         return state
 
     @override
@@ -740,6 +743,12 @@ class StandardMagenticManager(MagenticManagerBase):
                 self.task_ledger = _MagenticTaskLedger.from_dict(ledger)
             except Exception:  # pragma: no cover - defensive
                 logger.warning("Failed to restore manager task ledger from checkpoint state")
+        session_payload = state.get("agent_session")
+        if session_payload is not None:
+            try:
+                self._session = AgentSession.from_dict(session_payload)
+            except Exception:  # pragma: no cover - defensive
+                logger.warning("Failed to restore manager agent session from checkpoint state")
 
 
 # endregion Magentic Manager

python/packages/orchestrations/tests/test_magentic.py

@@ -1074,4 +1074,71 @@ def test_magentic_agent_factory_with_standard_manager_options():
     assert manager.final_answer_prompt == custom_final_prompt
 
 
+async def test_standard_manager_propagates_session_to_agent():
+    """Verify StandardMagenticManager passes a consistent session to the underlying agent.
+
+    Regression test for #4371: context providers (e.g. RedisHistoryProvider) configured on
+    the manager agent silently failed because no session was propagated.
+    """
+    captured_sessions: list[AgentSession | None] = []
+
+    class SessionCapturingAgent(BaseAgent):
+        """Agent that records the session passed to each run() call."""
+
+        def run(
+            self,
+            messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+            *,
+            stream: bool = False,
+            session: Any = None,
+            **kwargs: Any,
+        ) -> Awaitable[AgentResponse] | AsyncIterable[AgentResponseUpdate]:
+            captured_sessions.append(session)
+
+            async def _run() -> AgentResponse:
+                return AgentResponse(messages=[Message("assistant", ["ok"])])
+
+            return _run()
+
+    agent = SessionCapturingAgent()
+    mgr = StandardMagenticManager(agent=agent)
+    ctx = MagenticContext(task="task", participant_descriptions={"a": "desc"})
+
+    await mgr.plan(ctx.clone())
+
+    # plan() calls _complete twice (facts + plan), both should receive the same session
+    assert len(captured_sessions) == 2
+    assert all(s is not None for s in captured_sessions), "session must be passed to agent.run()"
+    assert captured_sessions[0] is captured_sessions[1], "same session instance must be reused across calls"
+    assert captured_sessions[0] is mgr._session
+
+
+def test_standard_manager_checkpoint_preserves_session():
+    """Verify that checkpoint save/restore preserves the manager's session identity."""
+    agent = StubManagerAgent()
+    mgr = StandardMagenticManager(agent=agent)
+    original_session_id = mgr._session.session_id
+
+    state = mgr.on_checkpoint_save()
+    assert "agent_session" in state
+
+    # Restore into a fresh manager and verify session_id is preserved
+    mgr2 = StandardMagenticManager(agent=agent)
+    assert mgr2._session.session_id != original_session_id
+    mgr2.on_checkpoint_restore(state)
+    assert mgr2._session.session_id == original_session_id
+
+
+def test_standard_manager_checkpoint_restore_empty_state():
+    """Verify that restoring from a state without agent_session leaves the session intact."""
+    agent = StubManagerAgent()
+    mgr = StandardMagenticManager(agent=agent)
+    original_session = mgr._session
+    original_session_id = original_session.session_id
+
+    mgr.on_checkpoint_restore({})
+    assert mgr._session is original_session
+    assert mgr._session.session_id == original_session_id
+
+
 # endregion

python/pyproject.toml

@@ -184,7 +184,6 @@ omit = [
 
 [tool.pyright]
 include = ["agent_framework*"]
-exclude = ["**/tests/**", "**/.venv/**", "packages/devui/frontend/**"]
 typeCheckingMode = "strict"
 reportUnnecessaryIsInstance = false
 reportMissingTypeStubs = false

python/samples/02-agents/context_providers/azure_ai_foundry_memory.py

@@ -61,7 +61,7 @@ async def main() -> None:
         print(f"Creating memory store '{memory_store_name}'...")
         try:
             # Create a memory store
-            memory_store = await project_client.memory_stores.create(
+            memory_store = await project_client.beta.memory_stores.create(
                 name=memory_store_name,
                 description="Memory store for Agent Framework with FoundryMemoryProvider",
                 definition=memory_store_definition,
@@ -126,7 +126,7 @@ async def main() -> None:
                 print(f"Agent: {result3}\n")
 
                 print(f"Stored memories from: {memory_store.name} ({memory_store.id})")
-                res = await project_client.memory_stores.search_memories(name=memory_store.name, scope="user_123")
+                res = await project_client.beta.memory_stores.search_memories(name=memory_store.name, scope="user_123")
                 for memory in res.memories:
                     print(f"Memory: {memory.memory_item.content}")
 
@@ -134,7 +134,7 @@ async def main() -> None:
                 print(f"An error occurred: {e}")
 
             finally:
-                await project_client.memory_stores.delete(memory_store_name)
+                await project_client.beta.memory_stores.delete(memory_store_name)
                 print("==========================================")
                 print("Memory store deleted")
 

python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py

@@ -8,7 +8,7 @@ from typing import Annotated
 from agent_framework import tool
 from agent_framework.azure import AzureAIProjectAgentProvider
 from azure.ai.projects.aio import AIProjectClient
-from azure.ai.projects.models import AgentReference, PromptAgentDefinition
+from azure.ai.projects.models import PromptAgentDefinition
 from azure.identity.aio import AzureCliCredential
 from dotenv import load_dotenv
 from pydantic import Field
@@ -116,7 +116,7 @@ async def get_agent_by_name_example() -> None:
 async def get_agent_by_reference_example() -> None:
     """Example of using provider.get_agent(reference=...) to retrieve a specific agent version.
 
-    This method fetches a specific version of an agent using an AgentReference.
+    This method fetches a specific version of an agent using a reference mapping.
     Use this when you need to use a particular version of an agent.
     """
     print("=== provider.get_agent(reference=...) Example ===")
@@ -136,9 +136,9 @@ async def get_agent_by_reference_example() -> None:
         )
 
         try:
-            # Get the agent using an AgentReference with specific version
+            # Get the agent using a reference mapping with specific version
             provider = AzureAIProjectAgentProvider(project_client=project_client)
-            reference = AgentReference(name=created_agent.name, version=created_agent.version)
+            reference = {"name": created_agent.name, "version": created_agent.version}
             agent = await provider.get_agent(reference=reference)
 
             print(f"Retrieved agent: {agent.name} (version via reference)")

python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py

@@ -43,7 +43,7 @@ async def main() -> None:
                 options=MemoryStoreDefaultOptions(user_profile_enabled=True, chat_summary_enabled=True),
             )
 
-            memory_store = await project_client.memory_stores.create(
+            memory_store = await project_client.beta.memory_stores.create(
                 name=memory_store_name,
                 description="Memory store for Agent Framework conversations",
                 definition=memory_store_definition,
@@ -57,7 +57,7 @@ async def main() -> None:
                 instructions="""You are a helpful assistant that remembers past conversations.
                 Use the memory search tool to recall relevant information from previous interactions.""",
                 tools={
-                    "type": "memory_search",
+                    "type": "memory_search_preview",
                     "memory_store_name": memory_store.name,
                     "scope": "user_123",
                     "update_delay": 1,  # Wait 1 second before updating memories (use higher value in production)
@@ -84,7 +84,7 @@ async def main() -> None:
 
         # Clean up - delete the memory store
         async with AIProjectClient(endpoint=endpoint, credential=credential) as project_client:
-            await project_client.memory_stores.delete(memory_store_name)
+            await project_client.beta.memory_stores.delete(memory_store_name)
             print("Memory store deleted")
 
 

python/uv.lock

@@ -401,7 +401,7 @@ requires-dist = [
     { name = "agent-framework-orchestrations", marker = "extra == 'all'", editable = "packages/orchestrations" },
     { name = "agent-framework-purview", marker = "extra == 'all'", editable = "packages/purview" },
     { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" },
-    { name = "azure-ai-projects", specifier = "==2.0.0b3" },
+    { name = "azure-ai-projects", specifier = "==2.0.0b4" },
     { name = "azure-identity", specifier = ">=1,<2" },
     { name = "mcp", extras = ["ws"], specifier = ">=1.24.0,<2" },
     { name = "openai", specifier = ">=1.99.0" },
@@ -1014,7 +1014,7 @@ wheels = [
 
 [[package]]
 name = "azure-ai-projects"
-version = "2.0.0b3"
+version = "2.0.0b4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "azure-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -1022,10 +1022,11 @@ dependencies = [
     { name = "azure-storage-blob", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
     { name = "isodate", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
     { name = "openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/24/e0/3512d3f07e9dd2eb4af684387c31598c435bd87833b6a81850972963cb9c/azure_ai_projects-2.0.0b3.tar.gz", hash = "sha256:6d09ad110086e450a47b991ee8a3644f1be97fa3085d5981d543f900d78f4505", size = 431749, upload-time = "2026-01-06T05:31:25.849Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/24/e9/1cb8e95a19fbf174cfd7b30368a011b3e17503928b7801b8d9129b7cc59b/azure_ai_projects-2.0.0b4.tar.gz", hash = "sha256:b6082eacf0a11db59ad4c48cb7962f5204b9a0391000bc22421236f229ff783a", size = 477764, upload-time = "2026-02-24T17:57:52.489Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/4e/b6/8fbd4786bb5c0dd19eaff86ddce0fbfb53a6f90d712038272161067a076a/azure_ai_projects-2.0.0b3-py3-none-any.whl", hash = "sha256:3b3048a3ba3904d556ba392b7bd20b6e84c93bb39df6d43a6470cdb0ad08af8c", size = 240717, upload-time = "2026-01-06T05:31:27.716Z" },
+    { url = "https://files.pythonhosted.org/packages/27/6e/6445d510a8cb6a54f57e4344c14d825c37c5146fa69ccf9d9d15a29d23e2/azure_ai_projects-2.0.0b4-py3-none-any.whl", hash = "sha256:f4cf1615bd815744ddce304b97eea9456b7f6f0bd8725547c4e54e3a67534635", size = 231920, upload-time = "2026-02-24T17:57:53.917Z" },
 ]
 
 [[package]]
@@ -1408,7 +1409,7 @@ name = "clr-loader"
 version = "0.2.10"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/18/24/c12faf3f61614b3131b5c98d3bf0d376b49c7feaa73edca559aeb2aee080/clr_loader-0.2.10.tar.gz", hash = "sha256:81f114afbc5005bafc5efe5af1341d400e22137e275b042a8979f3feb9fc9446", size = 83605, upload-time = "2026-01-03T23:13:06.984Z" }
 wheels = [
@@ -1887,7 +1888,7 @@ name = "exceptiongroup"
 version = "1.3.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform == 'darwin') or (python_full_version < '3.13' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform == 'win32')" },
+    { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" }
 wheels = [
@@ -4654,8 +4655,8 @@ name = "powerfx"
 version = "0.0.34"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
-    { name = "pythonnet", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
+    { name = "pythonnet", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/9f/fb/6c4bf87e0c74ca1c563921ce89ca1c5785b7576bca932f7255cdf81082a7/powerfx-0.0.34.tar.gz", hash = "sha256:956992e7afd272657ed16d80f4cad24ec95d9e4a79fb9dfa4a068a09e136af32", size = 3237555, upload-time = "2025-12-22T15:50:59.682Z" }
 wheels = [
@@ -5318,7 +5319,7 @@ name = "pythonnet"
 version = "3.0.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "clr-loader", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "clr-loader", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/9a/d6/1afd75edd932306ae9bd2c2d961d603dc2b52fcec51b04afea464f1f6646/pythonnet-3.0.5.tar.gz", hash = "sha256:48e43ca463941b3608b32b4e236db92d8d40db4c58a75ace902985f76dac21cf", size = 239212, upload-time = "2024-12-13T08:30:44.393Z" }
 wheels = [