sharepoint sample fix (#3108 )

fix: tool_choice parameter not being honored when passed to agent.run() (#3095 )
.NET: Fix message ordering inconsistency when using AIContextProvider (#2659 )
2026-06-16 21:04:09 +08:00 · 2026-01-06 22:57:54 +00:00 · 2026-01-06 22:51:20 +00:00 · 2026-01-06 15:39:42 +00:00 · 2026-01-06 14:55:29 +00:00 · 2026-01-06 10:44:13 +00:00
225 changed files with 6216 additions and 26555 deletions
@@ -1,8 +1,8 @@
-blank_issues_enabled: true
+blank_issues_enabled: false
 contact_links:
  - name: Documentation
    url: https://aka.ms/agent-framework
    about: Check out the official documentation for guides and API reference.
  - name: Discussions
    url: https://github.com/microsoft/agent-framework/discussions
-    about: Ask questions about Agent Framework.
+    about: Ask questions and share ideas in GitHub Discussions.
@@ -1,70 +0,0 @@
-name: .NET Bug Report
-description: Report a bug in the Agent Framework .NET SDK
-title: ".NET: [Bug]: "
-labels: ["bug", ".NET"]
-type: bug
-body:
-  - type: textarea
-    id: description
-    attributes:
-      label: Description
-      description: Please provide a clear and detailed description of the bug.
-      placeholder: |
-        - What happened?
-        - What did you expect to happen?
-        - Steps to reproduce the issue
-    validations:
-      required: true
-
-  - type: textarea
-    id: code-sample
-    attributes:
-      label: Code Sample
-      description: If applicable, provide a minimal code sample that demonstrates the issue.
-      placeholder: |
-        ```csharp
-        // Your code here
-        ```
-      render: markdown
-    validations:
-      required: false
-
-  - type: textarea
-    id: error-messages
-    attributes:
-      label: Error Messages / Stack Traces
-      description: Include any error messages or stack traces you received.
-      placeholder: |
-        ```
-        Paste error messages or stack traces here
-        ```
-      render: markdown
-    validations:
-      required: false
-
-  - type: input
-    id: dotnet-packages
-    attributes:
-      label: Package Versions
-      description: List the Microsoft.Agents.* packages and versions you are using
-      placeholder: "e.g., Microsoft.Agents.AI.Abstractions: 1.0.0, Microsoft.Agents.AI.OpenAI: 1.0.0"
-    validations:
-      required: true
-
-  - type: input
-    id: dotnet-version
-    attributes:
-      label: .NET Version
-      description: What version of .NET are you using?
-      placeholder: "e.g., .NET 8.0"
-    validations:
-      required: false
-
-  - type: textarea
-    id: additional-context
-    attributes:
-      label: Additional Context
-      description: Add any other context or screenshots that might be helpful.
-      placeholder: "Any additional information..."
-    validations:
-      required: false
@@ -1,51 +0,0 @@
-name: Feature Request
-description: Request a new feature for Microsoft Agent Framework
-title: "[Feature]: "
-type: feature
-body:
-
-  - type: textarea
-    id: description
-    attributes:
-      label: Description
-      description: Please describe the feature you'd like and why it would be useful.
-      placeholder: |
-        Describe the feature you're requesting:
-        - What problem does it solve?
-        - What would the expected behavior be?
-        - Are there any alternatives you've considered?
-    validations:
-      required: true
-
-  - type: textarea
-    id: code-sample
-    attributes:
-      label: Code Sample
-      description: If applicable, provide a code sample showing how you'd like to use this feature.
-      placeholder: |
-        ```python
-        # Your code here
-        ```
-
-        or
-
-        ```csharp
-        // Your code here
-        ```
-      render: markdown
-    validations:
-      required: false
-
-  - type: dropdown
-    id: language
-    attributes:
-      label: Language/SDK
-      description: Which language/SDK does this feature apply to?
-      options:
-        - Both
-        - .NET
-        - Python
-        - Other / Not Applicable
-      default: 0
-    validations:
-      required: false
@@ -0,0 +1,203 @@
+name: Issue Report
+description: Report a bug, request a feature, or ask a question about Microsoft Agent Framework
+title: "[Issue]: "
+labels: ["triage"]
+body:
+  - type: dropdown
+    id: language
+    attributes:
+      label: Language
+      description: Which language/SDK are you using?
+      options:
+        - .NET
+        - Python
+        - None / Not Applicable
+    validations:
+      required: true
+
+  - type: dropdown
+    id: issue-type
+    attributes:
+      label: Type of Issue
+      description: What type of issue is this?
+      options:
+        - Bug
+        - Feature Request
+        - Question
+    validations:
+      required: true
+
+  - type: markdown
+    attributes:
+      value: |
+        ## Version Information
+        Please provide the version of the package(s) you are using. Select the relevant packages below.
+
+  - type: markdown
+    attributes:
+      value: "### .NET Packages"
+
+  - type: input
+    id: dotnet-agents-ai
+    attributes:
+      label: Microsoft.Agents.AI
+      description: Version of Microsoft.Agents.AI (e.g., 1.0.0)
+      placeholder: "e.g., 1.0.0"
+    validations:
+      required: false
+
+  - type: input
+    id: dotnet-agents-ai-abstractions
+    attributes:
+      label: Microsoft.Agents.AI.Abstractions
+      description: Version of Microsoft.Agents.AI.Abstractions
+      placeholder: "e.g., 1.0.0"
+    validations:
+      required: false
+
+  - type: input
+    id: dotnet-agents-ai-openai
+    attributes:
+      label: Microsoft.Agents.AI.OpenAI
+      description: Version of Microsoft.Agents.AI.OpenAI
+      placeholder: "e.g., 1.0.0"
+    validations:
+      required: false
+
+  - type: input
+    id: dotnet-agents-ai-azureai
+    attributes:
+      label: Microsoft.Agents.AI.AzureAI
+      description: Version of Microsoft.Agents.AI.AzureAI
+      placeholder: "e.g., 1.0.0"
+    validations:
+      required: false
+
+  - type: input
+    id: dotnet-agents-ai-anthropic
+    attributes:
+      label: Microsoft.Agents.AI.Anthropic
+      description: Version of Microsoft.Agents.AI.Anthropic
+      placeholder: "e.g., 1.0.0"
+    validations:
+      required: false
+
+  - type: input
+    id: dotnet-agents-ai-hosting
+    attributes:
+      label: Microsoft.Agents.AI.Hosting
+      description: Version of Microsoft.Agents.AI.Hosting
+      placeholder: "e.g., 1.0.0"
+    validations:
+      required: false
+
+  - type: input
+    id: dotnet-agents-ai-workflows
+    attributes:
+      label: Microsoft.Agents.AI.Workflows
+      description: Version of Microsoft.Agents.AI.Workflows
+      placeholder: "e.g., 1.0.0"
+    validations:
+      required: false
+
+  - type: input
+    id: dotnet-other-packages
+    attributes:
+      label: Other .NET Packages
+      description: List any other Microsoft.Agents.* packages and versions you are using
+      placeholder: "e.g., Microsoft.Agents.AI.CopilotStudio: 1.0.0, Microsoft.Agents.AI.Purview: 1.0.0"
+    validations:
+      required: false
+
+  - type: markdown
+    attributes:
+      value: "### Python Packages"
+
+  - type: input
+    id: python-core
+    attributes:
+      label: agent-framework-core
+      description: Version of agent-framework-core
+      placeholder: "e.g., 1.0.0b1"
+    validations:
+      required: false
+
+  - type: input
+    id: python-azure-ai
+    attributes:
+      label: agent-framework-azure-ai
+      description: Version of agent-framework-azure-ai
+      placeholder: "e.g., 1.0.0b1"
+    validations:
+      required: false
+
+  - type: input
+    id: python-anthropic
+    attributes:
+      label: agent-framework-anthropic
+      description: Version of agent-framework-anthropic
+      placeholder: "e.g., 1.0.0b1"
+    validations:
+      required: false
+
+  - type: input
+    id: python-azurefunctions
+    attributes:
+      label: agent-framework-azurefunctions
+      description: Version of agent-framework-azurefunctions
+      placeholder: "e.g., 1.0.0b1"
+    validations:
+      required: false
+
+  - type: input
+    id: python-other-packages
+    attributes:
+      label: Other Python Packages
+      description: List any other agent-framework-* packages and versions you are using
+      placeholder: "e.g., agent-framework-mem0: 1.0.0b1, agent-framework-redis: 1.0.0b1"
+    validations:
+      required: false
+
+  - type: markdown
+    attributes:
+      value: "---"
+
+  - type: textarea
+    id: description
+    attributes:
+      label: Description
+      description: Please provide a clear and detailed description of the issue, feature request, or question.
+      placeholder: |
+        For bugs: Describe what happened, what you expected to happen, and steps to reproduce.
+        For features: Describe the feature you'd like and why it would be useful.
+        For questions: Describe what you're trying to accomplish.
+    validations:
+      required: true
+
+  - type: textarea
+    id: code-sample
+    attributes:
+      label: Code Sample
+      description: If applicable, provide a minimal code sample that demonstrates the issue or your use case.
+      placeholder: |
+        ```python
+        # Your code here
+        ```
+
+        or
+
+        ```csharp
+        // Your code here
+        ```
+      render: markdown
+    validations:
+      required: false
+
+  - type: textarea
+    id: additional-context
+    attributes:
+      label: Additional Context
+      description: Add any other context, screenshots, error messages, or stack traces that might be helpful.
+      placeholder: "Any additional information..."
+    validations:
+      required: false
@@ -1,70 +0,0 @@
-name: Python Bug Report
-description: Report a bug in the Agent Framework Python SDK
-title: "Python: [Bug]: "
-labels: ["bug", "Python"]
-type: bug
-body:
-  - type: textarea
-    id: description
-    attributes:
-      label: Description
-      description: Please provide a clear and detailed description of the bug.
-      placeholder: |
-        - What happened?
-        - What did you expect to happen?
-        - Steps to reproduce the issue
-    validations:
-      required: true
-
-  - type: textarea
-    id: code-sample
-    attributes:
-      label: Code Sample
-      description: If applicable, provide a minimal code sample that demonstrates the issue.
-      placeholder: |
-        ```python
-        # Your code here
-        ```
-      render: markdown
-    validations:
-      required: false
-
-  - type: textarea
-    id: error-messages
-    attributes:
-      label: Error Messages / Stack Traces
-      description: Include any error messages or stack traces you received.
-      placeholder: |
-        ```
-        Paste error messages or stack traces here
-        ```
-      render: markdown
-    validations:
-      required: false
-
-  - type: input
-    id: python-packages
-    attributes:
-      label: Package Versions
-      description: List the agent-framework-* packages and versions you are using
-      placeholder: "e.g., agent-framework-core: 1.0.0, agent-framework-azure-ai: 1.0.0"
-    validations:
-      required: true
-
-  - type: input
-    id: python-version
-    attributes:
-      label: Python Version
-      description: What version of Python are you using?
-      placeholder: "e.g., Python 3.11"
-    validations:
-      required: false
-
-  - type: textarea
-    id: additional-context
-    attributes:
-      label: Additional Context
-      description: Add any other context or screenshots that might be helpful.
-      placeholder: "Any additional information..."
-    validations:
-      required: false
@@ -208,14 +208,13 @@ WARP.md
 **/projectBrief.md

 # Azurite storage emulator files
-*/__azurite_db_blob__.json*
-*/__azurite_db_blob_extent__.json*
-*/__azurite_db_queue__.json*
-*/__azurite_db_queue_extent__.json*
-*/__azurite_db_table__.json*
+*/__azurite_db_blob__.json
+*/__azurite_db_blob_extent__.json
+*/__azurite_db_queue__.json
+*/__azurite_db_queue_extent__.json
+*/__azurite_db_table__.json
 */__blobstorage__/
 */__queuestorage__/
-*/AzuriteConfig

 # Azure Functions local settings
 local.settings.json
@@ -11,7 +11,7 @@
  </PropertyGroup>
  <ItemGroup>
    <!-- Aspire.* -->
-    <PackageVersion Include="Anthropic" Version="12.0.1" />
+    <PackageVersion Include="Anthropic" Version="12.0.0" />
    <PackageVersion Include="Anthropic.Foundry" Version="0.1.0" />
    <PackageVersion Include="Aspire.Azure.AI.OpenAI" Version="13.0.0-preview.1.25560.3" />
    <PackageVersion Include="Aspire.Hosting.AppHost" Version="$(AspireAppHostSdkVersion)" />
@@ -26,7 +26,7 @@
    <PackageVersion Include="Azure.Identity" Version="1.17.1" />
    <PackageVersion Include="Azure.Monitor.OpenTelemetry.Exporter" Version="1.4.0" />
    <!-- Google Gemini  -->
-    <PackageVersion Include="Google.GenAI" Version="0.9.0" />
+    <PackageVersion Include="Google.GenAI" Version="0.6.0" />
    <PackageVersion Include="Mscc.GenerativeAI.Microsoft" Version="2.9.3" />
    <!-- Microsoft.Azure.* -->
    <PackageVersion Include="Microsoft.Azure.Cosmos" Version="3.54.0" />
@@ -100,7 +100,7 @@
    <!-- MCP -->
    <PackageVersion Include="ModelContextProtocol" Version="0.4.0-preview.3" />
    <!-- Inference SDKs -->
-    <PackageVersion Include="AWSSDK.Extensions.Bedrock.MEAI" Version="4.0.5.1" />
+    <PackageVersion Include="AWSSDK.Extensions.Bedrock.MEAI" Version="4.0.5" />
    <PackageVersion Include="Microsoft.ML.OnnxRuntimeGenAI" Version="0.10.0" />
    <PackageVersion Include="OllamaSharp" Version="5.4.8" />
    <PackageVersion Include="OpenAI" Version="2.8.0" />
@@ -2,9 +2,9 @@
  <PropertyGroup>
    <!-- Central version prefix - applies to all nuget packages. -->
    <VersionPrefix>1.0.0</VersionPrefix>
-    <PackageVersion Condition="'$(VersionSuffix)' != ''">$(VersionPrefix)-$(VersionSuffix).260108.1</PackageVersion>
-    <PackageVersion Condition="'$(VersionSuffix)' == ''">$(VersionPrefix)-preview.260108.1</PackageVersion>
-    <GitTag>1.0.0-preview.260108.1</GitTag>
+    <PackageVersion Condition="'$(VersionSuffix)' != ''">$(VersionPrefix)-$(VersionSuffix).251219.1</PackageVersion>
+    <PackageVersion Condition="'$(VersionSuffix)' == ''">$(VersionPrefix)-preview.251219.1</PackageVersion>
+    <GitTag>1.0.0-preview.251219.1</GitTag>

    <Configurations>Debug;Release;Publish</Configurations>
    <IsPackable>true</IsPackable>
@@ -39,8 +39,8 @@
    <PackageReference Include="Azure.AI.AgentServer.AgentFramework" Version="1.0.0-beta.5" />
    <PackageReference Include="Azure.AI.OpenAI" Version="2.7.0-beta.2" />
    <PackageReference Include="Azure.Identity" Version="1.17.1" />
-    <PackageReference Include="Microsoft.Agents.AI.OpenAI" Version="1.0.0-preview.251219.1" />
-    <PackageReference Include="Microsoft.Extensions.AI.OpenAI" Version="10.1.1-preview.1.25612.2" />
+    <PackageReference Include="Microsoft.Agents.AI.OpenAI" Version="1.0.0-preview.251125.1" />
+    <PackageReference Include="Microsoft.Extensions.AI.OpenAI" Version="10.1.0-preview.1.25608.1" />
  </ItemGroup>

  <!-- Add analyzers with compatible versions -->
@@ -9,7 +9,7 @@ using Azure.AI.OpenAI;
 using Azure.Identity;
 using Microsoft.Agents.AI;
 using Microsoft.Extensions.AI;
-using OpenAI.Responses;
+using OpenAI;

 var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set.");
 var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-4o-mini";
@@ -25,7 +25,7 @@ AITool mcpTool = new HostedMcpServerTool(serverName: "microsoft_learn", serverAd
 AIAgent agent = new AzureOpenAIClient(
    new Uri(endpoint),
    new DefaultAzureCredential())
-    .GetResponsesClient(deploymentName)
+    .GetOpenAIResponseClient(deploymentName)
    .CreateAIAgent(
        instructions: "You answer questions by searching the Microsoft Learn content only.",
        name: "MicrosoftLearnAgent",
@@ -38,8 +38,8 @@
    <PackageReference Include="Azure.AI.AgentServer.AgentFramework" Version="1.0.0-beta.5" />
    <PackageReference Include="Azure.AI.OpenAI" Version="2.7.0-beta.2" />
    <PackageReference Include="Azure.Identity" Version="1.17.1" />
-    <PackageReference Include="Microsoft.Agents.AI.OpenAI" Version="1.0.0-preview.251219.1" />
-    <PackageReference Include="Microsoft.Extensions.AI.OpenAI" Version="10.1.1-preview.1.25612.2" />
+    <PackageReference Include="Microsoft.Agents.AI.OpenAI" Version="1.0.0-preview.251125.1" />
+    <PackageReference Include="Microsoft.Extensions.AI.OpenAI" Version="10.1.0-preview.1.25608.1" />
  </ItemGroup>

  <!-- Add analyzers with compatible versions -->
@@ -8,8 +8,9 @@ using Azure.AI.AgentServer.AgentFramework.Extensions;
 using Azure.AI.OpenAI;
 using Azure.Identity;
 using Microsoft.Agents.AI;
+using Microsoft.Agents.AI.Data;
 using Microsoft.Extensions.AI;
-using OpenAI.Chat;
+using OpenAI;

 var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set.");
 var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-4o-mini";
@@ -27,10 +28,7 @@ AIAgent agent = new AzureOpenAIClient(
    .GetChatClient(deploymentName)
    .CreateAIAgent(new ChatClientAgentOptions
    {
-        ChatOptions = new ChatOptions
-        {
-            Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available.",
-        },
+        Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available.",
        AIContextProviderFactory = ctx => new TextSearchProvider(MockSearchAsync, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions)
    });

@@ -38,7 +38,7 @@
    <PackageReference Include="Azure.AI.AgentServer.AgentFramework" Version="1.0.0-beta.5" />
    <PackageReference Include="Azure.AI.OpenAI" Version="2.7.0-beta.2" />
    <PackageReference Include="Azure.Identity" Version="1.17.1" />
-    <PackageReference Include="Microsoft.Agents.AI.Workflows" Version="1.0.0-preview.251219.1" />
+    <PackageReference Include="Microsoft.Agents.AI.Workflows" Version="1.0.0-preview.251125.1" />
    <PackageReference Include="Microsoft.Extensions.AI.OpenAI" Version="10.1.0-preview.1.25608.1" />
  </ItemGroup>

@@ -24,7 +24,7 @@ namespace Microsoft.Agents.AI.A2A;
 /// Support for tasks will be added later as part of the long-running
 /// executions work.
 /// </remarks>
-public sealed class A2AAgent : AIAgent
+internal sealed class A2AAgent : AIAgent
 {
    private readonly A2AClient _a2aClient;
    private readonly string? _id;
@@ -84,13 +84,9 @@ public sealed class A2AAgent : AIAgent
        }
        else
        {
-            MessageSendParams sendParams = new()
-            {
-                Message = CreateA2AMessage(typedThread, messages),
-                Metadata = options?.AdditionalProperties?.ToA2AMetadata()
-            };
+            var a2aMessage = CreateA2AMessage(typedThread, messages);

-            a2aResponse = await this._a2aClient.SendMessageAsync(sendParams, cancellationToken).ConfigureAwait(false);
+            a2aResponse = await this._a2aClient.SendMessageAsync(new MessageSendParams { Message = a2aMessage }, cancellationToken).ConfigureAwait(false);
        }

        this._logger.LogAgentChatClientInvokedAgent(nameof(RunAsync), this.Id, this.Name);
@@ -158,13 +154,9 @@ public sealed class A2AAgent : AIAgent
            // a2aSseEvents = this._a2aClient.SubscribeToTaskAsync(token.TaskId, cancellationToken).ConfigureAwait(false);  
        }

-        MessageSendParams sendParams = new()
-        {
-            Message = CreateA2AMessage(typedThread, messages),
-            Metadata = options?.AdditionalProperties?.ToA2AMetadata()
-        };
+        var a2aMessage = CreateA2AMessage(typedThread, messages);

-        a2aSseEvents = this._a2aClient.SendMessageStreamingAsync(sendParams, cancellationToken).ConfigureAwait(false);
+        a2aSseEvents = this._a2aClient.SendMessageStreamingAsync(new MessageSendParams { Message = a2aMessage }, cancellationToken).ConfigureAwait(false);

        this._logger.LogAgentChatClientInvokedAgent(nameof(RunStreamingAsync), this.Id, this.Name);

@@ -206,10 +198,10 @@ public sealed class A2AAgent : AIAgent
    protected override string? IdCore => this._id;

    /// <inheritdoc/>
-    public override string? Name => this._name;
+    public override string? Name => this._name ?? base.Name;

    /// <inheritdoc/>
-    public override string? Description => this._description;
+    public override string? Description => this._description ?? base.Description;

    private A2AAgentThread GetA2AThread(AgentThread? thread, AgentRunOptions? options)
    {
@@ -14,9 +14,6 @@ internal static class A2AMetadataExtensions
    /// <summary>
    /// Converts a dictionary of metadata to an <see cref="AdditionalPropertiesDictionary"/>.
    /// </summary>
-    /// <remarks>
-    /// This method can be replaced by the one from A2A SDK once it is public.
-    /// </remarks>
    /// <param name="metadata">The metadata dictionary to convert.</param>
    /// <returns>The converted <see cref="AdditionalPropertiesDictionary"/>, or null if the input is null or empty.</returns>
    internal static AdditionalPropertiesDictionary? ToAdditionalProperties(this Dictionary<string, JsonElement>? metadata)
@@ -1,44 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using Microsoft.Agents.AI;
-
-namespace Microsoft.Extensions.AI;
-
-/// <summary>
-/// Extension methods for AdditionalPropertiesDictionary.
-/// </summary>
-internal static class AdditionalPropertiesDictionaryExtensions
-{
-    /// <summary>
-    /// Converts an <see cref="AdditionalPropertiesDictionary"/> to a dictionary of <see cref="JsonElement"/> values suitable for A2A metadata.
-    /// </summary>
-    /// <remarks>
-    /// This method can be replaced by the one from A2A SDK once it is available.
-    /// </remarks>
-    /// <param name="additionalProperties">The additional properties dictionary to convert, or <c>null</c>.</param>
-    /// <returns>A dictionary of JSON elements representing the metadata, or <c>null</c> if the input is null or empty.</returns>
-    internal static Dictionary<string, JsonElement>? ToA2AMetadata(this AdditionalPropertiesDictionary? additionalProperties)
-    {
-        if (additionalProperties is not { Count: > 0 })
-        {
-            return null;
-        }
-
-        var metadata = new Dictionary<string, JsonElement>();
-
-        foreach (var kvp in additionalProperties)
-        {
-            if (kvp.Value is JsonElement)
-            {
-                metadata[kvp.Key] = (JsonElement)kvp.Value!;
-                continue;
-            }
-
-            metadata[kvp.Key] = JsonSerializer.SerializeToElement(kvp.Value, A2AJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object)));
-        }
-
-        return metadata;
-    }
-}
@@ -291,15 +291,6 @@ public class AgentRunResponse
        return updates;
    }

-    /// <summary>
-    /// Deserializes the response text into the given type.
-    /// </summary>
-    /// <typeparam name="T">The output type to deserialize into.</typeparam>
-    /// <returns>The result as the requested type.</returns>
-    /// <exception cref="InvalidOperationException">The result is not parsable into the requested type.</exception>
-    public T Deserialize<T>() =>
-        this.Deserialize<T>(AgentAbstractionsJsonUtilities.DefaultOptions);
-
    /// <summary>
    /// Deserializes the response text into the given type using the specified serializer options.
    /// </summary>
@@ -320,15 +311,6 @@ public class AgentRunResponse
        };
    }

-    /// <summary>
-    /// Tries to deserialize response text into the given type.
-    /// </summary>
-    /// <typeparam name="T">The output type to deserialize into.</typeparam>
-    /// <param name="structuredOutput">The parsed structured output.</param>
-    /// <returns><see langword="true" /> if parsing was successful; otherwise, <see langword="false" />.</returns>
-    public bool TryDeserialize<T>([NotNullWhen(true)] out T? structuredOutput) =>
-        this.TryDeserialize(AgentAbstractionsJsonUtilities.DefaultOptions, out structuredOutput);
-
    /// <summary>
    /// Tries to deserialize response text into the given type using the specified serializer options.
    /// </summary>
@@ -43,14 +43,10 @@ public static class AIAgentExtensions
        {
            var contextId = messageSendParams.Message.ContextId ?? Guid.NewGuid().ToString("N");
            var thread = await hostAgent.GetOrCreateThreadAsync(contextId, cancellationToken).ConfigureAwait(false);
-            var options = messageSendParams.Metadata is not { Count: > 0 }
-                ? null
-                : new AgentRunOptions { AdditionalProperties = messageSendParams.Metadata.ToAdditionalProperties() };

            var response = await hostAgent.RunAsync(
                messageSendParams.ToChatMessages(),
                thread: thread,
-                options: options,
                cancellationToken: cancellationToken).ConfigureAwait(false);

            await hostAgent.SaveThreadAsync(contextId, thread, cancellationToken).ConfigureAwait(false);
@@ -60,8 +56,7 @@ public static class AIAgentExtensions
                MessageId = response.ResponseId ?? Guid.NewGuid().ToString("N"),
                ContextId = contextId,
                Role = MessageRole.Agent,
-                Parts = parts,
-                Metadata = response.AdditionalProperties?.ToA2AMetadata()
+                Parts = parts
            };
        }
    }
@@ -1,36 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI.Hosting.A2A.Converters;
-
-/// <summary>
-/// Extension methods for A2A metadata dictionary.
-/// </summary>
-internal static class A2AMetadataExtensions
-{
-    /// <summary>
-    /// Converts a dictionary of metadata to an <see cref="AdditionalPropertiesDictionary"/>.
-    /// </summary>
-    /// <remarks>
-    /// This method can be replaced by the one from A2A SDK once it is public.
-    /// </remarks>
-    /// <param name="metadata">The metadata dictionary to convert.</param>
-    /// <returns>The converted <see cref="AdditionalPropertiesDictionary"/>, or null if the input is null or empty.</returns>
-    internal static AdditionalPropertiesDictionary? ToAdditionalProperties(this Dictionary<string, JsonElement>? metadata)
-    {
-        if (metadata is not { Count: > 0 })
-        {
-            return null;
-        }
-
-        var additionalProperties = new AdditionalPropertiesDictionary();
-        foreach (var kvp in metadata)
-        {
-            additionalProperties[kvp.Key] = kvp.Value;
-        }
-        return additionalProperties;
-    }
-}
@@ -1,45 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using A2A;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI.Hosting.A2A.Converters;
-
-/// <summary>
-/// Extension methods for AdditionalPropertiesDictionary.
-/// </summary>
-internal static class AdditionalPropertiesDictionaryExtensions
-{
-    /// <summary>
-    /// Converts an <see cref="AdditionalPropertiesDictionary"/> to a dictionary of <see cref="JsonElement"/> values suitable for A2A metadata.
-    /// </summary>
-    /// <remarks>
-    /// This method can be replaced by the one from A2A SDK once it is available.
-    /// </remarks>
-    /// <param name="additionalProperties">The additional properties dictionary to convert, or <c>null</c>.</param>
-    /// <returns>A dictionary of JSON elements representing the metadata, or <c>null</c> if the input is null or empty.</returns>
-    internal static Dictionary<string, JsonElement>? ToA2AMetadata(this AdditionalPropertiesDictionary? additionalProperties)
-    {
-        if (additionalProperties is not { Count: > 0 })
-        {
-            return null;
-        }
-
-        var metadata = new Dictionary<string, JsonElement>();
-
-        foreach (var kvp in additionalProperties)
-        {
-            if (kvp.Value is JsonElement)
-            {
-                metadata[kvp.Key] = (JsonElement)kvp.Value!;
-                continue;
-            }
-
-            metadata[kvp.Key] = JsonSerializer.SerializeToElement(kvp.Value, A2AJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object)));
-        }
-
-        return metadata;
-    }
-}
@@ -50,7 +50,7 @@ TokenCredential browserCredential = new InteractiveBrowserCredential(
 IChatClient client = new AzureOpenAIClient(
    new Uri(endpoint),
    new AzureCliCredential())
-    .GetResponsesClient(deploymentName)
+    .GetOpenAIResponseClient(deploymentName)
    .AsIChatClient()
    .AsBuilder()
    .WithPurview(browserCredential, new PurviewSettings("My Sample App"))
@@ -198,7 +198,7 @@ Use the chat middleware when you attach directly to a chat client (e.g. minimal
 IChatClient client = new AzureOpenAIClient(
    new Uri(endpoint),
    new AzureCliCredential())
-    .GetResponsesClient(deploymentName)
+    .GetOpenAIResponseClient(deploymentName)
    .AsIChatClient()
    .AsBuilder()
    .WithPurview(browserCredential, new PurviewSettings("Agent Framework Test App"))
@@ -10,10 +10,4 @@ namespace Microsoft.Agents.AI.Workflows;
 /// <param name="e">
 /// Optionally, the <see cref="Exception"/> representing the error.
 /// </param>
-public class WorkflowErrorEvent(Exception? e) : WorkflowEvent(e)
-{
-    /// <summary>
-    /// Gets the exception that caused the current operation to fail, if one occurred.
-    /// </summary>
-    public Exception? Exception => this.Data as Exception;
-}
+public class WorkflowErrorEvent(Exception? e) : WorkflowEvent(e);
@@ -18,12 +18,11 @@ internal sealed class WorkflowHostAgent : AIAgent
    private readonly string? _id;
    private readonly CheckpointManager? _checkpointManager;
    private readonly IWorkflowExecutionEnvironment _executionEnvironment;
-    private readonly bool _includeExceptionDetails;
    private readonly Task<ProtocolDescriptor> _describeTask;

    private readonly ConcurrentDictionary<string, string> _assignedRunIds = [];

-    public WorkflowHostAgent(Workflow workflow, string? id = null, string? name = null, string? description = null, CheckpointManager? checkpointManager = null, IWorkflowExecutionEnvironment? executionEnvironment = null, bool includeExceptionDetails = false)
+    public WorkflowHostAgent(Workflow workflow, string? id = null, string? name = null, string? description = null, CheckpointManager? checkpointManager = null, IWorkflowExecutionEnvironment? executionEnvironment = null)
    {
        this._workflow = Throw.IfNull(workflow);

@@ -31,7 +30,6 @@ internal sealed class WorkflowHostAgent : AIAgent
                                                              ? InProcessExecution.Concurrent
                                                              : InProcessExecution.OffThread);
        this._checkpointManager = checkpointManager;
-        this._includeExceptionDetails = includeExceptionDetails;

        this._id = id;
        this.Name = name;
@@ -63,10 +61,10 @@ internal sealed class WorkflowHostAgent : AIAgent
        protocol.ThrowIfNotChatProtocol();
    }

-    public override AgentThread GetNewThread() => new WorkflowThread(this._workflow, this.GenerateNewId(), this._executionEnvironment, this._checkpointManager, this._includeExceptionDetails);
+    public override AgentThread GetNewThread() => new WorkflowThread(this._workflow, this.GenerateNewId(), this._executionEnvironment, this._checkpointManager);

    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
-        => new WorkflowThread(this._workflow, serializedThread, this._executionEnvironment, this._checkpointManager, this._includeExceptionDetails, jsonSerializerOptions);
+        => new WorkflowThread(this._workflow, serializedThread, this._executionEnvironment, this._checkpointManager, jsonSerializerOptions);

    private ValueTask<WorkflowThread> UpdateThreadAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, CancellationToken cancellationToken = default)
    {
@@ -21,8 +21,6 @@ public static class WorkflowHostingExtensions
    /// <param name="executionEnvironment">Specify the execution environment to use when running the workflows. See
    /// <see cref="InProcessExecution.OffThread"/>, <see cref="InProcessExecution.Concurrent"/> and
    /// <see cref="InProcessExecution.Lockstep"/> for the in-process environments.</param>
-    /// <param name="includeExceptionDetails">If <see langword="true"/>, will include <see cref="System.Exception.Message"/>
-    /// in the <see cref="ErrorContent"/> representing the workflow error.</param>
    /// <returns></returns>
    public static AIAgent AsAgent(
        this Workflow workflow,
@@ -30,10 +28,9 @@ public static class WorkflowHostingExtensions
        string? name = null,
        string? description = null,
        CheckpointManager? checkpointManager = null,
-        IWorkflowExecutionEnvironment? executionEnvironment = null,
-        bool includeExceptionDetails = false)
+        IWorkflowExecutionEnvironment? executionEnvironment = null)
    {
-        return new WorkflowHostAgent(workflow, id, name, description, checkpointManager, executionEnvironment, includeExceptionDetails);
+        return new WorkflowHostAgent(workflow, id, name, description, checkpointManager, executionEnvironment);
    }

    internal static FunctionCallContent ToFunctionCall(this ExternalRequest request)
@@ -3,7 +3,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Reflection;
 using System.Runtime.CompilerServices;
 using System.Text.Json;
 using System.Threading;
@@ -18,16 +17,14 @@ internal sealed class WorkflowThread : AgentThread
 {
    private readonly Workflow _workflow;
    private readonly IWorkflowExecutionEnvironment _executionEnvironment;
-    private readonly bool _includeExceptionDetails;

    private readonly CheckpointManager _checkpointManager;
    private readonly InMemoryCheckpointManager? _inMemoryCheckpointManager;

-    public WorkflowThread(Workflow workflow, string runId, IWorkflowExecutionEnvironment executionEnvironment, CheckpointManager? checkpointManager = null, bool includeExceptionDetails = false)
+    public WorkflowThread(Workflow workflow, string runId, IWorkflowExecutionEnvironment executionEnvironment, CheckpointManager? checkpointManager = null)
    {
        this._workflow = Throw.IfNull(workflow);
        this._executionEnvironment = Throw.IfNull(executionEnvironment);
-        this._includeExceptionDetails = includeExceptionDetails;

        // If the user provided an external checkpoint manager, use that, otherwise rely on an in-memory one.
        // TODO: Implement persist-only-last functionality for in-memory checkpoint manager, to avoid unbounded
@@ -38,7 +35,7 @@ internal sealed class WorkflowThread : AgentThread
        this.MessageStore = new WorkflowMessageStore();
    }

-    public WorkflowThread(Workflow workflow, JsonElement serializedThread, IWorkflowExecutionEnvironment executionEnvironment, CheckpointManager? checkpointManager = null, bool includeExceptionDetails = false, JsonSerializerOptions? jsonSerializerOptions = null)
+    public WorkflowThread(Workflow workflow, JsonElement serializedThread, IWorkflowExecutionEnvironment executionEnvironment, CheckpointManager? checkpointManager = null, JsonSerializerOptions? jsonSerializerOptions = null)
    {
        this._workflow = Throw.IfNull(workflow);
        this._executionEnvironment = Throw.IfNull(executionEnvironment);
@@ -83,7 +80,7 @@ internal sealed class WorkflowThread : AgentThread
        return marshaller.Marshal(info);
    }

-    public AgentRunResponseUpdate CreateUpdate(string responseId, object raw, params AIContent[] parts)
+    public AgentRunResponseUpdate CreateUpdate(string responseId, params AIContent[] parts)
    {
        Throw.IfNullOrEmpty(parts);

@@ -92,8 +89,7 @@ internal sealed class WorkflowThread : AgentThread
            CreatedAt = DateTimeOffset.UtcNow,
            MessageId = Guid.NewGuid().ToString("N"),
            Role = ChatRole.Assistant,
-            ResponseId = responseId,
-            RawRepresentation = raw
+            ResponseId = responseId
        };

        this.MessageStore.AddMessages(update.ToChatMessage());
@@ -157,29 +153,10 @@ internal sealed class WorkflowThread : AgentThread

                    case RequestInfoEvent requestInfo:
                        FunctionCallContent fcContent = requestInfo.Request.ToFunctionCall();
-                        AgentRunResponseUpdate update = this.CreateUpdate(this.LastResponseId, evt, fcContent);
+                        AgentRunResponseUpdate update = this.CreateUpdate(this.LastResponseId, fcContent);
                        yield return update;
                        break;

-                    case WorkflowErrorEvent workflowError:
-                        Exception? exception = workflowError.Exception;
-                        if (exception is TargetInvocationException tie && tie.InnerException != null)
-                        {
-                            exception = tie.InnerException;
-                        }
-
-                        if (exception != null)
-                        {
-                            string message = this._includeExceptionDetails
-                                           ? exception.Message
-                                           : "An error occurred while executing the workflow.";
-
-                            ErrorContent errorContent = new(message);
-                            yield return this.CreateUpdate(this.LastResponseId, evt, errorContent);
-                        }
-
-                        break;
-
                    case SuperStepCompletedEvent stepCompleted:
                        this.LastCheckpoint = stepCompleted.CompletionInfo?.Checkpoint;
                        goto default;
@@ -162,10 +162,7 @@ public sealed partial class ChatClientAgent : AIAgent

        static AgentRunResponse CreateResponse(ChatResponse chatResponse)
        {
-            return new AgentRunResponse(chatResponse)
-            {
-                ContinuationToken = WrapContinuationToken(chatResponse.ContinuationToken)
-            };
+            return new AgentRunResponse(chatResponse);
        }

        return this.RunCoreAsync(GetResponseAsync, CreateResponse, messages, thread, options, cancellationToken);
@@ -204,14 +201,11 @@ public sealed partial class ChatClientAgent : AIAgent
    {
        var inputMessages = Throw.IfNull(messages) as IReadOnlyCollection<ChatMessage> ?? messages.ToList();

-        (ChatClientAgentThread safeThread,
-         ChatOptions? chatOptions,
-         List<ChatMessage> inputMessagesForChatClient,
-         IList<ChatMessage>? aiContextProviderMessages,
-         IList<ChatMessage>? chatMessageStoreMessages,
-         ChatClientAgentContinuationToken? continuationToken) =
+        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages, IList<ChatMessage>? chatMessageStoreMessages) =
            await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false);

+        ValidateStreamResumptionAllowed(chatOptions?.ContinuationToken, safeThread);
+
        var chatClient = this.ChatClient;

        chatClient = ApplyRunOptionsTransformations(options, chatClient);
@@ -220,7 +214,7 @@ public sealed partial class ChatClientAgent : AIAgent

        this._logger.LogAgentChatClientInvokingAgent(nameof(RunStreamingAsync), this.Id, loggingAgentName, this._chatClientType);

-        List<ChatResponseUpdate> responseUpdates = GetResponseUpdates(continuationToken);
+        List<ChatResponseUpdate> responseUpdates = [];

        IAsyncEnumerator<ChatResponseUpdate> responseUpdatesEnumerator;

@@ -231,8 +225,8 @@ public sealed partial class ChatClientAgent : AIAgent
        }
        catch (Exception ex)
        {
-            await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
-            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
            throw;
        }

@@ -246,8 +240,8 @@ public sealed partial class ChatClientAgent : AIAgent
        }
        catch (Exception ex)
        {
-            await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
-            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
            throw;
        }

@@ -259,12 +253,7 @@ public sealed partial class ChatClientAgent : AIAgent
                update.AuthorName ??= this.Name;

                responseUpdates.Add(update);
-
-                yield return new(update)
-                {
-                    AgentId = this.Id,
-                    ContinuationToken = WrapContinuationToken(update.ContinuationToken, GetInputMessages(inputMessages, continuationToken), responseUpdates)
-                };
+                yield return new(update) { AgentId = this.Id };
            }

            try
@@ -273,8 +262,8 @@ public sealed partial class ChatClientAgent : AIAgent
            }
            catch (Exception ex)
            {
-                await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
-                await NotifyAIContextProviderOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+                await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+                await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
                throw;
            }
        }
@@ -286,10 +275,10 @@ public sealed partial class ChatClientAgent : AIAgent
        this.UpdateThreadWithTypeAndConversationId(safeThread, chatResponse.ConversationId);

        // To avoid inconsistent state we only notify the thread of the input messages if no error occurs after the initial request.
-        await NotifyMessageStoreOfNewMessagesAsync(safeThread, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
+        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);

        // Notify the AIContextProvider of all new messages.
-        await NotifyAIContextProviderOfSuccessAsync(safeThread, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
+        await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc/>
@@ -393,12 +382,7 @@ public sealed partial class ChatClientAgent : AIAgent
    {
        var inputMessages = Throw.IfNull(messages) as IReadOnlyCollection<ChatMessage> ?? messages.ToList();

-        (ChatClientAgentThread safeThread,
-         ChatOptions? chatOptions,
-         List<ChatMessage> inputMessagesForChatClient,
-         IList<ChatMessage>? aiContextProviderMessages,
-         IList<ChatMessage>? chatMessageStoreMessages,
-         ChatClientAgentContinuationToken? _) =
+        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages, IList<ChatMessage>? chatMessageStoreMessages) =
            await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false);

        var chatClient = this.ChatClient;
@@ -490,20 +474,20 @@ public sealed partial class ChatClientAgent : AIAgent
    /// <param name="runOptions">Optional run options that may include specific chat configuration settings.</param>
    /// <returns>A <see cref="ChatOptions"/> object representing the merged chat configuration, or <see langword="null"/> if
    /// neither the run options nor the agent's chat options are available.</returns>
-    private (ChatOptions?, ChatClientAgentContinuationToken?) CreateConfiguredChatOptions(AgentRunOptions? runOptions)
+    private ChatOptions? CreateConfiguredChatOptions(AgentRunOptions? runOptions)
    {
        ChatOptions? requestChatOptions = (runOptions as ChatClientAgentRunOptions)?.ChatOptions?.Clone();

        // If no agent chat options were provided, return the request chat options as is.
        if (this._agentOptions?.ChatOptions is null)
        {
-            return GetContinuationTokenAndApplyBackgroundResponsesProperties(requestChatOptions, runOptions);
+            return ApplyBackgroundResponsesProperties(requestChatOptions, runOptions);
        }

        // If no request chat options were provided, use the agent's chat options clone.
        if (requestChatOptions is null)
        {
-            return GetContinuationTokenAndApplyBackgroundResponsesProperties(this._agentOptions?.ChatOptions.Clone(), runOptions);
+            return ApplyBackgroundResponsesProperties(this._agentOptions?.ChatOptions.Clone(), runOptions);
        }

        // If both are present, we need to merge them.
@@ -599,26 +583,19 @@ public sealed partial class ChatClientAgent : AIAgent
            }
        }

-        return GetContinuationTokenAndApplyBackgroundResponsesProperties(requestChatOptions, runOptions);
+        return ApplyBackgroundResponsesProperties(requestChatOptions, runOptions);

-        static (ChatOptions?, ChatClientAgentContinuationToken?) GetContinuationTokenAndApplyBackgroundResponsesProperties(ChatOptions? chatOptions, AgentRunOptions? agentRunOptions)
+        static ChatOptions? ApplyBackgroundResponsesProperties(ChatOptions? chatOptions, AgentRunOptions? agentRunOptions)
        {
-            if (agentRunOptions?.AllowBackgroundResponses is not null)
+            // If any of the background response properties are set in the run options, we should apply both to the chat options.
+            if (agentRunOptions?.AllowBackgroundResponses is not null || agentRunOptions?.ContinuationToken is not null)
            {
                chatOptions ??= new ChatOptions();
                chatOptions.AllowBackgroundResponses = agentRunOptions.AllowBackgroundResponses;
+                chatOptions.ContinuationToken = agentRunOptions.ContinuationToken;
            }

-            ChatClientAgentContinuationToken? agentContinuationToken = null;
-
-            if ((agentRunOptions?.ContinuationToken ?? chatOptions?.ContinuationToken) is { } continuationToken)
-            {
-                agentContinuationToken = ChatClientAgentContinuationToken.FromToken(continuationToken);
-                chatOptions ??= new ChatOptions();
-                chatOptions.ContinuationToken = agentContinuationToken!.InnerToken;
-            }
-
-            return (chatOptions, agentContinuationToken);
+            return chatOptions;
        }
    }

@@ -629,22 +606,21 @@ public sealed partial class ChatClientAgent : AIAgent
    /// <param name="inputMessages">The input messages to use.</param>
    /// <param name="runOptions">Optional parameters for agent invocation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A tuple containing the thread, chat options, messages and continuation token.</returns>
+    /// <returns>A tuple containing the thread, chat options, and thread messages.</returns>
    private async Task
        <(
            ChatClientAgentThread AgentThread,
            ChatOptions? ChatOptions,
            List<ChatMessage> InputMessagesForChatClient,
            IList<ChatMessage>? AIContextProviderMessages,
-            IList<ChatMessage>? ChatMessageStoreMessages,
-            ChatClientAgentContinuationToken? ContinuationToken
+            IList<ChatMessage>? ChatMessageStoreMessages
        )> PrepareThreadAndMessagesAsync(
        AgentThread? thread,
        IEnumerable<ChatMessage> inputMessages,
        AgentRunOptions? runOptions,
        CancellationToken cancellationToken)
    {
-        (ChatOptions? chatOptions, ChatClientAgentContinuationToken? continuationToken) = this.CreateConfiguredChatOptions(runOptions);
+        ChatOptions? chatOptions = this.CreateConfiguredChatOptions(runOptions);

        // Supplying a thread for background responses is required to prevent inconsistent experience
        // for callers if they forget to provide the thread for initial or follow-up runs.
@@ -665,6 +641,11 @@ public sealed partial class ChatClientAgent : AIAgent
            throw new InvalidOperationException("Input messages are not allowed when continuing a background response using a continuation token.");
        }

+        if (chatOptions?.ContinuationToken is not null && typedThread.ConversationId is null && typedThread.MessageStore is null)
+        {
+            throw new InvalidOperationException("Continuation tokens are not allowed to be used for initial runs.");
+        }
+
        List<ChatMessage> inputMessagesForChatClient = [];
        IList<ChatMessage>? aiContextProviderMessages = null;
        IList<ChatMessage>? chatMessageStoreMessages = null;
@@ -732,7 +713,7 @@ public sealed partial class ChatClientAgent : AIAgent
            chatOptions.ConversationId = typedThread.ConversationId;
        }

-        return (typedThread, chatOptions, inputMessagesForChatClient, aiContextProviderMessages, chatMessageStoreMessages, continuationToken);
+        return (typedThread, chatOptions, inputMessagesForChatClient, aiContextProviderMessages, chatMessageStoreMessages);
    }

    private void UpdateThreadWithTypeAndConversationId(ChatClientAgentThread thread, string? responseConversationId)
@@ -810,43 +791,26 @@ public sealed partial class ChatClientAgent : AIAgent
        return Task.CompletedTask;
    }

-    private static ChatClientAgentContinuationToken? WrapContinuationToken(ResponseContinuationToken? continuationToken, IEnumerable<ChatMessage>? inputMessages = null, List<ChatResponseUpdate>? responseUpdates = null)
+    private static void ValidateStreamResumptionAllowed(ResponseContinuationToken? continuationToken, ChatClientAgentThread safeThread)
    {
        if (continuationToken is null)
        {
-            return null;
+            return;
        }

-        return new(continuationToken)
+        // Streaming resumption is only supported with chat history managed by the agent service because, currently, there's no good solution
+        // to collect updates received in failed runs and pass them to the last successful run so it can store them to the message store.
+        if (safeThread.ConversationId is null)
        {
-            // Save input messages to the continuation token so they can be added to the thread and
-            // provided to the context provider in the last successful streaming resumption run.
-            // That's necessary for scenarios where initial streaming run is interrupted and streaming is resumed later.
-            InputMessages = inputMessages?.Any() is true ? inputMessages : null,
-
-            // Save all updates received so far to the continuation token so they can be provided to the
-            // message store and context provider in the last successful streaming resumption run.
-            // That's necessary for scenarios where a streaming run is interrupted after some updates were received.
-            ResponseUpdates = responseUpdates?.Count > 0 ? responseUpdates : null
-        };
-    }
-
-    private static IEnumerable<ChatMessage> GetInputMessages(IReadOnlyCollection<ChatMessage> inputMessages, ChatClientAgentContinuationToken? token)
-    {
-        // First, use input messages if provided.
-        if (inputMessages.Count > 0)
-        {
-            return inputMessages;
+            throw new NotSupportedException("Streaming resumption is only supported when chat history is stored and managed by the underlying AI service.");
        }

-        // Fallback to messages saved in the continuation token if available.
-        return token?.InputMessages ?? [];
-    }
-
-    private static List<ChatResponseUpdate> GetResponseUpdates(ChatClientAgentContinuationToken? token)
-    {
-        // Restore any previously received updates from the continuation token.
-        return token?.ResponseUpdates?.ToList() ?? [];
+        // Similarly, streaming resumption is not supported when a context provider is used because, currently, there's no good solution
+        // to collect updates received in failed runs and pass them to the last successful run so it can notify the context provider of the updates.
+        if (safeThread.AIContextProvider is not null)
+        {
+            throw new NotSupportedException("Using context provider with streaming resumption is not supported.");
+        }
    }

    private string GetLoggingAgentName() => this.Name ?? "UnnamedAgent";
@@ -1,170 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using System.IO;
-using System.Linq;
-using System.Text.Json;
-using System.Text.Json.Serialization;
-using Microsoft.Extensions.AI;
-using Microsoft.Shared.Diagnostics;
-
-namespace Microsoft.Agents.AI;
-
-/// <summary>
-/// Represents a continuation token for ChatClientAgent operations.
-/// </summary>
-internal class ChatClientAgentContinuationToken : ResponseContinuationToken
-{
-    private const string TokenTypeName = "chatClientAgentContinuationToken";
-    private const string TypeDiscriminator = "type";
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ChatClientAgentContinuationToken"/> class.
-    /// </summary>
-    /// <param name="innerToken">A continuation token provided by the underlying <see cref="IChatClient"/>.</param>
-    [JsonConstructor]
-    internal ChatClientAgentContinuationToken(ResponseContinuationToken innerToken)
-    {
-        this.InnerToken = innerToken;
-    }
-
-    public override ReadOnlyMemory<byte> ToBytes()
-    {
-        using MemoryStream stream = new();
-        using Utf8JsonWriter writer = new(stream);
-
-        writer.WriteStartObject();
-
-        // This property should be the first one written to identify the type during deserialization.
-        writer.WriteString(TypeDiscriminator, TokenTypeName);
-
-        writer.WriteString("innerToken", JsonSerializer.Serialize(this.InnerToken, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ResponseContinuationToken))));
-
-        if (this.InputMessages?.Any() is true)
-        {
-            writer.WriteString("inputMessages", JsonSerializer.Serialize(this.InputMessages, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(IEnumerable<ChatMessage>))));
-        }
-
-        if (this.ResponseUpdates?.Count > 0)
-        {
-            writer.WriteString("responseUpdates", JsonSerializer.Serialize(this.ResponseUpdates, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(IReadOnlyList<ChatResponseUpdate>))));
-        }
-
-        writer.WriteEndObject();
-
-        writer.Flush();
-
-        return stream.ToArray();
-    }
-
-    /// <summary>
-    /// Create a new instance of <see cref="ChatClientAgentContinuationToken"/> from the provided <paramref name="token"/>.
-    /// </summary>
-    /// <param name="token">The token to create the <see cref="ChatClientAgentContinuationToken"/> from.</param>
-    /// <returns>A <see cref="ChatClientAgentContinuationToken"/> equivalent of the provided <paramref name="token"/>.</returns>
-    internal static ChatClientAgentContinuationToken FromToken(ResponseContinuationToken token)
-    {
-        if (token is ChatClientAgentContinuationToken chatClientContinuationToken)
-        {
-            return chatClientContinuationToken;
-        }
-
-        ReadOnlyMemory<byte> data = token.ToBytes();
-
-        if (data.Length == 0)
-        {
-            Throw.ArgumentException(nameof(token), "Failed to create ChatClientAgentContinuationToken from provided token because it does not contain any data.");
-        }
-
-        Utf8JsonReader reader = new(data.Span);
-
-        // Move to the start object token.
-        _ = reader.Read();
-
-        // Validate that the token is of this type.
-        ValidateTokenType(reader, token);
-
-        ResponseContinuationToken? innerToken = null;
-        IEnumerable<ChatMessage>? inputMessages = null;
-        IReadOnlyList<ChatResponseUpdate>? responseUpdates = null;
-
-        while (reader.Read())
-        {
-            if (reader.TokenType == JsonTokenType.EndObject)
-            {
-                break;
-            }
-
-            if (reader.TokenType != JsonTokenType.PropertyName)
-            {
-                continue;
-            }
-            switch (reader.GetString())
-            {
-                case "innerToken":
-                    _ = reader.Read();
-                    var innerTokenJson = reader.GetString() ?? throw new ArgumentException("No content for innerToken property.", nameof(token));
-                    innerToken = (ResponseContinuationToken?)JsonSerializer.Deserialize(innerTokenJson, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ResponseContinuationToken)));
-                    break;
-                case "inputMessages":
-                    _ = reader.Read();
-                    var innerMessagesJson = reader.GetString() ?? throw new ArgumentException("No content for inputMessages property.", nameof(token));
-                    inputMessages = (IEnumerable<ChatMessage>?)JsonSerializer.Deserialize(innerMessagesJson, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(IEnumerable<ChatMessage>)));
-                    break;
-                case "responseUpdates":
-                    _ = reader.Read();
-                    var responseUpdatesJson = reader.GetString() ?? throw new ArgumentException("No content for responseUpdates property.", nameof(token));
-                    responseUpdates = (IReadOnlyList<ChatResponseUpdate>?)JsonSerializer.Deserialize(responseUpdatesJson, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(IReadOnlyList<ChatResponseUpdate>)));
-                    break;
-                default:
-                    break;
-            }
-        }
-
-        if (innerToken is null)
-        {
-            Throw.ArgumentException(nameof(token), "Failed to create ChatClientAgentContinuationToken from provided token because it does not contain an inner token.");
-        }
-
-        return new ChatClientAgentContinuationToken(innerToken)
-        {
-            InputMessages = inputMessages,
-            ResponseUpdates = responseUpdates
-        };
-    }
-
-    private static void ValidateTokenType(Utf8JsonReader reader, ResponseContinuationToken token)
-    {
-        try
-        {
-            // Move to the first property.
-            _ = reader.Read();
-
-            // If the first property name is not "type", or its value does not match this token type name, then we know its not this token type.
-            if (reader.GetString() != TypeDiscriminator || !reader.Read() || reader.GetString() != TokenTypeName)
-            {
-                Throw.ArgumentException(nameof(token), "Failed to create ChatClientAgentContinuationToken from provided token because it is not of the correct type.");
-            }
-        }
-        catch (JsonException ex)
-        {
-            Throw.ArgumentException(nameof(token), "Failed to create ChatClientAgentContinuationToken from provided token because it could not be parsed.", ex);
-        }
-    }
-
-    /// <summary>
-    /// Gets a continuation token provided by the underlying <see cref="IChatClient"/>.
-    /// </summary>
-    internal ResponseContinuationToken InnerToken { get; }
-
-    /// <summary>
-    /// Gets or sets the input messages used for streaming run.
-    /// </summary>
-    internal IEnumerable<ChatMessage>? InputMessages { get; set; }
-
-    /// <summary>
-    /// Gets or sets the response updates received so far.
-    /// </summary>
-    internal IReadOnlyList<ChatResponseUpdate>? ResponseUpdates { get; set; }
-}
@@ -1,253 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI;
-
-/// <summary>
-/// Provides extension methods for <see cref="ChatClientAgent"/> to enable discoverability of <see cref="ChatClientAgentRunOptions"/>.
-/// </summary>
-public partial class ChatClientAgent
-{
-    /// <summary>
-    /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread.
-    /// </summary>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with any response messages generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="AgentRunResponse"/> with the agent's output.</returns>
-    public Task<AgentRunResponse> RunAsync(
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync(thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent with a text message from the user.
-    /// </summary>
-    /// <param name="message">The user message to send to the agent.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input message and any response messages generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="AgentRunResponse"/> with the agent's output.</returns>
-    public Task<AgentRunResponse> RunAsync(
-        string message,
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync(message, thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent with a single chat message.
-    /// </summary>
-    /// <param name="message">The chat message to send to the agent.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input message and any response messages generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="AgentRunResponse"/> with the agent's output.</returns>
-    public Task<AgentRunResponse> RunAsync(
-        ChatMessage message,
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync(message, thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent with a collection of chat messages.
-    /// </summary>
-    /// <param name="messages">The collection of messages to send to the agent for processing.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input messages and any response messages generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="AgentRunResponse"/> with the agent's output.</returns>
-    public Task<AgentRunResponse> RunAsync(
-        IEnumerable<ChatMessage> messages,
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync(messages, thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent in streaming mode without providing new input messages, relying on existing context and instructions.
-    /// </summary>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with any response messages generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>An asynchronous enumerable of <see cref="AgentRunResponseUpdate"/> instances representing the streaming response.</returns>
-    public IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunStreamingAsync(thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent in streaming mode with a text message from the user.
-    /// </summary>
-    /// <param name="message">The user message to send to the agent.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input message and any response messages generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>An asynchronous enumerable of <see cref="AgentRunResponseUpdate"/> instances representing the streaming response.</returns>
-    public IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
-        string message,
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunStreamingAsync(message, thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent in streaming mode with a single chat message.
-    /// </summary>
-    /// <param name="message">The chat message to send to the agent.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input message and any response messages generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>An asynchronous enumerable of <see cref="AgentRunResponseUpdate"/> instances representing the streaming response.</returns>
-    public IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
-        ChatMessage message,
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunStreamingAsync(message, thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent in streaming mode with a collection of chat messages.
-    /// </summary>
-    /// <param name="messages">The collection of messages to send to the agent for processing.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input messages and any response updates generated during invocation.
-    /// </param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>An asynchronous enumerable of <see cref="AgentRunResponseUpdate"/> instances representing the streaming response.</returns>
-    public IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
-        IEnumerable<ChatMessage> messages,
-        AgentThread? thread,
-        ChatClientAgentRunOptions? options,
-        CancellationToken cancellationToken = default) =>
-        this.RunStreamingAsync(messages, thread, (AgentRunOptions?)options, cancellationToken);
-
-    /// <summary>
-    /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread, and requesting a response of the specified type <typeparamref name="T"/>.
-    /// </summary>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with any response messages generated during invocation.
-    /// </param>
-    /// <param name="serializerOptions">The JSON serialization options to use.</param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="useJsonSchemaResponseFormat">
-    /// <see langword="true" /> to set a JSON schema on the <see cref="ChatResponseFormat"/>; otherwise, <see langword="false" />. The default is <see langword="true" />.
-    /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it.
-    /// </param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="ChatClientAgentRunResponse{T}"/> with the agent's output.</returns>
-    public Task<ChatClientAgentRunResponse<T>> RunAsync<T>(
-        AgentThread? thread,
-        JsonSerializerOptions? serializerOptions,
-        ChatClientAgentRunOptions? options,
-        bool? useJsonSchemaResponseFormat = null,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync<T>(thread, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent with a text message from the user, requesting a response of the specified type <typeparamref name="T"/>.
-    /// </summary>
-    /// <param name="message">The user message to send to the agent.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input message and any response messages generated during invocation.
-    /// </param>
-    /// <param name="serializerOptions">The JSON serialization options to use.</param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="useJsonSchemaResponseFormat">
-    /// <see langword="true" /> to set a JSON schema on the <see cref="ChatResponseFormat"/>; otherwise, <see langword="false" />. The default is <see langword="true" />.
-    /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it.
-    /// </param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="ChatClientAgentRunResponse{T}"/> with the agent's output.</returns>
-    public Task<ChatClientAgentRunResponse<T>> RunAsync<T>(
-        string message,
-        AgentThread? thread,
-        JsonSerializerOptions? serializerOptions,
-        ChatClientAgentRunOptions? options,
-        bool? useJsonSchemaResponseFormat = null,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync<T>(message, thread, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent with a single chat message, requesting a response of the specified type <typeparamref name="T"/>.
-    /// </summary>
-    /// <param name="message">The chat message to send to the agent.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input message and any response messages generated during invocation.
-    /// </param>
-    /// <param name="serializerOptions">The JSON serialization options to use.</param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="useJsonSchemaResponseFormat">
-    /// <see langword="true" /> to set a JSON schema on the <see cref="ChatResponseFormat"/>; otherwise, <see langword="false" />. The default is <see langword="true" />.
-    /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it.
-    /// </param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="ChatClientAgentRunResponse{T}"/> with the agent's output.</returns>
-    public Task<ChatClientAgentRunResponse<T>> RunAsync<T>(
-        ChatMessage message,
-        AgentThread? thread,
-        JsonSerializerOptions? serializerOptions,
-        ChatClientAgentRunOptions? options,
-        bool? useJsonSchemaResponseFormat = null,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync<T>(message, thread, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken);
-
-    /// <summary>
-    /// Runs the agent with a collection of chat messages, requesting a response of the specified type <typeparamref name="T"/>.
-    /// </summary>
-    /// <param name="messages">The collection of messages to send to the agent for processing.</param>
-    /// <param name="thread">
-    /// The conversation thread to use for this invocation. If <see langword="null"/>, a new thread will be created.
-    /// The thread will be updated with the input messages and any response messages generated during invocation.
-    /// </param>
-    /// <param name="serializerOptions">The JSON serialization options to use.</param>
-    /// <param name="options">Configuration parameters for controlling the agent's invocation behavior.</param>
-    /// <param name="useJsonSchemaResponseFormat">
-    /// <see langword="true" /> to set a JSON schema on the <see cref="ChatResponseFormat"/>; otherwise, <see langword="false" />. The default is <see langword="true" />.
-    /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it.
-    /// </param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="ChatClientAgentRunResponse{T}"/> with the agent's output.</returns>
-    public Task<ChatClientAgentRunResponse<T>> RunAsync<T>(
-        IEnumerable<ChatMessage> messages,
-        AgentThread? thread,
-        JsonSerializerOptions? serializerOptions,
-        ChatClientAgentRunOptions? options,
-        bool? useJsonSchemaResponseFormat = null,
-        CancellationToken cancellationToken = default) =>
-        this.RunAsync<T>(messages, thread, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken);
-}
@@ -80,7 +80,7 @@ public sealed class ChatClientAgentOptions
    /// <summary>
    /// Context object passed to the <see cref="AIContextProviderFactory"/> to create a new instance of <see cref="AIContextProvider"/>.
    /// </summary>
-    public sealed class AIContextProviderFactoryContext
+    public class AIContextProviderFactoryContext
    {
        /// <summary>
        /// Gets or sets the serialized state of the <see cref="AIContextProvider"/>, if any.
@@ -97,7 +97,7 @@ public sealed class ChatClientAgentOptions
    /// <summary>
    /// Context object passed to the <see cref="ChatMessageStoreFactory"/> to create a new instance of <see cref="ChatMessageStore"/>.
    /// </summary>
-    public sealed class ChatMessageStoreFactoryContext
+    public class ChatMessageStoreFactoryContext
    {
        /// <summary>
        /// Gets or sets the serialized state of the chat message store, if any.
@@ -40,6 +40,7 @@ public sealed class ChatClientAgentRunResponse<T> : AgentRunResponse<T>
    /// </summary>
    /// <remarks>
    /// If the response did not contain JSON, or if deserialization fails, this property will throw.
+    /// To avoid exceptions, use <see cref="AgentRunResponse.TryDeserialize{T}"/> instead.
    /// </remarks>
    public override T Result => this._response.Result;
 }
@@ -154,10 +154,7 @@ public sealed partial class ChatClientAgent

        static ChatClientAgentRunResponse<T> CreateResponse(ChatResponse<T> chatResponse)
        {
-            return new ChatClientAgentRunResponse<T>(chatResponse)
-            {
-                ContinuationToken = WrapContinuationToken(chatResponse.ContinuationToken)
-            };
+            return new ChatClientAgentRunResponse<T>(chatResponse);
        }

        return this.RunCoreAsync(GetResponseAsync, CreateResponse, messages, thread, options, cancellationToken);
@@ -832,174 +832,6 @@ public sealed class A2AAgentTests : IDisposable
        });
    }

-    [Fact]
-    public async Task RunAsync_WithAgentMessageResponseMetadata_ReturnsMetadataAsAdditionalPropertiesAsync()
-    {
-        // Arrange
-        this._handler.ResponseToReturn = new AgentMessage
-        {
-            MessageId = "response-123",
-            Role = MessageRole.Agent,
-            Parts = [new TextPart { Text = "Response with metadata" }],
-            Metadata = new Dictionary<string, JsonElement>
-            {
-                { "responseKey1", JsonSerializer.SerializeToElement("responseValue1") },
-                { "responseCount", JsonSerializer.SerializeToElement(99) }
-            }
-        };
-
-        var inputMessages = new List<ChatMessage>
-        {
-            new(ChatRole.User, "Test message")
-        };
-
-        // Act
-        var result = await this._agent.RunAsync(inputMessages);
-
-        // Assert
-        Assert.NotNull(result.AdditionalProperties);
-        Assert.NotNull(result.AdditionalProperties["responseKey1"]);
-        Assert.Equal("responseValue1", ((JsonElement)result.AdditionalProperties["responseKey1"]!).GetString());
-        Assert.NotNull(result.AdditionalProperties["responseCount"]);
-        Assert.Equal(99, ((JsonElement)result.AdditionalProperties["responseCount"]!).GetInt32());
-    }
-
-    [Fact]
-    public async Task RunAsync_WithAdditionalProperties_PropagatesThemAsMetadataToMessageSendParamsAsync()
-    {
-        // Arrange
-        this._handler.ResponseToReturn = new AgentMessage
-        {
-            MessageId = "response-123",
-            Role = MessageRole.Agent,
-            Parts = [new TextPart { Text = "Response" }]
-        };
-
-        var inputMessages = new List<ChatMessage>
-        {
-            new(ChatRole.User, "Test message")
-        };
-
-        var options = new AgentRunOptions
-        {
-            AdditionalProperties = new()
-            {
-                { "key1", "value1" },
-                { "key2", 42 },
-                { "key3", true }
-            }
-        };
-
-        // Act
-        await this._agent.RunAsync(inputMessages, null, options);
-
-        // Assert
-        Assert.NotNull(this._handler.CapturedMessageSendParams);
-        Assert.NotNull(this._handler.CapturedMessageSendParams.Metadata);
-        Assert.Equal("value1", this._handler.CapturedMessageSendParams.Metadata["key1"].GetString());
-        Assert.Equal(42, this._handler.CapturedMessageSendParams.Metadata["key2"].GetInt32());
-        Assert.True(this._handler.CapturedMessageSendParams.Metadata["key3"].GetBoolean());
-    }
-
-    [Fact]
-    public async Task RunAsync_WithNullAdditionalProperties_DoesNotSetMetadataAsync()
-    {
-        // Arrange
-        this._handler.ResponseToReturn = new AgentMessage
-        {
-            MessageId = "response-123",
-            Role = MessageRole.Agent,
-            Parts = [new TextPart { Text = "Response" }]
-        };
-
-        var inputMessages = new List<ChatMessage>
-        {
-            new(ChatRole.User, "Test message")
-        };
-
-        var options = new AgentRunOptions
-        {
-            AdditionalProperties = null
-        };
-
-        // Act
-        await this._agent.RunAsync(inputMessages, null, options);
-
-        // Assert
-        Assert.NotNull(this._handler.CapturedMessageSendParams);
-        Assert.Null(this._handler.CapturedMessageSendParams.Metadata);
-    }
-
-    [Fact]
-    public async Task RunStreamingAsync_WithAdditionalProperties_PropagatesThemAsMetadataToMessageSendParamsAsync()
-    {
-        // Arrange
-        this._handler.StreamingResponseToReturn = new AgentMessage
-        {
-            MessageId = "stream-123",
-            Role = MessageRole.Agent,
-            Parts = [new TextPart { Text = "Streaming response" }]
-        };
-
-        var inputMessages = new List<ChatMessage>
-        {
-            new(ChatRole.User, "Test streaming message")
-        };
-
-        var options = new AgentRunOptions
-        {
-            AdditionalProperties = new()
-            {
-                { "streamKey1", "streamValue1" },
-                { "streamKey2", 100 },
-                { "streamKey3", false }
-            }
-        };
-
-        // Act
-        await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, null, options))
-        {
-        }
-
-        // Assert
-        Assert.NotNull(this._handler.CapturedMessageSendParams);
-        Assert.NotNull(this._handler.CapturedMessageSendParams.Metadata);
-        Assert.Equal("streamValue1", this._handler.CapturedMessageSendParams.Metadata["streamKey1"].GetString());
-        Assert.Equal(100, this._handler.CapturedMessageSendParams.Metadata["streamKey2"].GetInt32());
-        Assert.False(this._handler.CapturedMessageSendParams.Metadata["streamKey3"].GetBoolean());
-    }
-
-    [Fact]
-    public async Task RunStreamingAsync_WithNullAdditionalProperties_DoesNotSetMetadataAsync()
-    {
-        // Arrange
-        this._handler.StreamingResponseToReturn = new AgentMessage
-        {
-            MessageId = "stream-123",
-            Role = MessageRole.Agent,
-            Parts = [new TextPart { Text = "Streaming response" }]
-        };
-
-        var inputMessages = new List<ChatMessage>
-        {
-            new(ChatRole.User, "Test streaming message")
-        };
-
-        var options = new AgentRunOptions
-        {
-            AdditionalProperties = null
-        };
-
-        // Act
-        await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, null, options))
-        {
-        }
-
-        // Assert
-        Assert.NotNull(this._handler.CapturedMessageSendParams);
-        Assert.Null(this._handler.CapturedMessageSendParams.Metadata);
-    }
-
    [Fact]
    public async Task RunAsync_WithInvalidThreadType_ThrowsInvalidOperationExceptionAsync()
    {
@@ -1,186 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI.A2A.UnitTests;
-
-/// <summary>
-/// Unit tests for the <see cref="AdditionalPropertiesDictionaryExtensions"/> class.
-/// </summary>
-public sealed class AdditionalPropertiesDictionaryExtensionsTests
-{
-    [Fact]
-    public void ToA2AMetadata_WithNullAdditionalProperties_ReturnsNull()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary? additionalProperties = null;
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.Null(result);
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithEmptyAdditionalProperties_ReturnsNull()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = [];
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.Null(result);
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithStringValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "stringKey", "stringValue" }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("stringKey"));
-        Assert.Equal("stringValue", result["stringKey"].GetString());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithNumericValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "numberKey", 42 }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("numberKey"));
-        Assert.Equal(42, result["numberKey"].GetInt32());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithBooleanValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "booleanKey", true }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("booleanKey"));
-        Assert.True(result["booleanKey"].GetBoolean());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithMultipleProperties_ReturnsMetadataWithAllProperties()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "stringKey", "stringValue" },
-            { "numberKey", 42 },
-            { "booleanKey", true }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Equal(3, result.Count);
-
-        Assert.True(result.ContainsKey("stringKey"));
-        Assert.Equal("stringValue", result["stringKey"].GetString());
-
-        Assert.True(result.ContainsKey("numberKey"));
-        Assert.Equal(42, result["numberKey"].GetInt32());
-
-        Assert.True(result.ContainsKey("booleanKey"));
-        Assert.True(result["booleanKey"].GetBoolean());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithArrayValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        int[] arrayValue = [1, 2, 3];
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "arrayKey", arrayValue }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("arrayKey"));
-        Assert.Equal(JsonValueKind.Array, result["arrayKey"].ValueKind);
-        Assert.Equal(3, result["arrayKey"].GetArrayLength());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithNullValue_ReturnsMetadataWithNullJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "nullKey", null! }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("nullKey"));
-        Assert.Equal(JsonValueKind.Null, result["nullKey"].ValueKind);
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithJsonElementValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        JsonElement jsonElement = JsonSerializer.SerializeToElement(new { name = "test", value = 123 });
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "jsonElementKey", jsonElement }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("jsonElementKey"));
-        Assert.Equal(JsonValueKind.Object, result["jsonElementKey"].ValueKind);
-        Assert.Equal("test", result["jsonElementKey"].GetProperty("name").GetString());
-        Assert.Equal(123, result["jsonElementKey"].GetProperty("value").GetInt32());
-    }
-}
@@ -57,7 +57,7 @@ public class AgentRunResponseTests
            RawRepresentation = new object(),
            ResponseId = "responseId",
            Usage = new UsageDetails(),
-            ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 })
+            ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }),
        };

        AgentRunResponse response = new(chatResponse);
@@ -214,12 +214,6 @@ public class AgentRunResponseTests
        Assert.Equal(100, usageContent.Details.TotalTokenCount);
    }

-#if NETFRAMEWORK
-    /// <summary>
-    /// Since Json Serialization using reflection is disabled in .net core builds, and we are using a custom type here that wouldn't
-    /// be registered with the default source generated serializer, this test will only pass in .net framework builds where reflection-based
-    /// serialization is available.
-    /// </summary>
    [Fact]
    public void ParseAsStructuredOutputSuccess()
    {
@@ -227,24 +221,6 @@ public class AgentRunResponseTests
        var expectedResult = new Animal { Id = 1, FullName = "Tigger", Species = Species.Tiger };
        var response = new AgentRunResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal)));

-        // Act.
-        var animal = response.Deserialize<Animal>();
-
-        // Assert.
-        Assert.NotNull(animal);
-        Assert.Equal(expectedResult.Id, animal.Id);
-        Assert.Equal(expectedResult.FullName, animal.FullName);
-        Assert.Equal(expectedResult.Species, animal.Species);
-    }
-#endif
-
-    [Fact]
-    public void ParseAsStructuredOutputWithJSOSuccess()
-    {
-        // Arrange.
-        var expectedResult = new Animal { Id = 1, FullName = "Tigger", Species = Species.Tiger };
-        var response = new AgentRunResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal)));
-
        // Act.
        var animal = response.Deserialize<Animal>(TestJsonSerializerContext.Default.Options);

@@ -286,12 +262,6 @@ public class AgentRunResponseTests
        Assert.Throws<JsonException>(() => response.Deserialize<Animal>(TestJsonSerializerContext.Default.Options));
    }

-#if NETFRAMEWORK
-    /// <summary>
-    /// Since Json Serialization using reflection is disabled in .net core builds, and we are using a custom type here that wouldn't
-    /// be registered with the default source generated serializer, this test will only pass in .net framework builds where reflection-based
-    /// serialization is available.
-    /// </summary>
    [Fact]
    public void TryParseAsStructuredOutputSuccess()
    {
@@ -299,24 +269,6 @@ public class AgentRunResponseTests
        var expectedResult = new Animal { Id = 1, FullName = "Tigger", Species = Species.Tiger };
        var response = new AgentRunResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal)));

-        // Act.
-        response.TryDeserialize(out Animal? animal);
-
-        // Assert.
-        Assert.NotNull(animal);
-        Assert.Equal(expectedResult.Id, animal.Id);
-        Assert.Equal(expectedResult.FullName, animal.FullName);
-        Assert.Equal(expectedResult.Species, animal.Species);
-    }
-#endif
-
-    [Fact]
-    public void TryParseAsStructuredOutputWithJSOSuccess()
-    {
-        // Arrange.
-        var expectedResult = new Animal { Id = 1, FullName = "Tigger", Species = Species.Tiger };
-        var response = new AgentRunResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal)));
-
        // Act.
        response.TryDeserialize(TestJsonSerializerContext.Default.Options, out Animal? animal);

@@ -1,218 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
-using A2A;
-using Microsoft.Extensions.AI;
-using Moq;
-using Moq.Protected;
-
-namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests;
-
-/// <summary>
-/// Unit tests for the <see cref="AIAgentExtensions"/> class.
-/// </summary>
-public sealed class AIAgentExtensionsTests
-{
-    /// <summary>
-    /// Verifies that when messageSendParams.Metadata is null, the options passed to RunAsync are null.
-    /// </summary>
-    [Fact]
-    public async Task MapA2A_WhenMetadataIsNull_PassesNullOptionsToRunAsync()
-    {
-        // Arrange
-        AgentRunOptions? capturedOptions = null;
-        ITaskManager taskManager = CreateAgentMock(options => capturedOptions = options).Object.MapA2A();
-
-        // Act
-        await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
-        {
-            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] },
-            Metadata = null
-        });
-
-        // Assert
-        Assert.Null(capturedOptions);
-    }
-
-    /// <summary>
-    /// Verifies that when messageSendParams.Metadata has values, the options.AdditionalProperties contains the converted values.
-    /// </summary>
-    [Fact]
-    public async Task MapA2A_WhenMetadataHasValues_PassesOptionsWithAdditionalPropertiesToRunAsync()
-    {
-        // Arrange
-        AgentRunOptions? capturedOptions = null;
-        ITaskManager taskManager = CreateAgentMock(options => capturedOptions = options).Object.MapA2A();
-
-        // Act
-        await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
-        {
-            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] },
-            Metadata = new Dictionary<string, JsonElement>
-            {
-                ["key1"] = JsonSerializer.SerializeToElement("value1"),
-                ["key2"] = JsonSerializer.SerializeToElement(42)
-            }
-        });
-
-        // Assert
-        Assert.NotNull(capturedOptions);
-        Assert.NotNull(capturedOptions.AdditionalProperties);
-        Assert.Equal(2, capturedOptions.AdditionalProperties.Count);
-        Assert.True(capturedOptions.AdditionalProperties.ContainsKey("key1"));
-        Assert.True(capturedOptions.AdditionalProperties.ContainsKey("key2"));
-    }
-
-    /// <summary>
-    /// Verifies that when messageSendParams.Metadata is an empty dictionary, the options passed to RunAsync is null
-    /// because the ToAdditionalProperties extension method returns null for empty dictionaries.
-    /// </summary>
-    [Fact]
-    public async Task MapA2A_WhenMetadataIsEmptyDictionary_PassesNullOptionsToRunAsync()
-    {
-        // Arrange
-        AgentRunOptions? capturedOptions = null;
-        ITaskManager taskManager = CreateAgentMock(options => capturedOptions = options).Object.MapA2A();
-
-        // Act
-        await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
-        {
-            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] },
-            Metadata = []
-        });
-
-        // Assert
-        Assert.Null(capturedOptions);
-    }
-
-    /// <summary>
-    /// Verifies that when the agent response has AdditionalProperties, the returned AgentMessage.Metadata contains the converted values.
-    /// </summary>
-    [Fact]
-    public async Task MapA2A_WhenResponseHasAdditionalProperties_ReturnsAgentMessageWithMetadataAsync()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProps = new()
-        {
-            ["responseKey1"] = "responseValue1",
-            ["responseKey2"] = 123
-        };
-        AgentRunResponse response = new([new ChatMessage(ChatRole.Assistant, "Test response")])
-        {
-            AdditionalProperties = additionalProps
-        };
-        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
-
-        // Act
-        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
-        {
-            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
-        });
-
-        // Assert
-        AgentMessage agentMessage = Assert.IsType<AgentMessage>(a2aResponse);
-        Assert.NotNull(agentMessage.Metadata);
-        Assert.Equal(2, agentMessage.Metadata.Count);
-        Assert.True(agentMessage.Metadata.ContainsKey("responseKey1"));
-        Assert.True(agentMessage.Metadata.ContainsKey("responseKey2"));
-        Assert.Equal("responseValue1", agentMessage.Metadata["responseKey1"].GetString());
-        Assert.Equal(123, agentMessage.Metadata["responseKey2"].GetInt32());
-    }
-
-    /// <summary>
-    /// Verifies that when the agent response has null AdditionalProperties, the returned AgentMessage.Metadata is null.
-    /// </summary>
-    [Fact]
-    public async Task MapA2A_WhenResponseHasNullAdditionalProperties_ReturnsAgentMessageWithNullMetadataAsync()
-    {
-        // Arrange
-        AgentRunResponse response = new([new ChatMessage(ChatRole.Assistant, "Test response")])
-        {
-            AdditionalProperties = null
-        };
-        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
-
-        // Act
-        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
-        {
-            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
-        });
-
-        // Assert
-        AgentMessage agentMessage = Assert.IsType<AgentMessage>(a2aResponse);
-        Assert.Null(agentMessage.Metadata);
-    }
-
-    /// <summary>
-    /// Verifies that when the agent response has empty AdditionalProperties, the returned AgentMessage.Metadata is null.
-    /// </summary>
-    [Fact]
-    public async Task MapA2A_WhenResponseHasEmptyAdditionalProperties_ReturnsAgentMessageWithNullMetadataAsync()
-    {
-        // Arrange
-        AgentRunResponse response = new([new ChatMessage(ChatRole.Assistant, "Test response")])
-        {
-            AdditionalProperties = []
-        };
-        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
-
-        // Act
-        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
-        {
-            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
-        });
-
-        // Assert
-        AgentMessage agentMessage = Assert.IsType<AgentMessage>(a2aResponse);
-        Assert.Null(agentMessage.Metadata);
-    }
-
-    private static Mock<AIAgent> CreateAgentMock(Action<AgentRunOptions?> optionsCallback)
-    {
-        Mock<AIAgent> agentMock = new() { CallBase = true };
-        agentMock.SetupGet(x => x.Name).Returns("TestAgent");
-        agentMock.Setup(x => x.GetNewThread()).Returns(new TestAgentThread());
-        agentMock
-            .Protected()
-            .Setup<Task<AgentRunResponse>>("RunCoreAsync",
-                ItExpr.IsAny<IEnumerable<ChatMessage>>(),
-                ItExpr.IsAny<AgentThread?>(),
-                ItExpr.IsAny<AgentRunOptions?>(),
-                ItExpr.IsAny<CancellationToken>())
-            .Callback<IEnumerable<ChatMessage>, AgentThread?, AgentRunOptions?, CancellationToken>(
-                (_, _, options, _) => optionsCallback(options))
-            .ReturnsAsync(new AgentRunResponse([new ChatMessage(ChatRole.Assistant, "Test response")]));
-
-        return agentMock;
-    }
-
-    private static Mock<AIAgent> CreateAgentMockWithResponse(AgentRunResponse response)
-    {
-        Mock<AIAgent> agentMock = new() { CallBase = true };
-        agentMock.SetupGet(x => x.Name).Returns("TestAgent");
-        agentMock.Setup(x => x.GetNewThread()).Returns(new TestAgentThread());
-        agentMock
-            .Protected()
-            .Setup<Task<AgentRunResponse>>("RunCoreAsync",
-                ItExpr.IsAny<IEnumerable<ChatMessage>>(),
-                ItExpr.IsAny<AgentThread?>(),
-                ItExpr.IsAny<AgentRunOptions?>(),
-                ItExpr.IsAny<CancellationToken>())
-            .ReturnsAsync(response);
-
-        return agentMock;
-    }
-
-    private static async Task<A2AResponse> InvokeOnMessageReceivedAsync(ITaskManager taskManager, MessageSendParams messageSendParams)
-    {
-        Func<MessageSendParams, CancellationToken, Task<A2AResponse>>? handler = taskManager.OnMessageReceived;
-        Assert.NotNull(handler);
-        return await handler.Invoke(messageSendParams, CancellationToken.None);
-    }
-
-    private sealed class TestAgentThread : AgentThread;
-}
@@ -1,187 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using Microsoft.Agents.AI.Hosting.A2A.Converters;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests.Converters;
-
-/// <summary>
-/// Unit tests for the <see cref="AdditionalPropertiesDictionaryExtensions"/> class.
-/// </summary>
-public sealed class AdditionalPropertiesDictionaryExtensionsTests
-{
-    [Fact]
-    public void ToA2AMetadata_WithNullAdditionalProperties_ReturnsNull()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary? additionalProperties = null;
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.Null(result);
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithEmptyAdditionalProperties_ReturnsNull()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = [];
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.Null(result);
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithStringValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "stringKey", "stringValue" }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("stringKey"));
-        Assert.Equal("stringValue", result["stringKey"].GetString());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithNumericValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "numberKey", 42 }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("numberKey"));
-        Assert.Equal(42, result["numberKey"].GetInt32());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithBooleanValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "booleanKey", true }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("booleanKey"));
-        Assert.True(result["booleanKey"].GetBoolean());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithMultipleProperties_ReturnsMetadataWithAllProperties()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "stringKey", "stringValue" },
-            { "numberKey", 42 },
-            { "booleanKey", true }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Equal(3, result.Count);
-
-        Assert.True(result.ContainsKey("stringKey"));
-        Assert.Equal("stringValue", result["stringKey"].GetString());
-
-        Assert.True(result.ContainsKey("numberKey"));
-        Assert.Equal(42, result["numberKey"].GetInt32());
-
-        Assert.True(result.ContainsKey("booleanKey"));
-        Assert.True(result["booleanKey"].GetBoolean());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithArrayValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        int[] arrayValue = [1, 2, 3];
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "arrayKey", arrayValue }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("arrayKey"));
-        Assert.Equal(JsonValueKind.Array, result["arrayKey"].ValueKind);
-        Assert.Equal(3, result["arrayKey"].GetArrayLength());
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithNullValue_ReturnsMetadataWithNullJsonElement()
-    {
-        // Arrange
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "nullKey", null! }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("nullKey"));
-        Assert.Equal(JsonValueKind.Null, result["nullKey"].ValueKind);
-    }
-
-    [Fact]
-    public void ToA2AMetadata_WithJsonElementValue_ReturnsMetadataWithJsonElement()
-    {
-        // Arrange
-        JsonElement jsonElement = JsonSerializer.SerializeToElement(new { name = "test", value = 123 });
-        AdditionalPropertiesDictionary additionalProperties = new()
-        {
-            { "jsonElementKey", jsonElement }
-        };
-
-        // Act
-        Dictionary<string, JsonElement>? result = additionalProperties.ToA2AMetadata();
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result);
-        Assert.True(result.ContainsKey("jsonElementKey"));
-        Assert.Equal(JsonValueKind.Object, result["jsonElementKey"].ValueKind);
-        Assert.Equal("test", result["jsonElementKey"].GetProperty("name").GetString());
-        Assert.Equal(123, result["jsonElementKey"].GetProperty("value").GetInt32());
-    }
-}
@@ -1,128 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Linq;
-using System.Text.Json;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI.UnitTests.ChatClient;
-
-public class ChatClientAgentContinuationTokenTests
-{
-    [Fact]
-    public void ToBytes_Roundtrip()
-    {
-        // Arrange
-        ResponseContinuationToken originalToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3, 4, 5 });
-
-        ChatClientAgentContinuationToken chatClientToken = new(originalToken)
-        {
-            InputMessages =
-            [
-                new ChatMessage(ChatRole.User, "Hello!"),
-                new ChatMessage(ChatRole.User, "How are you?")
-            ],
-            ResponseUpdates =
-            [
-                new ChatResponseUpdate(ChatRole.Assistant, "I'm fine, thank you."),
-                new ChatResponseUpdate(ChatRole.Assistant, "How can I assist you today?")
-            ]
-        };
-
-        // Act
-        ReadOnlyMemory<byte> bytes = chatClientToken.ToBytes();
-
-        ChatClientAgentContinuationToken tokenFromBytes = ChatClientAgentContinuationToken.FromToken(ResponseContinuationToken.FromBytes(bytes));
-
-        // Assert
-        Assert.NotNull(tokenFromBytes);
-        Assert.Equal(chatClientToken.ToBytes().ToArray(), tokenFromBytes.ToBytes().ToArray());
-
-        // Verify InnerToken
-        Assert.Equal(chatClientToken.InnerToken.ToBytes().ToArray(), tokenFromBytes.InnerToken.ToBytes().ToArray());
-
-        // Verify InputMessages
-        Assert.NotNull(tokenFromBytes.InputMessages);
-        Assert.Equal(chatClientToken.InputMessages.Count(), tokenFromBytes.InputMessages.Count());
-        for (int i = 0; i < chatClientToken.InputMessages.Count(); i++)
-        {
-            Assert.Equal(chatClientToken.InputMessages.ElementAt(i).Role, tokenFromBytes.InputMessages.ElementAt(i).Role);
-            Assert.Equal(chatClientToken.InputMessages.ElementAt(i).Text, tokenFromBytes.InputMessages.ElementAt(i).Text);
-        }
-
-        // Verify ResponseUpdates
-        Assert.NotNull(tokenFromBytes.ResponseUpdates);
-        Assert.Equal(chatClientToken.ResponseUpdates.Count, tokenFromBytes.ResponseUpdates.Count);
-        for (int i = 0; i < chatClientToken.ResponseUpdates.Count; i++)
-        {
-            Assert.Equal(chatClientToken.ResponseUpdates.ElementAt(i).Role, tokenFromBytes.ResponseUpdates.ElementAt(i).Role);
-            Assert.Equal(chatClientToken.ResponseUpdates.ElementAt(i).Text, tokenFromBytes.ResponseUpdates.ElementAt(i).Text);
-        }
-    }
-
-    [Fact]
-    public void Serialization_Roundtrip()
-    {
-        // Arrange
-        ResponseContinuationToken originalToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3, 4, 5 });
-
-        ChatClientAgentContinuationToken chatClientToken = new(originalToken)
-        {
-            InputMessages =
-            [
-                new ChatMessage(ChatRole.User, "Hello!"),
-                new ChatMessage(ChatRole.User, "How are you?")
-            ],
-            ResponseUpdates =
-            [
-                new ChatResponseUpdate(ChatRole.Assistant, "I'm fine, thank you."),
-                new ChatResponseUpdate(ChatRole.Assistant, "How can I assist you today?")
-            ]
-        };
-
-        // Act
-        string json = JsonSerializer.Serialize(chatClientToken, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ResponseContinuationToken)));
-
-        ResponseContinuationToken? deserializedToken = (ResponseContinuationToken?)JsonSerializer.Deserialize(json, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ResponseContinuationToken)));
-
-        ChatClientAgentContinuationToken deserializedChatClientToken = ChatClientAgentContinuationToken.FromToken(deserializedToken!);
-
-        // Assert
-        Assert.NotNull(deserializedChatClientToken);
-        Assert.Equal(chatClientToken.ToBytes().ToArray(), deserializedChatClientToken.ToBytes().ToArray());
-
-        // Verify InnerToken
-        Assert.Equal(chatClientToken.InnerToken.ToBytes().ToArray(), deserializedChatClientToken.InnerToken.ToBytes().ToArray());
-
-        // Verify InputMessages
-        Assert.NotNull(deserializedChatClientToken.InputMessages);
-        Assert.Equal(chatClientToken.InputMessages.Count(), deserializedChatClientToken.InputMessages.Count());
-        for (int i = 0; i < chatClientToken.InputMessages.Count(); i++)
-        {
-            Assert.Equal(chatClientToken.InputMessages.ElementAt(i).Role, deserializedChatClientToken.InputMessages.ElementAt(i).Role);
-            Assert.Equal(chatClientToken.InputMessages.ElementAt(i).Text, deserializedChatClientToken.InputMessages.ElementAt(i).Text);
-        }
-
-        // Verify ResponseUpdates
-        Assert.NotNull(deserializedChatClientToken.ResponseUpdates);
-        Assert.Equal(chatClientToken.ResponseUpdates.Count, deserializedChatClientToken.ResponseUpdates.Count);
-        for (int i = 0; i < chatClientToken.ResponseUpdates.Count; i++)
-        {
-            Assert.Equal(chatClientToken.ResponseUpdates.ElementAt(i).Role, deserializedChatClientToken.ResponseUpdates.ElementAt(i).Role);
-            Assert.Equal(chatClientToken.ResponseUpdates.ElementAt(i).Text, deserializedChatClientToken.ResponseUpdates.ElementAt(i).Text);
-        }
-    }
-
-    [Fact]
-    public void FromToken_WithChatClientAgentContinuationToken_ReturnsSameInstance()
-    {
-        // Arrange
-        ChatClientAgentContinuationToken originalToken = new(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3, 4, 5 }));
-
-        // Act
-        ChatClientAgentContinuationToken fromToken = ChatClientAgentContinuationToken.FromToken(originalToken);
-
-        // Assert
-        Assert.Same(originalToken, fromToken);
-    }
-}
@@ -18,10 +18,10 @@ public class ChatClientAgent_BackgroundResponsesTests
    [Theory]
    [InlineData(true)]
    [InlineData(false)]
-    public async Task RunAsync_PropagatesBackgroundResponsesPropertiesToChatClientAsync(bool providePropsViaChatOptions)
+    public async Task RunAsyncPropagatesBackgroundResponsesPropertiesToChatClientAsync(bool providePropsViaChatOptions)
    {
        // Arrange
-        var continuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }));
+        var continuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
        ChatOptions? capturedChatOptions = null;
        Mock<IChatClient> mockChatClient = new();
        mockChatClient
@@ -63,15 +63,15 @@ public class ChatClientAgent_BackgroundResponsesTests
        // Assert
        Assert.NotNull(capturedChatOptions);
        Assert.True(capturedChatOptions.AllowBackgroundResponses);
-        Assert.Same(continuationToken.InnerToken, capturedChatOptions.ContinuationToken);
+        Assert.Same(continuationToken, capturedChatOptions.ContinuationToken);
    }

    [Fact]
-    public async Task RunAsync_WhenPropertiesSetInBothLocations_PrioritizesAgentRunOptionsOverChatOptionsAsync()
+    public async Task RunAsyncPrioritizesBackgroundResponsesPropertiesFromAgentRunOptionsOverOnesFromChatOptionsAsync()
    {
        // Arrange
-        var continuationToken1 = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }));
-        var continuationToken2 = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }));
+        var continuationToken1 = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
+        var continuationToken2 = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
        ChatOptions? capturedChatOptions = null;
        Mock<IChatClient> mockChatClient = new();
        mockChatClient
@@ -104,13 +104,13 @@ public class ChatClientAgent_BackgroundResponsesTests
        // Assert
        Assert.NotNull(capturedChatOptions);
        Assert.False(capturedChatOptions.AllowBackgroundResponses);
-        Assert.Same(continuationToken2.InnerToken, capturedChatOptions.ContinuationToken);
+        Assert.Same(continuationToken2, capturedChatOptions.ContinuationToken);
    }

    [Theory]
    [InlineData(true)]
    [InlineData(false)]
-    public async Task RunStreamingAsync_PropagatesBackgroundResponsesPropertiesToChatClientAsync(bool providePropsViaChatOptions)
+    public async Task RunStreamingAsyncPropagatesBackgroundResponsesPropertiesToChatClientAsync(bool providePropsViaChatOptions)
    {
        // Arrange
        ChatResponseUpdate[] returnUpdates =
@@ -119,7 +119,7 @@ public class ChatClientAgent_BackgroundResponsesTests
            new ChatResponseUpdate(role: ChatRole.Assistant, content: "at?") { ConversationId = "conversation-id" },
        ];

-        var continuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 })) { InputMessages = [new ChatMessage()] };
+        var continuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
        ChatOptions? capturedChatOptions = null;
        Mock<IChatClient> mockChatClient = new();
        mockChatClient
@@ -164,11 +164,11 @@ public class ChatClientAgent_BackgroundResponsesTests
        Assert.NotNull(capturedChatOptions);

        Assert.True(capturedChatOptions.AllowBackgroundResponses);
-        Assert.Same(continuationToken.InnerToken, capturedChatOptions.ContinuationToken);
+        Assert.Same(continuationToken, capturedChatOptions.ContinuationToken);
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenPropertiesSetInBothLocations_PrioritizesAgentRunOptionsOverChatOptionsAsync()
+    public async Task RunStreamingAsyncPrioritizesBackgroundResponsesPropertiesFromAgentRunOptionsOverOnesFromChatOptionsAsync()
    {
        // Arrange
        ChatResponseUpdate[] returnUpdates =
@@ -176,8 +176,8 @@ public class ChatClientAgent_BackgroundResponsesTests
            new ChatResponseUpdate(role: ChatRole.Assistant, content: "wh") { ConversationId = "conversation-id" },
        ];

-        var continuationToken1 = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 })) { InputMessages = [new ChatMessage()] };
-        var continuationToken2 = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 })) { InputMessages = [new ChatMessage()] };
+        var continuationToken1 = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
+        var continuationToken2 = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
        ChatOptions? capturedChatOptions = null;
        Mock<IChatClient> mockChatClient = new();
        mockChatClient
@@ -212,11 +212,11 @@ public class ChatClientAgent_BackgroundResponsesTests
        // Assert
        Assert.NotNull(capturedChatOptions);
        Assert.False(capturedChatOptions.AllowBackgroundResponses);
-        Assert.Same(continuationToken2.InnerToken, capturedChatOptions.ContinuationToken);
+        Assert.Same(continuationToken2, capturedChatOptions.ContinuationToken);
    }

    [Fact]
-    public async Task RunAsync_WhenContinuationTokenReceivedFromChatResponse_WrapsContinuationTokenAsync()
+    public async Task RunAsyncPropagatesContinuationTokenFromChatResponseToAgentRunResponseAsync()
    {
        // Arrange
        var continuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
@@ -237,11 +237,11 @@ public class ChatClientAgent_BackgroundResponsesTests
        var response = await agent.RunAsync([new(ChatRole.User, "hi")], thread, options: runOptions);

        // Assert
-        Assert.Same(continuationToken, (response.ContinuationToken as ChatClientAgentContinuationToken)?.InnerToken);
+        Assert.Same(continuationToken, response.ContinuationToken);
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenContinuationTokenReceived_WrapsContinuationTokenAsync()
+    public async Task RunStreamingAsyncPropagatesContinuationTokensFromUpdatesAsync()
    {
        // Arrange
        var token1 = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
@@ -272,19 +272,19 @@ public class ChatClientAgent_BackgroundResponsesTests

        // Assert
        Assert.Equal(2, actualUpdates.Count);
-        Assert.Same(token1, (actualUpdates[0].ContinuationToken as ChatClientAgentContinuationToken)?.InnerToken);
+        Assert.Same(token1, actualUpdates[0].ContinuationToken);
        Assert.Null(actualUpdates[1].ContinuationToken); // last update has null token
    }

    [Fact]
-    public async Task RunAsync_WhenMessagesProvidedWithContinuationToken_ThrowsInvalidOperationExceptionAsync()
+    public async Task RunAsyncThrowsWhenMessagesProvidedWithContinuationTokenAsync()
    {
        // Arrange
        Mock<IChatClient> mockChatClient = new();

        ChatClientAgent agent = new(mockChatClient.Object);

-        AgentRunOptions runOptions = new() { ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 })) };
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };

        IEnumerable<ChatMessage> inputMessages = [new ChatMessage(ChatRole.User, "test message")];

@@ -301,14 +301,14 @@ public class ChatClientAgent_BackgroundResponsesTests
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenMessagesProvidedWithContinuationToken_ThrowsInvalidOperationExceptionAsync()
+    public async Task RunStreamingAsyncThrowsWhenMessagesProvidedWithContinuationTokenAsync()
    {
        // Arrange
        Mock<IChatClient> mockChatClient = new();

        ChatClientAgent agent = new(mockChatClient.Object);

-        AgentRunOptions runOptions = new() { ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 })) };
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };

        IEnumerable<ChatMessage> inputMessages = [new ChatMessage(ChatRole.User, "test message")];

@@ -331,7 +331,7 @@ public class ChatClientAgent_BackgroundResponsesTests
    }

    [Fact]
-    public async Task RunAsync_WhenContinuationTokenProvided_SkipsThreadMessagePopulationAsync()
+    public async Task RunAsyncSkipsThreadMessagePopulationWithContinuationTokenAsync()
    {
        // Arrange
        List<ChatMessage> capturedMessages = [];
@@ -371,10 +371,7 @@ public class ChatClientAgent_BackgroundResponsesTests
            AIContextProvider = mockContextProvider.Object
        };

-        AgentRunOptions runOptions = new()
-        {
-            ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }))
-        };
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };

        // Act
        await agent.RunAsync([], thread, options: runOptions);
@@ -396,7 +393,7 @@ public class ChatClientAgent_BackgroundResponsesTests
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenContinuationTokenProvided_SkipsThreadMessagePopulationAsync()
+    public async Task RunStreamingAsyncSkipsThreadMessagePopulationWithContinuationTokenAsync()
    {
        // Arrange
        List<ChatMessage> capturedMessages = [];
@@ -436,15 +433,14 @@ public class ChatClientAgent_BackgroundResponsesTests
            AIContextProvider = mockContextProvider.Object
        };

-        AgentRunOptions runOptions = new()
-        {
-            ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 })) { InputMessages = [new ChatMessage()] }
-        };
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };

        // Act
-        await agent.RunStreamingAsync(thread, options: runOptions).ToListAsync();
+        var exception = await Assert.ThrowsAsync<NotSupportedException>(async () => await agent.RunStreamingAsync(thread, options: runOptions).ToListAsync());

        // Assert
+        Assert.Equal("Streaming resumption is only supported when chat history is stored and managed by the underlying AI service.", exception.Message);
+
        // With continuation token, thread message population should be skipped
        Assert.Empty(capturedMessages);

@@ -460,7 +456,7 @@ public class ChatClientAgent_BackgroundResponsesTests
    }

    [Fact]
-    public async Task RunAsync_WhenNoThreadProvidedForBackgroundResponses_ThrowsInvalidOperationExceptionAsync()
+    public async Task RunAsyncThrowsWhenNoThreadProvideForBackgroundResponsesAsync()
    {
        // Arrange
        Mock<IChatClient> mockChatClient = new();
@@ -484,7 +480,7 @@ public class ChatClientAgent_BackgroundResponsesTests
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenNoThreadProvidedForBackgroundResponses_ThrowsInvalidOperationExceptionAsync()
+    public async Task RunStreamingAsyncThrowsWhenNoThreadProvideForBackgroundResponsesAsync()
    {
        // Arrange
        Mock<IChatClient> mockChatClient = new();
@@ -514,287 +510,126 @@ public class ChatClientAgent_BackgroundResponsesTests
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenInputMessagesPresentInContinuationToken_ResumesStreamingAsync()
+    public async Task RunAsyncThrowsWhenContinuationTokenProvidedForInitialRunAsync()
    {
        // Arrange
-        ChatResponseUpdate[] returnUpdates =
-        [
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: "continuation") { ConversationId = "conversation-id" },
-        ];
-
        Mock<IChatClient> mockChatClient = new();
-        mockChatClient
-            .Setup(c => c.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()))
-            .Returns(ToAsyncEnumerableAsync(returnUpdates));

        ChatClientAgent agent = new(mockChatClient.Object);

-        ChatClientAgentThread thread = new() { ConversationId = "conversation-id" };
+        // Create a new thread with no ConversationId and no MessageStore (initial run state)
+        ChatClientAgentThread thread = new();

-        AgentRunOptions runOptions = new()
-        {
-            ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }))
-            {
-                InputMessages = [new ChatMessage(ChatRole.User, "previous message")]
-            }
-        };
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };

-        // Act
-        var updates = new List<AgentRunResponseUpdate>();
-        await foreach (var update in agent.RunStreamingAsync(thread, options: runOptions))
-        {
-            updates.Add(update);
-        }
+        // Act & Assert
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync(thread: thread, options: runOptions));
+        Assert.Equal("Continuation tokens are not allowed to be used for initial runs.", exception.Message);

-        // Assert
-        Assert.Single(updates);
+        // Verify that the IChatClient was never called due to early validation
+        mockChatClient.Verify(
+            c => c.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()),
+            Times.Never);
+    }

-        // Verify that the IChatClient was called
+    [Fact]
+    public async Task RunStreamingAsyncThrowsWhenContinuationTokenProvidedForInitialRunAsync()
+    {
+        // Arrange
+        Mock<IChatClient> mockChatClient = new();
+
+        ChatClientAgent agent = new(mockChatClient.Object);
+
+        // Create a new thread with no ConversationId and no MessageStore (initial run state)
+        ChatClientAgentThread thread = new();
+
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };
+
+        // Act & Assert
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await agent.RunStreamingAsync(thread: thread, options: runOptions).ToListAsync());
+        Assert.Equal("Continuation tokens are not allowed to be used for initial runs.", exception.Message);
+
+        // Verify that the IChatClient was never called due to early validation
        mockChatClient.Verify(
            c => c.GetStreamingResponseAsync(
                It.IsAny<IEnumerable<ChatMessage>>(),
                It.IsAny<ChatOptions>(),
                It.IsAny<CancellationToken>()),
-            Times.Once);
+            Times.Never);
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenResponseUpdatesPresentInContinuationToken_ResumesStreamingAsync()
+    public async Task RunStreamingAsyncThrowsWhenContinuationTokenUsedWithClientSideManagedChatHistoryAsync()
    {
        // Arrange
-        ChatResponseUpdate[] returnUpdates =
-        [
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: "continuation") { ConversationId = "conversation-id" },
-        ];
-
        Mock<IChatClient> mockChatClient = new();
-        mockChatClient
-            .Setup(c => c.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()))
-            .Returns(ToAsyncEnumerableAsync(returnUpdates));

        ChatClientAgent agent = new(mockChatClient.Object);

-        ChatClientAgentThread thread = new() { ConversationId = "conversation-id" };
-
-        AgentRunOptions runOptions = new()
+        // Create a thread with a MessageStore
+        ChatClientAgentThread thread = new()
        {
-            ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }))
-            {
-                ResponseUpdates = [new ChatResponseUpdate(ChatRole.Assistant, "previous update")]
-            }
+            MessageStore = new InMemoryChatMessageStore(), // Setting a message store to skip checking the continuation token in the initial run
+            ConversationId = null, // No conversation ID to simulate client-side managed chat history
        };

-        // Act
-        var updates = new List<AgentRunResponseUpdate>();
-        await foreach (var update in agent.RunStreamingAsync(thread, options: runOptions))
-        {
-            updates.Add(update);
-        }
+        // Create run options with a continuation token
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };

-        // Assert
-        Assert.Single(updates);
+        // Act & Assert
+        var exception = await Assert.ThrowsAsync<NotSupportedException>(async () => await agent.RunStreamingAsync(thread: thread, options: runOptions).ToListAsync());
+        Assert.Equal("Streaming resumption is only supported when chat history is stored and managed by the underlying AI service.", exception.Message);

-        // Verify that the IChatClient was called
+        // Verify that the IChatClient was never called due to early validation
        mockChatClient.Verify(
            c => c.GetStreamingResponseAsync(
                It.IsAny<IEnumerable<ChatMessage>>(),
                It.IsAny<ChatOptions>(),
                It.IsAny<CancellationToken>()),
-            Times.Once);
+            Times.Never);
    }

    [Fact]
-    public async Task RunStreamingAsync_WhenResumingStreaming_UsesUpdatesFromInitialRunForContextProviderAndMessageStoreAsync()
+    public async Task RunStreamingAsyncThrowsWhenContinuationTokenUsedWithAIContextProviderAsync()
    {
        // Arrange
-        ChatResponseUpdate[] returnUpdates =
-        [
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: "upon"),
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: " a"),
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: " time"),
-        ];
-
        Mock<IChatClient> mockChatClient = new();
-        mockChatClient
-            .Setup(c => c.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()))
-            .Returns(ToAsyncEnumerableAsync(returnUpdates));

        ChatClientAgent agent = new(mockChatClient.Object);

-        List<ChatMessage> capturedMessagesAddedToStore = [];
-        var mockMessageStore = new Mock<ChatMessageStore>();
-        mockMessageStore
-            .Setup(ms => ms.InvokedAsync(It.IsAny<ChatMessageStore.InvokedContext>(), It.IsAny<CancellationToken>()))
-            .Callback<ChatMessageStore.InvokedContext, CancellationToken>((ctx, ct) => capturedMessagesAddedToStore.AddRange(ctx.ResponseMessages ?? []))
-            .Returns(new ValueTask());
-
-        AIContextProvider.InvokedContext? capturedInvokedContext = null;
+        // Create a mock AIContextProvider
        var mockContextProvider = new Mock<AIContextProvider>();
        mockContextProvider
-            .Setup(cp => cp.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
-            .Callback<AIContextProvider.InvokedContext, CancellationToken>((context, ct) => capturedInvokedContext = context)
+            .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new AIContext());
+        mockContextProvider
+            .Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
            .Returns(new ValueTask());

+        // Create a thread with an AIContextProvider and conversation ID to simulate non-initial run
        ChatClientAgentThread thread = new()
        {
-            MessageStore = mockMessageStore.Object,
+            ConversationId = "existing-conversation-id",
            AIContextProvider = mockContextProvider.Object
        };

-        AgentRunOptions runOptions = new()
-        {
-            ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }))
-            {
-                ResponseUpdates = [new ChatResponseUpdate(ChatRole.Assistant, "once ")]
-            }
-        };
+        AgentRunOptions runOptions = new() { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) };

-        // Act
-        await agent.RunStreamingAsync(thread, options: runOptions).ToListAsync();
+        // Act & Assert
+        var exception = await Assert.ThrowsAsync<NotSupportedException>(async () => await agent.RunStreamingAsync(thread: thread, options: runOptions).ToListAsync());

-        // Assert
-        mockMessageStore.Verify(ms => ms.InvokedAsync(It.IsAny<ChatMessageStore.InvokedContext>(), It.IsAny<CancellationToken>()), Times.Once);
-        Assert.Single(capturedMessagesAddedToStore);
-        Assert.Contains("once upon a time", capturedMessagesAddedToStore[0].Text);
+        Assert.Equal("Using context provider with streaming resumption is not supported.", exception.Message);

-        mockContextProvider.Verify(cp => cp.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()), Times.Once);
-        Assert.NotNull(capturedInvokedContext?.ResponseMessages);
-        Assert.Single(capturedInvokedContext.ResponseMessages);
-        Assert.Contains("once upon a time", capturedInvokedContext.ResponseMessages.ElementAt(0).Text);
-    }
-
-    [Fact]
-    public async Task RunStreamingAsync_WhenResumingStreaming_UsesInputMessagesFromInitialRunForContextProviderAndMessageStoreAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient
-            .Setup(c => c.GetStreamingResponseAsync(
+        // Verify that the IChatClient was never called due to early validation
+        mockChatClient.Verify(
+            c => c.GetStreamingResponseAsync(
                It.IsAny<IEnumerable<ChatMessage>>(),
                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()))
-            .Returns(ToAsyncEnumerableAsync(Array.Empty<ChatResponseUpdate>()));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-
-        List<ChatMessage> capturedMessagesAddedToStore = [];
-        var mockMessageStore = new Mock<ChatMessageStore>();
-        mockMessageStore
-            .Setup(ms => ms.InvokedAsync(It.IsAny<ChatMessageStore.InvokedContext>(), It.IsAny<CancellationToken>()))
-            .Callback<ChatMessageStore.InvokedContext, CancellationToken>((ctx, ct) => capturedMessagesAddedToStore.AddRange(ctx.RequestMessages))
-            .Returns(new ValueTask());
-
-        AIContextProvider.InvokedContext? capturedInvokedContext = null;
-        var mockContextProvider = new Mock<AIContextProvider>();
-        mockContextProvider
-            .Setup(cp => cp.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
-            .Callback<AIContextProvider.InvokedContext, CancellationToken>((context, ct) => capturedInvokedContext = context)
-            .Returns(new ValueTask());
-
-        ChatClientAgentThread thread = new()
-        {
-            MessageStore = mockMessageStore.Object,
-            AIContextProvider = mockContextProvider.Object
-        };
-
-        AgentRunOptions runOptions = new()
-        {
-            ContinuationToken = new ChatClientAgentContinuationToken(ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }))
-            {
-                InputMessages = [new ChatMessage(ChatRole.User, "Tell me a story")],
-            }
-        };
-
-        // Act
-        await agent.RunStreamingAsync(thread, options: runOptions).ToListAsync();
-
-        // Assert
-        mockMessageStore.Verify(ms => ms.InvokedAsync(It.IsAny<ChatMessageStore.InvokedContext>(), It.IsAny<CancellationToken>()), Times.Once);
-        Assert.Single(capturedMessagesAddedToStore);
-        Assert.Contains("Tell me a story", capturedMessagesAddedToStore[0].Text);
-
-        mockContextProvider.Verify(cp => cp.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()), Times.Once);
-        Assert.NotNull(capturedInvokedContext?.RequestMessages);
-        Assert.Single(capturedInvokedContext.RequestMessages);
-        Assert.Contains("Tell me a story", capturedInvokedContext.RequestMessages.ElementAt(0).Text);
-    }
-
-    [Fact]
-    public async Task RunStreamingAsync_WhenResumingStreaming_SavesInputMessagesAndUpdatesInContinuationTokenAsync()
-    {
-        // Arrange
-        List<ChatResponseUpdate> returnUpdates =
-        [
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: "Once") { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) },
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: " upon") { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) },
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: " a") { ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) },
-            new ChatResponseUpdate(role: ChatRole.Assistant, content: " time"){ ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) },
-        ];
-
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient
-            .Setup(c => c.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()))
-            .Returns(ToAsyncEnumerableAsync(returnUpdates));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-
-        ChatClientAgentThread thread = new() { };
-
-        List<ChatClientAgentContinuationToken> capturedContinuationTokens = [];
-
-        ChatMessage userMessage = new(ChatRole.User, "Tell me a story");
-
-        // Act
-
-        // Do the initial run
-        await foreach (var update in agent.RunStreamingAsync(userMessage, thread))
-        {
-            capturedContinuationTokens.Add(Assert.IsType<ChatClientAgentContinuationToken>(update.ContinuationToken));
-            break;
-        }
-
-        // Now resume the run using the captured continuation token
-        returnUpdates.RemoveAt(0); // remove the first mock update as it was already processed
-        var options = new AgentRunOptions { ContinuationToken = capturedContinuationTokens[0] };
-        await foreach (var update in agent.RunStreamingAsync(thread, options: options))
-        {
-            capturedContinuationTokens.Add(Assert.IsType<ChatClientAgentContinuationToken>(update.ContinuationToken));
-        }
-
-        // Assert
-        Assert.Equal(4, capturedContinuationTokens.Count);
-
-        // Verify that the first continuation token has the initial input and first update
-        Assert.NotNull(capturedContinuationTokens[0].InputMessages);
-        Assert.Single(capturedContinuationTokens[0].InputMessages!);
-        Assert.Equal("Tell me a story", capturedContinuationTokens[0].InputMessages!.Last().Text);
-        Assert.NotNull(capturedContinuationTokens[0].ResponseUpdates);
-        Assert.Single(capturedContinuationTokens[0].ResponseUpdates!);
-        Assert.Equal("Once", capturedContinuationTokens[0].ResponseUpdates![0].Text);
-
-        // Verify the last continuation token has the input and all updates
-        var lastToken = capturedContinuationTokens[^1];
-        Assert.NotNull(lastToken.InputMessages);
-        Assert.Single(lastToken.InputMessages!);
-        Assert.Equal("Tell me a story", lastToken.InputMessages!.Last().Text);
-        Assert.NotNull(lastToken.ResponseUpdates);
-        Assert.Equal(4, lastToken.ResponseUpdates!.Count);
-        Assert.Equal("Once", lastToken.ResponseUpdates!.ElementAt(0).Text);
-        Assert.Equal(" upon", lastToken.ResponseUpdates!.ElementAt(1).Text);
-        Assert.Equal(" a", lastToken.ResponseUpdates!.ElementAt(2).Text);
-        Assert.Equal(" time", lastToken.ResponseUpdates!.ElementAt(3).Text);
+                It.IsAny<CancellationToken>()),
+            Times.Never);
    }

    private static async IAsyncEnumerable<T> ToAsyncEnumerableAsync<T>(IEnumerable<T> values)
@@ -1,456 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text.Json.Serialization;
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.Extensions.AI;
-using Moq;
-
-namespace Microsoft.Agents.AI.UnitTests;
-
-/// <summary>
-/// Tests for <see cref="ChatClientAgent"/> run methods with <see cref="ChatClientAgentRunOptions"/>.
-/// </summary>
-public sealed partial class ChatClientAgent_RunWithCustomOptionsTests
-{
-    #region RunAsync Tests
-
-    [Fact]
-    public async Task RunAsync_WithThreadAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "Response")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse result = await agent.RunAsync(thread, options);
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result.Messages);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunAsync_WithStringMessageAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "Response")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse result = await agent.RunAsync("Test message", thread, options);
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result.Messages);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Any(m => m.Text == "Test message")),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunAsync_WithChatMessageAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "Response")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatMessage message = new(ChatRole.User, "Test message");
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse result = await agent.RunAsync(message, thread, options);
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result.Messages);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Contains(message)),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunAsync_WithMessagesCollectionAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "Response")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        IEnumerable<ChatMessage> messages = [new(ChatRole.User, "Message 1"), new(ChatRole.User, "Message 2")];
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse result = await agent.RunAsync(messages, thread, options);
-
-        // Assert
-        Assert.NotNull(result);
-        Assert.Single(result.Messages);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunAsync_WithChatOptionsInRunOptions_UsesChatOptionsAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "Response")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        ChatClientAgentRunOptions options = new(new ChatOptions { Temperature = 0.5f });
-
-        // Act
-        AgentRunResponse result = await agent.RunAsync("Test", null, options);
-
-        // Assert
-        Assert.NotNull(result);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.Is<ChatOptions>(opts => opts.Temperature == 0.5f),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    #endregion
-
-    #region RunStreamingAsync Tests
-
-    [Fact]
-    public async Task RunStreamingAsync_WithThreadAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).Returns(GetAsyncUpdatesAsync());
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        var updates = new List<AgentRunResponseUpdate>();
-        await foreach (var update in agent.RunStreamingAsync(thread, options))
-        {
-            updates.Add(update);
-        }
-
-        // Assert
-        Assert.NotEmpty(updates);
-        mockChatClient.Verify(
-            x => x.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunStreamingAsync_WithStringMessageAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).Returns(GetAsyncUpdatesAsync());
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        var updates = new List<AgentRunResponseUpdate>();
-        await foreach (var update in agent.RunStreamingAsync("Test message", thread, options))
-        {
-            updates.Add(update);
-        }
-
-        // Assert
-        Assert.NotEmpty(updates);
-        mockChatClient.Verify(
-            x => x.GetStreamingResponseAsync(
-                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Any(m => m.Text == "Test message")),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunStreamingAsync_WithChatMessageAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).Returns(GetAsyncUpdatesAsync());
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatMessage message = new(ChatRole.User, "Test message");
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        var updates = new List<AgentRunResponseUpdate>();
-        await foreach (var update in agent.RunStreamingAsync(message, thread, options))
-        {
-            updates.Add(update);
-        }
-
-        // Assert
-        Assert.NotEmpty(updates);
-        mockChatClient.Verify(
-            x => x.GetStreamingResponseAsync(
-                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Contains(message)),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunStreamingAsync_WithMessagesCollectionAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).Returns(GetAsyncUpdatesAsync());
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        IEnumerable<ChatMessage> messages = [new ChatMessage(ChatRole.User, "Message 1"), new ChatMessage(ChatRole.User, "Message 2")];
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        var updates = new List<AgentRunResponseUpdate>();
-        await foreach (var update in agent.RunStreamingAsync(messages, thread, options))
-        {
-            updates.Add(update);
-        }
-
-        // Assert
-        Assert.NotEmpty(updates);
-        mockChatClient.Verify(
-            x => x.GetStreamingResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    #endregion
-
-    #region Helper Methods
-
-    private static async IAsyncEnumerable<ChatResponseUpdate> GetAsyncUpdatesAsync()
-    {
-        yield return new ChatResponseUpdate { Contents = new[] { new TextContent("Hello") } };
-        yield return new ChatResponseUpdate { Contents = new[] { new TextContent(" World") } };
-        await Task.CompletedTask;
-    }
-
-    #endregion
-
-    #region RunAsync{T} Tests
-
-    [Fact]
-    public async Task RunAsyncOfT_WithThreadAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, """{"id":2, "fullName":"Tigger", "species":"Tiger"}""")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse<Animal> agentRunResponse = await agent.RunAsync<Animal>(thread, JsonContext_WithCustomRunOptions.Default.Options, options);
-
-        // Assert
-        Assert.NotNull(agentRunResponse);
-        Assert.Single(agentRunResponse.Messages);
-        Assert.Equal("Tigger", agentRunResponse.Result.FullName);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunAsyncOfT_WithStringMessageAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, """{"id":2, "fullName":"Tigger", "species":"Tiger"}""")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse<Animal> agentRunResponse = await agent.RunAsync<Animal>("Test message", thread, JsonContext_WithCustomRunOptions.Default.Options, options);
-
-        // Assert
-        Assert.NotNull(agentRunResponse);
-        Assert.Single(agentRunResponse.Messages);
-        Assert.Equal("Tigger", agentRunResponse.Result.FullName);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Any(m => m.Text == "Test message")),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunAsyncOfT_WithChatMessageAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, """{"id":2, "fullName":"Tigger", "species":"Tiger"}""")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        ChatMessage message = new(ChatRole.User, "Test message");
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse<Animal> agentRunResponse = await agent.RunAsync<Animal>(message, thread, JsonContext_WithCustomRunOptions.Default.Options, options);
-
-        // Assert
-        Assert.NotNull(agentRunResponse);
-        Assert.Single(agentRunResponse.Messages);
-        Assert.Equal("Tigger", agentRunResponse.Result.FullName);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Contains(message)),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    [Fact]
-    public async Task RunAsyncOfT_WithMessagesCollectionAndOptions_CallsBaseMethodAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockChatClient = new();
-        mockChatClient.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, """{"id":2, "fullName":"Tigger", "species":"Tiger"}""")]));
-
-        ChatClientAgent agent = new(mockChatClient.Object);
-        AgentThread thread = agent.GetNewThread();
-        IEnumerable<ChatMessage> messages = [new(ChatRole.User, "Message 1"), new(ChatRole.User, "Message 2")];
-        ChatClientAgentRunOptions options = new();
-
-        // Act
-        AgentRunResponse<Animal> agentRunResponse = await agent.RunAsync<Animal>(messages, thread, JsonContext_WithCustomRunOptions.Default.Options, options);
-
-        // Assert
-        Assert.NotNull(agentRunResponse);
-        Assert.Single(agentRunResponse.Messages);
-        Assert.Equal("Tigger", agentRunResponse.Result.FullName);
-        mockChatClient.Verify(
-            x => x.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-    }
-
-    #endregion
-
-    private sealed class Animal
-    {
-        public int Id { get; set; }
-        public string? FullName { get; set; }
-        public Species Species { get; set; }
-    }
-
-    private enum Species
-    {
-        Bear,
-        Tiger,
-        Walrus,
-    }
-
-    [JsonSourceGenerationOptions(UseStringEnumConverter = true, PropertyNamingPolicy = JsonKnownNamingPolicy.CamelCase)]
-    [JsonSerializable(typeof(Animal))]
-    private sealed partial class JsonContext_WithCustomRunOptions : JsonSerializerContext;
-}
@@ -1,114 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Runtime.CompilerServices;
-using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
-using FluentAssertions;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI.Workflows.UnitTests;
-
-public sealed class ExpectedException : Exception
-{
-    public ExpectedException(string message)
-        : base(message)
-    {
-    }
-
-    public ExpectedException() : base()
-    {
-    }
-
-    public ExpectedException(string? message, Exception? innerException) : base(message, innerException)
-    {
-    }
-}
-
-public class WorkflowHostSmokeTests
-{
-    private sealed class AlwaysFailsAIAgent(bool failByThrowing) : AIAgent
-    {
-        private sealed class Thread : InMemoryAgentThread
-        {
-            public Thread() { }
-
-            public Thread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
-                : base(serializedThread, jsonSerializerOptions)
-            { }
-        }
-
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
-        {
-            return new Thread(serializedThread, jsonSerializerOptions);
-        }
-
-        public override AgentThread GetNewThread()
-        {
-            return new Thread();
-        }
-
-        protected override async Task<AgentRunResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
-        {
-            return await this.RunStreamingAsync(messages, thread, options, cancellationToken)
-                             .ToAgentRunResponseAsync(cancellationToken);
-        }
-
-        protected override async IAsyncEnumerable<AgentRunResponseUpdate> RunCoreStreamingAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
-        {
-            const string ErrorMessage = "Simulated agent failure.";
-            if (failByThrowing)
-            {
-                throw new ExpectedException(ErrorMessage);
-            }
-
-            yield return new AgentRunResponseUpdate(ChatRole.Assistant, [new ErrorContent(ErrorMessage)]);
-        }
-    }
-
-    private static Workflow CreateWorkflow(bool failByThrowing)
-    {
-        ExecutorBinding agent = new AlwaysFailsAIAgent(failByThrowing).BindAsExecutor(emitEvents: true);
-
-        return new WorkflowBuilder(agent).Build();
-    }
-
-    [Theory]
-    [InlineData(true, true)]
-    [InlineData(true, false)]
-    [InlineData(false, true)]
-    [InlineData(false, false)]
-    public async Task Test_AsAgent_ErrorContentStreamedOutAsync(bool includeExceptionDetails, bool failByThrowing)
-    {
-        string expectedMessage = !failByThrowing || includeExceptionDetails
-                               ? "Simulated agent failure."
-                               : "An error occurred while executing the workflow.";
-
-        // Arrange is done by the caller.
-        Workflow workflow = CreateWorkflow(failByThrowing);
-
-        // Act
-        List<AgentRunResponseUpdate> updates = await workflow.AsAgent("WorkflowAgent", includeExceptionDetails: includeExceptionDetails)
-                                                             .RunStreamingAsync(new ChatMessage(ChatRole.User, "Hello"))
-                                                             .ToListAsync();
-
-        // Assert
-        bool hadErrorContent = false;
-        foreach (AgentRunResponseUpdate update in updates)
-        {
-            if (update.Contents.Any())
-            {
-                // We should expect a single update which contains the error content.
-                update.Contents.Should().ContainSingle()
-                                        .Which.Should().BeOfType<ErrorContent>()
-                                        .Which.Message.Should().Be(expectedMessage);
-                hadErrorContent = true;
-            }
-        }
-
-        hadErrorContent.Should().BeTrue();
-    }
-}
@@ -7,41 +7,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

 ## [Unreleased]

-## [1.0.0b260107] - 2026-01-07
-
-### Added
-
- **agent-framework-devui**: Improve DevUI and add Context Inspector view as a new tab under traces ([#2742](https://github.com/microsoft/agent-framework/pull/2742)) by @victordibia
- **samples**: Add streaming sample for Azure Functions ([#3057](https://github.com/microsoft/agent-framework/pull/3057)) by @gavin-aguiar
-
-### Changed
-
- **repo**: Update templates ([#3106](https://github.com/microsoft/agent-framework/pull/3106)) by @eavanvalkenburg
-
-### Fixed
-
- **agent-framework-ag-ui**: Fix MCP tool result serialization for list[TextContent] ([#2523](https://github.com/microsoft/agent-framework/pull/2523)) by @claude89757
- **agent-framework-azure-ai**: Fix response_format handling for structured outputs ([#3114](https://github.com/microsoft/agent-framework/pull/3114)) by @moonbox3
-
-## [1.0.0b260106] - 2026-01-06
-
-### Added
-
- **repo**: Add issue template and additional labeling ([#3006](https://github.com/microsoft/agent-framework/pull/3006)) by @eavanvalkenburg
-
-### Changed
-
- None
-
-### Fixed
-
- **agent-framework-core**: Fix max tokens translation and add extra integer test ([#3037](https://github.com/microsoft/agent-framework/pull/3037)) by @eavanvalkenburg
- **agent-framework-azure-ai**: Fix failure when conversation history contains assistant messages ([#3076](https://github.com/microsoft/agent-framework/pull/3076)) by @moonbox3
- **agent-framework-core**: Use HTTP exporter for http/protobuf protocol ([#3070](https://github.com/microsoft/agent-framework/pull/3070)) by @takanori-terai
- **agent-framework-core**: Fix ExecutorInvokedEvent and ExecutorCompletedEvent observability data ([#3090](https://github.com/microsoft/agent-framework/pull/3090)) by @moonbox3
- **agent-framework-core**: Honor tool_choice parameter passed to agent.run() and chat client methods ([#3095](https://github.com/microsoft/agent-framework/pull/3095)) by @moonbox3
- **samples**: AzureAI SharePoint sample fix ([#3108](https://github.com/microsoft/agent-framework/pull/3108)) by @giles17
-
 ## [1.0.0b251223] - 2025-12-23

 ### Added
@@ -461,9 +426,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

 For more information, see the [announcement blog post](https://devblogs.microsoft.com/foundry/introducing-microsoft-agent-framework-the-open-source-engine-for-agentic-ai-apps/).

-[Unreleased]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b260107...HEAD
-[1.0.0b260107]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b260106...python-1.0.0b260107
-[1.0.0b260106]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b251223...python-1.0.0b260106
+[Unreleased]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b251223...HEAD
 [1.0.0b251223]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b251218...python-1.0.0b251223
 [1.0.0b251218]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b251216...python-1.0.0b251218
 [1.0.0b251216]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b251211...python-1.0.0b251216
@@ -4,7 +4,7 @@ description = "A2A integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -11,61 +11,21 @@ from typing import Any


 class ConfirmationStrategy(ABC):
-    """Strategy for generating confirmation messages during human-in-the-loop flows.
+    """Strategy for generating confirmation messages during human-in-the-loop flows."""

-    Subclasses must define the message properties. The methods use those properties
-    by default, but can be overridden for complete customization.
-    """
-
-    @property
    @abstractmethod
-    def approval_header(self) -> str:
-        """Header for approval accepted message. Must be overridden."""
-        ...
-
-    @property
-    @abstractmethod
-    def approval_footer(self) -> str:
-        """Footer for approval accepted message. Must be overridden."""
-        ...
-
-    @property
-    @abstractmethod
-    def rejection_message(self) -> str:
-        """Message when user rejects. Must be overridden."""
-        ...
-
-    @property
-    @abstractmethod
-    def state_confirmed_message(self) -> str:
-        """Message when state is confirmed. Must be overridden."""
-        ...
-
-    @property
-    @abstractmethod
-    def state_rejected_message(self) -> str:
-        """Message when state is rejected. Must be overridden."""
-        ...
-
    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
        """Generate message when user approves function execution.

-        Default implementation uses header/footer properties.
-        Override for complete customization.
-
        Args:
            steps: List of approved steps with 'description', 'status', etc.

        Returns:
            Message to display to user
        """
-        enabled_steps = [s for s in steps if s.get("status") == "enabled"]
-        message_parts = [self.approval_header.format(count=len(enabled_steps))]
-        for i, step in enumerate(enabled_steps, 1):
-            message_parts.append(f"{i}. {step['description']}\n")
-        message_parts.append(self.approval_footer)
-        return "".join(message_parts)
+        ...

+    @abstractmethod
    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
        """Generate message when user rejects function execution.

@@ -75,143 +35,141 @@ class ConfirmationStrategy(ABC):
        Returns:
            Message to display to user
        """
-        return self.rejection_message
+        ...

+    @abstractmethod
    def on_state_confirmed(self) -> str:
        """Generate message when user confirms predictive state changes.

        Returns:
            Message to display to user
        """
-        return self.state_confirmed_message
+        ...

+    @abstractmethod
    def on_state_rejected(self) -> str:
        """Generate message when user rejects predictive state changes.

        Returns:
            Message to display to user
        """
-        return self.state_rejected_message
+        ...


 class DefaultConfirmationStrategy(ConfirmationStrategy):
-    """Generic confirmation messages suitable for most agents."""
+    """Generic confirmation messages suitable for most agents.

-    @property
-    def approval_header(self) -> str:
-        return "Executing {count} approved steps:\n\n"
+    This preserves the original behavior from v1.
+    """

-    @property
-    def approval_footer(self) -> str:
-        return "\nAll steps completed successfully!"
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate generic approval message with step list."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]

-    @property
-    def rejection_message(self) -> str:
+        message_parts = [f"Executing {len(enabled_steps)} approved steps:\n\n"]
+
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nAll steps completed successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate generic rejection message."""
        return "No problem! What would you like me to change about the plan?"

-    @property
-    def state_confirmed_message(self) -> str:
+    def on_state_confirmed(self) -> str:
+        """Generate generic state confirmation message."""
        return "Changes confirmed and applied successfully!"

-    @property
-    def state_rejected_message(self) -> str:
+    def on_state_rejected(self) -> str:
+        """Generate generic state rejection message."""
        return "No problem! What would you like me to change?"


 class TaskPlannerConfirmationStrategy(ConfirmationStrategy):
    """Domain-specific confirmation messages for task planning agents."""

-    @property
-    def approval_header(self) -> str:
-        return "Executing your requested tasks:\n\n"
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate task-specific approval message."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]

-    @property
-    def approval_footer(self) -> str:
-        return "\nAll tasks completed successfully!"
+        message_parts = ["Executing your requested tasks:\n\n"]

-    @property
-    def rejection_message(self) -> str:
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nAll tasks completed successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate task-specific rejection message."""
        return "No problem! Let me revise the plan. What would you like me to change?"

-    @property
-    def state_confirmed_message(self) -> str:
+    def on_state_confirmed(self) -> str:
+        """Task planners typically don't use state confirmation."""
        return "Tasks confirmed and ready to execute!"

-    @property
-    def state_rejected_message(self) -> str:
+    def on_state_rejected(self) -> str:
+        """Task planners typically don't use state confirmation."""
        return "No problem! How should I adjust the task list?"


 class RecipeConfirmationStrategy(ConfirmationStrategy):
    """Domain-specific confirmation messages for recipe agents."""

-    @property
-    def approval_header(self) -> str:
-        return "Updating your recipe:\n\n"
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate recipe-specific approval message."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]

-    @property
-    def approval_footer(self) -> str:
-        return "\nRecipe updated successfully!"
+        message_parts = ["Updating your recipe:\n\n"]

-    @property
-    def rejection_message(self) -> str:
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nRecipe updated successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate recipe-specific rejection message."""
        return "No problem! What ingredients or steps should I change?"

-    @property
-    def state_confirmed_message(self) -> str:
+    def on_state_confirmed(self) -> str:
+        """Generate recipe-specific state confirmation message."""
        return "Recipe changes applied successfully!"

-    @property
-    def state_rejected_message(self) -> str:
+    def on_state_rejected(self) -> str:
+        """Generate recipe-specific state rejection message."""
        return "No problem! What would you like me to adjust in the recipe?"


 class DocumentWriterConfirmationStrategy(ConfirmationStrategy):
    """Domain-specific confirmation messages for document writing agents."""

-    @property
-    def approval_header(self) -> str:
-        return "Applying your edits:\n\n"
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate document-specific approval message."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]

-    @property
-    def approval_footer(self) -> str:
-        return "\nDocument updated successfully!"
+        message_parts = ["Applying your edits:\n\n"]

-    @property
-    def rejection_message(self) -> str:
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nDocument updated successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate document-specific rejection message."""
        return "No problem! Which changes should I keep or modify?"

-    @property
-    def state_confirmed_message(self) -> str:
+    def on_state_confirmed(self) -> str:
+        """Generate document-specific state confirmation message."""
        return "Document edits applied!"

-    @property
-    def state_rejected_message(self) -> str:
+    def on_state_rejected(self) -> str:
+        """Generate document-specific state rejection message."""
        return "No problem! What should I change about the document?"
-
-
-def apply_confirmation_strategy(
-    strategy: ConfirmationStrategy | None,
-    accepted: bool,
-    steps: list[dict[str, Any]],
-) -> str:
-    """Apply a confirmation strategy to generate a message.
-
-    This helper consolidates the pattern used in multiple orchestrators.
-
-    Args:
-        strategy: Strategy to use, or None for default
-        accepted: Whether the user approved
-        steps: List of steps (may be empty for state confirmations)
-
-    Returns:
-        Generated message string
-    """
-    if strategy is None:
-        strategy = DefaultConfirmationStrategy()
-
-    if not steps:
-        # State confirmation (no steps)
-        return strategy.on_state_confirmed() if accepted else strategy.on_state_rejected()
-    # Step-based approval
-    return strategy.on_approval_accepted(steps) if accepted else strategy.on_approval_rejected(steps)
@@ -11,6 +11,8 @@ from typing import Any
 from ag_ui.core import (
    BaseEvent,
    CustomEvent,
+    EventType,
+    MessagesSnapshotEvent,
    RunFinishedEvent,
    RunStartedEvent,
    StateDeltaEvent,
@@ -29,10 +31,9 @@ from agent_framework import (
    FunctionCallContent,
    FunctionResultContent,
    TextContent,
-    prepare_function_call_results,
 )

-from ._utils import extract_state_from_tool_args, generate_event_id, safe_json_parse
+from ._utils import generate_event_id

 logger = logging.getLogger(__name__)

@@ -47,8 +48,8 @@ class AgentFrameworkEventBridge:
        predict_state_config: dict[str, dict[str, str]] | None = None,
        current_state: dict[str, Any] | None = None,
        skip_text_content: bool = False,
+        input_messages: list[Any] | None = None,
        require_confirmation: bool = True,
-        approval_tool_name: str | None = None,
    ) -> None:
        """
        Initialize the event bridge.
@@ -60,6 +61,7 @@ class AgentFrameworkEventBridge:
                Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
            current_state: Reference to the current state dict for tracking updates.
            skip_text_content: If True, skip emitting TextMessageContentEvents (for structured outputs).
+            input_messages: The input messages from the conversation history.
            require_confirmation: Whether predictive state updates require user confirmation.
        """
        self.run_id = run_id
@@ -72,7 +74,6 @@ class AgentFrameworkEventBridge:
        self.pending_state_updates: dict[str, Any] = {}  # Track updates from tool calls
        self.skip_text_content = skip_text_content
        self.require_confirmation = require_confirmation
-        self.approval_tool_name = approval_tool_name

        # For predictive state updates: accumulate streaming arguments
        self.streaming_tool_args: str = ""  # Accumulated JSON string
@@ -81,6 +82,13 @@ class AgentFrameworkEventBridge:
        self.should_stop_after_confirm: bool = False  # Flag to stop run after confirm_changes
        self.suppressed_summary: str = ""  # Store LLM summary to show after confirmation

+        # For MessagesSnapshotEvent: track tool calls and results
+        self.input_messages = input_messages or []
+        self.pending_tool_calls: list[dict[str, Any]] = []  # Track tool calls for assistant message
+        self.tool_results: list[dict[str, Any]] = []  # Track tool results
+        self.tool_calls_ended: set[str] = set()  # Track which tool calls have had ToolCallEndEvent emitted
+        self.accumulated_text_content: str = ""  # Track accumulated text for final MessagesSnapshotEvent
+
    async def from_agent_run_update(self, update: AgentRunResponseUpdate) -> list[BaseEvent]:
        """
        Convert an AgentRunResponseUpdate to AG-UI events.
@@ -146,6 +154,7 @@ class AgentFrameworkEventBridge:
            message_id=self.current_message_id,
            delta=content.text,
        )
+        self.accumulated_text_content += content.text
        logger.info(f"  EMITTING TextMessageContentEvent with text_len={len(content.text)}")
        events.append(event)
        return events
@@ -160,10 +169,10 @@ class AgentFrameworkEventBridge:
            logger.warning(f"FunctionCallContent missing name and call_id. args_length={args_length}")

        tool_call_id = self._coalesce_tool_call_id(content)
-        # Only emit ToolCallStartEvent once per tool call (when it's a new tool call)
        if content.name and tool_call_id != self.current_tool_call_id:
            self.streaming_tool_args = ""
            self.state_delta_count = 0
+        if content.name:
            self.current_tool_call_id = tool_call_id
            self.current_tool_call_name = content.name

@@ -174,6 +183,17 @@ class AgentFrameworkEventBridge:
            )
            logger.info(f"Emitting ToolCallStartEvent with name='{content.name}', id='{tool_call_id}'")
            events.append(tool_start_event)
+
+            self.pending_tool_calls.append(
+                {
+                    "id": tool_call_id,
+                    "type": "function",
+                    "function": {
+                        "name": content.name,
+                        "arguments": "",
+                    },
+                }
+            )
        elif tool_call_id:
            self.current_tool_call_id = tool_call_id

@@ -186,7 +206,13 @@ class AgentFrameworkEventBridge:
            )
            events.append(args_event)

+            for tool_call in self.pending_tool_calls:
+                if tool_call["id"] == tool_call_id:
+                    tool_call["function"]["arguments"] += delta_str
+                    break
+
            events.extend(self._emit_predictive_state_deltas(delta_str))
+            events.extend(self._legacy_predictive_state(content))

        return events

@@ -209,8 +235,10 @@ class AgentFrameworkEventBridge:
            self.current_tool_call_name,
        )

-        parsed_args = safe_json_parse(self.streaming_tool_args)
-        if parsed_args is None:
+        parsed_args = None
+        try:
+            parsed_args = json.loads(self.streaming_tool_args)
+        except json.JSONDecodeError:
            for state_key, config in self.predict_state_config.items():
                if config["tool"] != self.current_tool_call_name:
                    continue
@@ -254,8 +282,11 @@ class AgentFrameworkEventBridge:
                    continue
                tool_arg_name = config["tool_argument"]

-                state_value = extract_state_from_tool_args(parsed_args, tool_arg_name)
-                if state_value is None:
+                if tool_arg_name == "*":
+                    state_value = parsed_args
+                elif tool_arg_name in parsed_args:
+                    state_value = parsed_args[tool_arg_name]
+                else:
                    continue

                if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != state_value:
@@ -286,6 +317,59 @@ class AgentFrameworkEventBridge:
                    self.pending_state_updates[state_key] = state_value
        return events

+    def _legacy_predictive_state(self, content: FunctionCallContent) -> list[BaseEvent]:
+        events: list[BaseEvent] = []
+        if not (content.name and content.arguments):
+            return events
+        parsed_args = content.parse_arguments()
+        if not parsed_args:
+            return events
+
+        logger.info(
+            "Checking predict_state_config keys: %s",
+            list(self.predict_state_config.keys()) if self.predict_state_config else "None",
+        )
+        for state_key, config in self.predict_state_config.items():
+            logger.info(f"Checking state_key='{state_key}'")
+            if config["tool"] != content.name:
+                continue
+            tool_arg_name = config["tool_argument"]
+            logger.info(f"MATCHED tool '{content.name}' for state key '{state_key}', arg='{tool_arg_name}'")
+
+            state_value: Any
+            if tool_arg_name == "*":
+                state_value = parsed_args
+                logger.info(f"Using all args as state value, keys: {list(state_value.keys())}")
+            elif tool_arg_name in parsed_args:
+                state_value = parsed_args[tool_arg_name]
+                logger.info(f"Using specific arg '{tool_arg_name}' as state value")
+            else:
+                logger.warning(f"Tool argument '{tool_arg_name}' not found in parsed args")
+                continue
+
+            previous_value = self.last_emitted_state.get(state_key, object())
+            if previous_value == state_value:
+                logger.info(
+                    "Skipping duplicate StateDeltaEvent for key '%s' - value unchanged",
+                    state_key,
+                )
+                continue
+
+            state_delta_event = StateDeltaEvent(
+                delta=[
+                    {
+                        "op": "replace",
+                        "path": f"/{state_key}",
+                        "value": state_value,
+                    }
+                ],
+            )
+            logger.info(f"Emitting StateDeltaEvent for key '{state_key}', value type: {type(state_value)}")  # type: ignore
+            events.append(state_delta_event)
+            self.pending_state_updates[state_key] = state_value
+            self.last_emitted_state[state_key] = state_value
+        return events
+
    def _handle_function_result_content(self, content: FunctionResultContent) -> list[BaseEvent]:
        events: list[BaseEvent] = []
        if content.call_id:
@@ -294,6 +378,7 @@ class AgentFrameworkEventBridge:
            )
            logger.info(f"Emitting ToolCallEndEvent for completed tool call '{content.call_id}'")
            events.append(end_event)
+            self.tool_calls_ended.add(content.call_id)

            if self.state_delta_count > 0:
                logger.info(
@@ -306,7 +391,12 @@ class AgentFrameworkEventBridge:
            self.state_delta_count = 0

        result_message_id = generate_event_id()
-        result_content = prepare_function_call_results(content.result)
+        if isinstance(content.result, dict):
+            result_content = json.dumps(content.result)  # type: ignore[arg-type]
+        elif content.result is not None:
+            result_content = str(content.result)
+        else:
+            result_content = ""

        result_event = ToolCallResultEvent(
            message_id=result_message_id,
@@ -315,10 +405,55 @@ class AgentFrameworkEventBridge:
            role="tool",
        )
        events.append(result_event)
+
+        self.tool_results.append(
+            {
+                "id": result_message_id,
+                "role": "tool",
+                "toolCallId": content.call_id,
+                "content": result_content,
+            }
+        )
+
+        events.extend(self._emit_snapshot_for_tool_result())
        events.extend(self._emit_state_snapshot_and_confirmation())

        return events

+    def _emit_snapshot_for_tool_result(self) -> list[BaseEvent]:
+        events: list[BaseEvent] = []
+        should_emit_snapshot = self.pending_tool_calls and self.tool_results
+
+        is_predictive_without_confirmation = False
+        if should_emit_snapshot and self.current_tool_call_name and self.predict_state_config:
+            for _, config in self.predict_state_config.items():
+                if config["tool"] == self.current_tool_call_name and not self.require_confirmation:
+                    is_predictive_without_confirmation = True
+                    logger.info(
+                        "Skipping intermediate MessagesSnapshotEvent for predictive tool '%s' - delaying until summary",
+                        self.current_tool_call_name,
+                    )
+                    break
+
+        if should_emit_snapshot and not is_predictive_without_confirmation:
+            from ._message_adapters import agent_framework_messages_to_agui
+
+            assistant_message = {
+                "id": generate_event_id(),
+                "role": "assistant",
+                "tool_calls": self.pending_tool_calls.copy(),
+            }
+            converted_input_messages = agent_framework_messages_to_agui(self.input_messages)
+            all_messages = converted_input_messages + [assistant_message] + self.tool_results.copy()
+
+            messages_snapshot_event = MessagesSnapshotEvent(
+                type=EventType.MESSAGES_SNAPSHOT,
+                messages=all_messages,  # type: ignore[arg-type]
+            )
+            logger.info(f"Emitting MessagesSnapshotEvent with {len(all_messages)} messages")
+            events.append(messages_snapshot_event)
+        return events
+
    def _emit_state_snapshot_and_confirmation(self) -> list[BaseEvent]:
        events: list[BaseEvent] = []
        if self.pending_state_updates:
@@ -367,46 +502,31 @@ class AgentFrameworkEventBridge:
            self.current_tool_call_name = None
        return events

-    def _emit_confirm_changes_tool_call(self, function_call: FunctionCallContent | None = None) -> list[BaseEvent]:
-        """Emit a confirm_changes tool call for Dojo UI compatibility.
-
-        Args:
-            function_call: Optional function call that needs confirmation.
-                If provided, includes function info in the confirm_changes args
-                so Dojo UI can display what's being confirmed.
-        """
+    def _emit_confirm_changes_tool_call(self) -> list[BaseEvent]:
        events: list[BaseEvent] = []
        confirm_call_id = generate_event_id()
        logger.info("Emitting confirm_changes tool call for predictive update")

+        self.pending_tool_calls.append(
+            {
+                "id": confirm_call_id,
+                "type": "function",
+                "function": {
+                    "name": "confirm_changes",
+                    "arguments": "{}",
+                },
+            }
+        )
+
        confirm_start = ToolCallStartEvent(
            tool_call_id=confirm_call_id,
            tool_call_name="confirm_changes",
-            parent_message_id=self.current_message_id,
        )
        events.append(confirm_start)

-        # Include function info if this is for a function approval
-        # This helps Dojo UI display meaningful confirmation info
-        if function_call:
-            args_dict = {
-                "function_name": function_call.name,
-                "function_call_id": function_call.call_id,
-                "function_arguments": function_call.parse_arguments() or {},
-                "steps": [
-                    {
-                        "description": f"Execute {function_call.name}",
-                        "status": "enabled",
-                    }
-                ],
-            }
-            args_json = json.dumps(args_dict)
-        else:
-            args_json = "{}"
-
        confirm_args = ToolCallArgsEvent(
            tool_call_id=confirm_call_id,
-            delta=args_json,
+            delta="{}",
        )
        events.append(confirm_args)

@@ -415,48 +535,23 @@ class AgentFrameworkEventBridge:
        )
        events.append(confirm_end)

-        self.should_stop_after_confirm = True
-        logger.info("Set flag to stop run after confirm_changes")
-        return events
+        from ._message_adapters import agent_framework_messages_to_agui

-    def _emit_function_approval_tool_call(self, function_call: FunctionCallContent) -> list[BaseEvent]:
-        """Emit a tool call that can drive UI approval for function requests."""
-        tool_call_name = "confirm_changes"
-        if self.approval_tool_name and self.approval_tool_name != function_call.name:
-            tool_call_name = self.approval_tool_name
-
-        tool_call_id = generate_event_id()
-        tool_start = ToolCallStartEvent(
-            tool_call_id=tool_call_id,
-            tool_call_name=tool_call_name,
-            parent_message_id=self.current_message_id,
-        )
-        events: list[BaseEvent] = [tool_start]
-
-        args_dict = {
-            "function_name": function_call.name,
-            "function_call_id": function_call.call_id,
-            "function_arguments": function_call.parse_arguments() or {},
-            "steps": [
-                {
-                    "description": f"Execute {function_call.name}",
-                    "status": "enabled",
-                }
-            ],
+        assistant_message = {
+            "id": generate_event_id(),
+            "role": "assistant",
+            "tool_calls": self.pending_tool_calls.copy(),
        }
-        args_json = json.dumps(args_dict)

-        events.append(
-            ToolCallArgsEvent(
-                tool_call_id=tool_call_id,
-                delta=args_json,
-            )
-        )
-        events.append(
-            ToolCallEndEvent(
-                tool_call_id=tool_call_id,
-            )
+        converted_input_messages = agent_framework_messages_to_agui(self.input_messages)
+        all_messages = converted_input_messages + [assistant_message] + self.tool_results.copy()
+
+        messages_snapshot_event = MessagesSnapshotEvent(
+            type=EventType.MESSAGES_SNAPSHOT,
+            messages=all_messages,  # type: ignore[arg-type]
        )
+        logger.info(f"Emitting MessagesSnapshotEvent for confirm_changes with {len(all_messages)} messages")
+        events.append(messages_snapshot_event)

        self.should_stop_after_confirm = True
        logger.info("Set flag to stop run after confirm_changes")
@@ -488,8 +583,12 @@ class AgentFrameworkEventBridge:
                    tool_arg_name,
                )

-                state_value = extract_state_from_tool_args(parsed_args, tool_arg_name)
-                if state_value is None:
+                state_value: Any
+                if tool_arg_name == "*":
+                    state_value = parsed_args
+                elif tool_arg_name in parsed_args:
+                    state_value = parsed_args[tool_arg_name]
+                else:
                    logger.warning(f"  Tool argument '{tool_arg_name}' not found in parsed args")
                    continue

@@ -506,8 +605,8 @@ class AgentFrameworkEventBridge:
            )
            logger.info(f"Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'")
            events.append(end_event)
+            self.tool_calls_ended.add(content.function_call.call_id)

-        # Emit the function_approval_request custom event for UI implementations that support it
        approval_event = CustomEvent(
            name="function_approval_request",
            value={
@@ -521,14 +620,6 @@ class AgentFrameworkEventBridge:
        )
        logger.info(f"Emitting function_approval_request custom event for '{content.function_call.name}'")
        events.append(approval_event)
-
-        # Emit a UI-friendly approval tool call for function approvals.
-        if self.require_confirmation:
-            events.extend(self._emit_function_approval_tool_call(content.function_call))
-
-        # Signal orchestrator to stop the run and wait for user approval response
-        self.should_stop_after_confirm = True
-        logger.info("Set flag to stop run - waiting for function approval response")
        return events

    def create_run_started_event(self) -> RunStartedEvent:
@@ -2,8 +2,6 @@

 """Message format conversion between AG-UI and Agent Framework."""

-import json
-import logging
 from typing import Any, cast

 from agent_framework import (
@@ -13,229 +11,20 @@ from agent_framework import (
    FunctionResultContent,
    Role,
    TextContent,
-    prepare_function_call_results,
 )

-from ._utils import (
-    AGUI_TO_FRAMEWORK_ROLE,
-    FRAMEWORK_TO_AGUI_ROLE,
-    get_role_value,
-    normalize_agui_role,
-    safe_json_parse,
-)
+# Role mapping constants
+_AGUI_TO_FRAMEWORK_ROLE = {
+    "user": Role.USER,
+    "assistant": Role.ASSISTANT,
+    "system": Role.SYSTEM,
+}

-logger = logging.getLogger(__name__)
-
-
-def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
-    """Normalize tool ordering and inject synthetic results for AG-UI edge cases."""
-    sanitized: list[ChatMessage] = []
-    pending_tool_call_ids: set[str] | None = None
-    pending_confirm_changes_id: str | None = None
-
-    for msg in messages:
-        role_value = get_role_value(msg)
-
-        if role_value == "assistant":
-            tool_ids = {
-                str(content.call_id)
-                for content in msg.contents or []
-                if isinstance(content, FunctionCallContent) and content.call_id
-            }
-            confirm_changes_call = None
-            for content in msg.contents or []:
-                if isinstance(content, FunctionCallContent) and content.name == "confirm_changes":
-                    confirm_changes_call = content
-                    break
-
-            sanitized.append(msg)
-            pending_tool_call_ids = tool_ids if tool_ids else None
-            pending_confirm_changes_id = (
-                str(confirm_changes_call.call_id) if confirm_changes_call and confirm_changes_call.call_id else None
-            )
-            continue
-
-        if role_value == "user":
-            approval_call_ids: set[str] = set()
-            approval_accepted: bool | None = None
-            for content in msg.contents or []:
-                if type(content) is FunctionApprovalResponseContent:
-                    if content.function_call and content.function_call.call_id:
-                        approval_call_ids.add(str(content.function_call.call_id))
-                    if approval_accepted is None:
-                        approval_accepted = bool(content.approved)
-                    else:
-                        approval_accepted = approval_accepted and bool(content.approved)
-
-            if approval_call_ids and pending_tool_call_ids:
-                pending_tool_call_ids -= approval_call_ids
-                logger.info(
-                    f"FunctionApprovalResponseContent found for call_ids={sorted(approval_call_ids)} - "
-                    "framework will handle execution"
-                )
-
-            if pending_confirm_changes_id and approval_accepted is not None:
-                logger.info(f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}")
-                synthetic_result = ChatMessage(
-                    role="tool",
-                    contents=[
-                        FunctionResultContent(
-                            call_id=pending_confirm_changes_id,
-                            result="Confirmed" if approval_accepted else "Rejected",
-                        )
-                    ],
-                )
-                sanitized.append(synthetic_result)
-                if pending_tool_call_ids:
-                    pending_tool_call_ids.discard(pending_confirm_changes_id)
-                pending_confirm_changes_id = None
-
-            if pending_confirm_changes_id:
-                user_text = ""
-                for content in msg.contents or []:
-                    if isinstance(content, TextContent):
-                        user_text = content.text
-                        break
-
-                try:
-                    parsed = json.loads(user_text)
-                    if "accepted" in parsed:
-                        logger.info(
-                            f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}"
-                        )
-                        synthetic_result = ChatMessage(
-                            role="tool",
-                            contents=[
-                                FunctionResultContent(
-                                    call_id=pending_confirm_changes_id,
-                                    result="Confirmed" if parsed.get("accepted") else "Rejected",
-                                )
-                            ],
-                        )
-                        sanitized.append(synthetic_result)
-                        if pending_tool_call_ids:
-                            pending_tool_call_ids.discard(pending_confirm_changes_id)
-                        pending_confirm_changes_id = None
-                        continue
-                except (json.JSONDecodeError, KeyError) as exc:
-                    logger.debug(f"Could not parse user message as confirm_changes response: {type(exc).__name__}")
-
-            if pending_tool_call_ids:
-                logger.info(
-                    f"User message arrived with {len(pending_tool_call_ids)} pending tool calls - "
-                    "injecting synthetic results"
-                )
-                for pending_call_id in pending_tool_call_ids:
-                    logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}")
-                    synthetic_result = ChatMessage(
-                        role="tool",
-                        contents=[
-                            FunctionResultContent(
-                                call_id=pending_call_id,
-                                result="Tool execution skipped - user provided follow-up message",
-                            )
-                        ],
-                    )
-                    sanitized.append(synthetic_result)
-                pending_tool_call_ids = None
-                pending_confirm_changes_id = None
-
-            sanitized.append(msg)
-            pending_confirm_changes_id = None
-            continue
-
-        if role_value == "tool":
-            if not pending_tool_call_ids:
-                continue
-            keep = False
-            for content in msg.contents or []:
-                if isinstance(content, FunctionResultContent):
-                    call_id = str(content.call_id)
-                    if call_id in pending_tool_call_ids:
-                        keep = True
-                        if call_id == pending_confirm_changes_id:
-                            pending_confirm_changes_id = None
-                        break
-            if keep:
-                sanitized.append(msg)
-            continue
-
-        sanitized.append(msg)
-        pending_tool_call_ids = None
-        pending_confirm_changes_id = None
-
-    return sanitized
-
-
-def _deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
-    """Remove duplicate messages while preserving order."""
-    seen_keys: dict[Any, int] = {}
-    unique_messages: list[ChatMessage] = []
-
-    for idx, msg in enumerate(messages):
-        role_value = get_role_value(msg)
-
-        if role_value == "tool" and msg.contents and isinstance(msg.contents[0], FunctionResultContent):
-            call_id = str(msg.contents[0].call_id)
-            key: Any = (role_value, call_id)
-
-            if key in seen_keys:
-                existing_idx = seen_keys[key]
-                existing_msg = unique_messages[existing_idx]
-
-                existing_result = None
-                if existing_msg.contents and isinstance(existing_msg.contents[0], FunctionResultContent):
-                    existing_result = existing_msg.contents[0].result
-                new_result = msg.contents[0].result
-
-                if (not existing_result or existing_result == "") and new_result:
-                    logger.info(f"Replacing empty tool result at index {existing_idx} with data from index {idx}")
-                    unique_messages[existing_idx] = msg
-                else:
-                    logger.info(f"Skipping duplicate tool result at index {idx}: call_id={call_id}")
-                continue
-
-            seen_keys[key] = len(unique_messages)
-            unique_messages.append(msg)
-
-        elif (
-            role_value == "assistant" and msg.contents and any(isinstance(c, FunctionCallContent) for c in msg.contents)
-        ):
-            tool_call_ids = tuple(
-                sorted(str(c.call_id) for c in msg.contents if isinstance(c, FunctionCallContent) and c.call_id)
-            )
-            key = (role_value, tool_call_ids)
-
-            if key in seen_keys:
-                logger.info(f"Skipping duplicate assistant tool call at index {idx}")
-                continue
-
-            seen_keys[key] = len(unique_messages)
-            unique_messages.append(msg)
-
-        else:
-            content_str = str([str(c) for c in msg.contents]) if msg.contents else ""
-            key = (role_value, hash(content_str))
-
-            if key in seen_keys:
-                logger.info(f"Skipping duplicate message at index {idx}: role={role_value}")
-                continue
-
-            seen_keys[key] = len(unique_messages)
-            unique_messages.append(msg)
-
-    return unique_messages
-
-
-def normalize_agui_input_messages(
-    messages: list[dict[str, Any]],
-) -> tuple[list[ChatMessage], list[dict[str, Any]]]:
-    """Normalize raw AG-UI messages into provider and snapshot formats."""
-    provider_messages = agui_messages_to_agent_framework(messages)
-    provider_messages = _sanitize_tool_history(provider_messages)
-    provider_messages = _deduplicate_messages(provider_messages)
-    snapshot_messages = agui_messages_to_snapshot_format(messages)
-    return provider_messages, snapshot_messages
+_FRAMEWORK_TO_AGUI_ROLE = {
+    Role.USER: "user",
+    Role.ASSISTANT: "assistant",
+    Role.SYSTEM: "system",
+}


 def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[ChatMessage]:
@@ -247,108 +36,11 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
    Returns:
        List of Agent Framework ChatMessage objects
    """
-
-    def _update_tool_call_arguments(
-        raw_messages: list[dict[str, Any]],
-        tool_call_id: str,
-        modified_args: dict[str, Any],
-    ) -> None:
-        for raw_msg in raw_messages:
-            tool_calls = raw_msg.get("tool_calls") or raw_msg.get("toolCalls")
-            if not isinstance(tool_calls, list):
-                continue
-            tool_calls_list = cast(list[Any], tool_calls)
-            for tool_call in tool_calls_list:
-                if not isinstance(tool_call, dict):
-                    continue
-                tool_call_dict = cast(dict[str, Any], tool_call)
-                if str(tool_call_dict.get("id", "")) != tool_call_id:
-                    continue
-                function_payload = tool_call_dict.get("function")
-                if not isinstance(function_payload, dict):
-                    return
-                function_payload_dict = cast(dict[str, Any], function_payload)
-                existing_args = function_payload_dict.get("arguments")
-                if isinstance(existing_args, str):
-                    function_payload_dict["arguments"] = json.dumps(modified_args)
-                else:
-                    function_payload_dict["arguments"] = modified_args
-                return
-
-    def _find_matching_func_call(call_id: str) -> FunctionCallContent | None:
-        for prev_msg in result:
-            role_val = prev_msg.role.value if hasattr(prev_msg.role, "value") else str(prev_msg.role)
-            if role_val != "assistant":
-                continue
-            for content in prev_msg.contents or []:
-                if isinstance(content, FunctionCallContent):
-                    if content.call_id == call_id and content.name != "confirm_changes":
-                        return content
-        return None
-
-    def _parse_arguments(arguments: Any) -> dict[str, Any] | None:
-        return safe_json_parse(arguments)
-
-    def _resolve_approval_call_id(tool_call_id: str, parsed_payload: dict[str, Any] | None) -> str | None:
-        if parsed_payload:
-            explicit_call_id = parsed_payload.get("function_call_id")
-            if explicit_call_id:
-                return str(explicit_call_id)
-
-        for prev_msg in result:
-            role_val = prev_msg.role.value if hasattr(prev_msg.role, "value") else str(prev_msg.role)
-            if role_val != "assistant":
-                continue
-            direct_call = None
-            confirm_call = None
-            sibling_calls: list[FunctionCallContent] = []
-            for content in prev_msg.contents or []:
-                if not isinstance(content, FunctionCallContent):
-                    continue
-                if content.call_id == tool_call_id:
-                    direct_call = content
-                if content.name == "confirm_changes" and content.call_id == tool_call_id:
-                    confirm_call = content
-                elif content.name != "confirm_changes":
-                    sibling_calls.append(content)
-
-            if direct_call:
-                direct_args = direct_call.parse_arguments() or {}
-                if isinstance(direct_args, dict):
-                    explicit_call_id = direct_args.get("function_call_id")
-                    if explicit_call_id:
-                        return str(explicit_call_id)
-
-            if not confirm_call:
-                continue
-
-            confirm_args = confirm_call.parse_arguments() or {}
-            if isinstance(confirm_args, dict):
-                explicit_call_id = confirm_args.get("function_call_id")
-                if explicit_call_id:
-                    return str(explicit_call_id)
-
-            if len(sibling_calls) == 1 and sibling_calls[0].call_id:
-                return str(sibling_calls[0].call_id)
-
-        return None
-
-    def _filter_modified_args(
-        modified_args: dict[str, Any],
-        original_args: dict[str, Any] | None,
-    ) -> dict[str, Any]:
-        if not modified_args:
-            return {}
-        if not isinstance(original_args, dict) or not original_args:
-            return {}
-        allowed_keys = set(original_args.keys())
-        return {key: value for key, value in modified_args.items() if key in allowed_keys}
-
    result: list[ChatMessage] = []
    for msg in messages:
        # Handle standard tool result messages early (role="tool") to preserve provider invariants
        # This path maps AG‑UI tool messages to FunctionResultContent with the correct tool_call_id
-        role_str = normalize_agui_role(msg.get("role", "user"))
+        role_str = msg.get("role", "user")
        if role_str == "tool":
            # Prefer explicit tool_call_id fields; fall back to backend fields only if necessary
            tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
@@ -365,153 +57,31 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
                result_content = msg.get("result", "")

            # Distinguish approval payloads from actual tool results
-            parsed: dict[str, Any] | None = None
+            is_approval = False
            if isinstance(result_content, str) and result_content:
-                try:
-                    parsed_candidate = json.loads(result_content)
-                except Exception:
-                    parsed_candidate = None
-                if isinstance(parsed_candidate, dict):
-                    parsed = cast(dict[str, Any], parsed_candidate)
-            elif isinstance(result_content, dict):
-                parsed = cast(dict[str, Any], result_content)
+                import json as _json

-            is_approval = parsed is not None and "accepted" in parsed
+                try:
+                    parsed = _json.loads(result_content)
+                    is_approval = isinstance(parsed, dict) and "accepted" in parsed
+                except Exception:
+                    is_approval = False

            if is_approval:
-                # Look for the matching function call in previous messages to create
-                # a proper FunctionApprovalResponseContent. This enables the agent framework
-                # to execute the approved tool (fix for GitHub issue #3034).
-                accepted = parsed.get("accepted", False) if parsed is not None else False
-                approval_payload_text = result_content if isinstance(result_content, str) else json.dumps(parsed)
-
-                # Log the full approval payload to debug modified arguments
-                import logging
-
-                logger = logging.getLogger(__name__)
-                logger.info(f"Approval payload received: {parsed}")
-
-                approval_call_id = tool_call_id
-                resolved_call_id = _resolve_approval_call_id(tool_call_id, parsed)
-                if resolved_call_id:
-                    approval_call_id = resolved_call_id
-                matching_func_call = _find_matching_func_call(approval_call_id)
-
-                if matching_func_call:
-                    # Remove any existing tool result for this call_id since the framework
-                    # will re-execute the tool after approval. Keeping old results causes
-                    # OpenAI API errors ("tool message must follow assistant with tool_calls").
-                    result = [
-                        m
-                        for m in result
-                        if not (
-                            (m.role.value if hasattr(m.role, "value") else str(m.role)) == "tool"
-                            and any(
-                                isinstance(c, FunctionResultContent) and c.call_id == approval_call_id
-                                for c in (m.contents or [])
-                            )
-                        )
-                    ]
-
-                    # Check if the approval payload contains modified arguments
-                    # The UI sends back the modified state (e.g., deselected steps) in the approval payload
-                    modified_args = {k: v for k, v in parsed.items() if k != "accepted"} if parsed else {}
-                    original_args = matching_func_call.parse_arguments()
-                    filtered_args = _filter_modified_args(modified_args, original_args)
-                    state_args: dict[str, Any] | None = None
-                    if filtered_args:
-                        original_args = original_args or {}
-                        merged_args: dict[str, Any]
-                        if isinstance(original_args, dict) and original_args:
-                            merged_args = {**original_args, **filtered_args}
-                        else:
-                            merged_args = dict(filtered_args)
-
-                        if isinstance(filtered_args.get("steps"), list):
-                            original_steps = original_args.get("steps") if isinstance(original_args, dict) else None
-                            if isinstance(original_steps, list):
-                                approved_steps_list = list(filtered_args.get("steps") or [])
-                                approved_by_description: dict[str, dict[str, Any]] = {}
-                                for step_item in approved_steps_list:
-                                    if isinstance(step_item, dict):
-                                        step_item_dict = cast(dict[str, Any], step_item)
-                                        desc = step_item_dict.get("description")
-                                        if desc:
-                                            approved_by_description[str(desc)] = step_item_dict
-                                merged_steps: list[Any] = []
-                                original_steps_list = cast(list[Any], original_steps)
-                                for orig_step in original_steps_list:
-                                    if not isinstance(orig_step, dict):
-                                        merged_steps.append(orig_step)
-                                        continue
-                                    orig_step_dict = cast(dict[str, Any], orig_step)
-                                    description = str(orig_step_dict.get("description", ""))
-                                    approved_step = approved_by_description.get(description)
-                                    status: str = (
-                                        str(approved_step.get("status"))
-                                        if approved_step is not None and approved_step.get("status")
-                                        else "disabled"
-                                    )
-                                    updated_step: dict[str, Any] = orig_step_dict.copy()
-                                    updated_step["status"] = status
-                                    merged_steps.append(updated_step)
-                                merged_args["steps"] = merged_steps
-                        state_args = merged_args
-
-                        # Keep the original tool call and AG-UI snapshot in sync with approved args.
-                        updated_args = (
-                            json.dumps(merged_args) if isinstance(matching_func_call.arguments, str) else merged_args
-                        )
-                        matching_func_call.arguments = updated_args
-                        _update_tool_call_arguments(messages, str(approval_call_id), merged_args)
-                        # Create a new FunctionCallContent with the modified arguments
-                        func_call_for_approval = FunctionCallContent(
-                            call_id=matching_func_call.call_id,
-                            name=matching_func_call.name,
-                            arguments=json.dumps(filtered_args),
-                        )
-                        logger.info(f"Using modified arguments from approval: {filtered_args}")
-                    else:
-                        # No modified arguments - use the original function call
-                        func_call_for_approval = matching_func_call
-
-                    # Create FunctionApprovalResponseContent for the agent framework
-                    approval_response = FunctionApprovalResponseContent(
-                        approved=accepted,
-                        id=str(approval_call_id),
-                        function_call=func_call_for_approval,
-                        additional_properties={"ag_ui_state_args": state_args} if state_args else None,
-                    )
-                    chat_msg = ChatMessage(
-                        role=Role.USER,
-                        contents=[approval_response],
-                    )
-                else:
-                    # No matching function call found - this is likely a confirm_changes approval
-                    # Keep the old behavior for backwards compatibility
-                    chat_msg = ChatMessage(
-                        role=Role.USER,
-                        contents=[TextContent(text=approval_payload_text)],
-                        additional_properties={"is_tool_result": True, "tool_call_id": str(tool_call_id or "")},
-                    )
+                # Approval responses should be treated as user messages to trigger human-in-the-loop flow
+                chat_msg = ChatMessage(
+                    role=Role.USER,
+                    contents=[TextContent(text=str(result_content))],
+                    additional_properties={"is_tool_result": True, "tool_call_id": str(tool_call_id or "")},
+                )
                if "id" in msg:
                    chat_msg.message_id = msg["id"]
                result.append(chat_msg)
                continue

-            # Cast result_content to acceptable type for FunctionResultContent
-            func_result: str | dict[str, Any] | list[Any]
-            if isinstance(result_content, str):
-                func_result = result_content
-            elif isinstance(result_content, dict):
-                func_result = cast(dict[str, Any], result_content)
-            elif isinstance(result_content, list):
-                func_result = cast(list[Any], result_content)
-            else:
-                func_result = str(result_content)
            chat_msg = ChatMessage(
                role=Role.TOOL,
-                contents=[FunctionResultContent(call_id=str(tool_call_id), result=func_result)],
+                contents=[FunctionResultContent(call_id=str(tool_call_id), result=result_content)],
            )
            if "id" in msg:
                chat_msg.message_id = msg["id"]
@@ -572,7 +142,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha

        # No special handling required for assistant/plain messages here

-        role = AGUI_TO_FRAMEWORK_ROLE.get(role_str, Role.USER)
+        role = _AGUI_TO_FRAMEWORK_ROLE.get(role_str, Role.USER)

        # Check if this message contains function approvals
        if "function_approvals" in msg and msg["function_approvals"]:
@@ -628,7 +198,6 @@ def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str
        if isinstance(msg, dict):
            # Always work on a copy to avoid mutating input
            normalized_msg = msg.copy()
-            normalized_msg["role"] = normalize_agui_role(normalized_msg.get("role"))
            # Ensure ID exists
            if "id" not in normalized_msg:
                normalized_msg["id"] = generate_event_id()
@@ -645,7 +214,7 @@ def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str
            continue

        # Convert ChatMessage to AG-UI format
-        role = FRAMEWORK_TO_AGUI_ROLE.get(msg.role, "user")
+        role = _FRAMEWORK_TO_AGUI_ROLE.get(msg.role, "user")

        content_text = ""
        tool_calls: list[dict[str, Any]] = []
@@ -668,8 +237,13 @@ def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str
            elif isinstance(content, FunctionResultContent):
                # Tool result content - extract call_id and result
                tool_result_call_id = content.call_id
-                # Serialize result to string using core utility
-                content_text = prepare_function_call_results(content.result)
+                # Serialize result to string
+                if isinstance(content.result, dict):
+                    import json
+
+                    content_text = json.dumps(content.result)  # type: ignore
+                elif content.result is not None:
+                    content_text = str(content.result)

        agui_msg: dict[str, Any] = {
            "id": msg.message_id if msg.message_id else generate_event_id(),  # Always include id
@@ -734,44 +308,22 @@ def agui_messages_to_snapshot_format(messages: list[dict[str, Any]]) -> list[dic
        content = normalized_msg.get("content")
        if isinstance(content, list):
            # Convert content array format to simple string
-            text_parts: list[str] = []
-            content_list = cast(list[Any], content)
-            for item in content_list:
+            text_parts = []
+            for item in content:
                if isinstance(item, dict):
-                    item_dict = cast(dict[str, Any], item)
                    # Convert 'input_text' to 'text' type
-                    if item_dict.get("type") == "input_text":
-                        text_parts.append(str(item_dict.get("text", "")))
-                    elif item_dict.get("type") == "text":
-                        text_parts.append(str(item_dict.get("text", "")))
+                    if item.get("type") == "input_text":
+                        text_parts.append(item.get("text", ""))
+                    elif item.get("type") == "text":
+                        text_parts.append(item.get("text", ""))
                    else:
                        # Other types - just extract text field if present
-                        text_parts.append(str(item_dict.get("text", "")))
+                        text_parts.append(item.get("text", ""))
            normalized_msg["content"] = "".join(text_parts)
        elif content is None:
            normalized_msg["content"] = ""

-        tool_calls = normalized_msg.get("tool_calls") or normalized_msg.get("toolCalls")
-        if isinstance(tool_calls, list):
-            tool_calls_list = cast(list[Any], tool_calls)
-            for tool_call in tool_calls_list:
-                if not isinstance(tool_call, dict):
-                    continue
-                tool_call_dict = cast(dict[str, Any], tool_call)
-                function_payload = tool_call_dict.get("function")
-                if not isinstance(function_payload, dict):
-                    continue
-                function_payload_dict = cast(dict[str, Any], function_payload)
-                if "arguments" not in function_payload_dict:
-                    continue
-                arguments = function_payload_dict.get("arguments")
-                if arguments is None:
-                    function_payload_dict["arguments"] = ""
-                elif not isinstance(arguments, str):
-                    function_payload_dict["arguments"] = json.dumps(arguments)
-
        # Normalize tool_call_id to toolCallId for tool messages
-        normalized_msg["role"] = normalize_agui_role(normalized_msg.get("role"))
        if normalized_msg.get("role") == "tool":
            if "tool_call_id" in normalized_msg:
                normalized_msg["toolCallId"] = normalized_msg["tool_call_id"]
@@ -1,391 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Helper functions for orchestration logic."""
-
-import json
-import logging
-from typing import TYPE_CHECKING, Any
-
-from ag_ui.core import StateSnapshotEvent
-from agent_framework import (
-    ChatMessage,
-    FunctionApprovalResponseContent,
-    FunctionCallContent,
-    FunctionResultContent,
-    TextContent,
-)
-
-from .._utils import get_role_value, safe_json_parse
-
-if TYPE_CHECKING:
-    from .._events import AgentFrameworkEventBridge
-    from ._state_manager import StateManager
-
-logger = logging.getLogger(__name__)
-
-
-def pending_tool_call_ids(messages: list[ChatMessage]) -> set[str]:
-    """Get IDs of tool calls without corresponding results.
-
-    Args:
-        messages: List of messages to scan
-
-    Returns:
-        Set of pending tool call IDs
-    """
-    pending_ids: set[str] = set()
-    resolved_ids: set[str] = set()
-    for msg in messages:
-        for content in msg.contents:
-            if isinstance(content, FunctionCallContent) and content.call_id:
-                pending_ids.add(str(content.call_id))
-            elif isinstance(content, FunctionResultContent) and content.call_id:
-                resolved_ids.add(str(content.call_id))
-    return pending_ids - resolved_ids
-
-
-def is_state_context_message(message: ChatMessage) -> bool:
-    """Check if a message is a state context system message.
-
-    Args:
-        message: Message to check
-
-    Returns:
-        True if this is a state context message
-    """
-    if get_role_value(message) != "system":
-        return False
-    for content in message.contents:
-        if isinstance(content, TextContent) and content.text.startswith("Current state of the application:"):
-            return True
-    return False
-
-
-def ensure_tool_call_entry(
-    tool_call_id: str,
-    tool_calls_by_id: dict[str, dict[str, Any]],
-    pending_tool_calls: list[dict[str, Any]],
-) -> dict[str, Any]:
-    """Get or create a tool call entry in the tracking dicts.
-
-    Args:
-        tool_call_id: The tool call ID
-        tool_calls_by_id: Dict mapping IDs to tool call entries
-        pending_tool_calls: List of pending tool calls
-
-    Returns:
-        The tool call entry dict
-    """
-    entry = tool_calls_by_id.get(tool_call_id)
-    if entry is None:
-        entry = {
-            "id": tool_call_id,
-            "type": "function",
-            "function": {
-                "name": "",
-                "arguments": "",
-            },
-        }
-        tool_calls_by_id[tool_call_id] = entry
-        pending_tool_calls.append(entry)
-    return entry
-
-
-def tool_name_for_call_id(
-    tool_calls_by_id: dict[str, dict[str, Any]],
-    tool_call_id: str,
-) -> str | None:
-    """Get the tool name for a given call ID.
-
-    Args:
-        tool_calls_by_id: Dict mapping IDs to tool call entries
-        tool_call_id: The tool call ID to look up
-
-    Returns:
-        Tool name or None if not found
-    """
-    entry = tool_calls_by_id.get(tool_call_id)
-    if not entry:
-        return None
-    function = entry.get("function")
-    if not isinstance(function, dict):
-        return None
-    name = function.get("name")
-    return str(name) if name else None
-
-
-def tool_calls_match_state(
-    provider_messages: list[ChatMessage],
-    state_manager: "StateManager",
-) -> bool:
-    """Check if tool calls in messages match current state.
-
-    Args:
-        provider_messages: Messages to check
-        state_manager: State manager with config and current state
-
-    Returns:
-        True if tool calls match state configuration
-    """
-    if not state_manager.predict_state_config or not state_manager.current_state:
-        return False
-
-    for state_key, config in state_manager.predict_state_config.items():
-        tool_name = config["tool"]
-        tool_arg_name = config["tool_argument"]
-        tool_args: dict[str, Any] | None = None
-
-        for msg in reversed(provider_messages):
-            if get_role_value(msg) != "assistant":
-                continue
-            for content in msg.contents:
-                if isinstance(content, FunctionCallContent) and content.name == tool_name:
-                    tool_args = safe_json_parse(content.arguments)
-                    break
-            if tool_args is not None:
-                break
-
-        if not tool_args:
-            return False
-
-        if tool_arg_name == "*":
-            state_value = tool_args
-        elif tool_arg_name in tool_args:
-            state_value = tool_args[tool_arg_name]
-        else:
-            return False
-
-        if state_manager.current_state.get(state_key) != state_value:
-            return False
-
-    return True
-
-
-def schema_has_steps(schema: Any) -> bool:
-    """Check if a schema has a steps array property.
-
-    Args:
-        schema: JSON schema to check
-
-    Returns:
-        True if schema has steps array
-    """
-    if not isinstance(schema, dict):
-        return False
-    properties = schema.get("properties")
-    if not isinstance(properties, dict):
-        return False
-    steps_schema = properties.get("steps")
-    if not isinstance(steps_schema, dict):
-        return False
-    return steps_schema.get("type") == "array"
-
-
-def select_approval_tool_name(client_tools: list[Any] | None) -> str | None:
-    """Select appropriate approval tool from client tools.
-
-    Args:
-        client_tools: List of client tool definitions
-
-    Returns:
-        Name of approval tool, or None if not found
-    """
-    if not client_tools:
-        return None
-    for tool in client_tools:
-        tool_name = getattr(tool, "name", None)
-        if not tool_name:
-            continue
-        params_fn = getattr(tool, "parameters", None)
-        if not callable(params_fn):
-            continue
-        schema = params_fn()
-        if schema_has_steps(schema):
-            return str(tool_name)
-    return None
-
-
-def select_messages_to_run(
-    provider_messages: list[ChatMessage],
-    state_manager: "StateManager",
-) -> list[ChatMessage]:
-    """Select and prepare messages for agent execution.
-
-    Injects state context message when appropriate.
-
-    Args:
-        provider_messages: Original messages from client
-        state_manager: State manager instance
-
-    Returns:
-        Messages ready for agent execution
-    """
-    if not provider_messages:
-        return []
-
-    is_new_user_turn = get_role_value(provider_messages[-1]) == "user"
-    conversation_has_tool_calls = tool_calls_match_state(provider_messages, state_manager)
-    state_context_msg = state_manager.state_context_message(
-        is_new_user_turn=is_new_user_turn, conversation_has_tool_calls=conversation_has_tool_calls
-    )
-    if not state_context_msg:
-        return list(provider_messages)
-
-    messages_to_run = [msg for msg in provider_messages if not is_state_context_message(msg)]
-    if pending_tool_call_ids(messages_to_run):
-        return messages_to_run
-
-    insert_index = len(messages_to_run) - 1 if is_new_user_turn else len(messages_to_run)
-    if insert_index < 0:
-        insert_index = 0
-    messages_to_run.insert(insert_index, state_context_msg)
-    return messages_to_run
-
-
-def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any]:
-    """Build metadata dict with truncated string values.
-
-    Args:
-        thread_metadata: Raw metadata dict
-
-    Returns:
-        Metadata with string values truncated to 512 chars
-    """
-    if not thread_metadata:
-        return {}
-    safe_metadata: dict[str, Any] = {}
-    for key, value in thread_metadata.items():
-        value_str = value if isinstance(value, str) else json.dumps(value)
-        if len(value_str) > 512:
-            value_str = value_str[:512]
-        safe_metadata[key] = value_str
-    return safe_metadata
-
-
-def collect_approved_state_snapshots(
-    provider_messages: list[ChatMessage],
-    predict_state_config: dict[str, dict[str, str]] | None,
-    current_state: dict[str, Any],
-    event_bridge: "AgentFrameworkEventBridge",
-) -> list[StateSnapshotEvent]:
-    """Collect state snapshots from approved function calls.
-
-    Args:
-        provider_messages: Messages containing approvals
-        predict_state_config: Predictive state configuration
-        current_state: Current state dict (will be mutated)
-        event_bridge: Event bridge for creating events
-
-    Returns:
-        List of state snapshot events
-    """
-    if not predict_state_config:
-        return []
-
-    events: list[StateSnapshotEvent] = []
-    for msg in provider_messages:
-        if get_role_value(msg) != "user":
-            continue
-        for content in msg.contents:
-            if type(content) is FunctionApprovalResponseContent:
-                if not content.function_call or not content.approved:
-                    continue
-                parsed_args = content.function_call.parse_arguments()
-                state_args = None
-                if content.additional_properties:
-                    state_args = content.additional_properties.get("ag_ui_state_args")
-                if not isinstance(state_args, dict):
-                    state_args = parsed_args
-                if not state_args:
-                    continue
-                for state_key, config in predict_state_config.items():
-                    if config["tool"] != content.function_call.name:
-                        continue
-                    tool_arg_name = config["tool_argument"]
-                    if tool_arg_name == "*":
-                        state_value = state_args
-                    elif isinstance(state_args, dict) and tool_arg_name in state_args:
-                        state_value = state_args[tool_arg_name]
-                    else:
-                        continue
-                    current_state[state_key] = state_value
-                    event_bridge.current_state[state_key] = state_value
-                    logger.info(
-                        f"Emitting StateSnapshotEvent for approved state key '{state_key}' "
-                        f"with {len(state_value) if isinstance(state_value, list) else 'N/A'} items"
-                    )
-                    events.append(StateSnapshotEvent(snapshot=current_state))
-                    break
-    return events
-
-
-def latest_approval_response(messages: list[ChatMessage]) -> FunctionApprovalResponseContent | None:
-    """Get the latest approval response from messages.
-
-    Args:
-        messages: Messages to search
-
-    Returns:
-        Latest approval response or None
-    """
-    if not messages:
-        return None
-    last_message = messages[-1]
-    for content in last_message.contents:
-        if type(content) is FunctionApprovalResponseContent:
-            return content
-    return None
-
-
-def approval_steps(approval: FunctionApprovalResponseContent) -> list[Any]:
-    """Extract steps from an approval response.
-
-    Args:
-        approval: Approval response content
-
-    Returns:
-        List of steps, or empty list if none
-    """
-    state_args: Any | None = None
-    if approval.additional_properties:
-        state_args = approval.additional_properties.get("ag_ui_state_args")
-    if isinstance(state_args, dict):
-        steps = state_args.get("steps")
-        if isinstance(steps, list):
-            return steps
-
-    if approval.function_call:
-        parsed_args = approval.function_call.parse_arguments()
-        if isinstance(parsed_args, dict):
-            steps = parsed_args.get("steps")
-            if isinstance(steps, list):
-                return steps
-
-    return []
-
-
-def is_step_based_approval(
-    approval: FunctionApprovalResponseContent,
-    predict_state_config: dict[str, dict[str, str]] | None,
-) -> bool:
-    """Check if an approval is step-based.
-
-    Args:
-        approval: Approval response to check
-        predict_state_config: Predictive state configuration
-
-    Returns:
-        True if this is a step-based approval
-    """
-    steps = approval_steps(approval)
-    if steps:
-        return True
-    if not approval.function_call:
-        return False
-    if not predict_state_config:
-        return False
-    tool_name = approval.function_call.name
-    for config in predict_state_config.values():
-        if config.get("tool") == tool_name and config.get("tool_argument") == "steps":
-            return True
-    return False
@@ -0,0 +1,176 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Message hygiene utilities for orchestrators."""
+
+import json
+import logging
+from typing import Any
+
+from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent
+
+logger = logging.getLogger(__name__)
+
+
+def sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
+    """Normalize tool ordering and inject synthetic results for AG-UI edge cases."""
+    sanitized: list[ChatMessage] = []
+    pending_tool_call_ids: set[str] | None = None
+    pending_confirm_changes_id: str | None = None
+
+    for msg in messages:
+        role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+
+        if role_value == "assistant":
+            tool_ids = {
+                str(content.call_id)
+                for content in msg.contents or []
+                if isinstance(content, FunctionCallContent) and content.call_id
+            }
+            confirm_changes_call = None
+            for content in msg.contents or []:
+                if isinstance(content, FunctionCallContent) and content.name == "confirm_changes":
+                    confirm_changes_call = content
+                    break
+
+            sanitized.append(msg)
+            pending_tool_call_ids = tool_ids if tool_ids else None
+            pending_confirm_changes_id = (
+                str(confirm_changes_call.call_id) if confirm_changes_call and confirm_changes_call.call_id else None
+            )
+            continue
+
+        if role_value == "user":
+            if pending_confirm_changes_id:
+                user_text = ""
+                for content in msg.contents or []:
+                    if isinstance(content, TextContent):
+                        user_text = content.text
+                        break
+
+                try:
+                    parsed = json.loads(user_text)
+                    if "accepted" in parsed:
+                        logger.info(
+                            f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}"
+                        )
+                        synthetic_result = ChatMessage(
+                            role="tool",
+                            contents=[
+                                FunctionResultContent(
+                                    call_id=pending_confirm_changes_id,
+                                    result="Confirmed" if parsed.get("accepted") else "Rejected",
+                                )
+                            ],
+                        )
+                        sanitized.append(synthetic_result)
+                        if pending_tool_call_ids:
+                            pending_tool_call_ids.discard(pending_confirm_changes_id)
+                        pending_confirm_changes_id = None
+                        continue
+                except (json.JSONDecodeError, KeyError) as exc:
+                    logger.debug("Could not parse user message as confirm_changes response: %s", type(exc).__name__)
+
+            if pending_tool_call_ids:
+                logger.info(
+                    f"User message arrived with {len(pending_tool_call_ids)} pending tool calls - injecting synthetic results"
+                )
+                for pending_call_id in pending_tool_call_ids:
+                    logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}")
+                    synthetic_result = ChatMessage(
+                        role="tool",
+                        contents=[
+                            FunctionResultContent(
+                                call_id=pending_call_id,
+                                result="Tool execution skipped - user provided follow-up message",
+                            )
+                        ],
+                    )
+                    sanitized.append(synthetic_result)
+                pending_tool_call_ids = None
+                pending_confirm_changes_id = None
+
+            sanitized.append(msg)
+            pending_confirm_changes_id = None
+            continue
+
+        if role_value == "tool":
+            if not pending_tool_call_ids:
+                continue
+            keep = False
+            for content in msg.contents or []:
+                if isinstance(content, FunctionResultContent):
+                    call_id = str(content.call_id)
+                    if call_id in pending_tool_call_ids:
+                        keep = True
+                        if call_id == pending_confirm_changes_id:
+                            pending_confirm_changes_id = None
+                        break
+            if keep:
+                sanitized.append(msg)
+            continue
+
+        sanitized.append(msg)
+        pending_tool_call_ids = None
+        pending_confirm_changes_id = None
+
+    return sanitized
+
+
+def deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
+    """Remove duplicate messages while preserving order."""
+    seen_keys: dict[Any, int] = {}
+    unique_messages: list[ChatMessage] = []
+
+    for idx, msg in enumerate(messages):
+        role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+
+        if role_value == "tool" and msg.contents and isinstance(msg.contents[0], FunctionResultContent):
+            call_id = str(msg.contents[0].call_id)
+            key: Any = (role_value, call_id)
+
+            if key in seen_keys:
+                existing_idx = seen_keys[key]
+                existing_msg = unique_messages[existing_idx]
+
+                existing_result = None
+                if existing_msg.contents and isinstance(existing_msg.contents[0], FunctionResultContent):
+                    existing_result = existing_msg.contents[0].result
+                new_result = msg.contents[0].result
+
+                if (not existing_result or existing_result == "") and new_result:
+                    logger.info(f"Replacing empty tool result at index {existing_idx} with data from index {idx}")
+                    unique_messages[existing_idx] = msg
+                else:
+                    logger.info(f"Skipping duplicate tool result at index {idx}: call_id={call_id}")
+                continue
+
+            seen_keys[key] = len(unique_messages)
+            unique_messages.append(msg)
+
+        elif (
+            role_value == "assistant" and msg.contents and any(isinstance(c, FunctionCallContent) for c in msg.contents)
+        ):
+            tool_call_ids = tuple(
+                sorted(str(c.call_id) for c in msg.contents if isinstance(c, FunctionCallContent) and c.call_id)
+            )
+            key = (role_value, tool_call_ids)
+
+            if key in seen_keys:
+                logger.info(f"Skipping duplicate assistant tool call at index {idx}")
+                continue
+
+            seen_keys[key] = len(unique_messages)
+            unique_messages.append(msg)
+
+        else:
+            content_str = str([str(c) for c in msg.contents]) if msg.contents else ""
+            key = (role_value, hash(content_str))
+
+            if key in seen_keys:
+                logger.info(f"Skipping duplicate message at index {idx}: role={role_value}")
+                continue
+
+            seen_keys[key] = len(unique_messages)
+            unique_messages.append(msg)
+
+    return unique_messages
@@ -1,230 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Predictive state handling utilities."""
-
-import json
-import logging
-import re
-from typing import Any
-
-from ag_ui.core import StateDeltaEvent
-
-from .._utils import safe_json_parse
-
-logger = logging.getLogger(__name__)
-
-
-class PredictiveStateHandler:
-    """Handles predictive state updates from streaming tool calls."""
-
-    def __init__(
-        self,
-        predict_state_config: dict[str, dict[str, str]] | None = None,
-        current_state: dict[str, Any] | None = None,
-    ) -> None:
-        """Initialize the handler.
-
-        Args:
-            predict_state_config: Configuration mapping state keys to tool/argument pairs
-            current_state: Reference to current state dict
-        """
-        self.predict_state_config = predict_state_config or {}
-        self.current_state = current_state or {}
-        self.streaming_tool_args: str = ""
-        self.last_emitted_state: dict[str, Any] = {}
-        self.state_delta_count: int = 0
-        self.pending_state_updates: dict[str, Any] = {}
-
-    def reset_streaming(self) -> None:
-        """Reset streaming state for a new tool call."""
-        self.streaming_tool_args = ""
-        self.state_delta_count = 0
-
-    def extract_state_value(
-        self,
-        tool_name: str,
-        args: dict[str, Any] | str | None,
-    ) -> tuple[str, Any] | None:
-        """Extract state value from tool arguments based on config.
-
-        Args:
-            tool_name: Name of the tool being called
-            args: Tool arguments (dict or JSON string)
-
-        Returns:
-            Tuple of (state_key, state_value) or None if no match
-        """
-        if not self.predict_state_config:
-            return None
-
-        parsed_args = safe_json_parse(args) if isinstance(args, str) else args
-        if not parsed_args:
-            return None
-
-        for state_key, config in self.predict_state_config.items():
-            if config["tool"] != tool_name:
-                continue
-            tool_arg_name = config["tool_argument"]
-            if tool_arg_name == "*":
-                return (state_key, parsed_args)
-            if tool_arg_name in parsed_args:
-                return (state_key, parsed_args[tool_arg_name])
-
-        return None
-
-    def is_predictive_tool(self, tool_name: str | None) -> bool:
-        """Check if a tool is configured for predictive state.
-
-        Args:
-            tool_name: Name of the tool to check
-
-        Returns:
-            True if tool is in predictive state config
-        """
-        if not tool_name or not self.predict_state_config:
-            return False
-        for config in self.predict_state_config.values():
-            if config["tool"] == tool_name:
-                return True
-        return False
-
-    def emit_streaming_deltas(
-        self,
-        tool_name: str | None,
-        argument_chunk: str,
-    ) -> list[StateDeltaEvent]:
-        """Process streaming argument chunk and emit state deltas.
-
-        Args:
-            tool_name: Name of the current tool
-            argument_chunk: New chunk of JSON arguments
-
-        Returns:
-            List of state delta events to emit
-        """
-        events: list[StateDeltaEvent] = []
-        if not tool_name or not self.predict_state_config:
-            return events
-
-        self.streaming_tool_args += argument_chunk
-        logger.debug(
-            "Predictive state: accumulated %s chars for tool '%s'",
-            len(self.streaming_tool_args),
-            tool_name,
-        )
-
-        # Try to parse complete JSON first
-        parsed_args = None
-        try:
-            parsed_args = json.loads(self.streaming_tool_args)
-        except json.JSONDecodeError:
-            # Fall back to regex matching for partial JSON
-            events.extend(self._emit_partial_deltas(tool_name))
-
-        if parsed_args:
-            events.extend(self._emit_complete_deltas(tool_name, parsed_args))
-
-        return events
-
-    def _emit_partial_deltas(self, tool_name: str) -> list[StateDeltaEvent]:
-        """Emit deltas from partial JSON using regex matching.
-
-        Args:
-            tool_name: Name of the current tool
-
-        Returns:
-            List of state delta events
-        """
-        events: list[StateDeltaEvent] = []
-
-        for state_key, config in self.predict_state_config.items():
-            if config["tool"] != tool_name:
-                continue
-            tool_arg_name = config["tool_argument"]
-            pattern = rf'"{re.escape(tool_arg_name)}":\s*"([^"]*)'
-            match = re.search(pattern, self.streaming_tool_args)
-
-            if match:
-                partial_value = match.group(1).replace("\\n", "\n").replace('\\"', '"').replace("\\\\", "\\")
-
-                if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != partial_value:
-                    event = self._create_delta_event(state_key, partial_value)
-                    events.append(event)
-                    self.last_emitted_state[state_key] = partial_value
-                    self.pending_state_updates[state_key] = partial_value
-
-        return events
-
-    def _emit_complete_deltas(
-        self,
-        tool_name: str,
-        parsed_args: dict[str, Any],
-    ) -> list[StateDeltaEvent]:
-        """Emit deltas from complete parsed JSON.
-
-        Args:
-            tool_name: Name of the current tool
-            parsed_args: Fully parsed arguments dict
-
-        Returns:
-            List of state delta events
-        """
-        events: list[StateDeltaEvent] = []
-
-        for state_key, config in self.predict_state_config.items():
-            if config["tool"] != tool_name:
-                continue
-            tool_arg_name = config["tool_argument"]
-
-            if tool_arg_name == "*":
-                state_value = parsed_args
-            elif tool_arg_name in parsed_args:
-                state_value = parsed_args[tool_arg_name]
-            else:
-                continue
-
-            if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != state_value:
-                event = self._create_delta_event(state_key, state_value)
-                events.append(event)
-                self.last_emitted_state[state_key] = state_value
-                self.pending_state_updates[state_key] = state_value
-
-        return events
-
-    def _create_delta_event(self, state_key: str, value: Any) -> StateDeltaEvent:
-        """Create a state delta event with logging.
-
-        Args:
-            state_key: The state key being updated
-            value: The new value
-
-        Returns:
-            StateDeltaEvent instance
-        """
-        self.state_delta_count += 1
-        if self.state_delta_count % 10 == 1:
-            logger.info(
-                "StateDeltaEvent #%s for '%s': op=replace, path=/%s, value_length=%s",
-                self.state_delta_count,
-                state_key,
-                state_key,
-                len(str(value)),
-            )
-        elif self.state_delta_count % 100 == 0:
-            logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
-
-        return StateDeltaEvent(
-            delta=[
-                {
-                    "op": "replace",
-                    "path": f"/{state_key}",
-                    "value": value,
-                }
-            ],
-        )
-
-    def apply_pending_updates(self) -> None:
-        """Apply pending updates to current state and clear them."""
-        for key, value in self.pending_state_updates.items():
-            self.current_state[key] = value
-        self.pending_state_updates.clear()
@@ -22,11 +22,9 @@ class StateManager:
        self.predict_state_config = predict_state_config or {}
        self.require_confirmation = require_confirmation
        self.current_state: dict[str, Any] = {}
-        self._state_from_input: bool = False

    def initialize(self, initial_state: dict[str, Any] | None) -> dict[str, Any]:
        """Initialize state with schema defaults."""
-        self._state_from_input = initial_state is not None
        self.current_state = (initial_state or {}).copy()
        self._apply_schema_defaults()
        return self.current_state
@@ -62,9 +60,7 @@ class StateManager:
        """Inject state context only when starting a new user turn."""
        if not self.current_state or not self.state_schema:
            return None
-        if not is_new_user_turn:
-            return None
-        if conversation_has_tool_calls and not self._state_from_input:
+        if not is_new_user_turn or conversation_has_tool_calls:
            return None

        state_json = json.dumps(self.current_state, indent=2)
@@ -16,10 +16,6 @@ from ag_ui.core import (
    TextMessageContentEvent,
    TextMessageEndEvent,
    TextMessageStartEvent,
-    ToolCallArgsEvent,
-    ToolCallEndEvent,
-    ToolCallResultEvent,
-    ToolCallStartEvent,
 )
 from agent_framework import (
    AgentProtocol,
@@ -29,31 +25,8 @@ from agent_framework import (
    FunctionResultContent,
    TextContent,
 )
-from agent_framework._middleware import extract_and_merge_function_middleware
-from agent_framework._tools import (
-    FunctionInvocationConfiguration,
-    _collect_approval_responses,  # type: ignore
-    _replace_approval_contents_with_results,  # type: ignore
-    _try_execute_function_calls,  # type: ignore
-)

-from ._orchestration._helpers import (
-    approval_steps,
-    build_safe_metadata,
-    collect_approved_state_snapshots,
-    ensure_tool_call_entry,
-    is_step_based_approval,
-    latest_approval_response,
-    select_approval_tool_name,
-    select_messages_to_run,
-    tool_name_for_call_id,
-)
-from ._orchestration._tooling import (
-    collect_server_tools,
-    merge_tools,
-    register_additional_client_tools,
-)
-from ._utils import convert_agui_tools_to_agent_framework, generate_event_id, get_role_value
+from ._utils import convert_agui_tools_to_agent_framework, generate_event_id

 if TYPE_CHECKING:
    from ._agent import AgentConfig
@@ -88,7 +61,6 @@ class ExecutionContext:

        # Lazy-loaded properties
        self._messages = None
-        self._snapshot_messages = None
        self._last_message = None
        self._run_id: str | None = None
        self._thread_id: str | None = None
@@ -97,27 +69,12 @@ class ExecutionContext:
    def messages(self):
        """Get converted Agent Framework messages (lazy loaded)."""
        if self._messages is None:
-            from ._message_adapters import normalize_agui_input_messages
+            from ._message_adapters import agui_messages_to_agent_framework

            raw = self.input_data.get("messages", [])
-            if not isinstance(raw, list):
-                raw = []
-            self._messages, self._snapshot_messages = normalize_agui_input_messages(raw)
+            self._messages = agui_messages_to_agent_framework(raw)
        return self._messages

-    @property
-    def snapshot_messages(self) -> list[dict[str, Any]]:
-        """Get normalized AG-UI snapshot messages (lazy loaded)."""
-        if self._snapshot_messages is None:
-            if self._messages is None:
-                _ = self.messages
-            else:
-                from ._message_adapters import agent_framework_messages_to_agui, agui_messages_to_snapshot_format
-
-                raw_snapshot = agent_framework_messages_to_agui(self._messages)
-                self._snapshot_messages = agui_messages_to_snapshot_format(raw_snapshot)
-        return self._snapshot_messages or []
-
    @property
    def last_message(self):
        """Get the last message in the conversation (lazy loaded)."""
@@ -313,7 +270,14 @@ class DefaultOrchestrator(Orchestrator):
            AG-UI events
        """
        from ._events import AgentFrameworkEventBridge
+        from ._message_adapters import agui_messages_to_snapshot_format
+        from ._orchestration._message_hygiene import deduplicate_messages, sanitize_tool_history
        from ._orchestration._state_manager import StateManager
+        from ._orchestration._tooling import (
+            collect_server_tools,
+            merge_tools,
+            register_additional_client_tools,
+        )

        logger.info(f"Starting default agent run for thread_id={context.thread_id}, run_id={context.run_id}")

@@ -322,15 +286,12 @@ class DefaultOrchestrator(Orchestrator):
            response_format = context.agent.chat_options.response_format
        skip_text_content = response_format is not None

-        client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
-        approval_tool_name = select_approval_tool_name(client_tools)
-
        state_manager = StateManager(
            state_schema=context.config.state_schema,
            predict_state_config=context.config.predict_state_config,
            require_confirmation=context.config.require_confirmation,
        )
-        current_state = state_manager.initialize(context.input_data.get("state"))
+        current_state = state_manager.initialize(context.input_data.get("state", {}))

        event_bridge = AgentFrameworkEventBridge(
            run_id=context.run_id,
@@ -338,8 +299,8 @@ class DefaultOrchestrator(Orchestrator):
            predict_state_config=context.config.predict_state_config,
            current_state=current_state,
            skip_text_content=skip_text_content,
+            input_messages=context.input_data.get("messages", []),
            require_confirmation=context.config.require_confirmation,
-            approval_tool_name=approval_tool_name,
        )

        yield event_bridge.create_run_started_event()
@@ -360,18 +321,17 @@ class DefaultOrchestrator(Orchestrator):
        if current_state:
            thread.metadata["current_state"] = current_state  # type: ignore[attr-defined]

-        provider_messages = context.messages or []
-        snapshot_messages = context.snapshot_messages
-        if not provider_messages:
+        raw_messages = context.messages or []
+        if not raw_messages:
            logger.warning("No messages provided in AG-UI input")
            yield event_bridge.create_run_finished_event()
            return

-        logger.info(f"Received {len(provider_messages)} provider messages from client")
-        for i, msg in enumerate(provider_messages):
-            role = get_role_value(msg)
+        logger.info(f"Received {len(raw_messages)} raw messages from client")
+        for i, msg in enumerate(raw_messages):
+            role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
            msg_id = getattr(msg, "message_id", None)
-            logger.info(f"  Message {i}: role={role}, id={msg_id}")
+            logger.info(f"  Raw message {i}: role={role}, id={msg_id}")
            if hasattr(msg, "contents") and msg.contents:
                for j, content in enumerate(msg.contents):
                    content_type = type(content).__name__
@@ -394,26 +354,62 @@ class DefaultOrchestrator(Orchestrator):
                    else:
                        logger.debug(f"    Content {j}: {content_type}")

-        pending_tool_calls: list[dict[str, Any]] = []
-        tool_calls_by_id: dict[str, dict[str, Any]] = {}
-        tool_results: list[dict[str, Any]] = []
-        tool_calls_ended: set[str] = set()
-        messages_snapshot_emitted = False
-        accumulated_text_content = ""
-        active_message_id: str | None = None
+        sanitized_messages = sanitize_tool_history(raw_messages)
+        provider_messages = deduplicate_messages(sanitized_messages)

-        # Check for FunctionApprovalResponseContent and emit updated state snapshot
-        # This ensures the UI shows the approved state (e.g., 2 steps) not the original (3 steps)
-        for snapshot_evt in collect_approved_state_snapshots(
-            provider_messages,
-            context.config.predict_state_config,
-            current_state,
-            event_bridge,
-        ):
-            yield snapshot_evt
+        if not provider_messages:
+            logger.info("No provider-eligible messages after filtering; finishing run without invoking agent.")
+            yield event_bridge.create_run_finished_event()
+            return

-        messages_to_run = select_messages_to_run(provider_messages, state_manager)
+        logger.info(f"Processing {len(provider_messages)} provider messages after sanitization/deduplication")
+        for i, msg in enumerate(provider_messages):
+            role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+            logger.info(f"  Message {i}: role={role}")
+            if hasattr(msg, "contents") and msg.contents:
+                for j, content in enumerate(msg.contents):
+                    content_type = type(content).__name__
+                    if isinstance(content, TextContent):
+                        logger.info(f"    Content {j}: {content_type} - text_length={len(content.text)}")
+                    elif isinstance(content, FunctionCallContent):
+                        arg_length = len(str(content.arguments)) if content.arguments else 0
+                        logger.info("    Content %s: %s - %s args_length=%s", j, content_type, content.name, arg_length)
+                    elif isinstance(content, FunctionResultContent):
+                        result_preview = type(content.result).__name__ if content.result is not None else "None"
+                        logger.info(
+                            "    Content %s: %s - call_id=%s, result_type=%s",
+                            j,
+                            content_type,
+                            content.call_id,
+                            result_preview,
+                        )
+                    else:
+                        logger.info(f"    Content {j}: {content_type}")

+        messages_to_run: list[Any] = []
+        is_new_user_turn = False
+        if provider_messages:
+            last_msg = provider_messages[-1]
+            role_value = last_msg.role.value if hasattr(last_msg.role, "value") else str(last_msg.role)
+            is_new_user_turn = role_value == "user"
+
+        conversation_has_tool_calls = False
+        for msg in provider_messages:
+            role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+            if role_value == "assistant" and hasattr(msg, "contents") and msg.contents:
+                if any(isinstance(content, FunctionCallContent) for content in msg.contents):
+                    conversation_has_tool_calls = True
+                    break
+
+        state_context_msg = state_manager.state_context_message(
+            is_new_user_turn=is_new_user_turn, conversation_has_tool_calls=conversation_has_tool_calls
+        )
+        if state_context_msg:
+            messages_to_run.append(state_context_msg)
+
+        messages_to_run.extend(provider_messages)
+
+        client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
        logger.info(f"[TOOLS] Client sent {len(client_tools) if client_tools else 0} tools")
        if client_tools:
            for tool in client_tools:
@@ -425,11 +421,17 @@ class DefaultOrchestrator(Orchestrator):
        register_additional_client_tools(context.agent, client_tools)
        tools_param = merge_tools(server_tools, client_tools)

-        collect_updates = response_format is not None
-        all_updates: list[Any] | None = [] if collect_updates else None
+        all_updates: list[Any] = []
        update_count = 0
        # Prepare metadata for chat client (Azure requires string values)
-        safe_metadata = build_safe_metadata(getattr(thread, "metadata", None))
+        safe_metadata: dict[str, Any] = {}
+        thread_metadata = getattr(thread, "metadata", None)
+        if thread_metadata:
+            for key, value in thread_metadata.items():
+                value_str = value if isinstance(value, str) else json.dumps(value)
+                if len(value_str) > 512:
+                    value_str = value_str[:512]
+                safe_metadata[key] = value_str

        run_kwargs: dict[str, Any] = {
            "thread": thread,
@@ -439,200 +441,27 @@ class DefaultOrchestrator(Orchestrator):
        if safe_metadata:
            run_kwargs["store"] = True

-        async def _resolve_approval_responses(
-            messages: list[Any],
-            tools_for_execution: list[Any],
-        ) -> None:
-            fcc_todo = _collect_approval_responses(messages)
-            if not fcc_todo:
-                return
-
-            approved_responses = [resp for resp in fcc_todo.values() if resp.approved]
-            approved_function_results: list[Any] = []
-            if approved_responses and tools_for_execution:
-                chat_client = getattr(context.agent, "chat_client", None)
-                config = (
-                    getattr(chat_client, "function_invocation_configuration", None) or FunctionInvocationConfiguration()
-                )
-                middleware_pipeline = extract_and_merge_function_middleware(chat_client, run_kwargs)
-                try:
-                    results, _ = await _try_execute_function_calls(
-                        custom_args=run_kwargs,
-                        attempt_idx=0,
-                        function_calls=approved_responses,
-                        tools=tools_for_execution,
-                        middleware_pipeline=middleware_pipeline,
-                        config=config,
-                    )
-                    approved_function_results = list(results)
-                except Exception:
-                    logger.error("Failed to execute approved tool calls; injecting error results.")
-                    approved_function_results = []
-
-            normalized_results: list[FunctionResultContent] = []
-            for idx, approval in enumerate(approved_responses):
-                if idx < len(approved_function_results) and isinstance(
-                    approved_function_results[idx], FunctionResultContent
-                ):
-                    normalized_results.append(approved_function_results[idx])
-                    continue
-                call_id = approval.function_call.call_id or approval.id
-                normalized_results.append(
-                    FunctionResultContent(call_id=call_id, result="Error: Tool call invocation failed.")
-                )
-
-            _replace_approval_contents_with_results(messages, fcc_todo, normalized_results)  # type: ignore
-
-        def _should_emit_tool_snapshot(tool_name: str | None) -> bool:
-            if not pending_tool_calls or not tool_results:
-                return False
-            if tool_name and context.config.predict_state_config and not context.config.require_confirmation:
-                for config in context.config.predict_state_config.values():
-                    if config["tool"] == tool_name:
-                        logger.info(
-                            f"Skipping intermediate MessagesSnapshotEvent for predictive tool '{tool_name}' "
-                            " - delaying until summary"
-                        )
-                        return False
-            return True
-
-        def _build_messages_snapshot(tool_message_id: str | None = None) -> MessagesSnapshotEvent:
-            has_text_content = bool(accumulated_text_content)
-            all_messages = snapshot_messages.copy()
-
-            if pending_tool_calls:
-                if tool_message_id and not has_text_content:
-                    tool_call_message_id = tool_message_id
-                else:
-                    tool_call_message_id = (
-                        active_message_id if not has_text_content and active_message_id else generate_event_id()
-                    )
-                tool_call_message = {
-                    "id": tool_call_message_id,
-                    "role": "assistant",
-                    "tool_calls": pending_tool_calls.copy(),
-                }
-                all_messages.append(tool_call_message)
-
-            all_messages.extend(tool_results)
-
-            if has_text_content and active_message_id:
-                assistant_text_message = {
-                    "id": active_message_id,
-                    "role": "assistant",
-                    "content": accumulated_text_content,
-                }
-                all_messages.append(assistant_text_message)
-
-            return MessagesSnapshotEvent(
-                messages=all_messages,  # type: ignore[arg-type]
-            )
-
-        # Use tools_param if available (includes client tools), otherwise fall back to server_tools
-        # This ensures both server tools AND client tools can be executed after approval
-        tools_for_approval = tools_param if tools_param is not None else server_tools
-        latest_approval = latest_approval_response(messages_to_run)
-        await _resolve_approval_responses(messages_to_run, tools_for_approval)
-
-        if latest_approval and is_step_based_approval(latest_approval, context.config.predict_state_config):
-            from ._confirmation_strategies import DefaultConfirmationStrategy
-
-            strategy = context.confirmation_strategy
-            if strategy is None:
-                strategy = DefaultConfirmationStrategy()
-
-            steps = approval_steps(latest_approval)
-            if steps:
-                if latest_approval.approved:
-                    confirmation_message = strategy.on_approval_accepted(steps)
-                else:
-                    confirmation_message = strategy.on_approval_rejected(steps)
-            else:
-                if latest_approval.approved:
-                    confirmation_message = strategy.on_state_confirmed()
-                else:
-                    confirmation_message = strategy.on_state_rejected()
-
-            message_id = generate_event_id()
-            yield TextMessageStartEvent(message_id=message_id, role="assistant")
-            yield TextMessageContentEvent(message_id=message_id, delta=confirmation_message)
-            yield TextMessageEndEvent(message_id=message_id)
-            yield event_bridge.create_run_finished_event()
-            return
-
        async for update in context.agent.run_stream(messages_to_run, **run_kwargs):
            update_count += 1
            logger.info(f"[STREAM] Received update #{update_count} from agent")
-            if all_updates is not None:
-                all_updates.append(update)
-            if event_bridge.current_message_id is None and update.contents:
-                has_tool_call = any(isinstance(content, FunctionCallContent) for content in update.contents)
-                has_text = any(isinstance(content, TextContent) for content in update.contents)
-                if has_tool_call and not has_text:
-                    tool_message_id = generate_event_id()
-                    event_bridge.current_message_id = tool_message_id
-                    active_message_id = tool_message_id
-                    accumulated_text_content = ""
-                    logger.info(
-                        "[STREAM] Emitting TextMessageStartEvent for tool-only response message_id=%s",
-                        tool_message_id,
-                    )
-                    yield TextMessageStartEvent(message_id=tool_message_id, role="assistant")
+            all_updates.append(update)
            events = await event_bridge.from_agent_run_update(update)
            logger.info(f"[STREAM] Update #{update_count} produced {len(events)} events")
            for event in events:
-                if isinstance(event, TextMessageStartEvent):
-                    active_message_id = event.message_id
-                    accumulated_text_content = ""
-                elif isinstance(event, TextMessageContentEvent):
-                    accumulated_text_content += event.delta
-                elif isinstance(event, ToolCallStartEvent):
-                    tool_call_entry = ensure_tool_call_entry(event.tool_call_id, tool_calls_by_id, pending_tool_calls)
-                    tool_call_entry["function"]["name"] = event.tool_call_name
-                elif isinstance(event, ToolCallArgsEvent):
-                    tool_call_entry = ensure_tool_call_entry(event.tool_call_id, tool_calls_by_id, pending_tool_calls)
-                    tool_call_entry["function"]["arguments"] += event.delta
-                elif isinstance(event, ToolCallEndEvent):
-                    tool_calls_ended.add(event.tool_call_id)
-                elif isinstance(event, ToolCallResultEvent):
-                    tool_results.append(
-                        {
-                            "id": event.message_id,
-                            "role": "tool",
-                            "toolCallId": event.tool_call_id,
-                            "content": event.content,
-                        }
-                    )
                logger.info(f"[STREAM] Yielding event: {type(event).__name__}")
                yield event
-                if isinstance(event, ToolCallResultEvent):
-                    tool_name = tool_name_for_call_id(tool_calls_by_id, event.tool_call_id)
-                    if _should_emit_tool_snapshot(tool_name):
-                        messages_snapshot_emitted = True
-                        messages_snapshot = _build_messages_snapshot()
-                        logger.info(f"[STREAM] Yielding event: {type(messages_snapshot).__name__}")
-                        yield messages_snapshot
-                elif isinstance(event, ToolCallEndEvent):
-                    tool_name = tool_name_for_call_id(tool_calls_by_id, event.tool_call_id)
-                    if tool_name == "confirm_changes":
-                        messages_snapshot_emitted = True
-                        messages_snapshot = _build_messages_snapshot()
-                        logger.info(f"[STREAM] Yielding event: {type(messages_snapshot).__name__}")
-                        yield messages_snapshot

        logger.info(f"[STREAM] Agent stream completed. Total updates: {update_count}")

        if event_bridge.should_stop_after_confirm:
-            logger.info("Stopping run - waiting for user approval/confirmation response")
-            if event_bridge.current_message_id:
-                logger.info(f"[CONFIRM] Emitting TextMessageEndEvent for message_id={event_bridge.current_message_id}")
-                yield event_bridge.create_message_end_event(event_bridge.current_message_id)
-                event_bridge.current_message_id = None
+            logger.info("Stopping run after confirm_changes - waiting for user response")
            yield event_bridge.create_run_finished_event()
            return

-        if pending_tool_calls:
-            pending_without_end = [tc for tc in pending_tool_calls if tc.get("id") not in tool_calls_ended]
+        if event_bridge.pending_tool_calls:
+            pending_without_end = [
+                tc for tc in event_bridge.pending_tool_calls if tc.get("id") not in event_bridge.tool_calls_ended
+            ]
            if pending_without_end:
                logger.info(
                    "Found %s pending tool calls without end event - emitting ToolCallEndEvent",
@@ -641,11 +470,13 @@ class DefaultOrchestrator(Orchestrator):
                for tool_call in pending_without_end:
                    tool_call_id = tool_call.get("id")
                    if tool_call_id:
+                        from ag_ui.core import ToolCallEndEvent
+
                        end_event = ToolCallEndEvent(tool_call_id=tool_call_id)
                        logger.info(f"Emitting ToolCallEndEvent for declaration-only tool call '{tool_call_id}'")
                        yield end_event

-        if response_format and all_updates:
+        if all_updates and response_format:
            from agent_framework import AgentRunResponse
            from pydantic import BaseModel

@@ -677,22 +508,37 @@ class DefaultOrchestrator(Orchestrator):
            logger.info(f"[FINALIZE] Emitting TextMessageEndEvent for message_id={event_bridge.current_message_id}")
            yield event_bridge.create_message_end_event(event_bridge.current_message_id)

-            messages_snapshot = _build_messages_snapshot(tool_message_id=event_bridge.current_message_id)
-            messages_snapshot_emitted = True
+            assistant_text_message = {
+                "id": event_bridge.current_message_id,
+                "role": "assistant",
+                "content": event_bridge.accumulated_text_content,
+            }
+
+            converted_input_messages = agui_messages_to_snapshot_format(event_bridge.input_messages)
+            all_messages = converted_input_messages.copy()
+
+            if event_bridge.pending_tool_calls:
+                tool_call_message = {
+                    "id": generate_event_id(),
+                    "role": "assistant",
+                    "tool_calls": event_bridge.pending_tool_calls.copy(),
+                }
+                all_messages.append(tool_call_message)
+
+            all_messages.extend(event_bridge.tool_results.copy())
+            all_messages.append(assistant_text_message)
+
+            messages_snapshot = MessagesSnapshotEvent(
+                messages=all_messages,  # type: ignore[arg-type]
+            )
            logger.info(
-                f"[FINALIZE] Emitting MessagesSnapshotEvent with {len(messages_snapshot.messages)} messages "
-                f"(text content length: {len(accumulated_text_content)})"
+                "[FINALIZE] Emitting MessagesSnapshotEvent with %s messages (text content length: %s)",
+                len(all_messages),
+                len(event_bridge.accumulated_text_content),
            )
            yield messages_snapshot
        else:
            logger.info("[FINALIZE] No current_message_id - skipping TextMessageEndEvent")
-            if not messages_snapshot_emitted and (pending_tool_calls or tool_results):
-                messages_snapshot = _build_messages_snapshot()
-                messages_snapshot_emitted = True
-                logger.info(
-                    f"[FINALIZE] Emitting MessagesSnapshotEvent with {len(messages_snapshot.messages)} messages"
-                )
-                yield messages_snapshot

        logger.info("[FINALIZE] Emitting RUN_FINISHED event")
        yield event_bridge.create_run_finished_event()
@@ -3,29 +3,13 @@
 """Utility functions for AG-UI integration."""

 import copy
-import json
 import uuid
 from collections.abc import Callable, MutableMapping, Sequence
 from dataclasses import asdict, is_dataclass
 from datetime import date, datetime
 from typing import Any

-from agent_framework import AIFunction, Role, ToolProtocol
-
-# Role mapping constants
-AGUI_TO_FRAMEWORK_ROLE: dict[str, Role] = {
-    "user": Role.USER,
-    "assistant": Role.ASSISTANT,
-    "system": Role.SYSTEM,
-}
-
-FRAMEWORK_TO_AGUI_ROLE: dict[Role, str] = {
-    Role.USER: "user",
-    Role.ASSISTANT: "assistant",
-    Role.SYSTEM: "system",
-}
-
-ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool"}
+from agent_framework import AIFunction, ToolProtocol


 def generate_event_id() -> str:
@@ -33,85 +17,6 @@ def generate_event_id() -> str:
    return str(uuid.uuid4())


-def safe_json_parse(value: Any) -> dict[str, Any] | None:
-    """Safely parse a value as JSON dict.
-
-    Args:
-        value: String or dict to parse
-
-    Returns:
-        Parsed dict or None if parsing fails
-    """
-    if isinstance(value, dict):
-        return value
-    if isinstance(value, str):
-        try:
-            parsed = json.loads(value)
-            if isinstance(parsed, dict):
-                return parsed
-        except json.JSONDecodeError:
-            pass
-    return None
-
-
-def get_role_value(message: Any) -> str:
-    """Extract role string from a message object.
-
-    Handles both enum roles (with .value) and string roles.
-
-    Args:
-        message: Message object with role attribute
-
-    Returns:
-        Role as lowercase string, or empty string if not found
-    """
-    role = getattr(message, "role", None)
-    if role is None:
-        return ""
-    if hasattr(role, "value"):
-        return str(role.value)
-    return str(role)
-
-
-def normalize_agui_role(raw_role: Any) -> str:
-    """Normalize an AG-UI role to a standard role string.
-
-    Args:
-        raw_role: Raw role value from AG-UI message
-
-    Returns:
-        Normalized role string (user, assistant, system, or tool)
-    """
-    if not isinstance(raw_role, str):
-        return "user"
-    role = raw_role.lower()
-    if role == "developer":
-        return "system"
-    if role in ALLOWED_AGUI_ROLES:
-        return role
-    return "user"
-
-
-def extract_state_from_tool_args(
-    args: dict[str, Any] | None,
-    tool_arg_name: str,
-) -> Any:
-    """Extract state value from tool arguments based on config.
-
-    Args:
-        args: Parsed tool arguments dict
-        tool_arg_name: Name of the argument to extract, or "*" for entire args
-
-    Returns:
-        Extracted state value, or None if not found
-    """
-    if not args:
-        return None
-    if tool_arg_name == "*":
-        return args
-    return args.get(tool_arg_name)
-
-
 def merge_state(current: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]:
    """Merge state updates.

@@ -75,10 +75,8 @@ def human_in_the_loop_agent(chat_client: ChatClientProtocol) -> ChatAgent:
    9. "Calibrate systems"
    10. "Final testing"

-    IMPORTANT: When you call generate_task_steps, the user will be shown the steps and asked to approve.
-    Do NOT output any text along with the function call - just call the function.
-    After the user approves and the function executes, THEN provide a brief acknowledgment like:
-    "The plan has been created with X steps selected."
+    After calling the function, provide a brief acknowledgment like:
+    "I've created a plan with 10 steps. You can customize which steps to enable before I proceed."
    """,
        chat_client=chat_client,
        tools=[generate_task_steps],
@@ -1,6 +1,6 @@
 [project]
 name = "agent-framework-ag-ui"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 description = "AG-UI protocol integration for Agent Framework"
 readme = "README.md"
 license-files = ["LICENSE"]
@@ -630,179 +630,3 @@ async def test_suppressed_summary_with_document_state():
    # Should contain some reference to the document
    full_text = "".join(e.delta for e in text_events)
    assert "written" in full_text.lower() or "document" in full_text.lower()
-
-
-async def test_function_approval_mode_executes_tool():
-    """Test that function approval with approval_mode='always_require' sends the correct messages."""
-    from agent_framework import FunctionResultContent, ai_function
-    from agent_framework.ag_ui import AgentFrameworkAgent
-
-    messages_received: list[Any] = []
-
-    @ai_function(
-        name="get_datetime",
-        description="Get the current date and time",
-        approval_mode="always_require",
-    )
-    def get_datetime() -> str:
-        return "2025/12/01 12:00:00"
-
-    async def stream_fn(
-        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
-    ) -> AsyncIterator[ChatResponseUpdate]:
-        # Capture the messages received by the chat client
-        messages_received.clear()
-        messages_received.extend(messages)
-        yield ChatResponseUpdate(contents=[TextContent(text="Processing completed")])
-
-    agent = ChatAgent(
-        name="test_agent",
-        instructions="Test",
-        chat_client=StreamingChatClientStub(stream_fn),
-        tools=[get_datetime],
-    )
-    wrapper = AgentFrameworkAgent(agent=agent)
-
-    # Simulate the conversation history with:
-    # 1. User message asking for time
-    # 2. Assistant message with the function call that needs approval
-    # 3. Tool approval message from user
-    tool_result: dict[str, Any] = {"accepted": True}
-    input_data: dict[str, Any] = {
-        "messages": [
-            {
-                "role": "user",
-                "content": "What time is it?",
-            },
-            {
-                "role": "assistant",
-                "content": "",
-                "tool_calls": [
-                    {
-                        "id": "call_get_datetime_123",
-                        "type": "function",
-                        "function": {
-                            "name": "get_datetime",
-                            "arguments": "{}",
-                        },
-                    }
-                ],
-            },
-            {
-                "role": "tool",
-                "content": json.dumps(tool_result),
-                "toolCallId": "call_get_datetime_123",
-            },
-        ],
-    }
-
-    events: list[Any] = []
-    async for event in wrapper.run_agent(input_data):
-        events.append(event)
-
-    # Verify the run completed successfully
-    run_started = [e for e in events if e.type == "RUN_STARTED"]
-    run_finished = [e for e in events if e.type == "RUN_FINISHED"]
-    assert len(run_started) == 1
-    assert len(run_finished) == 1
-
-    # Verify that a FunctionResultContent was created and sent to the agent
-    # Approved tool calls are resolved before the model run.
-    tool_result_found = False
-    for msg in messages_received:
-        for content in msg.contents:
-            if isinstance(content, FunctionResultContent):
-                tool_result_found = True
-                assert content.call_id == "call_get_datetime_123"
-                assert content.result == "2025/12/01 12:00:00"
-                break
-
-    assert tool_result_found, (
-        "FunctionResultContent should be included in messages sent to agent. "
-        "This is required for the model to see the approved tool execution result."
-    )
-
-
-async def test_function_approval_mode_rejection():
-    """Test that function approval rejection creates a rejection response."""
-    from agent_framework import FunctionResultContent, ai_function
-    from agent_framework.ag_ui import AgentFrameworkAgent
-
-    messages_received: list[Any] = []
-
-    @ai_function(
-        name="delete_all_data",
-        description="Delete all user data",
-        approval_mode="always_require",
-    )
-    def delete_all_data() -> str:
-        return "All data deleted"
-
-    async def stream_fn(
-        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
-    ) -> AsyncIterator[ChatResponseUpdate]:
-        # Capture the messages received by the chat client
-        messages_received.clear()
-        messages_received.extend(messages)
-        yield ChatResponseUpdate(contents=[TextContent(text="Operation cancelled")])
-
-    agent = ChatAgent(
-        name="test_agent",
-        instructions="Test",
-        chat_client=StreamingChatClientStub(stream_fn),
-        tools=[delete_all_data],
-    )
-    wrapper = AgentFrameworkAgent(agent=agent)
-
-    # Simulate rejection
-    tool_result: dict[str, Any] = {"accepted": False}
-    input_data: dict[str, Any] = {
-        "messages": [
-            {
-                "role": "user",
-                "content": "Delete all my data",
-            },
-            {
-                "role": "assistant",
-                "content": "",
-                "tool_calls": [
-                    {
-                        "id": "call_delete_123",
-                        "type": "function",
-                        "function": {
-                            "name": "delete_all_data",
-                            "arguments": "{}",
-                        },
-                    }
-                ],
-            },
-            {
-                "role": "tool",
-                "content": json.dumps(tool_result),
-                "toolCallId": "call_delete_123",
-            },
-        ],
-    }
-
-    events: list[Any] = []
-    async for event in wrapper.run_agent(input_data):
-        events.append(event)
-
-    # Verify the run completed
-    run_finished = [e for e in events if e.type == "RUN_FINISHED"]
-    assert len(run_finished) == 1
-
-    # Verify that a FunctionResultContent with rejection payload was created
-    rejection_found = False
-    for msg in messages_received:
-        for content in msg.contents:
-            if isinstance(content, FunctionResultContent):
-                rejection_found = True
-                assert content.call_id == "call_delete_123"
-                assert content.result == "Error: Tool call invocation was rejected by user."
-                break
-
-    assert rejection_found, (
-        "FunctionResultContent with rejection details should be included in messages sent to agent. "
-        "This tells the model that the tool was rejected."
-    )
@@ -52,8 +52,8 @@ async def test_tool_call_flow():
    update2 = AgentRunResponseUpdate(contents=[tool_result])
    events2 = await bridge.from_agent_run_update(update2)

-    # Should have: ToolCallEndEvent, ToolCallResultEvent
-    assert len(events2) == 2
+    # Should have: ToolCallEndEvent, ToolCallResultEvent, MessagesSnapshotEvent
+    assert len(events2) == 3
    assert isinstance(events2[0], ToolCallEndEvent)
    assert isinstance(events2[1], ToolCallResultEvent)

@@ -201,8 +201,7 @@ async def test_tool_result_with_none():
    assert len(events) == 2
    assert events[0].type == "TOOL_CALL_END"
    assert events[1].type == "TOOL_CALL_RESULT"
-    # prepare_function_call_results serializes None as JSON "null"
-    assert events[1].content == "null"
+    assert events[1].content == ""


 async def test_multiple_tool_results_in_sequence():
@@ -231,12 +230,7 @@ async def test_function_approval_request_basic():
    """Test FunctionApprovalRequestContent conversion."""
    from agent_framework_ag_ui._events import AgentFrameworkEventBridge

-    # Set require_confirmation=False to test just the function_approval_request event
-    bridge = AgentFrameworkEventBridge(
-        run_id="test_run",
-        thread_id="test_thread",
-        require_confirmation=False,
-    )
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")

    func_call = FunctionCallContent(
        call_id="call_123",
@@ -289,12 +283,14 @@ async def test_empty_predict_state_config():
    assert "STATE_DELTA" not in event_types
    assert "STATE_SNAPSHOT" not in event_types

-    # Should have: ToolCallStart, ToolCallArgs, ToolCallEnd, ToolCallResult
+    # Should have: ToolCallStart, ToolCallArgs, ToolCallEnd, ToolCallResult, MessagesSnapshot
+    # MessagesSnapshotEvent is emitted after tool results to track the conversation
    assert event_types == [
        "TOOL_CALL_START",
        "TOOL_CALL_ARGS",
        "TOOL_CALL_END",
        "TOOL_CALL_RESULT",
+        "MESSAGES_SNAPSHOT",
    ]


@@ -692,97 +688,3 @@ async def test_state_delta_count_logging():

    # State delta count should have incremented (one per unique state update)
    assert bridge.state_delta_count >= 1
-
-
-# Tests for list type tool results (MCP tool serialization)
-
-
-async def test_tool_result_with_empty_list():
-    """Test FunctionResultContent with empty list result."""
-    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
-
-    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
-
-    update = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_123", result=[])])
-    events = await bridge.from_agent_run_update(update)
-
-    assert len(events) == 2
-    assert events[0].type == "TOOL_CALL_END"
-    assert events[1].type == "TOOL_CALL_RESULT"
-    # Empty list serializes as JSON empty array
-    assert events[1].content == "[]"
-
-
-async def test_tool_result_with_single_text_content():
-    """Test FunctionResultContent with single TextContent-like item (MCP tool result)."""
-    from dataclasses import dataclass
-
-    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
-
-    @dataclass
-    class MockTextContent:
-        text: str
-
-    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
-
-    update = AgentRunResponseUpdate(
-        contents=[FunctionResultContent(call_id="call_123", result=[MockTextContent("Hello from MCP tool!")])]
-    )
-    events = await bridge.from_agent_run_update(update)
-
-    assert len(events) == 2
-    assert events[0].type == "TOOL_CALL_END"
-    assert events[1].type == "TOOL_CALL_RESULT"
-    # TextContent text is extracted and serialized as JSON array
-    assert events[1].content == '["Hello from MCP tool!"]'
-
-
-async def test_tool_result_with_multiple_text_contents():
-    """Test FunctionResultContent with multiple TextContent-like items (MCP tool result)."""
-    from dataclasses import dataclass
-
-    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
-
-    @dataclass
-    class MockTextContent:
-        text: str
-
-    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
-
-    update = AgentRunResponseUpdate(
-        contents=[
-            FunctionResultContent(
-                call_id="call_123",
-                result=[MockTextContent("First result"), MockTextContent("Second result")],
-            )
-        ]
-    )
-    events = await bridge.from_agent_run_update(update)
-
-    assert len(events) == 2
-    assert events[0].type == "TOOL_CALL_END"
-    assert events[1].type == "TOOL_CALL_RESULT"
-    # Multiple TextContent items should return JSON array
-    assert events[1].content == '["First result", "Second result"]'
-
-
-async def test_tool_result_with_model_dump_objects():
-    """Test FunctionResultContent with Pydantic BaseModel objects."""
-    from pydantic import BaseModel
-
-    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
-
-    class MockModel(BaseModel):
-        value: int
-
-    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
-
-    update = AgentRunResponseUpdate(
-        contents=[FunctionResultContent(call_id="call_123", result=[MockModel(value=1), MockModel(value=2)])]
-    )
-    events = await bridge.from_agent_run_update(update)
-
-    assert len(events) == 2
-    assert events[1].type == "TOOL_CALL_RESULT"
-    # Should be properly serialized JSON array without double escaping
-    assert events[1].content == '[{"value": 1}, {"value": 2}]'
@@ -18,7 +18,6 @@ from agent_framework import (
 from agent_framework._clients import BaseChatClient
 from agent_framework._types import ChatResponse, ChatResponseUpdate

-from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
 from agent_framework_ag_ui._orchestrators import ExecutionContext

 StreamFn = Callable[..., AsyncIterator[ChatResponseUpdate]]
@@ -135,9 +134,5 @@ class StubAgent(AgentProtocol):
 class TestExecutionContext(ExecutionContext):
    """ExecutionContext helper that allows setting messages for tests."""

-    def set_messages(self, messages: list[ChatMessage], *, normalize: bool = True) -> None:
-        if normalize:
-            self._messages = _deduplicate_messages(_sanitize_tool_history(messages))
-        else:
-            self._messages = messages
-        self._snapshot_messages = None
+    def set_messages(self, messages: list[ChatMessage]) -> None:
+        self._messages = messages
@@ -10,11 +10,9 @@ from agent_framework_ag_ui._events import AgentFrameworkEventBridge

 async def test_function_approval_request_emission():
    """Test that CustomEvent is emitted for FunctionApprovalRequestContent."""
-    # Set require_confirmation=False to test just the function_approval_request event
    bridge = AgentFrameworkEventBridge(
        run_id="test_run",
        thread_id="test_thread",
-        require_confirmation=False,
    )

    # Create approval request
@@ -49,65 +47,11 @@ async def test_function_approval_request_emission():
    assert event.value["function_call"]["arguments"]["subject"] == "Test"


-async def test_function_approval_request_with_confirm_changes():
-    """Test that confirm_changes is also emitted when require_confirmation=True."""
-    bridge = AgentFrameworkEventBridge(
-        run_id="test_run",
-        thread_id="test_thread",
-        require_confirmation=True,
-    )
-
-    func_call = FunctionCallContent(
-        call_id="call_456",
-        name="delete_file",
-        arguments={"path": "/tmp/test.txt"},
-    )
-    approval_request = FunctionApprovalRequestContent(
-        id="approval_002",
-        function_call=func_call,
-    )
-
-    update = AgentRunResponseUpdate(contents=[approval_request])
-    events = await bridge.from_agent_run_update(update)
-
-    # Should emit: ToolCallEndEvent, CustomEvent, and confirm_changes (Start, Args, End) = 5 events
-    assert len(events) == 5
-
-    # Check ToolCallEndEvent
-    assert events[0].type == "TOOL_CALL_END"
-    assert events[0].tool_call_id == "call_456"
-
-    # Check function_approval_request CustomEvent
-    assert events[1].type == "CUSTOM"
-    assert events[1].name == "function_approval_request"
-
-    # Check confirm_changes tool call events
-    assert events[2].type == "TOOL_CALL_START"
-    assert events[2].tool_call_name == "confirm_changes"
-    assert events[3].type == "TOOL_CALL_ARGS"
-    # Verify confirm_changes includes function info for Dojo UI
-    import json
-
-    args = json.loads(events[3].delta)
-    assert args["function_name"] == "delete_file"
-    assert args["function_call_id"] == "call_456"
-    assert args["function_arguments"] == {"path": "/tmp/test.txt"}
-    assert args["steps"] == [
-        {
-            "description": "Execute delete_file",
-            "status": "enabled",
-        }
-    ]
-    assert events[4].type == "TOOL_CALL_END"
-
-
 async def test_multiple_approval_requests():
    """Test handling multiple approval requests in one update."""
-    # Set require_confirmation=False to simplify the test
    bridge = AgentFrameworkEventBridge(
        run_id="test_run",
        thread_id="test_thread",
-        require_confirmation=False,
    )

    func_call_1 = FunctionCallContent(
@@ -150,32 +94,3 @@ async def test_multiple_approval_requests():
    assert events[3].type == "CUSTOM"
    assert events[3].name == "function_approval_request"
    assert events[3].value["id"] == "approval_2"
-
-
-async def test_function_approval_request_sets_stop_flag():
-    """Test that function approval request sets should_stop_after_confirm flag.
-
-    This ensures the orchestrator stops the run after emitting the approval request,
-    allowing the UI to send back an approval response.
-    """
-    bridge = AgentFrameworkEventBridge(
-        run_id="test_run",
-        thread_id="test_thread",
-    )
-
-    assert bridge.should_stop_after_confirm is False
-
-    func_call = FunctionCallContent(
-        call_id="call_stop_test",
-        name="get_datetime",
-        arguments={},
-    )
-    approval_request = FunctionApprovalRequestContent(
-        id="approval_stop_test",
-        function_call=func_call,
-    )
-
-    update = AgentRunResponseUpdate(contents=[approval_request])
-    await bridge.from_agent_run_update(update)
-
-    assert bridge.should_stop_after_confirm is True
@@ -2,15 +2,12 @@

 """Tests for message adapters."""

-import json
-
 import pytest
-from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, Role, TextContent
+from agent_framework import ChatMessage, FunctionCallContent, Role, TextContent

 from agent_framework_ag_ui._message_adapters import (
    agent_framework_messages_to_agui,
    agui_messages_to_agent_framework,
-    agui_messages_to_snapshot_format,
    extract_text_from_contents,
 )

@@ -46,32 +43,6 @@ def test_agent_framework_to_agui_basic(sample_agent_framework_message):
    assert messages[0]["id"] == "msg-123"


-def test_agent_framework_to_agui_normalizes_dict_roles():
-    """Dict inputs normalize unknown roles for UI compatibility."""
-    messages = [
-        {"role": "developer", "content": "policy"},
-        {"role": "weird_role", "content": "payload"},
-    ]
-
-    converted = agent_framework_messages_to_agui(messages)
-
-    assert converted[0]["role"] == "system"
-    assert converted[1]["role"] == "user"
-
-
-def test_agui_snapshot_format_normalizes_roles():
-    """Snapshot normalization coerces roles into supported AG-UI values."""
-    messages = [
-        {"role": "Developer", "content": "policy"},
-        {"role": "unknown", "content": "payload"},
-    ]
-
-    normalized = agui_messages_to_snapshot_format(messages)
-
-    assert normalized[0]["role"] == "system"
-    assert normalized[1]["role"] == "user"
-
-
 def test_agui_tool_result_to_agent_framework():
    """Test converting AG-UI tool result message to Agent Framework."""
    tool_result_message = {
@@ -97,237 +68,6 @@ def test_agui_tool_result_to_agent_framework():
    assert message.additional_properties.get("tool_call_id") == "call_123"


-def test_agui_tool_approval_updates_tool_call_arguments():
-    """Tool approval updates matching tool call arguments for snapshots and agent context."""
-    messages_input = [
-        {
-            "role": "assistant",
-            "content": "",
-            "tool_calls": [
-                {
-                    "id": "call_123",
-                    "type": "function",
-                    "function": {
-                        "name": "generate_task_steps",
-                        "arguments": {
-                            "steps": [
-                                {"description": "Boil water", "status": "enabled"},
-                                {"description": "Brew coffee", "status": "enabled"},
-                                {"description": "Serve coffee", "status": "enabled"},
-                            ]
-                        },
-                    },
-                }
-            ],
-            "id": "msg_1",
-        },
-        {
-            "role": "tool",
-            "content": json.dumps(
-                {
-                    "accepted": True,
-                    "steps": [
-                        {"description": "Boil water", "status": "enabled"},
-                        {"description": "Serve coffee", "status": "enabled"},
-                    ],
-                }
-            ),
-            "toolCallId": "call_123",
-            "id": "msg_2",
-        },
-    ]
-
-    messages = agui_messages_to_agent_framework(messages_input)
-
-    assert len(messages) == 2
-    assistant_msg = messages[0]
-    func_call = next(content for content in assistant_msg.contents if isinstance(content, FunctionCallContent))
-    assert func_call.arguments == {
-        "steps": [
-            {"description": "Boil water", "status": "enabled"},
-            {"description": "Brew coffee", "status": "disabled"},
-            {"description": "Serve coffee", "status": "enabled"},
-        ]
-    }
-    assert messages_input[0]["tool_calls"][0]["function"]["arguments"] == {
-        "steps": [
-            {"description": "Boil water", "status": "enabled"},
-            {"description": "Brew coffee", "status": "disabled"},
-            {"description": "Serve coffee", "status": "enabled"},
-        ]
-    }
-
-    from agent_framework import FunctionApprovalResponseContent
-
-    approval_msg = messages[1]
-    approval_content = next(
-        content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
-    )
-    assert approval_content.function_call.parse_arguments() == {
-        "steps": [
-            {"description": "Boil water", "status": "enabled"},
-            {"description": "Serve coffee", "status": "enabled"},
-        ]
-    }
-    assert approval_content.additional_properties is not None
-    assert approval_content.additional_properties.get("ag_ui_state_args") == {
-        "steps": [
-            {"description": "Boil water", "status": "enabled"},
-            {"description": "Brew coffee", "status": "disabled"},
-            {"description": "Serve coffee", "status": "enabled"},
-        ]
-    }
-
-
-def test_agui_tool_approval_from_confirm_changes_maps_to_function_call():
-    """Confirm_changes approvals map back to the original tool call when metadata is present."""
-    messages_input = [
-        {
-            "role": "assistant",
-            "content": "",
-            "tool_calls": [
-                {
-                    "id": "call_tool",
-                    "type": "function",
-                    "function": {"name": "get_datetime", "arguments": {}},
-                },
-                {
-                    "id": "call_confirm",
-                    "type": "function",
-                    "function": {
-                        "name": "confirm_changes",
-                        "arguments": {"function_call_id": "call_tool"},
-                    },
-                },
-            ],
-            "id": "msg_1",
-        },
-        {
-            "role": "tool",
-            "content": json.dumps({"accepted": True, "function_call_id": "call_tool"}),
-            "toolCallId": "call_confirm",
-            "id": "msg_2",
-        },
-    ]
-
-    messages = agui_messages_to_agent_framework(messages_input)
-
-    from agent_framework import FunctionApprovalResponseContent
-
-    approval_msg = messages[1]
-    approval_content = next(
-        content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
-    )
-
-    assert approval_content.function_call.call_id == "call_tool"
-    assert approval_content.function_call.name == "get_datetime"
-    assert approval_content.function_call.parse_arguments() == {}
-    assert messages_input[0]["tool_calls"][0]["function"]["arguments"] == {}
-
-
-def test_agui_tool_approval_from_confirm_changes_falls_back_to_sibling_call():
-    """Confirm_changes approvals map to the only sibling tool call when metadata is missing."""
-    messages_input = [
-        {
-            "role": "assistant",
-            "content": "",
-            "tool_calls": [
-                {
-                    "id": "call_tool",
-                    "type": "function",
-                    "function": {"name": "get_datetime", "arguments": {}},
-                },
-                {
-                    "id": "call_confirm",
-                    "type": "function",
-                    "function": {"name": "confirm_changes", "arguments": {}},
-                },
-            ],
-            "id": "msg_1",
-        },
-        {
-            "role": "tool",
-            "content": json.dumps(
-                {
-                    "accepted": True,
-                    "steps": [{"description": "Approve get_datetime", "status": "enabled"}],
-                }
-            ),
-            "toolCallId": "call_confirm",
-            "id": "msg_2",
-        },
-    ]
-
-    messages = agui_messages_to_agent_framework(messages_input)
-
-    from agent_framework import FunctionApprovalResponseContent
-
-    approval_msg = messages[1]
-    approval_content = next(
-        content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
-    )
-
-    assert approval_content.function_call.call_id == "call_tool"
-    assert approval_content.function_call.name == "get_datetime"
-    assert approval_content.function_call.parse_arguments() == {}
-    assert messages_input[0]["tool_calls"][0]["function"]["arguments"] == {}
-
-
-def test_agui_tool_approval_from_generate_task_steps_maps_to_function_call():
-    """Approval tool payloads map to the referenced function call when function_call_id is present."""
-    messages_input = [
-        {
-            "role": "assistant",
-            "content": "",
-            "tool_calls": [
-                {
-                    "id": "call_tool",
-                    "type": "function",
-                    "function": {"name": "get_datetime", "arguments": {}},
-                },
-                {
-                    "id": "call_steps",
-                    "type": "function",
-                    "function": {
-                        "name": "generate_task_steps",
-                        "arguments": {
-                            "function_name": "get_datetime",
-                            "function_call_id": "call_tool",
-                            "function_arguments": {},
-                            "steps": [{"description": "Execute get_datetime", "status": "enabled"}],
-                        },
-                    },
-                },
-            ],
-            "id": "msg_1",
-        },
-        {
-            "role": "tool",
-            "content": json.dumps(
-                {
-                    "accepted": True,
-                    "steps": [{"description": "Execute get_datetime", "status": "enabled"}],
-                }
-            ),
-            "toolCallId": "call_steps",
-            "id": "msg_2",
-        },
-    ]
-
-    messages = agui_messages_to_agent_framework(messages_input)
-
-    from agent_framework import FunctionApprovalResponseContent
-
-    approval_msg = messages[1]
-    approval_content = next(
-        content for content in approval_msg.contents if isinstance(content, FunctionApprovalResponseContent)
-    )
-
-    assert approval_content.function_call.call_id == "call_tool"
-    assert approval_content.function_call.name == "get_datetime"
-    assert approval_content.function_call.parse_arguments() == {}
-
-
 def test_agui_multiple_messages_to_agent_framework():
    """Test converting multiple AG-UI messages."""
    messages_input = [
@@ -538,119 +278,3 @@ def test_extract_text_from_custom_contents():
    result = extract_text_from_contents(contents)

    assert result == "Custom Mixed"
-
-
-# Tests for FunctionResultContent serialization in agent_framework_messages_to_agui
-
-
-def test_agent_framework_to_agui_function_result_dict():
-    """Test converting FunctionResultContent with dict result to AG-UI."""
-    msg = ChatMessage(
-        role=Role.TOOL,
-        contents=[FunctionResultContent(call_id="call-123", result={"key": "value", "count": 42})],
-        message_id="msg-789",
-    )
-
-    messages = agent_framework_messages_to_agui([msg])
-
-    assert len(messages) == 1
-    agui_msg = messages[0]
-    assert agui_msg["role"] == "tool"
-    assert agui_msg["toolCallId"] == "call-123"
-    assert agui_msg["content"] == '{"key": "value", "count": 42}'
-
-
-def test_agent_framework_to_agui_function_result_none():
-    """Test converting FunctionResultContent with None result to AG-UI."""
-    msg = ChatMessage(
-        role=Role.TOOL,
-        contents=[FunctionResultContent(call_id="call-123", result=None)],
-        message_id="msg-789",
-    )
-
-    messages = agent_framework_messages_to_agui([msg])
-
-    assert len(messages) == 1
-    agui_msg = messages[0]
-    # None serializes as JSON null
-    assert agui_msg["content"] == "null"
-
-
-def test_agent_framework_to_agui_function_result_string():
-    """Test converting FunctionResultContent with string result to AG-UI."""
-    msg = ChatMessage(
-        role=Role.TOOL,
-        contents=[FunctionResultContent(call_id="call-123", result="plain text result")],
-        message_id="msg-789",
-    )
-
-    messages = agent_framework_messages_to_agui([msg])
-
-    assert len(messages) == 1
-    agui_msg = messages[0]
-    assert agui_msg["content"] == "plain text result"
-
-
-def test_agent_framework_to_agui_function_result_empty_list():
-    """Test converting FunctionResultContent with empty list result to AG-UI."""
-    msg = ChatMessage(
-        role=Role.TOOL,
-        contents=[FunctionResultContent(call_id="call-123", result=[])],
-        message_id="msg-789",
-    )
-
-    messages = agent_framework_messages_to_agui([msg])
-
-    assert len(messages) == 1
-    agui_msg = messages[0]
-    # Empty list serializes as JSON empty array
-    assert agui_msg["content"] == "[]"
-
-
-def test_agent_framework_to_agui_function_result_single_text_content():
-    """Test converting FunctionResultContent with single TextContent-like item."""
-    from dataclasses import dataclass
-
-    @dataclass
-    class MockTextContent:
-        text: str
-
-    msg = ChatMessage(
-        role=Role.TOOL,
-        contents=[FunctionResultContent(call_id="call-123", result=[MockTextContent("Hello from MCP!")])],
-        message_id="msg-789",
-    )
-
-    messages = agent_framework_messages_to_agui([msg])
-
-    assert len(messages) == 1
-    agui_msg = messages[0]
-    # TextContent text is extracted and serialized as JSON array
-    assert agui_msg["content"] == '["Hello from MCP!"]'
-
-
-def test_agent_framework_to_agui_function_result_multiple_text_contents():
-    """Test converting FunctionResultContent with multiple TextContent-like items."""
-    from dataclasses import dataclass
-
-    @dataclass
-    class MockTextContent:
-        text: str
-
-    msg = ChatMessage(
-        role=Role.TOOL,
-        contents=[
-            FunctionResultContent(
-                call_id="call-123",
-                result=[MockTextContent("First result"), MockTextContent("Second result")],
-            )
-        ],
-        message_id="msg-789",
-    )
-
-    messages = agent_framework_messages_to_agui([msg])
-
-    assert len(messages) == 1
-    agui_msg = messages[0]
-    # Multiple items should return JSON array
-    assert agui_msg["content"] == '["First result", "Second result"]'
@@ -2,7 +2,10 @@

 from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent

-from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
+from agent_framework_ag_ui._orchestration._message_hygiene import (
+    deduplicate_messages,
+    sanitize_tool_history,
+)


 def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
@@ -23,7 +26,7 @@ def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
        ),
    ]

-    sanitized = _sanitize_tool_history(messages)
+    sanitized = sanitize_tool_history(messages)

    tool_messages = [
        msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
@@ -45,6 +48,6 @@ def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
        ),
    ]

-    deduped = _deduplicate_messages(messages)
+    deduped = deduplicate_messages(messages)
    assert len(deduped) == 1
    assert deduped[0].contents[0].result == "result data"
@@ -42,29 +42,6 @@ class DummyAgent:
        yield AgentRunResponseUpdate(contents=[TextContent(text="ok")], role="assistant")


-class RecordingAgent:
-    """Agent stub that captures messages passed to run_stream."""
-
-    def __init__(self) -> None:
-        self.chat_options = SimpleNamespace(tools=[], response_format=None)
-        self.tools: list[Any] = []
-        self.chat_client = SimpleNamespace(
-            function_invocation_configuration=FunctionInvocationConfiguration(),
-        )
-        self.seen_messages: list[Any] | None = None
-
-    async def run_stream(
-        self,
-        messages: list[Any],
-        *,
-        thread: Any,
-        tools: list[Any] | None = None,
-        **kwargs: Any,
-    ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
-        self.seen_messages = messages
-        yield AgentRunResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
-
-
 async def test_default_orchestrator_merges_client_tools() -> None:
    """Client tool declarations are merged with server tools before running agent."""

@@ -174,104 +151,3 @@ async def test_default_orchestrator_with_snake_case_ids() -> None:
    last_event = events[-1]
    assert last_event.run_id == "test-snakecase-runid"
    assert last_event.thread_id == "test-snakecase-threadid"
-
-
-async def test_state_context_injected_when_tool_call_state_mismatch() -> None:
-    """State context should be injected when current state differs from tool call args."""
-
-    agent = RecordingAgent()
-    orchestrator = DefaultOrchestrator()
-
-    tool_recipe = {"title": "Salad", "special_preferences": []}
-    current_recipe = {"title": "Salad", "special_preferences": ["Vegetarian"]}
-
-    input_data = {
-        "state": {"recipe": current_recipe},
-        "messages": [
-            {"role": "system", "content": "Instructions"},
-            {
-                "role": "assistant",
-                "tool_calls": [
-                    {
-                        "id": "call_1",
-                        "type": "function",
-                        "function": {"name": "update_recipe", "arguments": {"recipe": tool_recipe}},
-                    }
-                ],
-            },
-            {"role": "user", "content": "What are the dietary preferences?"},
-        ],
-    }
-
-    context = ExecutionContext(
-        input_data=input_data,
-        agent=agent,
-        config=AgentConfig(
-            state_schema={"recipe": {"type": "object"}},
-            predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
-            require_confirmation=False,
-        ),
-    )
-
-    async for _event in orchestrator.run(context):
-        pass
-
-    assert agent.seen_messages is not None
-    state_messages = []
-    for msg in agent.seen_messages:
-        role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
-        if role_value != "system":
-            continue
-        for content in msg.contents or []:
-            if isinstance(content, TextContent) and content.text.startswith("Current state of the application:"):
-                state_messages.append(content.text)
-    assert state_messages
-    assert "Vegetarian" in state_messages[0]
-
-
-async def test_state_context_not_injected_when_tool_call_matches_state() -> None:
-    """State context should be skipped when tool call args match current state."""
-
-    agent = RecordingAgent()
-    orchestrator = DefaultOrchestrator()
-
-    input_data = {
-        "messages": [
-            {"role": "system", "content": "Instructions"},
-            {
-                "role": "assistant",
-                "tool_calls": [
-                    {
-                        "id": "call_1",
-                        "type": "function",
-                        "function": {"name": "update_recipe", "arguments": {"recipe": {}}},
-                    }
-                ],
-            },
-            {"role": "user", "content": "What are the dietary preferences?"},
-        ],
-    }
-
-    context = ExecutionContext(
-        input_data=input_data,
-        agent=agent,
-        config=AgentConfig(
-            state_schema={"recipe": {"type": "object"}},
-            predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
-            require_confirmation=False,
-        ),
-    )
-
-    async for _event in orchestrator.run(context):
-        pass
-
-    assert agent.seen_messages is not None
-    state_messages = []
-    for msg in agent.seen_messages:
-        role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
-        if role_value != "system":
-            continue
-        for content in msg.contents or []:
-            if isinstance(content, TextContent) and content.text.startswith("Current state of the application:"):
-                state_messages.append(content.text)
-    assert not state_messages
@@ -62,7 +62,7 @@ async def test_human_in_the_loop_json_decode_error() -> None:
        agent=agent,
        config=AgentConfig(),
    )
-    context.set_messages(messages, normalize=False)
+    context.set_messages(messages)

    assert orchestrator.can_handle(context)

@@ -385,8 +385,8 @@ async def test_state_context_injection() -> None:
    assert "banana" in system_messages[0].contents[0].text


-async def test_state_context_injection_with_tool_calls_and_input_state() -> None:
-    """Test state context is injected when state is provided, even with tool calls."""
+async def test_no_state_context_injection_with_tool_calls() -> None:
+    """Test state context is NOT injected if conversation has tool calls."""
    from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent

    messages = [
@@ -420,13 +420,13 @@ async def test_state_context_injection_with_tool_calls_and_input_state() -> None
    async for event in orchestrator.run(context):
        events.append(event)

-    # Should inject state context system message because input state is provided
+    # Should NOT inject state context system message since conversation has tool calls
    system_messages = [
        msg
        for msg in agent.messages_received
        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
    ]
-    assert len(system_messages) == 1
+    assert len(system_messages) == 0


 async def test_structured_output_processing() -> None:
@@ -685,54 +685,6 @@ async def test_confirm_changes_with_invalid_json_fallback() -> None:
    assert len(user_messages) == 1


-async def test_confirm_changes_closes_active_message_before_finish() -> None:
-    """Confirm-changes flow closes any active text message before run finishes."""
-    from ag_ui.core import TextMessageEndEvent, TextMessageStartEvent
-    from agent_framework import FunctionCallContent, FunctionResultContent
-
-    updates = [
-        AgentRunResponseUpdate(
-            contents=[
-                FunctionCallContent(
-                    name="write_document_local",
-                    call_id="call_1",
-                    arguments='{"document": "Draft"}',
-                )
-            ]
-        ),
-        AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")]),
-    ]
-
-    orchestrator = DefaultOrchestrator()
-    input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Start"}]}
-    agent = StubAgent(
-        chat_options=DEFAULT_CHAT_OPTIONS,
-        updates=updates,
-    )
-    context = TestExecutionContext(
-        input_data=input_data,
-        agent=agent,
-        config=AgentConfig(
-            predict_state_config={"document": {"tool": "write_document_local", "tool_argument": "document"}},
-            require_confirmation=True,
-        ),
-    )
-
-    events: list[Any] = []
-    async for event in orchestrator.run(context):
-        events.append(event)
-
-    start_events = [e for e in events if isinstance(e, TextMessageStartEvent)]
-    end_events = [e for e in events if isinstance(e, TextMessageEndEvent)]
-    assert len(start_events) == 1
-    assert len(end_events) == 1
-    assert end_events[0].message_id == start_events[0].message_id
-
-    end_index = events.index(end_events[0])
-    finished_index = events.index([e for e in events if e.type == "RUN_FINISHED"][0])
-    assert end_index < finished_index
-
-
 async def test_tool_result_kept_when_call_id_matches() -> None:
    """Test tool result is kept when call_id matches pending tool calls."""
    from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent
@@ -5,11 +5,7 @@
 from dataclasses import dataclass
 from datetime import date, datetime

-from agent_framework_ag_ui._utils import (
-    generate_event_id,
-    make_json_safe,
-    merge_state,
-)
+from agent_framework_ag_ui._utils import generate_event_id, make_json_safe, merge_state


 def test_generate_event_id():
@@ -1,4 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
+
 from collections.abc import AsyncIterable, MutableMapping, MutableSequence, Sequence
 from typing import Any, ClassVar, Final, TypeVar

@@ -12,10 +13,7 @@ from agent_framework import (
    ChatResponse,
    ChatResponseUpdate,
    CitationAnnotation,
-    CodeInterpreterToolCallContent,
-    CodeInterpreterToolResultContent,
    Contents,
-    ErrorContent,
    FinishReason,
    FunctionCallContent,
    FunctionResultContent,
@@ -23,8 +21,6 @@ from agent_framework import (
    HostedFileContent,
    HostedMCPTool,
    HostedWebSearchTool,
-    MCPServerToolCallContent,
-    MCPServerToolResultContent,
    Role,
    TextContent,
    TextReasoningContent,
@@ -49,8 +45,6 @@ from anthropic.types.beta import (
    BetaTextBlock,
    BetaUsage,
 )
-from anthropic.types.beta.beta_bash_code_execution_tool_result_error import BetaBashCodeExecutionToolResultError
-from anthropic.types.beta.beta_code_execution_tool_result_error import BetaCodeExecutionToolResultError
 from pydantic import SecretStr, ValidationError

 logger = get_logger("agent_framework.anthropic")
@@ -511,7 +505,7 @@ class AnthropicClient(BaseChatClient):
            usage_details=self._parse_usage_from_anthropic(message.usage),
            model_id=message.model,
            finish_reason=FINISH_REASON_MAP.get(message.stop_reason) if message.stop_reason else None,
-            raw_representation=message,
+            raw_response=message,
        )

    def _process_stream_event(self, event: BetaRawMessageStreamEvent) -> ChatResponseUpdate | None:
@@ -536,14 +530,13 @@ class AnthropicClient(BaseChatClient):
                    finish_reason=FINISH_REASON_MAP.get(event.message.stop_reason)
                    if event.message.stop_reason
                    else None,
-                    raw_representation=event,
+                    raw_response=event,
                )
            case "message_delta":
                usage = self._parse_usage_from_anthropic(event.usage)
                return ChatResponseUpdate(
                    contents=[UsageContent(details=usage, raw_representation=event.usage)] if usage else [],
-                    finish_reason=FINISH_REASON_MAP.get(event.delta.stop_reason) if event.delta.stop_reason else None,
-                    raw_representation=event,
+                    raw_response=event,
                )
            case "message_stop":
                logger.debug("Received message_stop event; no content to process.")
@@ -551,13 +544,13 @@ class AnthropicClient(BaseChatClient):
                contents = self._parse_contents_from_anthropic([event.content_block])
                return ChatResponseUpdate(
                    contents=contents,
-                    raw_representation=event,
+                    raw_response=event,
                )
            case "content_block_delta":
                contents = self._parse_contents_from_anthropic([event.delta])
                return ChatResponseUpdate(
                    contents=contents,
-                    raw_representation=event,
+                    raw_response=event,
                )
            case "content_block_stop":
                logger.debug("Received content_block_stop event; no content to process.")
@@ -595,49 +588,23 @@ class AnthropicClient(BaseChatClient):
                    )
                case "tool_use" | "mcp_tool_use" | "server_tool_use":
                    self._last_call_id_name = (content_block.id, content_block.name)
-                    if content_block.type == "mcp_tool_use":
-                        contents.append(
-                            MCPServerToolCallContent(
-                                call_id=content_block.id,
-                                tool_name=content_block.name,
-                                server_name=None,
-                                arguments=content_block.input,
-                                raw_representation=content_block,
-                            )
-                        )
-                    elif "code_execution" in (content_block.name or ""):
-                        contents.append(
-                            CodeInterpreterToolCallContent(
-                                call_id=content_block.id,
-                                inputs=[TextContent(text=str(content_block.input), raw_representation=content_block)],
-                                raw_representation=content_block,
-                            )
-                        )
-                    else:
-                        contents.append(
-                            FunctionCallContent(
-                                call_id=content_block.id,
-                                name=content_block.name,
-                                arguments=content_block.input,
-                                raw_representation=content_block,
-                            )
+                    contents.append(
+                        FunctionCallContent(
+                            call_id=content_block.id,
+                            name=content_block.name,
+                            arguments=content_block.input,
+                            raw_representation=content_block,
                        )
+                    )
                case "mcp_tool_result":
                    call_id, name = self._last_call_id_name or (None, None)
-                    parsed_output: list[Contents] | None = None
-                    if content_block.content:
-                        if isinstance(content_block.content, list):
-                            parsed_output = self._parse_contents_from_anthropic(content_block.content)
-                        elif isinstance(content_block.content, (str, bytes)):
-                            parsed_output = [
-                                TextContent(text=str(content_block.content), raw_representation=content_block)
-                            ]
-                        else:
-                            parsed_output = self._parse_contents_from_anthropic([content_block.content])
                    contents.append(
-                        MCPServerToolResultContent(
+                        FunctionResultContent(
                            call_id=content_block.tool_use_id,
-                            output=parsed_output,
+                            name=name if name and call_id == content_block.tool_use_id else "mcp_tool",
+                            result=self._parse_contents_from_anthropic(content_block.content)
+                            if isinstance(content_block.content, list)
+                            else content_block.content,
                            raw_representation=content_block,
                        )
                    )
@@ -651,183 +618,30 @@ class AnthropicClient(BaseChatClient):
                            raw_representation=content_block,
                        )
                    )
-                case "code_execution_tool_result":
-                    code_outputs: list[Contents] = []
-                    if content_block.content:
-                        if isinstance(content_block.content, BetaCodeExecutionToolResultError):
-                            code_outputs.append(
-                                ErrorContent(
-                                    message=content_block.content.error_code,
-                                    raw_representation=content_block.content,
-                                )
-                            )
-                        else:
-                            if content_block.content.stdout:
-                                code_outputs.append(
-                                    TextContent(
-                                        text=content_block.content.stdout,
-                                        raw_representation=content_block.content,
-                                    )
-                                )
-                            if content_block.content.stderr:
-                                code_outputs.append(
-                                    ErrorContent(
-                                        message=content_block.content.stderr,
-                                        raw_representation=content_block.content,
-                                    )
-                                )
-                            for code_file_content in content_block.content.content:
-                                code_outputs.append(
-                                    HostedFileContent(
-                                        file_id=code_file_content.file_id, raw_representation=code_file_content
-                                    )
-                                )
-                    contents.append(
-                        CodeInterpreterToolResultContent(
-                            call_id=content_block.tool_use_id,
-                            raw_representation=content_block,
-                            outputs=code_outputs,
+                case (
+                    "code_execution_tool_result"
+                    | "bash_code_execution_tool_result"
+                    | "text_editor_code_execution_tool_result"
+                ):
+                    call_id, name = self._last_call_id_name or (None, None)
+                    if (
+                        content_block.content
+                        and (
+                            content_block.content.type == "bash_code_execution_result"
+                            or content_block.content.type == "code_execution_result"
                        )
-                    )
-                case "bash_code_execution_tool_result":
-                    bash_outputs: list[Contents] = []
-                    if content_block.content:
-                        if isinstance(
-                            content_block.content,
-                            BetaBashCodeExecutionToolResultError,
-                        ):
-                            bash_outputs.append(
-                                ErrorContent(
-                                    message=content_block.content.error_code,
-                                    raw_representation=content_block.content,
-                                )
-                            )
-                        else:
-                            if content_block.content.stdout:
-                                bash_outputs.append(
-                                    TextContent(
-                                        text=content_block.content.stdout,
-                                        raw_representation=content_block.content,
-                                    )
-                                )
-                            if content_block.content.stderr:
-                                bash_outputs.append(
-                                    ErrorContent(
-                                        message=content_block.content.stderr,
-                                        raw_representation=content_block.content,
-                                    )
-                                )
-                            for bash_file_content in content_block.content.content:
+                        and content_block.content.content
+                    ):
+                        for result_content in content_block.content.content:
+                            if hasattr(result_content, "file_id"):
                                contents.append(
-                                    HostedFileContent(
-                                        file_id=bash_file_content.file_id, raw_representation=bash_file_content
-                                    )
+                                    HostedFileContent(file_id=result_content.file_id, raw_representation=result_content)
                                )
                    contents.append(
                        FunctionResultContent(
                            call_id=content_block.tool_use_id,
-                            name=content_block.type,
-                            result=bash_outputs,
-                            raw_representation=content_block,
-                        )
-                    )
-                case "text_editor_code_execution_tool_result":
-                    text_editor_outputs: list[Contents] = []
-                    match content_block.content.type:
-                        case "text_editor_code_execution_tool_result_error":
-                            text_editor_outputs.append(
-                                ErrorContent(
-                                    message=content_block.content.error_code
-                                    and getattr(content_block.content, "error_message", ""),
-                                    raw_representation=content_block.content,
-                                )
-                            )
-                        case "text_editor_code_execution_view_result":
-                            annotations = (
-                                [
-                                    CitationAnnotation(
-                                        raw_representation=content_block.content,
-                                        annotated_regions=[
-                                            TextSpanRegion(
-                                                start_index=content_block.content.start_line,
-                                                end_index=content_block.content.start_line
-                                                + (content_block.content.num_lines or 0),
-                                            )
-                                        ],
-                                    )
-                                ]
-                                if content_block.content.num_lines is not None
-                                and content_block.content.start_line is not None
-                                else None
-                            )
-                            text_editor_outputs.append(
-                                TextContent(
-                                    text=content_block.content.content,
-                                    annotations=annotations,
-                                    raw_representation=content_block.content,
-                                )
-                            )
-                        case "text_editor_code_execution_str_replace_result":
-                            old_annotation = (
-                                CitationAnnotation(
-                                    raw_representation=content_block.content,
-                                    annotated_regions=[
-                                        TextSpanRegion(
-                                            start_index=content_block.content.old_start or 0,
-                                            end_index=(
-                                                (content_block.content.old_start or 0)
-                                                + (content_block.content.old_lines or 0)
-                                            ),
-                                        )
-                                    ],
-                                )
-                                if content_block.content.old_lines is not None
-                                and content_block.content.old_start is not None
-                                else None
-                            )
-                            new_annotation = (
-                                CitationAnnotation(
-                                    raw_representation=content_block.content,
-                                    snippet="\n".join(content_block.content.lines)
-                                    if content_block.content.lines
-                                    else None,
-                                    annotated_regions=[
-                                        TextSpanRegion(
-                                            start_index=content_block.content.new_start or 0,
-                                            end_index=(
-                                                (content_block.content.new_start or 0)
-                                                + (content_block.content.new_lines or 0)
-                                            ),
-                                        )
-                                    ],
-                                )
-                                if content_block.content.new_lines is not None
-                                and content_block.content.new_start is not None
-                                else None
-                            )
-                            annotations = [ann for ann in [old_annotation, new_annotation] if ann is not None]
-
-                            text_editor_outputs.append(
-                                TextContent(
-                                    text=(
-                                        "\n".join(content_block.content.lines) if content_block.content.lines else ""
-                                    ),
-                                    annotations=annotations or None,
-                                    raw_representation=content_block.content,
-                                )
-                            )
-                        case "text_editor_code_execution_create_result":
-                            text_editor_outputs.append(
-                                TextContent(
-                                    text=f"File update: {content_block.content.is_file_update}",
-                                    raw_representation=content_block.content,
-                                )
-                            )
-                    contents.append(
-                        FunctionResultContent(
-                            call_id=content_block.tool_use_id,
-                            name=content_block.type,
-                            result=text_editor_outputs,
+                            name=name if name and call_id == content_block.tool_use_id else "code_execution_tool",
+                            result=content_block.content,
                            raw_representation=content_block,
                        )
                    )
@@ -4,7 +4,7 @@ description = "Anthropic integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -4,7 +4,7 @@ description = "Azure AI Search integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -300,26 +300,13 @@ class AzureAIClient(OpenAIBaseResponsesClient):
        raise ServiceInvalidRequestError("response_format must be a Pydantic model or mapping.")

    async def _get_agent_reference_or_create(
-        self,
-        run_options: dict[str, Any],
-        messages_instructions: str | None,
-        chat_options: ChatOptions | None = None,
+        self, run_options: dict[str, Any], messages_instructions: str | None
    ) -> dict[str, str]:
        """Determine which agent to use and create if needed.

-        Args:
-            run_options: The prepared options for the API call.
-            messages_instructions: Instructions extracted from messages.
-            chat_options: The chat options containing response_format and other settings.
-
        Returns:
            dict[str, str]: The agent reference to use.
        """
-        # chat_options is needed separately because the base class excludes response_format
-        # from run_options (transforming it to text/text_format for OpenAI). Azure's agent
-        # creation API requires the original response_format to build its own config format.
-        if chat_options is None:
-            chat_options = ChatOptions()
        # Agent name must be explicitly provided by the user.
        if self.agent_name is None:
            raise ServiceInitializationError(
@@ -354,14 +341,8 @@ class AzureAIClient(OpenAIBaseResponsesClient):
            if "top_p" in run_options:
                args["top_p"] = run_options["top_p"]

-            # response_format is accessed from chat_options or additional_properties
-            # since the base class excludes it from run_options
-            response_format: Any = (
-                chat_options.response_format
-                if chat_options.response_format is not None
-                else chat_options.additional_properties.get("response_format")
-            )
-            if response_format:
+            if "response_format" in run_options:
+                response_format = run_options["response_format"]
                args["text"] = PromptAgentDefinitionText(format=self._create_text_format_config(response_format))

            # Combine instructions from messages and options
@@ -409,12 +390,12 @@ class AzureAIClient(OpenAIBaseResponsesClient):

        if not self._is_application_endpoint:
            # Application-scoped response APIs do not support "agent" property.
-            agent_reference = await self._get_agent_reference_or_create(run_options, instructions, chat_options)
+            agent_reference = await self._get_agent_reference_or_create(run_options, instructions)
            run_options["extra_body"] = {"agent": agent_reference}

        # Remove properties that are not supported on request level
        # but were configured on agent level
-        exclude = ["model", "tools", "response_format", "temperature", "top_p", "text", "text_format"]
+        exclude = ["model", "tools", "response_format", "temperature", "top_p"]

        for property in exclude:
            run_options.pop(property, None)
@@ -4,7 +4,7 @@ description = "Azure AI Foundry integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -723,10 +723,9 @@ async def test_azure_ai_client_agent_creation_with_response_format(
    mock_agent.version = "1.0"
    mock_project_client.agents.create_version = AsyncMock(return_value=mock_agent)

-    run_options = {"model": "test-model"}
-    chat_options = ChatOptions(response_format=ResponseFormatModel)
+    run_options = {"model": "test-model", "response_format": ResponseFormatModel}

-    await client._get_agent_reference_or_create(run_options, None, chat_options)  # type: ignore
+    await client._get_agent_reference_or_create(run_options, None)  # type: ignore

    # Verify agent was created with response format configuration
    call_args = mock_project_client.agents.create_version.call_args
@@ -777,18 +776,19 @@ async def test_azure_ai_client_agent_creation_with_mapping_response_format(
        "additionalProperties": False,
    }

-    run_options = {"model": "test-model"}
-    response_format_mapping = {
-        "type": "json_schema",
-        "json_schema": {
-            "name": runtime_schema["title"],
-            "strict": True,
-            "schema": runtime_schema,
+    run_options = {
+        "model": "test-model",
+        "response_format": {
+            "type": "json_schema",
+            "json_schema": {
+                "name": runtime_schema["title"],
+                "strict": True,
+                "schema": runtime_schema,
+            },
        },
    }
-    chat_options = ChatOptions(response_format=response_format_mapping)  # type: ignore

-    await client._get_agent_reference_or_create(run_options, None, chat_options)  # type: ignore
+    await client._get_agent_reference_or_create(run_options, None)  # type: ignore

    call_args = mock_project_client.agents.create_version.call_args
    created_definition = call_args[1]["definition"]
@@ -805,7 +805,7 @@ async def test_azure_ai_client_agent_creation_with_mapping_response_format(
 async def test_azure_ai_client_prepare_options_excludes_response_format(
    mock_project_client: MagicMock,
 ) -> None:
-    """Test that prepare_options excludes response_format, text, and text_format from final run options."""
+    """Test that prepare_options excludes response_format from final run options."""
    client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")

    messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
@@ -815,12 +815,7 @@ async def test_azure_ai_client_prepare_options_excludes_response_format(
        patch.object(
            client.__class__.__bases__[0],
            "_prepare_options",
-            return_value={
-                "model": "test-model",
-                "response_format": ResponseFormatModel,
-                "text": {"format": {"type": "json_schema", "name": "test"}},
-                "text_format": ResponseFormatModel,
-            },
+            return_value={"model": "test-model", "response_format": ResponseFormatModel},
        ),
        patch.object(
            client,
@@ -830,11 +825,8 @@ async def test_azure_ai_client_prepare_options_excludes_response_format(
    ):
        run_options = await client._prepare_options(messages, chat_options)

-        # response_format, text, and text_format should be excluded from final run options
-        # because they are configured at agent level, not request level
+        # response_format should be excluded from final run options
        assert "response_format" not in run_options
-        assert "text" not in run_options
-        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"
@@ -1017,91 +1009,3 @@ async def test_azure_ai_chat_client_agent_with_tools() -> None:
        assert response.text is not None
        assert len(response.text) > 0
        assert any(word in response.text.lower() for word in ["sunny", "25"])
-
-
-class ReleaseBrief(BaseModel):
-    """Structured output model for release brief."""
-
-    title: str = Field(description="A short title for the release.")
-    summary: str = Field(description="A brief summary of what was released.")
-    highlights: list[str] = Field(description="Key highlights from the release.")
-    model_config = ConfigDict(extra="forbid")
-
-
-@pytest.mark.flaky
-@skip_if_azure_ai_integration_tests_disabled
-async def test_azure_ai_chat_client_agent_with_response_format() -> None:
-    """Test ChatAgent with response_format (structured output) using AzureAIClient."""
-    async with (
-        temporary_chat_client(agent_name="ResponseFormatAgent") as chat_client,
-        ChatAgent(chat_client=chat_client) as agent,
-    ):
-        response = await agent.run(
-            "Summarize the following release notes into a ReleaseBrief:\n\n"
-            "Version 2.0 Release Notes:\n"
-            "- Added new streaming API for real-time responses\n"
-            "- Improved error handling with detailed messages\n"
-            "- Performance boost of 50% in batch processing\n"
-            "- Fixed memory leak in connection pooling",
-            response_format=ReleaseBrief,
-        )
-
-        # Validate response
-        assert isinstance(response, AgentRunResponse)
-        assert response.value is not None
-        assert isinstance(response.value, ReleaseBrief)
-
-        # Validate structured output fields
-        brief = response.value
-        assert len(brief.title) > 0
-        assert len(brief.summary) > 0
-        assert len(brief.highlights) > 0
-
-
-@pytest.mark.flaky
-@skip_if_azure_ai_integration_tests_disabled
-async def test_azure_ai_chat_client_agent_with_runtime_json_schema() -> None:
-    """Test ChatAgent with runtime JSON schema (structured output) using AzureAIClient."""
-    runtime_schema = {
-        "title": "WeatherDigest",
-        "type": "object",
-        "properties": {
-            "location": {"type": "string"},
-            "conditions": {"type": "string"},
-            "temperature_c": {"type": "number"},
-            "advisory": {"type": "string"},
-        },
-        "required": ["location", "conditions", "temperature_c", "advisory"],
-        "additionalProperties": False,
-    }
-
-    async with (
-        temporary_chat_client(agent_name="RuntimeSchemaAgent") as chat_client,
-        ChatAgent(chat_client=chat_client) as agent,
-    ):
-        response = await agent.run(
-            "Give a brief weather digest for Seattle.",
-            additional_chat_options={
-                "response_format": {
-                    "type": "json_schema",
-                    "json_schema": {
-                        "name": runtime_schema["title"],
-                        "strict": True,
-                        "schema": runtime_schema,
-                    },
-                },
-            },
-        )
-
-        # Validate response
-        assert isinstance(response, AgentRunResponse)
-        assert response.text is not None
-
-        # Parse JSON and validate structure
-        import json
-
-        parsed = json.loads(response.text)
-        assert "location" in parsed
-        assert "conditions" in parsed
-        assert "temperature_c" in parsed
-        assert "advisory" in parsed
@@ -2,9 +2,9 @@

 import importlib.metadata

-from agent_framework_durabletask import AgentCallbackContext, AgentResponseCallbackProtocol, DurableAIAgent
-
 from ._app import AgentFunctionApp
+from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
+from ._orchestration import DurableAIAgent

 try:
    __version__ = importlib.metadata.version(__name__)
@@ -8,16 +8,16 @@ with Azure Durable Entities, enabling stateful and durable AI agent execution.

 import json
 import re
-import uuid
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, TypeVar, cast

 import azure.durable_functions as df
 import azure.functions as func
 from agent_framework import AgentProtocol, get_logger
-from agent_framework_durabletask import (
+
+from ._callbacks import AgentResponseCallbackProtocol
+from ._constants import (
    DEFAULT_MAX_POLL_RETRIES,
    DEFAULT_POLL_INTERVAL_SECONDS,
    MIMETYPE_APPLICATION_JSON,
@@ -28,17 +28,12 @@ from agent_framework_durabletask import (
    THREAD_ID_HEADER,
    WAIT_FOR_RESPONSE_FIELD,
    WAIT_FOR_RESPONSE_HEADER,
-    AgentResponseCallbackProtocol,
-    AgentSessionId,
-    ApiResponseFields,
-    DurableAgentState,
-    DurableAIAgent,
-    RunRequest,
 )
-
+from ._durable_agent_state import DurableAgentState
 from ._entities import create_agent_entity
 from ._errors import IncomingRequestError
-from ._orchestration import AgentOrchestrationContextType, AgentTask, AzureFunctionsAgentExecutor
+from ._models import AgentSessionId, RunRequest
+from ._orchestration import AgentOrchestrationContextType, DurableAIAgent

 logger = get_logger("agent_framework.azurefunctions")

@@ -299,7 +294,7 @@ class AgentFunctionApp(DFAppBase):
        self,
        context: AgentOrchestrationContextType,
        agent_name: str,
-    ) -> DurableAIAgent[AgentTask]:
+    ) -> DurableAIAgent:
        """Return a DurableAIAgent proxy for a registered agent.

        Args:
@@ -310,15 +305,14 @@ class AgentFunctionApp(DFAppBase):
            ValueError: If the requested agent has not been registered.

        Returns:
-            DurableAIAgent[AgentTask] wrapper bound to the orchestration context.
+            DurableAIAgent wrapper bound to the orchestration context.
        """
        normalized_name = str(agent_name)

        if normalized_name not in self._agent_metadata:
            raise ValueError(f"Agent '{normalized_name}' is not registered with this app.")

-        executor = AzureFunctionsAgentExecutor(context)
-        return DurableAIAgent(executor, normalized_name)
+        return DurableAIAgent(context, normalized_name)

    def _setup_agent_functions(
        self,
@@ -381,6 +375,8 @@ class AgentFunctionApp(DFAppBase):
                "enable_tool_calls": true|false (optional, default: true)
            }
            """
+            logger.debug(f"[HTTP Trigger] Received request on route: /api/agents/{agent_name}/run")
+
            request_response_format: str = REQUEST_RESPONSE_FORMAT_JSON
            thread_id: str | None = None

@@ -389,9 +385,9 @@ class AgentFunctionApp(DFAppBase):
                thread_id = self._resolve_thread_id(req=req, req_body=req_body)
                wait_for_response = self._should_wait_for_response(req=req, req_body=req_body)

-                logger.debug(
-                    f"[HTTP Trigger] Message: {message}, Thread ID: {thread_id}, wait_for_response: {wait_for_response}"
-                )
+                logger.debug(f"[HTTP Trigger] Message: {message}")
+                logger.debug(f"[HTTP Trigger] Thread ID: {thread_id}")
+                logger.debug(f"[HTTP Trigger] wait_for_response: {wait_for_response}")

                if not message:
                    logger.warning("[HTTP Trigger] Request rejected: Missing message")
@@ -405,18 +401,15 @@ class AgentFunctionApp(DFAppBase):
                session_id = self._create_session_id(agent_name, thread_id)
                correlation_id = self._generate_unique_id()

-                logger.debug(
-                    f"[HTTP Trigger] Calling entity to run agent using session ID: {session_id} "
-                    f"and correlation ID: {correlation_id}"
-                )
+                logger.debug(f"[HTTP Trigger] Using session ID: {session_id}")
+                logger.debug(f"[HTTP Trigger] Generated correlation ID: {correlation_id}")
+                logger.debug("[HTTP Trigger] Calling entity to run agent...")

-                entity_instance_id = df.EntityId(
-                    name=session_id.entity_name,
-                    key=session_id.key,
-                )
+                entity_instance_id = session_id.to_entity_id()
                run_request = self._build_request_data(
                    req_body,
                    message,
+                    thread_id,
                    correlation_id,
                    request_response_format,
                )
@@ -629,16 +622,14 @@ class AgentFunctionApp(DFAppBase):
            session_id = AgentSessionId.with_random_key(agent_name)

        # Build entity instance ID
-        entity_instance_id = df.EntityId(
-            name=session_id.entity_name,
-            key=session_id.key,
-        )
+        entity_instance_id = session_id.to_entity_id()

        # Create run request
        correlation_id = self._generate_unique_id()
        run_request = self._build_request_data(
            req_body={"message": query, "role": "user"},
            message=query,
+            thread_id=str(session_id),
            correlation_id=correlation_id,
            request_response_format=REQUEST_RESPONSE_FORMAT_TEXT,
        )
@@ -790,7 +781,7 @@ class AgentFunctionApp(DFAppBase):
            agent_response = state.try_get_agent_response(correlation_id)
            if agent_response:
                result = self._build_success_result(
-                    response_message=agent_response.text,
+                    response_data=agent_response,
                    message=message,
                    thread_id=thread_id,
                    correlation_id=correlation_id,
@@ -836,22 +827,23 @@ class AgentFunctionApp(DFAppBase):
        )

    def _build_success_result(
-        self, response_message: str, message: str, thread_id: str, correlation_id: str, state: DurableAgentState
+        self, response_data: dict[str, Any], message: str, thread_id: str, correlation_id: str, state: DurableAgentState
    ) -> dict[str, Any]:
        """Build the success result returned to the HTTP caller."""
        return self._build_response_payload(
-            response=response_message,
+            response=response_data.get("content"),
            message=message,
            thread_id=thread_id,
            status="success",
            correlation_id=correlation_id,
-            extra_fields={ApiResponseFields.MESSAGE_COUNT: state.message_count},
+            extra_fields={"message_count": response_data.get("message_count", state.message_count)},
        )

    def _build_request_data(
        self,
        req_body: dict[str, Any],
        message: str,
+        thread_id: str,
        correlation_id: str,
        request_response_format: str,
    ) -> dict[str, Any]:
@@ -865,8 +857,8 @@ class AgentFunctionApp(DFAppBase):
            request_response_format=request_response_format,
            response_format=req_body.get("response_format"),
            enable_tool_calls=enable_tool_calls,
+            thread_id=thread_id,
            correlation_id=correlation_id,
-            created_at=datetime.now(timezone.utc),
        ).to_dict()

    def _build_accepted_response(self, message: str, thread_id: str, correlation_id: str) -> dict[str, Any]:
@@ -918,13 +910,15 @@ class AgentFunctionApp(DFAppBase):

    def _generate_unique_id(self) -> str:
        """Generate a new unique identifier."""
+        import uuid
+
        return uuid.uuid4().hex

-    def _create_session_id(self, agent_name: str, thread_id: str | None) -> AgentSessionId:
+    def _create_session_id(self, func_name: str, thread_id: str | None) -> AgentSessionId:
        """Create a session identifier using the provided thread id or a random value."""
        if thread_id:
-            return AgentSessionId(name=agent_name, key=thread_id)
-        return AgentSessionId.with_random_key(name=agent_name)
+            return AgentSessionId(name=func_name, key=thread_id)
+        return AgentSessionId.with_random_key(name=func_name)

    def _resolve_thread_id(self, req: func.HttpRequest, req_body: dict[str, Any]) -> str:
        """Retrieve the thread identifier from request body or query parameters."""
@@ -28,10 +28,6 @@ WAIT_FOR_RESPONSE_HEADER: str = "x-ms-wait-for-response"
 DEFAULT_MAX_POLL_RETRIES: int = 30
 DEFAULT_POLL_INTERVAL_SECONDS: float = 1.0

-# TTL configuration
-DEFAULT_TIME_TO_LIVE_DAYS: int = 14
-MINIMUM_TTL_SIGNAL_DELAY_MINUTES: int = 5
-

 # =============================================================================
 # JSON Field Name Constants for Durable Agent State Serialization
@@ -102,9 +98,6 @@ class DurableStateFields:
    # History field
    CONVERSATION_HISTORY: Final[str] = "conversationHistory"

-    # TTL field
-    EXPIRATION_TIME_UTC: Final[str] = "expirationTimeUtc"
-

 class ContentTypes:
    """Content type discriminator values for the $type field.
@@ -53,10 +53,10 @@ from agent_framework import (
 )
 from dateutil import parser as date_parser

-from ._constants import ContentTypes, DurableStateFields
+from ._constants import ApiResponseFields, ContentTypes, DurableStateFields
 from ._models import RunRequest, serialize_response_format

-logger = get_logger("agent_framework.durabletask.durable_agent_state")
+logger = get_logger("agent_framework.azurefunctions.durable_agent_state")


 class DurableAgentStateEntryJsonType(str, Enum):
@@ -82,10 +82,7 @@ def _parse_created_at(value: Any) -> datetime:
        except (ValueError, TypeError):
            pass

-    logger.warning(
-        f"Invalid or missing created_at value in durable agent state; defaulting to current UTC time, {value}",
-        stack_info=True,
-    )
+    logger.warning("Invalid or missing created_at value in durable agent state; defaulting to current UTC time.")
    return datetime.now(tz=timezone.utc)


@@ -332,29 +329,24 @@ class DurableAgentStateData:
    Attributes:
        conversation_history: Ordered list of conversation entries (requests and responses)
        extension_data: Optional dictionary for custom metadata (not part of core schema)
-        expiration_time_utc: Optional UTC datetime when the entity should expire (TTL feature)
    """

    conversation_history: list[DurableAgentStateEntry]
    extension_data: dict[str, Any] | None
-    expiration_time_utc: datetime | None

    def __init__(
        self,
        conversation_history: list[DurableAgentStateEntry] | None = None,
        extension_data: dict[str, Any] | None = None,
-        expiration_time_utc: datetime | None = None,
    ) -> None:
        """Initialize the data container.

        Args:
            conversation_history: Initial conversation history (defaults to empty list)
            extension_data: Optional custom metadata
-            expiration_time_utc: Optional UTC datetime when the entity should expire
        """
        self.conversation_history = conversation_history or []
        self.extension_data = extension_data
-        self.expiration_time_utc = expiration_time_utc

    def to_dict(self) -> dict[str, Any]:
        result: dict[str, Any] = {
@@ -362,19 +354,13 @@ class DurableAgentStateData:
        }
        if self.extension_data is not None:
            result[DurableStateFields.EXTENSION_DATA] = self.extension_data
-        if self.expiration_time_utc is not None:
-            result[DurableStateFields.EXPIRATION_TIME_UTC] = self.expiration_time_utc.isoformat()
        return result

    @classmethod
    def from_dict(cls, data_dict: dict[str, Any]) -> DurableAgentStateData:
-        expiration_time_raw = data_dict.get(DurableStateFields.EXPIRATION_TIME_UTC)
-        expiration_time_utc = _parse_created_at(expiration_time_raw) if expiration_time_raw else None
-
        return cls(
            conversation_history=_parse_history_entries(data_dict),
            extension_data=data_dict.get(DurableStateFields.EXTENSION_DATA),
-            expiration_time_utc=expiration_time_utc,
        )


@@ -463,7 +449,7 @@ class DurableAgentState:
        """Get the count of conversation entries (requests + responses)."""
        return len(self.data.conversation_history)

-    def try_get_agent_response(self, correlation_id: str) -> AgentRunResponse | None:
+    def try_get_agent_response(self, correlation_id: str) -> dict[str, Any] | None:
        """Try to get an agent response by correlation ID.

        This method searches the conversation history for a response entry matching the given
@@ -485,8 +471,14 @@ class DurableAgentState:
        for entry in self.data.conversation_history:
            if entry.correlation_id == correlation_id and isinstance(entry, DurableAgentStateResponse):
                # Found the entry, extract response data
-                return DurableAgentStateResponse.to_run_response(entry)
+                # Get the text content from assistant messages only
+                content = "\n".join(message.text for message in entry.messages if message.text)

+                return {
+                    ApiResponseFields.CONTENT: content,
+                    ApiResponseFields.MESSAGE_COUNT: self.message_count,
+                    ApiResponseFields.CORRELATION_ID: correlation_id,
+                }
        return None


@@ -710,21 +702,6 @@ class DurableAgentStateResponse(DurableAgentStateEntry):
            usage=DurableAgentStateUsage.from_usage(response.usage_details),
        )

-    @staticmethod
-    def to_run_response(
-        response_entry: DurableAgentStateResponse,
-    ) -> AgentRunResponse:
-        """Converts a DurableAgentStateResponse back to an AgentRunResponse."""
-        messages = [m.to_chat_message() for m in response_entry.messages]
-
-        usage_details = response_entry.usage.to_usage_details() if response_entry.usage is not None else UsageDetails()
-
-        return AgentRunResponse(
-            created_at=response_entry.created_at.isoformat(),
-            messages=messages,
-            usage_details=usage_details,
-        )
-

 class DurableAgentStateMessage:
    """Represents a message within a conversation history entry.
@@ -1234,24 +1211,14 @@ class DurableAgentStateUsage:
            input_token_count=usage.input_token_count,
            output_token_count=usage.output_token_count,
            total_token_count=usage.total_token_count,
-            extensionData=usage.additional_counts,
        )

    def to_usage_details(self) -> UsageDetails:
        # Convert back to AI SDK UsageDetails
-        extension_data: dict[str, int] = {}
-        if self.extensionData is not None:
-            for k, v in self.extensionData.items():
-                try:
-                    extension_data[k] = int(v)
-                except (ValueError, TypeError):
-                    continue
-
        return UsageDetails(
            input_token_count=self.input_token_count,
            output_token_count=self.output_token_count,
            total_token_count=self.total_token_count,
-            **extension_data,
        )


@@ -8,41 +8,346 @@ allows for long-running agent conversations.
 """

 import asyncio
-from collections.abc import Callable
+import inspect
+from collections.abc import AsyncIterable, Callable
 from typing import Any, cast

 import azure.durable_functions as df
-from agent_framework import AgentProtocol, get_logger
-from agent_framework_durabletask import (
-    AgentEntity,
-    AgentEntityStateProviderMixin,
-    AgentResponseCallbackProtocol,
+from agent_framework import (
+    AgentProtocol,
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    ChatMessage,
+    ErrorContent,
+    Role,
+    get_logger,
 )

+from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
+from ._durable_agent_state import (
+    DurableAgentState,
+    DurableAgentStateData,
+    DurableAgentStateEntry,
+    DurableAgentStateRequest,
+    DurableAgentStateResponse,
+)
+from ._models import RunRequest
+
 logger = get_logger("agent_framework.azurefunctions.entities")


-class AzureFunctionEntityStateProvider(AgentEntityStateProviderMixin):
-    """Azure Functions Durable Entity state provider for AgentEntity.
+class AgentEntity:
+    """Durable entity that manages agent execution and conversation state.

-    This class utilizes the Durable Entity context from `azure-functions-durable` package
-    to get and set the state of the agent entity.
+    This entity:
+    - Maintains conversation history
+    - Executes agent with messages
+    - Stores agent responses
+    - Handles tool execution
+
+    Operations:
+    - run: Execute the agent with a message
+    - run_agent: (Deprecated) Execute the agent with a message
+    - reset: Clear conversation history
+
+    Attributes:
+        agent: The AgentProtocol instance
+        state: The DurableAgentState managing conversation history
    """

-    def __init__(self, context: df.DurableEntityContext) -> None:
-        self._context = context
+    agent: AgentProtocol
+    state: DurableAgentState

-    def _get_state_dict(self) -> dict[str, Any]:
-        raw_state = self._context.get_state(lambda: {})
-        if not isinstance(raw_state, dict):
-            return {}
-        return cast(dict[str, Any], raw_state)
+    def __init__(
+        self,
+        agent: AgentProtocol,
+        callback: AgentResponseCallbackProtocol | None = None,
+    ):
+        """Initialize the agent entity.

-    def _set_state_dict(self, state: dict[str, Any]) -> None:
-        self._context.set_state(state)
+        Args:
+            agent: The Microsoft Agent Framework agent instance (must implement AgentProtocol)
+            callback: Optional callback invoked during streaming updates and final responses
+        """
+        self.agent = agent
+        self.state = DurableAgentState()
+        self.callback = callback

-    def _get_thread_id_from_entity(self) -> str:
-        return self._context.entity_key
+        logger.debug(f"[AgentEntity] Initialized with agent type: {type(agent).__name__}")
+
+    def _is_error_response(self, entry: DurableAgentStateEntry) -> bool:
+        """Check if a conversation history entry is an error response.
+
+        Error responses should be kept in history for tracking but not sent to the agent
+        since Azure OpenAI doesn't support 'error' content type.
+
+        Args:
+            entry: A conversation history entry (DurableAgentStateEntry or dict)
+
+        Returns:
+            True if the entry is a response containing error content, False otherwise
+        """
+        if isinstance(entry, DurableAgentStateResponse):
+            return entry.is_error
+        return False
+
+    async def run_agent(
+        self,
+        context: df.DurableEntityContext,
+        request: RunRequest | dict[str, Any] | str,
+    ) -> AgentRunResponse:
+        """(Deprecated) Execute the agent with a message directly in the entity.
+
+        Args:
+            context: Entity context
+            request: RunRequest object, dict, or string message (for backward compatibility)
+
+        Returns:
+            AgentRunResponse enriched with execution metadata.
+        """
+        return await self.run(context, request)
+
+    async def run(
+        self,
+        context: df.DurableEntityContext,
+        request: RunRequest | dict[str, Any] | str,
+    ) -> AgentRunResponse:
+        """Execute the agent with a message directly in the entity.
+
+        Args:
+            context: Entity context
+            request: RunRequest object, dict, or string message (for backward compatibility)
+
+        Returns:
+            AgentRunResponse enriched with execution metadata.
+        """
+        if isinstance(request, str):
+            run_request = RunRequest(message=request, role=Role.USER)
+        elif isinstance(request, dict):
+            run_request = RunRequest.from_dict(request)
+        else:
+            run_request = request
+
+        message = run_request.message
+        thread_id = run_request.thread_id
+        correlation_id = run_request.correlation_id
+        if not thread_id:
+            raise ValueError("RunRequest must include a thread_id")
+        if not correlation_id:
+            raise ValueError("RunRequest must include a correlation_id")
+        response_format = run_request.response_format
+        enable_tool_calls = run_request.enable_tool_calls
+
+        state_request = DurableAgentStateRequest.from_run_request(run_request)
+        self.state.data.conversation_history.append(state_request)
+
+        logger.debug(f"[AgentEntity.run] Received Message: {state_request}")
+
+        try:
+            # Build messages from conversation history, excluding error responses
+            # Error responses are kept in history for tracking but not sent to the agent
+            chat_messages: list[ChatMessage] = [
+                m.to_chat_message()
+                for entry in self.state.data.conversation_history
+                if not self._is_error_response(entry)
+                for m in entry.messages
+            ]
+
+            run_kwargs: dict[str, Any] = {"messages": chat_messages}
+            if not enable_tool_calls:
+                run_kwargs["tools"] = None
+            if response_format:
+                run_kwargs["response_format"] = response_format
+
+            agent_run_response: AgentRunResponse = await self._invoke_agent(
+                run_kwargs=run_kwargs,
+                correlation_id=correlation_id,
+                thread_id=thread_id,
+                request_message=message,
+            )
+
+            logger.debug(
+                "[AgentEntity.run] Agent invocation completed - response type: %s",
+                type(agent_run_response).__name__,
+            )
+
+            try:
+                response_text = agent_run_response.text if agent_run_response.text else "No response"
+                logger.debug(f"Response: {response_text[:100]}...")
+            except Exception as extraction_error:
+                logger.error(
+                    "Error extracting response text: %s",
+                    extraction_error,
+                    exc_info=True,
+                )
+
+            state_response = DurableAgentStateResponse.from_run_response(correlation_id, agent_run_response)
+            self.state.data.conversation_history.append(state_response)
+
+            logger.debug("[AgentEntity.run] AgentRunResponse stored in conversation history")
+
+            return agent_run_response
+
+        except Exception as exc:
+            logger.exception("[AgentEntity.run] Agent execution failed.")
+
+            # Create error message
+            error_message = ChatMessage(
+                role=Role.ASSISTANT, contents=[ErrorContent(message=str(exc), error_code=type(exc).__name__)]
+            )
+
+            error_response = AgentRunResponse(messages=[error_message])
+
+            # Create and store error response in conversation history
+            error_state_response = DurableAgentStateResponse.from_run_response(correlation_id, error_response)
+            error_state_response.is_error = True
+            self.state.data.conversation_history.append(error_state_response)
+
+            return error_response
+
+    async def _invoke_agent(
+        self,
+        run_kwargs: dict[str, Any],
+        correlation_id: str,
+        thread_id: str,
+        request_message: str,
+    ) -> AgentRunResponse:
+        """Execute the agent, preferring streaming when available."""
+        callback_context: AgentCallbackContext | None = None
+        if self.callback is not None:
+            callback_context = self._build_callback_context(
+                correlation_id=correlation_id,
+                thread_id=thread_id,
+                request_message=request_message,
+            )
+
+        run_stream_callable = getattr(self.agent, "run_stream", None)
+        if callable(run_stream_callable):
+            try:
+                stream_candidate = run_stream_callable(**run_kwargs)
+                if inspect.isawaitable(stream_candidate):
+                    stream_candidate = await stream_candidate
+
+                return await self._consume_stream(
+                    stream=cast(AsyncIterable[AgentRunResponseUpdate], stream_candidate),
+                    callback_context=callback_context,
+                )
+            except TypeError as type_error:
+                if "__aiter__" not in str(type_error):
+                    raise
+                logger.debug(
+                    "run_stream returned a non-async result; falling back to run(): %s",
+                    type_error,
+                )
+            except Exception as stream_error:
+                logger.warning(
+                    "run_stream failed; falling back to run(): %s",
+                    stream_error,
+                    exc_info=True,
+                )
+        else:
+            logger.debug("Agent does not expose run_stream; falling back to run().")
+
+        agent_run_response = await self._invoke_non_stream(run_kwargs)
+        await self._notify_final_response(agent_run_response, callback_context)
+        return agent_run_response
+
+    async def _consume_stream(
+        self,
+        stream: AsyncIterable[AgentRunResponseUpdate],
+        callback_context: AgentCallbackContext | None = None,
+    ) -> AgentRunResponse:
+        """Consume streaming responses and build the final AgentRunResponse."""
+        updates: list[AgentRunResponseUpdate] = []
+
+        async for update in stream:
+            updates.append(update)
+            await self._notify_stream_update(update, callback_context)
+
+        if updates:
+            response = AgentRunResponse.from_agent_run_response_updates(updates)
+        else:
+            logger.debug("[AgentEntity] No streaming updates received; creating empty response")
+            response = AgentRunResponse(messages=[])
+
+        await self._notify_final_response(response, callback_context)
+        return response
+
+    async def _invoke_non_stream(self, run_kwargs: dict[str, Any]) -> AgentRunResponse:
+        """Invoke the agent without streaming support."""
+        run_callable = getattr(self.agent, "run", None)
+        if run_callable is None or not callable(run_callable):
+            raise AttributeError("Agent does not implement run() method")
+
+        result = run_callable(**run_kwargs)
+        if inspect.isawaitable(result):
+            result = await result
+
+        if not isinstance(result, AgentRunResponse):
+            raise TypeError(f"Agent run() must return an AgentRunResponse instance; received {type(result).__name__}")
+
+        return result
+
+    async def _notify_stream_update(
+        self,
+        update: AgentRunResponseUpdate,
+        context: AgentCallbackContext | None,
+    ) -> None:
+        """Invoke the streaming callback if one is registered."""
+        if self.callback is None or context is None:
+            return
+
+        try:
+            callback_result = self.callback.on_streaming_response_update(update, context)
+            if inspect.isawaitable(callback_result):
+                await callback_result
+        except Exception as exc:
+            logger.warning(
+                "[AgentEntity] Streaming callback raised an exception: %s",
+                exc,
+                exc_info=True,
+            )
+
+    async def _notify_final_response(
+        self,
+        response: AgentRunResponse,
+        context: AgentCallbackContext | None,
+    ) -> None:
+        """Invoke the final response callback if one is registered."""
+        if self.callback is None or context is None:
+            return
+
+        try:
+            callback_result = self.callback.on_agent_response(response, context)
+            if inspect.isawaitable(callback_result):
+                await callback_result
+        except Exception as exc:
+            logger.warning(
+                "[AgentEntity] Response callback raised an exception: %s",
+                exc,
+                exc_info=True,
+            )
+
+    def _build_callback_context(
+        self,
+        correlation_id: str,
+        thread_id: str,
+        request_message: str,
+    ) -> AgentCallbackContext:
+        """Create the callback context provided to consumers."""
+        agent_name = getattr(self.agent, "name", None) or type(self.agent).__name__
+        return AgentCallbackContext(
+            agent_name=agent_name,
+            correlation_id=correlation_id,
+            thread_id=thread_id,
+            request_message=request_message,
+        )
+
+    def reset(self, context: df.DurableEntityContext) -> None:
+        """Reset the entity state (clear conversation history)."""
+        logger.debug("[AgentEntity.reset] Resetting entity state")
+        self.state.data = DurableAgentStateData(conversation_history=[])
+        logger.debug("[AgentEntity.reset] State reset complete")


 def create_agent_entity(
@@ -63,10 +368,19 @@ def create_agent_entity(
        """Async handler that executes the entity operations."""
        try:
            logger.debug("[entity_function] Entity triggered")
-            logger.debug("[entity_function] Operation: %s", context.operation_name)
+            logger.debug(f"[entity_function] Operation: {context.operation_name}")

-            state_provider = AzureFunctionEntityStateProvider(context)
-            entity = AgentEntity(agent, callback, state_provider=state_provider)
+            current_state = context.get_state(lambda: None)
+            logger.debug("Retrieved state: %s", str(current_state)[:100])
+            entity = AgentEntity(agent, callback)
+
+            if current_state is not None:
+                entity.state = DurableAgentState.from_dict(current_state)
+                logger.debug(
+                    "[entity_function] Restored entity from state (message_count: %s)", entity.state.message_count
+                )
+            else:
+                logger.debug("[entity_function] Created new entity instance")

            operation = context.operation_name

@@ -80,18 +394,21 @@ def create_agent_entity(
                    # Fall back to treating input as message string
                    request = "" if input_data is None else str(cast(object, input_data))

-                result = await entity.run(request)
+                result = await entity.run(context, request)
                context.set_result(result.to_dict())

            elif operation == "reset":
-                entity.reset()
+                entity.reset(context)
                context.set_result({"status": "reset"})

            else:
                logger.error("[entity_function] Unknown operation: %s", operation)
                context.set_result({"error": f"Unknown operation: {operation}"})

-            logger.info("[entity_function] Operation %s completed successfully", operation)
+            serialized_state = entity.state.to_dict()
+            logger.debug("State dict: %s", serialized_state)
+            context.set_state(serialized_state)
+            logger.info(f"[entity_function] Operation {operation} completed successfully")

        except Exception as exc:
            logger.exception("[entity_function] Error executing entity operation %s", exc)
@@ -8,14 +8,13 @@ This module defines the request and response models used by the framework.
 from __future__ import annotations

 import inspect
-import json
 import uuid
 from collections.abc import MutableMapping
 from dataclasses import dataclass
-from datetime import datetime, timezone
 from importlib import import_module
 from typing import TYPE_CHECKING, Any, cast

+import azure.durable_functions as df
 from agent_framework import AgentThread, Role

 from ._constants import REQUEST_RESPONSE_FORMAT_TEXT
@@ -33,6 +32,187 @@ else:
    _PydanticBaseModel = _RuntimeBaseModel


+@dataclass
+class AgentSessionId:
+    """Represents an agent session ID, which is used to identify a long-running agent session.
+
+    Attributes:
+        name: The name of the agent that owns the session (case-insensitive)
+        key: The unique key of the agent session (case-sensitive)
+    """
+
+    name: str
+    key: str
+
+    ENTITY_NAME_PREFIX: str = "dafx-"
+
+    @staticmethod
+    def to_entity_name(name: str) -> str:
+        """Converts an agent name to an entity name by adding the DAFx prefix.
+
+        Args:
+            name: The agent name
+
+        Returns:
+            The entity name with the dafx- prefix
+        """
+        return f"{AgentSessionId.ENTITY_NAME_PREFIX}{name}"
+
+    @staticmethod
+    def with_random_key(name: str) -> AgentSessionId:
+        """Creates a new AgentSessionId with the specified name and a randomly generated key.
+
+        Args:
+            name: The name of the agent that owns the session
+
+        Returns:
+            A new AgentSessionId with the specified name and a random GUID key
+        """
+        return AgentSessionId(name=name, key=uuid.uuid4().hex)
+
+    def to_entity_id(self) -> df.EntityId:
+        """Converts this AgentSessionId to a Durable Functions EntityId.
+
+        Returns:
+            EntityId for use with Durable Functions APIs
+        """
+        return df.EntityId(self.to_entity_name(self.name), self.key)
+
+    @staticmethod
+    def from_entity_id(entity_id: df.EntityId) -> AgentSessionId:
+        """Creates an AgentSessionId from a Durable Functions EntityId.
+
+        Args:
+            entity_id: The EntityId to convert
+
+        Returns:
+            AgentSessionId instance
+
+        Raises:
+            ValueError: If the entity ID does not have the expected prefix
+        """
+        if not entity_id.name.startswith(AgentSessionId.ENTITY_NAME_PREFIX):
+            raise ValueError(
+                f"'{entity_id}' is not a valid agent session ID. "
+                f"Expected entity name to start with '{AgentSessionId.ENTITY_NAME_PREFIX}'"
+            )
+
+        agent_name = entity_id.name[len(AgentSessionId.ENTITY_NAME_PREFIX) :]
+        return AgentSessionId(name=agent_name, key=entity_id.key)
+
+    def __str__(self) -> str:
+        """Returns a string representation in the form @name@key."""
+        return f"@{self.name}@{self.key}"
+
+    def __repr__(self) -> str:
+        """Returns a detailed string representation."""
+        return f"AgentSessionId(name='{self.name}', key='{self.key}')"
+
+    @staticmethod
+    def parse(session_id_string: str) -> AgentSessionId:
+        """Parses a string representation of an agent session ID.
+
+        Args:
+            session_id_string: A string in the form @name@key
+
+        Returns:
+            AgentSessionId instance
+
+        Raises:
+            ValueError: If the string format is invalid
+        """
+        if not session_id_string.startswith("@"):
+            raise ValueError(f"Invalid agent session ID format: {session_id_string}")
+
+        parts = session_id_string[1:].split("@", 1)
+        if len(parts) != 2:
+            raise ValueError(f"Invalid agent session ID format: {session_id_string}")
+
+        return AgentSessionId(name=parts[0], key=parts[1])
+
+
+class DurableAgentThread(AgentThread):
+    """Durable agent thread that tracks the owning :class:`AgentSessionId`."""
+
+    _SERIALIZED_SESSION_ID_KEY = "durable_session_id"
+
+    def __init__(
+        self,
+        *,
+        session_id: AgentSessionId | None = None,
+        service_thread_id: str | None = None,
+        message_store: Any = None,
+        context_provider: Any = None,
+    ) -> None:
+        super().__init__(
+            service_thread_id=service_thread_id,
+            message_store=message_store,
+            context_provider=context_provider,
+        )
+        self._session_id: AgentSessionId | None = session_id
+
+    @property
+    def session_id(self) -> AgentSessionId | None:
+        """Returns the durable agent session identifier for this thread."""
+        return self._session_id
+
+    def attach_session(self, session_id: AgentSessionId) -> None:
+        """Associates the thread with the provided :class:`AgentSessionId`."""
+        self._session_id = session_id
+
+    @classmethod
+    def from_session_id(
+        cls,
+        session_id: AgentSessionId,
+        *,
+        service_thread_id: str | None = None,
+        message_store: Any = None,
+        context_provider: Any = None,
+    ) -> DurableAgentThread:
+        """Creates a durable thread pre-associated with the supplied session ID."""
+        return cls(
+            session_id=session_id,
+            service_thread_id=service_thread_id,
+            message_store=message_store,
+            context_provider=context_provider,
+        )
+
+    async def serialize(self, **kwargs: Any) -> dict[str, Any]:
+        """Serializes thread state including the durable session identifier."""
+        state = await super().serialize(**kwargs)
+        if self._session_id is not None:
+            state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id)
+        return state
+
+    @classmethod
+    async def deserialize(
+        cls,
+        serialized_thread_state: MutableMapping[str, Any],
+        *,
+        message_store: Any = None,
+        **kwargs: Any,
+    ) -> DurableAgentThread:
+        """Restores a durable thread, rehydrating the stored session identifier."""
+        state_payload = dict(serialized_thread_state)
+        session_id_value = state_payload.pop(cls._SERIALIZED_SESSION_ID_KEY, None)
+        thread = await super().deserialize(
+            state_payload,
+            message_store=message_store,
+            **kwargs,
+        )
+        if not isinstance(thread, DurableAgentThread):
+            raise TypeError("Deserialized thread is not a DurableAgentThread instance")
+
+        if session_id_value is None:
+            return thread
+
+        if not isinstance(session_id_value, str):
+            raise ValueError("durable_session_id must be a string when present in serialized state")
+
+        thread.attach_session(AgentSessionId.parse(session_id_value))
+        return thread
+
+
 def serialize_response_format(response_format: type[BaseModel] | None) -> Any:
    """Serialize response format for transport across durable function boundaries."""
    if response_format is None:
@@ -104,43 +284,42 @@ class RunRequest:
        role: The role of the message sender (user, system, or assistant)
        response_format: Optional Pydantic BaseModel type describing the structured response format
        enable_tool_calls: Whether to enable tool calls for this request
-        wait_for_response: If True (default), caller will wait for agent response. If False,
-                          returns immediately after signaling (fire-and-forget mode)
-        correlation_id: Correlation ID for tracking the response to this specific request
+        thread_id: Optional thread ID for tracking
+        correlation_id: Optional correlation ID for tracking the response to this specific request
        created_at: Optional timestamp when the request was created
        orchestration_id: Optional ID of the orchestration that initiated this request
    """

    message: str
    request_response_format: str
-    correlation_id: str
    role: Role = Role.USER
    response_format: type[BaseModel] | None = None
    enable_tool_calls: bool = True
-    wait_for_response: bool = True
-    created_at: datetime | None = None
+    thread_id: str | None = None
+    correlation_id: str | None = None
+    created_at: str | None = None
    orchestration_id: str | None = None

    def __init__(
        self,
        message: str,
-        correlation_id: str,
        request_response_format: str = REQUEST_RESPONSE_FORMAT_TEXT,
        role: Role | str | None = Role.USER,
        response_format: type[BaseModel] | None = None,
        enable_tool_calls: bool = True,
-        wait_for_response: bool = True,
-        created_at: datetime | None = None,
+        thread_id: str | None = None,
+        correlation_id: str | None = None,
+        created_at: str | None = None,
        orchestration_id: str | None = None,
    ) -> None:
        self.message = message
-        self.correlation_id = correlation_id
        self.role = self.coerce_role(role)
        self.response_format = response_format
        self.request_response_format = request_response_format
        self.enable_tool_calls = enable_tool_calls
-        self.wait_for_response = wait_for_response
-        self.created_at = created_at if created_at is not None else datetime.now(tz=timezone.utc)
+        self.thread_id = thread_id
+        self.correlation_id = correlation_id
+        self.created_at = created_at
        self.orchestration_id = orchestration_id

    @staticmethod
@@ -160,155 +339,33 @@ class RunRequest:
        result = {
            "message": self.message,
            "enable_tool_calls": self.enable_tool_calls,
-            "wait_for_response": self.wait_for_response,
            "role": self.role.value,
            "request_response_format": self.request_response_format,
-            "correlationId": self.correlation_id,
        }
        if self.response_format:
            result["response_format"] = serialize_response_format(self.response_format)
+        if self.thread_id:
+            result["thread_id"] = self.thread_id
+        if self.correlation_id:
+            result["correlationId"] = self.correlation_id
        if self.created_at:
-            result["created_at"] = self.created_at.isoformat()
+            result["created_at"] = self.created_at
        if self.orchestration_id:
            result["orchestrationId"] = self.orchestration_id

        return result

-    @classmethod
-    def from_json(cls, data: str) -> RunRequest:
-        """Create RunRequest from JSON string."""
-        try:
-            dict_data = json.loads(data)
-        except json.JSONDecodeError as e:
-            raise ValueError("The durable agent state is not valid JSON.") from e
-
-        return cls.from_dict(dict_data)
-
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> RunRequest:
        """Create RunRequest from dictionary."""
-        created_at = data.get("created_at")
-        if isinstance(created_at, str):
-            try:
-                created_at = datetime.fromisoformat(created_at)
-            except ValueError:
-                created_at = None
-
-        correlation_id = data.get("correlationId")
-        if not correlation_id:
-            raise ValueError("correlationId is required in RunRequest data")
-
        return cls(
            message=data.get("message", ""),
-            correlation_id=correlation_id,
            request_response_format=data.get("request_response_format", REQUEST_RESPONSE_FORMAT_TEXT),
            role=cls.coerce_role(data.get("role")),
            response_format=_deserialize_response_format(data.get("response_format")),
-            wait_for_response=data.get("wait_for_response", True),
            enable_tool_calls=data.get("enable_tool_calls", True),
-            created_at=created_at,
+            thread_id=data.get("thread_id"),
+            correlation_id=data.get("correlationId"),
+            created_at=data.get("created_at"),
            orchestration_id=data.get("orchestrationId"),
        )
-
-
-@dataclass
-class AgentSessionId:
-    """Represents an agent session identifier (name + key)."""
-
-    name: str
-    key: str
-
-    ENTITY_NAME_PREFIX: str = "dafx-"
-
-    @staticmethod
-    def to_entity_name(name: str) -> str:
-        return f"{AgentSessionId.ENTITY_NAME_PREFIX}{name}"
-
-    @staticmethod
-    def with_random_key(name: str) -> AgentSessionId:
-        return AgentSessionId(name=name, key=uuid.uuid4().hex)
-
-    @property
-    def entity_name(self) -> str:
-        return self.to_entity_name(self.name)
-
-    def __str__(self) -> str:
-        return f"@{self.name}@{self.key}"
-
-    def __repr__(self) -> str:
-        return f"AgentSessionId(name='{self.name}', key='{self.key}')"
-
-    @staticmethod
-    def parse(session_id_string: str) -> AgentSessionId:
-        if not session_id_string.startswith("@"):
-            raise ValueError(f"Invalid agent session ID format: {session_id_string}")
-
-        parts = session_id_string[1:].split("@", 1)
-        if len(parts) != 2:
-            raise ValueError(f"Invalid agent session ID format: {session_id_string}")
-
-        return AgentSessionId(name=parts[0], key=parts[1])
-
-
-class DurableAgentThread(AgentThread):
-    """Durable agent thread that tracks the owning :class:`AgentSessionId`."""
-
-    _SERIALIZED_SESSION_ID_KEY = "durable_session_id"
-
-    def __init__(
-        self,
-        *,
-        session_id: AgentSessionId | None = None,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(**kwargs)
-        self._session_id: AgentSessionId | None = session_id
-
-    @property
-    def session_id(self) -> AgentSessionId | None:
-        return self._session_id
-
-    @session_id.setter
-    def session_id(self, value: AgentSessionId | None) -> None:
-        self._session_id = value
-
-    @classmethod
-    def from_session_id(
-        cls,
-        session_id: AgentSessionId,
-        **kwargs: Any,
-    ) -> DurableAgentThread:
-        return cls(session_id=session_id, **kwargs)
-
-    async def serialize(self, **kwargs: Any) -> dict[str, Any]:
-        state = await super().serialize(**kwargs)
-        if self._session_id is not None:
-            state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id)
-        return state
-
-    @classmethod
-    async def deserialize(
-        cls,
-        serialized_thread_state: MutableMapping[str, Any],
-        *,
-        message_store: Any = None,
-        **kwargs: Any,
-    ) -> DurableAgentThread:
-        state_payload = dict(serialized_thread_state)
-        session_id_value = state_payload.pop(cls._SERIALIZED_SESSION_ID_KEY, None)
-        thread = await super().deserialize(
-            state_payload,
-            message_store=message_store,
-            **kwargs,
-        )
-        if not isinstance(thread, DurableAgentThread):
-            raise TypeError("Deserialized thread is not a DurableAgentThread instance")
-
-        if session_id_value is None:
-            return thread
-
-        if not isinstance(session_id_value, str):
-            raise ValueError("durable_session_id must be a string when present in serialized state")
-
-        thread.session_id = AgentSessionId.parse(session_id_value)
-        return thread
@@ -5,22 +5,24 @@
 This module provides support for using agents inside Durable Function orchestrations.
 """

-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, TypeAlias
+import uuid
+from collections.abc import AsyncIterator, Callable
+from typing import TYPE_CHECKING, Any, TypeAlias, cast

-import azure.durable_functions as df
-from agent_framework import AgentThread, get_logger
-from agent_framework_durabletask import (
-    DurableAgentExecutor,
-    RunRequest,
-    ensure_response_format,
-    load_agent_response,
+from agent_framework import (
+    AgentProtocol,
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    ChatMessage,
+    get_logger,
 )
 from azure.durable_functions.models import TaskBase
-from azure.durable_functions.models.actions.NoOpAction import NoOpAction
 from azure.durable_functions.models.Task import CompoundTask, TaskState
 from pydantic import BaseModel

+from ._models import AgentSessionId, DurableAgentThread, RunRequest
+
 logger = get_logger("agent_framework.azurefunctions.orchestration")

 CompoundActionConstructor: TypeAlias = Callable[[list[Any]], Any] | None
@@ -43,25 +45,6 @@ else:
    _TypedCompoundTask = CompoundTask


-class PreCompletedTask(TaskBase):
-    """A simple task that is already completed with a result.
-
-    Used for fire-and-forget mode where we want to return immediately
-    with an acceptance response without waiting for entity processing.
-    """
-
-    def __init__(self, result: Any):
-        """Initialize with a completed result.
-
-        Args:
-            result: The result value for this completed task
-        """
-        # Initialize with a NoOp action since we don't need actual orchestration actions
-        super().__init__(-1, NoOpAction())
-        # Immediately mark as completed with the result
-        self.set_value(is_error=False, value=result)
-
-
 class AgentTask(_TypedCompoundTask):
    """A custom Task that wraps entity calls and provides typed AgentRunResponse results.

@@ -82,13 +65,10 @@ class AgentTask(_TypedCompoundTask):
            response_format: Optional Pydantic model for response parsing
            correlation_id: Correlation ID for logging
        """
-        # Set instance variables BEFORE calling super().__init__
-        # because super().__init__ may trigger try_set_value for pre-completed tasks
+        super().__init__([entity_task])
        self._response_format = response_format
        self._correlation_id = correlation_id

-        super().__init__([entity_task])
-
        # Override action_repr to expose the inner task's action directly
        # This ensures compatibility with ReplaySchema V3 which expects Action objects.
        self.action_repr = entity_task.action_repr
@@ -115,10 +95,10 @@ class AgentTask(_TypedCompoundTask):
                )

                try:
-                    response = load_agent_response(raw_result)
+                    response = self._load_agent_response(raw_result)

                    if self._response_format is not None:
-                        ensure_response_format(
+                        self._ensure_response_format(
                            self._response_format,
                            self._correlation_id,
                            response,
@@ -138,88 +118,249 @@ class AgentTask(_TypedCompoundTask):
                self._first_error = child.result
                self.set_value(is_error=True, value=self._first_error)

+    def _load_agent_response(self, agent_response: AgentRunResponse | dict[str, Any] | None) -> AgentRunResponse:
+        """Convert raw payloads into AgentRunResponse instance."""
+        if agent_response is None:
+            raise ValueError("agent_response cannot be None")

-class AzureFunctionsAgentExecutor(DurableAgentExecutor[AgentTask]):
-    """Executor that executes durable agents inside Azure Functions orchestrations."""
+        logger.debug("[load_agent_response] Loading agent response of type: %s", type(agent_response))

-    def __init__(self, context: AgentOrchestrationContextType):
-        self.context = context
+        if isinstance(agent_response, AgentRunResponse):
+            return agent_response
+        if isinstance(agent_response, dict):
+            logger.debug("[load_agent_response] Converting dict payload using AgentRunResponse.from_dict")
+            return AgentRunResponse.from_dict(agent_response)

-    def generate_unique_id(self) -> str:
-        return str(self.context.new_uuid())
+        raise TypeError(f"Unsupported type for agent_response: {type(agent_response)}")

-    def get_run_request(
+    def _ensure_response_format(
        self,
-        message: str,
        response_format: type[BaseModel] | None,
-        enable_tool_calls: bool,
-        wait_for_response: bool = True,
-    ) -> RunRequest:
-        """Get the current run request from the orchestration context.
+        correlation_id: str,
+        response: AgentRunResponse,
+    ) -> None:
+        """Ensure the AgentRunResponse value is parsed into the expected response_format."""
+        if response_format is not None and not isinstance(response.value, response_format):
+            response.try_parse_value(response_format)
+
+            logger.debug(
+                "[DurableAIAgent] Loaded AgentRunResponse.value for correlation_id %s with type: %s",
+                correlation_id,
+                type(response.value).__name__,
+            )
+
+
+class DurableAIAgent(AgentProtocol):
+    """A durable agent implementation that uses entity methods to interact with agent entities.
+
+    This class implements AgentProtocol and provides methods to work with Azure Durable Functions
+    orchestrations, which use generators and yield instead of async/await.
+
+    Key methods:
+    - get_new_thread(): Create a new conversation thread
+    - run(): Execute the agent and return a Task for yielding in orchestrations
+
+    Note: The run() method is NOT async. It returns a Task directly that must be
+    yielded in orchestrations to wait for the entity call to complete.
+
+    Example usage in orchestration:
+        writer = app.get_agent(context, "WriterAgent")
+        thread = writer.get_new_thread()  # NOT yielded - returns immediately
+
+        response = yield writer.run(  # Yielded - waits for entity call
+            message="Write a haiku about coding",
+            thread=thread
+        )
+    """
+
+    def __init__(self, context: AgentOrchestrationContextType, agent_name: str):
+        """Initialize the DurableAIAgent.

        Args:
-            message: The message to send to the agent
+            context: The orchestration context
+            agent_name: Name of the agent (used to construct entity ID)
+        """
+        self.context = context
+        self.agent_name = agent_name
+        self._id = str(uuid.uuid4())
+        self._name = agent_name
+        self._display_name = agent_name
+        self._description = f"Durable agent proxy for {agent_name}"
+        logger.debug("[DurableAIAgent] Initialized for agent: %s", agent_name)
+
+    @property
+    def id(self) -> str:
+        """Get the unique identifier for this agent."""
+        return self._id
+
+    @property
+    def name(self) -> str | None:
+        """Get the name of the agent."""
+        return self._name
+
+    @property
+    def display_name(self) -> str:
+        """Get the display name of the agent."""
+        return self._display_name
+
+    @property
+    def description(self) -> str | None:
+        """Get the description of the agent."""
+        return self._description
+
+    # We return an AgentTask here which is a TaskBase subclass.
+    # This is an intentional deviation from AgentProtocol which defines run() as async.
+    # The AgentTask can be yielded in Durable Functions orchestrations and will provide
+    # a typed AgentRunResponse result.
+    def run(  # type: ignore[override]
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        response_format: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> AgentTask:
+        """Execute the agent with messages and return an AgentTask for orchestrations.
+
+        This method implements AgentProtocol and returns an AgentTask (subclass of TaskBase)
+        that can be yielded in Durable Functions orchestrations. The task's result will be
+        a typed AgentRunResponse.
+
+        Args:
+            messages: The message(s) to send to the agent
+            thread: Optional agent thread for conversation context
            response_format: Optional Pydantic model for response parsing
-            enable_tool_calls: Whether to enable tool calls
-            wait_for_response: Must be True for orchestration contexts
+            **kwargs: Additional arguments (enable_tool_calls)

        Returns:
-            RunRequest: The current run request
+            An AgentTask that resolves to an AgentRunResponse when yielded

-        Raises:
-            ValueError: If wait_for_response=False (not supported in orchestrations)
+        Example:
+            @app.orchestration_trigger(context_name="context")
+            def my_orchestration(context):
+                agent = app.get_agent(context, "MyAgent")
+                thread = agent.get_new_thread()
+                response = yield agent.run("Hello", thread=thread)
+                # response is typed as AgentRunResponse
        """
-        request = super().get_run_request(
-            message,
-            response_format,
-            enable_tool_calls,
-            wait_for_response,
-        )
-        request.orchestration_id = self.context.instance_id
-        return request
+        message_str = self._normalize_messages(messages)

-    def run_durable_agent(
-        self,
-        agent_name: str,
-        run_request: RunRequest,
-        thread: AgentThread | None = None,
-    ) -> AgentTask:
+        # Extract optional parameters from kwargs
+        enable_tool_calls = kwargs.get("enable_tool_calls", True)

-        # Resolve session
-        session_id = self._create_session_id(agent_name, thread)
+        # Get the session ID for the entity
+        if isinstance(thread, DurableAgentThread) and thread.session_id is not None:
+            session_id = thread.session_id
+        else:
+            # Create a unique session ID for each call when no thread is provided
+            # This ensures each call gets its own conversation context
+            session_key = str(self.context.new_uuid())
+            session_id = AgentSessionId(name=self.agent_name, key=session_key)
+            logger.debug("[DurableAIAgent] No thread provided, created unique session_id: %s", session_id)

-        entity_id = df.EntityId(
-            name=session_id.entity_name,
-            key=session_id.key,
-        )
+        # Create entity ID from session ID
+        entity_id = session_id.to_entity_id()

+        # Generate a deterministic correlation ID for this call
+        # This is required by the entity and must be unique per call
+        correlation_id = str(self.context.new_uuid())
        logger.debug(
-            "[AzureFunctionsAgentProvider] correlation_id: %s entity_id: %s session_id: %s",
-            run_request.correlation_id,
+            "[DurableAIAgent] Using correlation_id: %s for entity_id: %s for session_id: %s",
+            correlation_id,
            entity_id,
            session_id,
        )

-        # Branch based on wait_for_response
-        if not run_request.wait_for_response:
-            # Fire-and-forget mode: signal entity and return pre-completed task
-            logger.debug(
-                "[AzureFunctionsAgentExecutor] Fire-and-forget mode: signaling entity (correlation: %s)",
-                run_request.correlation_id,
-            )
-            self.context.signal_entity(entity_id, "run", run_request.to_dict())
-
-            # Create acceptance response using base class helper
-            acceptance_response = self._create_acceptance_response(run_request.correlation_id)
-
-            # Create a pre-completed task with the acceptance response
-            entity_task = PreCompletedTask(acceptance_response)
-        else:
-            # Blocking mode: call entity and wait for response
-            entity_task = self.context.call_entity(entity_id, "run", run_request.to_dict())
-
-        return AgentTask(
-            entity_task=entity_task,
-            response_format=run_request.response_format,
-            correlation_id=run_request.correlation_id,
+        # Prepare the request using RunRequest model
+        # Include the orchestration's instance_id so it can be stored in the agent's entity state
+        run_request = RunRequest(
+            message=message_str,
+            enable_tool_calls=enable_tool_calls,
+            correlation_id=correlation_id,
+            thread_id=session_id.key,
+            response_format=response_format,
+            orchestration_id=self.context.instance_id,
        )
+
+        logger.debug("[DurableAIAgent] Calling entity %s with message: %s", entity_id, message_str[:100])
+
+        # Call the entity to get the underlying task
+        entity_task = self.context.call_entity(entity_id, "run", run_request.to_dict())
+
+        # Wrap it in an AgentTask that will convert the result to AgentRunResponse
+        agent_task = AgentTask(
+            entity_task=entity_task,
+            response_format=response_format,
+            correlation_id=correlation_id,
+        )
+
+        logger.debug(
+            "[DurableAIAgent] Created AgentTask for correlation_id %s",
+            correlation_id,
+        )
+
+        return agent_task
+
+    def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[AgentRunResponseUpdate]:
+        """Run the agent with streaming (not supported for durable agents).
+
+        Raises:
+            NotImplementedError: Streaming is not supported for durable agents.
+        """
+        raise NotImplementedError("Streaming is not supported for durable agents in orchestrations.")
+
+    def get_new_thread(self, **kwargs: Any) -> AgentThread:
+        """Create a new agent thread for this orchestration instance.
+
+        Each call creates a unique thread with its own conversation context.
+        The session ID is deterministic (uses context.new_uuid()) to ensure
+        orchestration replay works correctly.
+
+        Returns:
+            A new AgentThread instance with a unique session ID
+        """
+        # Generate a deterministic unique key for this thread
+        # Using context.new_uuid() ensures the same GUID is generated during replay
+        session_key = str(self.context.new_uuid())
+
+        # Create AgentSessionId with agent name and session key
+        session_id = AgentSessionId(name=self.agent_name, key=session_key)
+
+        thread = DurableAgentThread.from_session_id(session_id, **kwargs)
+
+        logger.debug("[DurableAIAgent] Created new thread with session_id: %s", session_id)
+        return thread
+
+    def _messages_to_string(self, messages: list[ChatMessage]) -> str:
+        """Convert a list of ChatMessage objects to a single string.
+
+        Args:
+            messages: List of ChatMessage objects
+
+        Returns:
+            Concatenated string of message contents
+        """
+        return "\n".join([msg.text or "" for msg in messages])
+
+    def _normalize_messages(self, messages: str | ChatMessage | list[str] | list[ChatMessage] | None) -> str:
+        """Convert supported message inputs to a single string."""
+        if messages is None:
+            return ""
+        if isinstance(messages, str):
+            return messages
+        if isinstance(messages, ChatMessage):
+            return messages.text or ""
+        if isinstance(messages, list):
+            if not messages:
+                return ""
+            first_item = messages[0]
+            if isinstance(first_item, str):
+                return "\n".join(cast(list[str], messages))
+            return self._messages_to_string(cast(list[ChatMessage], messages))
+        return str(messages)
@@ -4,7 +4,7 @@ description = "Azure Functions integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -23,7 +23,6 @@ classifiers = [
 ]
 dependencies = [
    "agent-framework-core",
-    "agent-framework-durabletask",
    "azure-functions",
    "azure-functions-durable",
 ]
@@ -15,7 +15,8 @@ Usage:
 """

 import pytest
-from agent_framework_durabletask import THREAD_ID_HEADER
+
+from agent_framework_azurefunctions._constants import THREAD_ID_HEADER

 from .testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled

@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for Callbacks Sample
+
+Tests the callbacks sample for event tracking and management.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite or Azure Storage account configured
+
+Usage:
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_03_callbacks.py -v
+"""
+
+from typing import Any
+
+import pytest
+import requests
+
+from .testutils import (
+    TIMEOUT,
+    SampleTestHelper,
+    skip_if_azure_functions_integration_tests_disabled,
+)
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.sample("03_callbacks"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+class TestSampleCallbacks:
+    """Tests for 03_callbacks sample."""
+
+    @pytest.fixture(autouse=True)
+    def _set_base_url(self, base_url: str) -> None:
+        """Provide the callback agent base URL for each test."""
+        self.base_url = f"{base_url}/api/agents/CallbackAgent"
+
+    @staticmethod
+    def _wait_for_callback_events(base_url: str, thread_id: str) -> list[dict[str, Any]]:
+        events: list[dict[str, Any]] = []
+        response = SampleTestHelper.get(f"{base_url}/callbacks/{thread_id}")
+        if response.status_code == 200:
+            events = response.json()
+        return events
+
+    def test_agent_with_callbacks(self) -> None:
+        """Test agent execution with callback tracking."""
+        thread_id = "test-callback"
+
+        response = SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Tell me about Python", "thread_id": thread_id},
+        )
+        assert response.status_code == 200
+        data = response.json()
+
+        assert data["status"] == "success"
+
+        events = self._wait_for_callback_events(self.base_url, thread_id)
+
+        assert events
+        assert any(event.get("event_type") == "final" for event in events)
+
+    def test_get_callbacks(self) -> None:
+        """Test retrieving callback events."""
+        thread_id = "test-callback-retrieve"
+
+        # Send a message first
+        SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Hello", "thread_id": thread_id, "wait_for_response": False},
+        )
+
+        # Get callbacks
+        response = SampleTestHelper.get(f"{self.base_url}/callbacks/{thread_id}")
+        assert response.status_code == 200
+        data = response.json()
+        assert isinstance(data, list)
+
+    def test_delete_callbacks(self) -> None:
+        """Test clearing callback events."""
+        thread_id = "test-callback-delete"
+
+        # Send a message first
+        SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Test", "thread_id": thread_id, "wait_for_response": False},
+        )
+
+        # Delete callbacks
+        response = requests.delete(f"{self.base_url}/callbacks/{thread_id}", timeout=TIMEOUT)
+        assert response.status_code == 204
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -1,125 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-"""
-Integration Tests for Reliable Streaming Sample
-
-Tests the reliable streaming sample using Redis Streams for persistent message delivery.
-
-The function app is automatically started by the test fixture.
-
-Prerequisites:
- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
- Azurite or Azure Storage account configured
- Redis running (docker run -d --name redis -p 6379:6379 redis:latest)
-
-Usage:
-    uv run pytest packages/azurefunctions/tests/integration_tests/test_03_reliable_streaming.py -v
-"""
-
-import time
-
-import pytest
-import requests
-
-from .testutils import (
-    SampleTestHelper,
-    skip_if_azure_functions_integration_tests_disabled,
-)
-
-# Module-level markers - applied to all tests in this file
-pytestmark = [
-    pytest.mark.sample("03_reliable_streaming"),
-    pytest.mark.usefixtures("function_app_for_test"),
-    skip_if_azure_functions_integration_tests_disabled,
-]
-
-
-class TestSampleReliableStreaming:
-    """Tests for 03_reliable_streaming sample."""
-
-    @pytest.fixture(autouse=True)
-    def _set_base_url(self, base_url: str) -> None:
-        """Provide the base URL for each test."""
-        self.base_url = base_url
-        self.agent_url = f"{base_url}/api/agents/TravelPlanner"
-        self.stream_url = f"{base_url}/api/agent/stream"
-
-    def test_agent_run_and_stream(self) -> None:
-        """Test agent execution with Redis streaming."""
-        # Start agent run
-        response = SampleTestHelper.post_json(
-            f"{self.agent_url}/run",
-            {"message": "Plan a 1-day trip to Seattle in 1 sentence", "wait_for_response": False},
-        )
-        assert response.status_code == 202
-        data = response.json()
-
-        thread_id = data.get("thread_id")
-
-        # Wait a moment for the agent to start writing to Redis
-        time.sleep(2)
-
-        # Stream response from Redis with shorter timeout
-        # Note: We use text/plain to avoid SSE parsing complexity
-        stream_response = requests.get(
-            f"{self.stream_url}/{thread_id}",
-            headers={"Accept": "text/plain"},
-            timeout=30,  # Shorter timeout for test
-        )
-        assert stream_response.status_code == 200
-
-    def test_stream_with_sse_format(self) -> None:
-        """Test streaming with Server-Sent Events format."""
-        # Start agent run
-        response = SampleTestHelper.post_json(
-            f"{self.agent_url}/run",
-            {"message": "What's the weather like?", "wait_for_response": False},
-        )
-        assert response.status_code == 202
-        data = response.json()
-        thread_id = data.get("thread_id")
-
-        # Wait for agent to start writing
-        time.sleep(2)
-
-        # Stream with SSE format
-        stream_response = requests.get(
-            f"{self.stream_url}/{thread_id}",
-            headers={"Accept": "text/event-stream"},
-            timeout=30,  # Shorter timeout
-        )
-        assert stream_response.status_code == 200
-        content_type = stream_response.headers.get("content-type", "")
-        assert "text/event-stream" in content_type
-
-        # Check for SSE event markers if we got content
-        content = stream_response.text
-        if content:
-            assert "event:" in content or "data:" in content
-
-    def test_stream_nonexistent_conversation(self) -> None:
-        """Test streaming from a non-existent conversation.
-
-        The endpoint will wait for data in Redis, but since the conversation
-        doesn't exist, it will timeout. This is expected behavior.
-        """
-        fake_id = "nonexistent-conversation-12345"
-
-        # Should timeout since the conversation doesn't exist
-        with pytest.raises(requests.exceptions.ReadTimeout):
-            requests.get(
-                f"{self.stream_url}/{fake_id}",
-                headers={"Accept": "text/plain"},
-                timeout=10,  # Short timeout for non-existent ID
-            )
-
-    def test_health_endpoint(self) -> None:
-        """Test health check endpoint."""
-        response = SampleTestHelper.get(f"{self.base_url}/api/health")
-        assert response.status_code == 200
-        data = response.json()
-        assert data["status"] == "healthy"
-        assert "agents" in data
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])
@@ -2,8 +2,6 @@

 """Unit tests for AgentFunctionApp."""

-# pyright: reportPrivateUsage=false
-
 import json
 from collections.abc import Awaitable, Callable
 from typing import Any, TypeVar
@@ -13,42 +11,20 @@ import azure.durable_functions as df
 import azure.functions as func
 import pytest
 from agent_framework import AgentRunResponse, ChatMessage, ErrorContent
-from agent_framework_durabletask import (
+
+from agent_framework_azurefunctions import AgentFunctionApp
+from agent_framework_azurefunctions._app import WAIT_FOR_RESPONSE_FIELD, WAIT_FOR_RESPONSE_HEADER
+from agent_framework_azurefunctions._constants import (
    MIMETYPE_APPLICATION_JSON,
    MIMETYPE_TEXT_PLAIN,
    THREAD_ID_HEADER,
-    WAIT_FOR_RESPONSE_FIELD,
-    WAIT_FOR_RESPONSE_HEADER,
-    AgentEntity,
-    AgentEntityStateProviderMixin,
-    DurableAgentState,
 )
-
-from agent_framework_azurefunctions import AgentFunctionApp
-from agent_framework_azurefunctions._entities import create_agent_entity
+from agent_framework_azurefunctions._durable_agent_state import DurableAgentState
+from agent_framework_azurefunctions._entities import AgentEntity, create_agent_entity

 TFunc = TypeVar("TFunc", bound=Callable[..., Any])


-def _identity_decorator(func: TFunc) -> TFunc:
-    return func
-
-
-class _InMemoryStateProvider(AgentEntityStateProviderMixin):
-    def __init__(self, *, thread_id: str = "test-thread", initial_state: dict[str, Any] | None = None) -> None:
-        self._thread_id = thread_id
-        self._state_dict: dict[str, Any] = initial_state or {}
-
-    def _get_state_dict(self) -> dict[str, Any]:
-        return self._state_dict
-
-    def _set_state_dict(self, state: dict[str, Any]) -> None:
-        self._state_dict = state
-
-    def _get_thread_id_from_entity(self) -> str:
-        return self._thread_id
-
-
 class TestAgentFunctionAppInit:
    """Test suite for AgentFunctionApp initialization."""

@@ -112,7 +88,7 @@ class TestAgentFunctionAppInit:
            app.add_agent(mock_agent, callback=specific_callback)

        setup_mock.assert_called_once()
-        _, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
+        _, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
        assert passed_callback is specific_callback
        assert enable_http_endpoint is True

@@ -128,7 +104,7 @@ class TestAgentFunctionAppInit:
            app.add_agent(mock_agent)

        setup_mock.assert_called_once()
-        _, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
+        _, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
        assert passed_callback is default_callback
        assert enable_http_endpoint is True

@@ -143,7 +119,7 @@ class TestAgentFunctionAppInit:
            AgentFunctionApp(agents=[mock_agent], default_callback=default_callback)

        setup_mock.assert_called_once()
-        _, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
+        _, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
        assert passed_callback is default_callback
        assert enable_http_endpoint is True

@@ -359,12 +335,13 @@ class TestAgentEntityOperations:
            return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Test response")])
        )

-        entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="test-conv-123"))
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()

-        result = await entity.run({
-            "message": "Test message",
-            "correlationId": "corr-app-entity-1",
-        })
+        result = await entity.run(
+            mock_context,
+            {"message": "Test message", "thread_id": "test-conv-123", "correlationId": "corr-app-entity-1"},
+        )

        assert isinstance(result, AgentRunResponse)
        assert result.text == "Test response"
@@ -377,17 +354,22 @@ class TestAgentEntityOperations:
            return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response 1")])
        )

-        entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()

        # Send first message
-        await entity.run({"message": "Message 1", "correlationId": "corr-app-entity-2"})
+        await entity.run(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-app-entity-2"}
+        )

        # Each conversation turn creates 2 entries: request and response
        history = entity.state.data.conversation_history[0].messages  # Request entry
        assert len(history) == 1  # Just the user message

        # Send second message
-        await entity.run({"message": "Message 2", "correlationId": "corr-app-entity-2b"})
+        await entity.run(
+            mock_context, {"message": "Message 2", "thread_id": "conv-2", "correlationId": "corr-app-entity-2b"}
+        )

        # Now we have 4 entries total (2 requests + 2 responses)
        # Access the first request entry
@@ -411,26 +393,32 @@ class TestAgentEntityOperations:
            return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response")])
        )

-        entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()

        assert len(entity.state.data.conversation_history) == 0

-        await entity.run({"message": "Message 1", "correlationId": "corr-app-entity-3a"})
+        await entity.run(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-app-entity-3a"}
+        )
        assert len(entity.state.data.conversation_history) == 2

-        await entity.run({"message": "Message 2", "correlationId": "corr-app-entity-3b"})
+        await entity.run(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-app-entity-3b"}
+        )
        assert len(entity.state.data.conversation_history) == 4

    def test_entity_reset(self) -> None:
        """Test that entity reset clears state."""
        mock_agent = Mock()
-        entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider())
+        entity = AgentEntity(mock_agent)

        # Set some state
        entity.state = DurableAgentState()

        # Reset
-        entity.reset()
+        mock_context = Mock()
+        entity.reset(mock_context)

        assert len(entity.state.data.conversation_history) == 0

@@ -459,6 +447,7 @@ class TestAgentEntityFactory:
        mock_context.operation_name = "run"
        mock_context.get_input.return_value = {
            "message": "Test message",
+            "thread_id": "conv-123",
            "correlationId": "corr-app-factory-1",
        }
        mock_context.get_state.return_value = None
@@ -486,6 +475,7 @@ class TestAgentEntityFactory:
        mock_context.operation_name = "run_agent"
        mock_context.get_input.return_value = {
            "message": "Test message",
+            "thread_id": "conv-123",
            "correlationId": "corr-app-factory-1",
        }
        mock_context.get_state.return_value = None
@@ -605,11 +595,7 @@ class TestAgentEntityFactory:
        }

        mock_context = Mock()
-        mock_context.operation_name = "run"
-        mock_context.get_input.return_value = {
-            "message": "Test message",
-            "correlationId": "corr-restore-1",
-        }
+        mock_context.operation_name = "reset"
        mock_context.get_state.return_value = existing_state

        with patch.object(DurableAgentState, "from_dict", wraps=DurableAgentState.from_dict) as from_dict_mock:
@@ -626,12 +612,12 @@ class TestErrorHandling:
        mock_agent = Mock()
        mock_agent.run = AsyncMock(side_effect=Exception("Agent error"))

-        entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()

-        result = await entity.run({
-            "message": "Test message",
-            "correlationId": "corr-app-error-1",
-        })
+        result = await entity.run(
+            mock_context, {"message": "Test message", "thread_id": "conv-1", "correlationId": "corr-app-error-1"}
+        )

        assert isinstance(result, AgentRunResponse)
        assert len(result.messages) == 1
@@ -724,7 +710,7 @@ class TestIncomingRequestParsing:

        request = Mock()
        request.params = {"thread_id": "query-thread"}
-        req_body: dict[str, Any] = {}
+        req_body = {}

        thread_id = app._resolve_thread_id(request, req_body)

@@ -791,7 +777,7 @@ class TestHttpRunRoute:

        assert run_request["message"] == "Plain text via HTTP"
        assert run_request["role"] == "user"
-        assert "thread_id" not in run_request
+        assert "thread_id" in run_request

    async def test_http_run_accept_header_returns_json(self) -> None:
        """Test that Accept header requesting JSON results in JSON response."""
@@ -927,9 +913,9 @@ class TestMCPToolEndpoint:
            patch.object(app, "durable_client_input") as client_mock,
        ):
            # Setup mock decorator chain
-            func_name_mock.return_value = _identity_decorator
-            mcp_trigger_mock.return_value = _identity_decorator
-            client_mock.return_value = _identity_decorator
+            func_name_mock.return_value = lambda f: f
+            mcp_trigger_mock.return_value = lambda f: f
+            client_mock.return_value = lambda f: f

            app._setup_mcp_tool_trigger(mock_agent.name, mock_agent.description)

@@ -952,11 +938,11 @@ class TestMCPToolEndpoint:
        app = AgentFunctionApp()

        with (
-            patch.object(app, "function_name", return_value=_identity_decorator),
+            patch.object(app, "function_name", return_value=lambda f: f),
            patch.object(app, "mcp_tool_trigger") as mcp_trigger_mock,
-            patch.object(app, "durable_client_input", return_value=_identity_decorator),
+            patch.object(app, "durable_client_input", return_value=lambda f: f),
        ):
-            mcp_trigger_mock.return_value = _identity_decorator
+            mcp_trigger_mock.return_value = lambda f: f

            app._setup_mcp_tool_trigger(mock_agent.name, None)

@@ -1078,10 +1064,10 @@ class TestMCPToolEndpoint:
        app = AgentFunctionApp(agents=[mock_agent], enable_mcp_tool_trigger=True)

        # Capture the health check handler function
-        captured_handler: Callable[[func.HttpRequest], func.HttpResponse] | None = None
+        captured_handler = None

-        def capture_decorator(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
-            def decorator(func: TFunc) -> TFunc:
+        def capture_decorator(*args, **kwargs):
+            def decorator(func):
                nonlocal captured_handler
                captured_handler = func
                return func
@@ -1,22 +1,42 @@
 # Copyright (c) Microsoft. All rights reserved.

-"""Unit tests for create_agent_entity factory function.
+"""Unit tests for AgentEntity and entity operations.

 Run with: pytest tests/test_entities.py -v
 """

-from collections.abc import Callable
+import asyncio
+from collections.abc import AsyncIterator, Callable
+from datetime import datetime
 from typing import Any, TypeVar
-from unittest.mock import AsyncMock, Mock
+from unittest.mock import AsyncMock, Mock, patch

 import pytest
-from agent_framework import AgentRunResponse, ChatMessage
+from agent_framework import AgentRunResponse, AgentRunResponseUpdate, ChatMessage, ErrorContent, Role
+from pydantic import BaseModel

-from agent_framework_azurefunctions._entities import create_agent_entity
+from agent_framework_azurefunctions._durable_agent_state import (
+    DurableAgentState,
+    DurableAgentStateData,
+    DurableAgentStateMessage,
+    DurableAgentStateRequest,
+    DurableAgentStateTextContent,
+)
+from agent_framework_azurefunctions._entities import AgentEntity, create_agent_entity
+from agent_framework_azurefunctions._models import RunRequest

 TFunc = TypeVar("TFunc", bound=Callable[..., Any])


+def _role_value(chat_message: DurableAgentStateMessage) -> str:
+    """Helper to extract the string role from a ChatMessage."""
+    role = getattr(chat_message, "role", None)
+    role_value = getattr(role, "value", role)
+    if role_value is None:
+        return ""
+    return str(role_value)
+
+
 def _agent_response(text: str | None) -> AgentRunResponse:
    """Create an AgentRunResponse with a single assistant message."""
    message = (
@@ -25,6 +45,379 @@ def _agent_response(text: str | None) -> AgentRunResponse:
    return AgentRunResponse(messages=[message])


+class RecordingCallback:
+    """Callback implementation capturing streaming and final responses for assertions."""
+
+    def __init__(self):
+        self.stream_mock = AsyncMock()
+        self.response_mock = AsyncMock()
+
+    async def on_streaming_response_update(
+        self,
+        update: AgentRunResponseUpdate,
+        context: Any,
+    ) -> None:
+        await self.stream_mock(update, context)
+
+    async def on_agent_response(self, response: AgentRunResponse, context: Any) -> None:
+        await self.response_mock(response, context)
+
+
+class EntityStructuredResponse(BaseModel):
+    answer: float
+
+
+class TestAgentEntityInit:
+    """Test suite for AgentEntity initialization."""
+
+    def test_init_creates_entity(self) -> None:
+        """Test that AgentEntity initializes correctly."""
+        mock_agent = Mock()
+
+        entity = AgentEntity(mock_agent)
+
+        assert entity.agent == mock_agent
+        assert len(entity.state.data.conversation_history) == 0
+        assert entity.state.data.extension_data is None
+        assert entity.state.schema_version == DurableAgentState.SCHEMA_VERSION
+
+    def test_init_stores_agent_reference(self) -> None:
+        """Test that the agent reference is stored correctly."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        entity = AgentEntity(mock_agent)
+
+        assert entity.agent.name == "TestAgent"
+
+    def test_init_with_different_agent_types(self) -> None:
+        """Test initialization with different agent types."""
+        agent1 = Mock()
+        agent1.__class__.__name__ = "AzureOpenAIAgent"
+
+        agent2 = Mock()
+        agent2.__class__.__name__ = "CustomAgent"
+
+        entity1 = AgentEntity(agent1)
+        entity2 = AgentEntity(agent2)
+
+        assert entity1.agent.__class__.__name__ == "AzureOpenAIAgent"
+        assert entity2.agent.__class__.__name__ == "CustomAgent"
+
+
+class TestAgentEntityRunAgent:
+    """Test suite for the run_agent operation."""
+
+    async def test_run_executes_agent(self) -> None:
+        """Test that run executes the agent."""
+        mock_agent = Mock()
+        mock_response = _agent_response("Test response")
+        mock_agent.run = AsyncMock(return_value=mock_response)
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run(
+            mock_context, {"message": "Test message", "thread_id": "conv-123", "correlationId": "corr-entity-1"}
+        )
+
+        # Verify agent.run was called
+        mock_agent.run.assert_called_once()
+        _, kwargs = mock_agent.run.call_args
+        sent_messages: list[Any] = kwargs.get("messages")
+        assert len(sent_messages) == 1
+        sent_message = sent_messages[0]
+        assert isinstance(sent_message, ChatMessage)
+        assert getattr(sent_message, "text", None) == "Test message"
+        assert getattr(sent_message.role, "value", sent_message.role) == "user"
+
+        # Verify result
+        assert isinstance(result, AgentRunResponse)
+        assert result.text == "Test response"
+
+    async def test_run_agent_executes_agent(self) -> None:
+        """Test that run_agent executes the agent."""
+        mock_agent = Mock()
+        mock_response = _agent_response("Test response")
+        mock_agent.run = AsyncMock(return_value=mock_response)
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Test message", "thread_id": "conv-123", "correlationId": "corr-entity-1"}
+        )
+
+        # Verify agent.run was called
+        mock_agent.run.assert_called_once()
+        _, kwargs = mock_agent.run.call_args
+        sent_messages: list[Any] = kwargs.get("messages")
+        assert len(sent_messages) == 1
+        sent_message = sent_messages[0]
+        assert isinstance(sent_message, ChatMessage)
+        assert getattr(sent_message, "text", None) == "Test message"
+        assert getattr(sent_message.role, "value", sent_message.role) == "user"
+
+        # Verify result
+        assert isinstance(result, AgentRunResponse)
+        assert result.text == "Test response"
+
+    async def test_run_agent_streaming_callbacks_invoked(self) -> None:
+        """Ensure streaming updates trigger callbacks and run() is not used."""
+
+        updates = [
+            AgentRunResponseUpdate(text="Hello"),
+            AgentRunResponseUpdate(text=" world"),
+        ]
+
+        async def update_generator() -> AsyncIterator[AgentRunResponseUpdate]:
+            for update in updates:
+                yield update
+
+        mock_agent = Mock()
+        mock_agent.name = "StreamingAgent"
+        mock_agent.run_stream = Mock(return_value=update_generator())
+        mock_agent.run = AsyncMock(side_effect=AssertionError("run() should not be called when streaming succeeds"))
+
+        callback = RecordingCallback()
+        entity = AgentEntity(mock_agent, callback=callback)
+        mock_context = Mock()
+
+        result = await entity.run(
+            mock_context,
+            {
+                "message": "Tell me something",
+                "thread_id": "session-1",
+                "correlationId": "corr-stream-1",
+            },
+        )
+
+        assert isinstance(result, AgentRunResponse)
+        assert "Hello" in result.text
+        assert callback.stream_mock.await_count == len(updates)
+        assert callback.response_mock.await_count == 1
+        mock_agent.run.assert_not_called()
+
+        # Validate callback arguments
+        stream_calls = callback.stream_mock.await_args_list
+        for expected_update, recorded_call in zip(updates, stream_calls, strict=True):
+            assert recorded_call.args[0] is expected_update
+            context = recorded_call.args[1]
+            assert context.agent_name == "StreamingAgent"
+            assert context.correlation_id == "corr-stream-1"
+            assert context.thread_id == "session-1"
+            assert context.request_message == "Tell me something"
+
+        final_call = callback.response_mock.await_args
+        assert final_call is not None
+        final_response, final_context = final_call.args
+        assert final_context.agent_name == "StreamingAgent"
+        assert final_context.correlation_id == "corr-stream-1"
+        assert final_context.thread_id == "session-1"
+        assert final_context.request_message == "Tell me something"
+        assert getattr(final_response, "text", "").strip()
+
+    async def test_run_agent_final_callback_without_streaming(self) -> None:
+        """Ensure the final callback fires even when streaming is unavailable."""
+
+        mock_agent = Mock()
+        mock_agent.name = "NonStreamingAgent"
+        mock_agent.run_stream = None
+        agent_response = _agent_response("Final response")
+        mock_agent.run = AsyncMock(return_value=agent_response)
+
+        callback = RecordingCallback()
+        entity = AgentEntity(mock_agent, callback=callback)
+        mock_context = Mock()
+
+        result = await entity.run(
+            mock_context,
+            {
+                "message": "Hi",
+                "thread_id": "session-2",
+                "correlationId": "corr-final-1",
+            },
+        )
+
+        assert isinstance(result, AgentRunResponse)
+        assert result.text == "Final response"
+        assert callback.stream_mock.await_count == 0
+        assert callback.response_mock.await_count == 1
+
+        final_call = callback.response_mock.await_args
+        assert final_call is not None
+        assert final_call.args[0] is agent_response
+        final_context = final_call.args[1]
+        assert final_context.agent_name == "NonStreamingAgent"
+        assert final_context.correlation_id == "corr-final-1"
+        assert final_context.thread_id == "session-2"
+        assert final_context.request_message == "Hi"
+
+    async def test_run_agent_updates_conversation_history(self) -> None:
+        """Test that run_agent updates the conversation history."""
+        mock_agent = Mock()
+        mock_response = _agent_response("Agent response")
+        mock_agent.run = AsyncMock(return_value=mock_response)
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        await entity.run(
+            mock_context, {"message": "User message", "thread_id": "conv-1", "correlationId": "corr-entity-2"}
+        )
+
+        # Should have 1 entry: user message + assistant response
+        user_history = entity.state.data.conversation_history[0].messages
+        assistant_history = entity.state.data.conversation_history[1].messages
+
+        assert len(user_history) == 1
+
+        user_msg = user_history[0]
+        assert _role_value(user_msg) == "user"
+        assert user_msg.text == "User message"
+
+        assistant_msg = assistant_history[0]
+        assert _role_value(assistant_msg) == "assistant"
+        assert assistant_msg.text == "Agent response"
+
+    async def test_run_agent_increments_message_count(self) -> None:
+        """Test that run_agent increments the message count."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        assert len(entity.state.data.conversation_history) == 0
+
+        await entity.run(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-3a"}
+        )
+        assert len(entity.state.data.conversation_history) == 2
+
+        await entity.run(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-3b"}
+        )
+        assert len(entity.state.data.conversation_history) == 4
+
+        await entity.run(
+            mock_context, {"message": "Message 3", "thread_id": "conv-1", "correlationId": "corr-entity-3c"}
+        )
+        assert len(entity.state.data.conversation_history) == 6
+
+    async def test_run_agent_with_none_thread_id(self) -> None:
+        """Test run_agent with a None thread identifier."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        with pytest.raises(ValueError, match="thread_id"):
+            await entity.run(mock_context, {"message": "Message", "thread_id": None, "correlationId": "corr-entity-5"})
+
+    async def test_run_agent_multiple_conversations(self) -> None:
+        """Test that run_agent maintains history across multiple messages."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Send multiple messages
+        await entity.run(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-8a"}
+        )
+        await entity.run(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-8b"}
+        )
+        await entity.run(
+            mock_context, {"message": "Message 3", "thread_id": "conv-1", "correlationId": "corr-entity-8c"}
+        )
+
+        history = entity.state.data.conversation_history
+        assert len(history) == 6
+        assert entity.state.message_count == 6
+
+
+class TestAgentEntityReset:
+    """Test suite for the reset operation."""
+
+    def test_reset_clears_conversation_history(self) -> None:
+        """Test that reset clears the conversation history."""
+        mock_agent = Mock()
+        entity = AgentEntity(mock_agent)
+
+        # Add some history with proper DurableAgentStateEntry objects
+        entity.state.data.conversation_history = [
+            DurableAgentStateRequest(
+                correlation_id="test-1",
+                created_at=datetime.now(),
+                messages=[
+                    DurableAgentStateMessage(
+                        role="user",
+                        contents=[DurableAgentStateTextContent(text="msg1")],
+                    )
+                ],
+            ),
+        ]
+
+        mock_context = Mock()
+        entity.reset(mock_context)
+
+        assert entity.state.data.conversation_history == []
+
+    def test_reset_with_extension_data(self) -> None:
+        """Test that reset works when entity has extension data."""
+        mock_agent = Mock()
+        entity = AgentEntity(mock_agent)
+
+        # Set up some initial state with conversation history
+        entity.state.data = DurableAgentStateData(conversation_history=[], extension_data={"some_key": "some_value"})
+
+        mock_context = Mock()
+        entity.reset(mock_context)
+
+        assert len(entity.state.data.conversation_history) == 0
+
+    def test_reset_clears_message_count(self) -> None:
+        """Test that reset clears the message count."""
+        mock_agent = Mock()
+        entity = AgentEntity(mock_agent)
+
+        mock_context = Mock()
+        entity.reset(mock_context)
+
+        assert len(entity.state.data.conversation_history) == 0
+
+    async def test_reset_after_conversation(self) -> None:
+        """Test reset after a full conversation."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Have a conversation
+        await entity.run(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-10a"}
+        )
+        await entity.run(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-10b"}
+        )
+
+        # Verify state before reset
+        assert entity.state.message_count == 4
+        assert len(entity.state.data.conversation_history) == 4
+
+        # Reset
+        entity.reset(mock_context)
+
+        # Verify state after reset
+        assert entity.state.message_count == 0
+        assert len(entity.state.data.conversation_history) == 0
+
+
 class TestCreateAgentEntity:
    """Test suite for the create_agent_entity factory function."""

@@ -46,9 +439,9 @@ class TestCreateAgentEntity:
        # Mock context
        mock_context = Mock()
        mock_context.operation_name = "run"
-        mock_context.entity_key = "conv-123"
        mock_context.get_input.return_value = {
            "message": "Test message",
+            "thread_id": "conv-123",
            "correlationId": "corr-entity-factory",
        }
        mock_context.get_state.return_value = None
@@ -142,7 +535,7 @@ class TestCreateAgentEntity:
        assert state["data"] == {"conversationHistory": []}

    def test_entity_function_restores_existing_state(self) -> None:
-        """Test that the entity function can operate when existing state is present."""
+        """Test that the entity function restores existing state."""
        mock_agent = Mock()

        entity_function = create_agent_entity(mock_agent)
@@ -191,14 +584,482 @@ class TestCreateAgentEntity:
        mock_context.operation_name = "reset"
        mock_context.get_state.return_value = existing_state

+        with patch.object(DurableAgentState, "from_dict", wraps=DurableAgentState.from_dict) as from_dict_mock:
+            entity_function(mock_context)
+
+        from_dict_mock.assert_called_once_with(existing_state)
+
+
+class TestErrorHandling:
+    """Test suite for error handling in entities."""
+
+    async def test_run_agent_handles_agent_exception(self) -> None:
+        """Test that run_agent handles agent exceptions."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=Exception("Agent failed"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-1"}
+        )
+
+        assert isinstance(result, AgentRunResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, ErrorContent)
+        assert "Agent failed" in (content.message or "")
+        assert content.error_code == "Exception"
+
+    async def test_run_agent_handles_value_error(self) -> None:
+        """Test that run_agent handles ValueError instances."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=ValueError("Invalid input"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-2"}
+        )
+
+        assert isinstance(result, AgentRunResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, ErrorContent)
+        assert content.error_code == "ValueError"
+        assert "Invalid input" in str(content.message)
+
+    async def test_run_agent_handles_timeout_error(self) -> None:
+        """Test that run_agent handles TimeoutError instances."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=TimeoutError("Request timeout"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-3"}
+        )
+
+        assert isinstance(result, AgentRunResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, ErrorContent)
+        assert content.error_code == "TimeoutError"
+
+    def test_entity_function_handles_exception_in_operation(self) -> None:
+        """Test that the entity function handles exceptions gracefully."""
+        mock_agent = Mock()
+
+        entity_function = create_agent_entity(mock_agent)
+
+        mock_context = Mock()
+        mock_context.operation_name = "run"
+        mock_context.get_input.side_effect = Exception("Input error")
+        mock_context.get_state.return_value = None
+
+        # Execute - should not raise
        entity_function(mock_context)

+        # Verify error was set
        assert mock_context.set_result.called
+        result = mock_context.set_result.call_args[0][0]
+        assert "error" in result

-        # Reset should clear history and persist via set_state
-        assert mock_context.set_state.called
-        persisted_state = mock_context.set_state.call_args[0][0]
-        assert persisted_state["data"]["conversationHistory"] == []
+    async def test_run_agent_preserves_message_on_error(self) -> None:
+        """Test that run_agent preserves message information on error."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=Exception("Error"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run(
+            mock_context,
+            {"message": "Test message", "thread_id": "conv-123", "correlationId": "corr-entity-error-4"},
+        )
+
+        # Even on error, message info should be preserved
+        assert isinstance(result, AgentRunResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, ErrorContent)
+
+
+class TestConversationHistory:
+    """Test suite for conversation history tracking."""
+
+    async def test_conversation_history_has_timestamps(self) -> None:
+        """Test that conversation history entries include timestamps."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        await entity.run(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-history-1"}
+        )
+
+        # Check both user and assistant messages have timestamps
+        for entry in entity.state.data.conversation_history:
+            timestamp = entry.created_at
+            assert timestamp is not None
+            # Verify timestamp is in ISO format
+            datetime.fromisoformat(str(timestamp))
+
+    async def test_conversation_history_ordering(self) -> None:
+        """Test that conversation history maintains the correct order."""
+        mock_agent = Mock()
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Send multiple messages with different responses
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 1"))
+        await entity.run(
+            mock_context,
+            {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-history-2a"},
+        )
+
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 2"))
+        await entity.run(
+            mock_context,
+            {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-history-2b"},
+        )
+
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 3"))
+        await entity.run(
+            mock_context,
+            {"message": "Message 3", "thread_id": "conv-1", "correlationId": "corr-entity-history-2c"},
+        )
+
+        # Verify order
+        history = entity.state.data.conversation_history
+        # Each conversation turn creates 2 entries: request and response
+        assert history[0].messages[0].text == "Message 1"  # Request 1
+        assert history[1].messages[0].text == "Response 1"  # Response 1
+        assert history[2].messages[0].text == "Message 2"  # Request 2
+        assert history[3].messages[0].text == "Response 2"  # Response 2
+        assert history[4].messages[0].text == "Message 3"  # Request 3
+        assert history[5].messages[0].text == "Response 3"  # Response 3
+
+    async def test_conversation_history_role_alternation(self) -> None:
+        """Test that conversation history alternates between user and assistant roles."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        await entity.run(
+            mock_context,
+            {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-history-3a"},
+        )
+        await entity.run(
+            mock_context,
+            {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-history-3b"},
+        )
+
+        # Check role alternation
+        history = entity.state.data.conversation_history
+        # Each conversation turn creates 2 entries: request and response
+        assert history[0].messages[0].role == "user"  # Request 1
+        assert history[1].messages[0].role == "assistant"  # Response 1
+        assert history[2].messages[0].role == "user"  # Request 2
+        assert history[3].messages[0].role == "assistant"  # Response 2
+
+
+class TestRunRequestSupport:
+    """Test suite for RunRequest support in entities."""
+
+    async def test_run_agent_with_run_request_object(self) -> None:
+        """Test run_agent with a RunRequest object."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request = RunRequest(
+            message="Test message",
+            thread_id="conv-123",
+            role=Role.USER,
+            enable_tool_calls=True,
+            correlation_id="corr-runreq-1",
+        )
+
+        result = await entity.run(mock_context, request)
+
+        assert isinstance(result, AgentRunResponse)
+        assert result.text == "Response"
+
+    async def test_run_agent_with_dict_request(self) -> None:
+        """Test run_agent with a dictionary request."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request_dict = {
+            "message": "Test message",
+            "thread_id": "conv-456",
+            "role": "system",
+            "enable_tool_calls": False,
+            "correlationId": "corr-runreq-2",
+        }
+
+        result = await entity.run(mock_context, request_dict)
+
+        assert isinstance(result, AgentRunResponse)
+        assert result.text == "Response"
+
+    async def test_run_agent_with_string_raises_without_correlation(self) -> None:
+        """Test that run_agent rejects legacy string input without correlation ID."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        with pytest.raises(ValueError):
+            await entity.run(mock_context, "Simple message")
+
+    async def test_run_agent_stores_role_in_history(self) -> None:
+        """Test that run_agent stores the role in conversation history."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Send as system role
+        request = RunRequest(
+            message="System message",
+            thread_id="conv-runreq-3",
+            role=Role.SYSTEM,
+            correlation_id="corr-runreq-3",
+        )
+
+        await entity.run(mock_context, request)
+
+        # Check that system role was stored
+        history = entity.state.data.conversation_history
+        assert history[0].messages[0].role == "system"
+        assert history[0].messages[0].text == "System message"
+
+    async def test_run_agent_with_response_format(self) -> None:
+        """Test run_agent with a JSON response format."""
+        mock_agent = Mock()
+        # Return JSON response
+        mock_agent.run = AsyncMock(return_value=_agent_response('{"answer": 42}'))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request = RunRequest(
+            message="What is the answer?",
+            thread_id="conv-runreq-4",
+            response_format=EntityStructuredResponse,
+            correlation_id="corr-runreq-4",
+        )
+
+        result = await entity.run(mock_context, request)
+
+        assert isinstance(result, AgentRunResponse)
+        assert result.text == '{"answer": 42}'
+        assert result.value is None
+
+    async def test_run_agent_disable_tool_calls(self) -> None:
+        """Test run_agent with tool calls disabled."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request = RunRequest(
+            message="Test", thread_id="conv-runreq-5", enable_tool_calls=False, correlation_id="corr-runreq-5"
+        )
+
+        result = await entity.run(mock_context, request)
+
+        assert isinstance(result, AgentRunResponse)
+        # Agent should have been called (tool disabling is framework-dependent)
+        mock_agent.run.assert_called_once()
+
+    async def test_entity_function_with_run_request_dict(self) -> None:
+        """Test that the entity function handles the RunRequest dict format."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity_function = create_agent_entity(mock_agent)
+
+        mock_context = Mock()
+        mock_context.operation_name = "run"
+        mock_context.get_input.return_value = {
+            "message": "Test message",
+            "thread_id": "conv-789",
+            "role": "user",
+            "enable_tool_calls": True,
+            "correlationId": "corr-runreq-6",
+        }
+        mock_context.get_state.return_value = None
+
+        await asyncio.to_thread(entity_function, mock_context)
+
+        # Verify result was set
+        assert mock_context.set_result.called
+        result = mock_context.set_result.call_args[0][0]
+        assert isinstance(result, dict)
+
+        # Check if messages are present
+        assert "messages" in result
+        assert len(result["messages"]) > 0
+        message = result["messages"][0]
+
+        # Check for text in various possible locations
+        text_found = False
+        if "text" in message and message["text"] == "Response":
+            text_found = True
+        elif "contents" in message:
+            for content in message["contents"]:
+                if isinstance(content, dict) and content.get("text") == "Response":
+                    text_found = True
+                    break
+
+        assert text_found, f"Response text not found in message: {message}"
+
+
+class TestDurableAgentStateRequestOrchestrationId:
+    """Test suite for DurableAgentStateRequest orchestration_id field."""
+
+    def test_request_with_orchestration_id(self) -> None:
+        """Test creating a request with an orchestration_id."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+            orchestration_id="orch-456",
+        )
+
+        assert request.orchestration_id == "orch-456"
+
+    def test_request_to_dict_includes_orchestration_id(self) -> None:
+        """Test that to_dict includes orchestrationId when set."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+            orchestration_id="orch-789",
+        )
+
+        data = request.to_dict()
+
+        assert "orchestrationId" in data
+        assert data["orchestrationId"] == "orch-789"
+
+    def test_request_to_dict_excludes_orchestration_id_when_none(self) -> None:
+        """Test that to_dict excludes orchestrationId when not set."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+        )
+
+        data = request.to_dict()
+
+        assert "orchestrationId" not in data
+
+    def test_request_from_dict_with_orchestration_id(self) -> None:
+        """Test from_dict correctly parses orchestrationId."""
+        data = {
+            "$type": "request",
+            "correlationId": "corr-123",
+            "createdAt": "2024-01-01T00:00:00Z",
+            "messages": [{"role": "user", "contents": [{"$type": "text", "text": "test"}]}],
+            "orchestrationId": "orch-from-dict",
+        }
+
+        request = DurableAgentStateRequest.from_dict(data)
+
+        assert request.orchestration_id == "orch-from-dict"
+
+    def test_request_from_run_request_with_orchestration_id(self) -> None:
+        """Test from_run_request correctly transfers orchestration_id."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+            orchestration_id="orch-from-run-request",
+        )
+
+        durable_request = DurableAgentStateRequest.from_run_request(run_request)
+
+        assert durable_request.orchestration_id == "orch-from-run-request"
+
+    def test_request_from_run_request_without_orchestration_id(self) -> None:
+        """Test from_run_request correctly handles missing orchestration_id."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+        )
+
+        durable_request = DurableAgentStateRequest.from_run_request(run_request)
+
+        assert durable_request.orchestration_id is None
+
+
+class TestDurableAgentStateMessageCreatedAt:
+    """Test suite for DurableAgentStateMessage created_at field handling."""
+
+    def test_message_from_run_request_without_created_at_preserves_none(self) -> None:
+        """Test from_run_request preserves None created_at instead of defaulting to current time.
+
+        When a RunRequest has no created_at value, the resulting DurableAgentStateMessage
+        should also have None for created_at, not default to current UTC time.
+        """
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+            created_at=None,  # Explicitly None
+        )
+
+        durable_message = DurableAgentStateMessage.from_run_request(run_request)
+
+        assert durable_message.created_at is None
+
+    def test_message_from_run_request_with_created_at_parses_correctly(self) -> None:
+        """Test from_run_request correctly parses a valid created_at timestamp."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+            created_at="2024-01-15T10:30:00Z",
+        )
+
+        durable_message = DurableAgentStateMessage.from_run_request(run_request)
+
+        assert durable_message.created_at is not None
+        assert durable_message.created_at.year == 2024
+        assert durable_message.created_at.month == 1
+        assert durable_message.created_at.day == 15


 if __name__ == "__main__":
@@ -1,115 +1,230 @@
 # Copyright (c) Microsoft. All rights reserved.

-"""Unit tests for data models (RunRequest)."""
+"""Unit tests for data models (AgentSessionId, RunRequest, AgentResponse)."""

+import azure.durable_functions as df
 import pytest
 from agent_framework import Role
 from pydantic import BaseModel

-from agent_framework_durabletask._models import RunRequest
+from agent_framework_azurefunctions._models import AgentSessionId, RunRequest


 class ModuleStructuredResponse(BaseModel):
    value: int


+class TestAgentSessionId:
+    """Test suite for AgentSessionId."""
+
+    def test_init_creates_session_id(self) -> None:
+        """Test that AgentSessionId initializes correctly."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key-123"
+
+    def test_with_random_key_generates_guid(self) -> None:
+        """Test that with_random_key generates a GUID."""
+        session_id = AgentSessionId.with_random_key(name="AgentEntity")
+
+        assert session_id.name == "AgentEntity"
+        assert len(session_id.key) == 32  # UUID hex is 32 chars
+        # Verify it's a valid hex string
+        int(session_id.key, 16)
+
+    def test_with_random_key_unique_keys(self) -> None:
+        """Test that with_random_key generates unique keys."""
+        session_id1 = AgentSessionId.with_random_key(name="AgentEntity")
+        session_id2 = AgentSessionId.with_random_key(name="AgentEntity")
+
+        assert session_id1.key != session_id2.key
+
+    def test_to_entity_id_conversion(self) -> None:
+        """Test conversion to EntityId."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key")
+        entity_id = session_id.to_entity_id()
+
+        assert isinstance(entity_id, df.EntityId)
+        assert entity_id.name == "dafx-AgentEntity"
+        assert entity_id.key == "test-key"
+
+    def test_from_entity_id_conversion(self) -> None:
+        """Test creation from EntityId."""
+        entity_id = df.EntityId(name="dafx-AgentEntity", key="test-key")
+        session_id = AgentSessionId.from_entity_id(entity_id)
+
+        assert isinstance(session_id, AgentSessionId)
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key"
+
+    def test_round_trip_entity_id_conversion(self) -> None:
+        """Test round-trip conversion to and from EntityId."""
+        original = AgentSessionId(name="AgentEntity", key="test-key")
+        entity_id = original.to_entity_id()
+        restored = AgentSessionId.from_entity_id(entity_id)
+
+        assert restored.name == original.name
+        assert restored.key == original.key
+
+    def test_str_representation(self) -> None:
+        """Test string representation."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
+        str_repr = str(session_id)
+
+        assert str_repr == "@AgentEntity@test-key-123"
+
+    def test_repr_representation(self) -> None:
+        """Test repr representation."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key")
+        repr_str = repr(session_id)
+
+        assert "AgentSessionId" in repr_str
+        assert "AgentEntity" in repr_str
+        assert "test-key" in repr_str
+
+    def test_parse_valid_session_id(self) -> None:
+        """Test parsing valid session ID string."""
+        session_id = AgentSessionId.parse("@AgentEntity@test-key-123")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key-123"
+
+    def test_parse_invalid_format_no_prefix(self) -> None:
+        """Test parsing invalid format without @ prefix."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("AgentEntity@test-key")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
+    def test_parse_invalid_format_single_part(self) -> None:
+        """Test parsing invalid format with single part."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("@AgentEntity")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
+    def test_parse_with_multiple_at_signs_in_key(self) -> None:
+        """Test parsing with @ signs in the key."""
+        session_id = AgentSessionId.parse("@AgentEntity@key-with@symbols")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "key-with@symbols"
+
+    def test_parse_round_trip(self) -> None:
+        """Test round-trip parse and string conversion."""
+        original = AgentSessionId(name="AgentEntity", key="test-key")
+        str_repr = str(original)
+        parsed = AgentSessionId.parse(str_repr)
+
+        assert parsed.name == original.name
+        assert parsed.key == original.key
+
+    def test_to_entity_name_adds_prefix(self) -> None:
+        """Test that to_entity_name adds the dafx- prefix."""
+        entity_name = AgentSessionId.to_entity_name("TestAgent")
+        assert entity_name == "dafx-TestAgent"
+
+    def test_from_entity_id_strips_prefix(self) -> None:
+        """Test that from_entity_id strips the dafx- prefix."""
+        entity_id = df.EntityId(name="dafx-TestAgent", key="key123")
+        session_id = AgentSessionId.from_entity_id(entity_id)
+
+        assert session_id.name == "TestAgent"
+        assert session_id.key == "key123"
+
+    def test_from_entity_id_raises_without_prefix(self) -> None:
+        """Test that from_entity_id raises ValueError when entity name lacks the prefix."""
+        entity_id = df.EntityId(name="TestAgent", key="key123")
+
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.from_entity_id(entity_id)
+
+        assert "not a valid agent session ID" in str(exc_info.value)
+        assert "dafx-" in str(exc_info.value)
+
+
 class TestRunRequest:
    """Test suite for RunRequest."""

    def test_init_with_defaults(self) -> None:
        """Test RunRequest initialization with defaults."""
-        request = RunRequest(message="Hello", correlation_id="corr-001")
+        request = RunRequest(message="Hello", thread_id="thread-default")

        assert request.message == "Hello"
-        assert request.correlation_id == "corr-001"
        assert request.role == Role.USER
        assert request.response_format is None
        assert request.enable_tool_calls is True
-        assert request.wait_for_response is True
+        assert request.thread_id == "thread-default"

    def test_init_with_all_fields(self) -> None:
        """Test RunRequest initialization with all fields."""
        schema = ModuleStructuredResponse
        request = RunRequest(
            message="Hello",
-            correlation_id="corr-002",
+            thread_id="thread-123",
            role=Role.SYSTEM,
            response_format=schema,
            enable_tool_calls=False,
-            wait_for_response=False,
        )

        assert request.message == "Hello"
-        assert request.correlation_id == "corr-002"
        assert request.role == Role.SYSTEM
        assert request.response_format is schema
        assert request.enable_tool_calls is False
-        assert request.wait_for_response is False
+        assert request.thread_id == "thread-123"

    def test_init_coerces_string_role(self) -> None:
        """Ensure string role values are coerced into Role instances."""
-        request = RunRequest(message="Hello", correlation_id="corr-003", role="system")  # type: ignore[arg-type]
+        request = RunRequest(message="Hello", thread_id="thread-str-role", role="system")  # type: ignore[arg-type]

        assert request.role == Role.SYSTEM

    def test_to_dict_with_defaults(self) -> None:
        """Test to_dict with default values."""
-        request = RunRequest(message="Test message", correlation_id="corr-004")
+        request = RunRequest(message="Test message", thread_id="thread-to-dict")
        data = request.to_dict()

        assert data["message"] == "Test message"
        assert data["enable_tool_calls"] is True
-        assert data["wait_for_response"] is True
        assert data["role"] == "user"
-        assert data["correlationId"] == "corr-004"
        assert "response_format" not in data or data["response_format"] is None
-        assert "thread_id" not in data
+        assert data["thread_id"] == "thread-to-dict"

    def test_to_dict_with_all_fields(self) -> None:
        """Test to_dict with all fields."""
        schema = ModuleStructuredResponse
        request = RunRequest(
            message="Hello",
-            correlation_id="corr-005",
+            thread_id="thread-456",
            role=Role.ASSISTANT,
            response_format=schema,
            enable_tool_calls=False,
-            wait_for_response=False,
        )
        data = request.to_dict()

        assert data["message"] == "Hello"
-        assert data["correlationId"] == "corr-005"
        assert data["role"] == "assistant"
        assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
        assert data["response_format"]["module"] == schema.__module__
        assert data["response_format"]["qualname"] == schema.__qualname__
        assert data["enable_tool_calls"] is False
-        assert data["wait_for_response"] is False
-        assert "thread_id" not in data
+        assert data["thread_id"] == "thread-456"

    def test_from_dict_with_defaults(self) -> None:
        """Test from_dict with minimal data."""
-        data = {"message": "Hello", "correlationId": "corr-006"}
+        data = {"message": "Hello", "thread_id": "thread-from-dict"}
        request = RunRequest.from_dict(data)

        assert request.message == "Hello"
-        assert request.correlation_id == "corr-006"
        assert request.role == Role.USER
        assert request.enable_tool_calls is True
-        assert request.wait_for_response is True
-
-    def test_from_dict_ignores_thread_id_field(self) -> None:
-        """Ensure legacy thread_id input does not break RunRequest parsing."""
-        request = RunRequest.from_dict({"message": "Hello", "correlationId": "corr-007", "thread_id": "ignored"})
-
-        assert request.message == "Hello"
+        assert request.thread_id == "thread-from-dict"

    def test_from_dict_with_all_fields(self) -> None:
        """Test from_dict with all fields."""
        data = {
            "message": "Test",
-            "correlationId": "corr-008",
            "role": "system",
            "response_format": {
                "__response_schema_type__": "pydantic_model",
@@ -117,18 +232,19 @@ class TestRunRequest:
                "qualname": ModuleStructuredResponse.__qualname__,
            },
            "enable_tool_calls": False,
+            "thread_id": "thread-789",
        }
        request = RunRequest.from_dict(data)

        assert request.message == "Test"
-        assert request.correlation_id == "corr-008"
        assert request.role == Role.SYSTEM
        assert request.response_format is ModuleStructuredResponse
        assert request.enable_tool_calls is False
+        assert request.thread_id == "thread-789"

-    def test_from_dict_unknown_role_preserves_value(self) -> None:
+    def test_from_dict_with_unknown_role_preserves_value(self) -> None:
        """Test from_dict keeps custom roles intact."""
-        data = {"message": "Test", "correlationId": "corr-009", "role": "reviewer"}
+        data = {"message": "Test", "role": "reviewer", "thread_id": "thread-with-custom-role"}
        request = RunRequest.from_dict(data)

        assert request.role.value == "reviewer"
@@ -136,22 +252,18 @@ class TestRunRequest:

    def test_from_dict_empty_message(self) -> None:
        """Test from_dict with empty message."""
-        request = RunRequest.from_dict({"correlationId": "corr-010"})
+        data = {"thread_id": "thread-empty"}
+        request = RunRequest.from_dict(data)

        assert request.message == ""
-        assert request.correlation_id == "corr-010"
        assert request.role == Role.USER
-
-    def test_from_dict_missing_correlation_id_raises(self) -> None:
-        """Test from_dict raises when correlationId is missing."""
-        with pytest.raises(ValueError, match="correlationId is required"):
-            RunRequest.from_dict({"message": "Test"})
+        assert request.thread_id == "thread-empty"

    def test_round_trip_dict_conversion(self) -> None:
        """Test round-trip to_dict and from_dict."""
        original = RunRequest(
            message="Test message",
-            correlation_id="corr-011",
+            thread_id="thread-123",
            role=Role.SYSTEM,
            response_format=ModuleStructuredResponse,
            enable_tool_calls=False,
@@ -161,16 +273,16 @@ class TestRunRequest:
        restored = RunRequest.from_dict(data)

        assert restored.message == original.message
-        assert restored.correlation_id == original.correlation_id
        assert restored.role == original.role
        assert restored.response_format is ModuleStructuredResponse
        assert restored.enable_tool_calls == original.enable_tool_calls
+        assert restored.thread_id == original.thread_id

    def test_round_trip_with_pydantic_response_format(self) -> None:
        """Ensure Pydantic response formats serialize and deserialize properly."""
        original = RunRequest(
            message="Structured",
-            correlation_id="corr-012",
+            thread_id="thread-pydantic",
            response_format=ModuleStructuredResponse,
        )

@@ -185,14 +297,14 @@ class TestRunRequest:

    def test_init_with_correlationId(self) -> None:
        """Test RunRequest initialization with correlationId."""
-        request = RunRequest(message="Test message", correlation_id="corr-123")
+        request = RunRequest(message="Test message", thread_id="thread-corr-init", correlation_id="corr-123")

        assert request.message == "Test message"
        assert request.correlation_id == "corr-123"

    def test_to_dict_with_correlationId(self) -> None:
        """Test to_dict includes correlationId."""
-        request = RunRequest(message="Test", correlation_id="corr-456")
+        request = RunRequest(message="Test", thread_id="thread-corr-to-dict", correlation_id="corr-456")
        data = request.to_dict()

        assert data["message"] == "Test"
@@ -200,18 +312,20 @@ class TestRunRequest:

    def test_from_dict_with_correlationId(self) -> None:
        """Test from_dict with correlationId."""
-        data = {"message": "Test", "correlationId": "corr-789"}
+        data = {"message": "Test", "correlationId": "corr-789", "thread_id": "thread-corr-from-dict"}
        request = RunRequest.from_dict(data)

        assert request.message == "Test"
        assert request.correlation_id == "corr-789"
+        assert request.thread_id == "thread-corr-from-dict"

    def test_round_trip_with_correlationId(self) -> None:
        """Test round-trip to_dict and from_dict with correlationId."""
        original = RunRequest(
            message="Test message",
+            thread_id="thread-123",
            role=Role.SYSTEM,
-            correlation_id="corr-124",
+            correlation_id="corr-123",
        )

        data = original.to_dict()
@@ -220,12 +334,13 @@ class TestRunRequest:
        assert restored.message == original.message
        assert restored.role == original.role
        assert restored.correlation_id == original.correlation_id
+        assert restored.thread_id == original.thread_id

    def test_init_with_orchestration_id(self) -> None:
        """Test RunRequest initialization with orchestration_id."""
        request = RunRequest(
            message="Test message",
-            correlation_id="corr-125",
+            thread_id="thread-orch-init",
            orchestration_id="orch-123",
        )

@@ -236,7 +351,7 @@ class TestRunRequest:
        """Test to_dict includes orchestrationId."""
        request = RunRequest(
            message="Test",
-            correlation_id="corr-126",
+            thread_id="thread-orch-to-dict",
            orchestration_id="orch-456",
        )
        data = request.to_dict()
@@ -248,7 +363,7 @@ class TestRunRequest:
        """Test to_dict excludes orchestrationId when not set."""
        request = RunRequest(
            message="Test",
-            correlation_id="corr-127",
+            thread_id="thread-orch-none",
        )
        data = request.to_dict()

@@ -258,20 +373,22 @@ class TestRunRequest:
        """Test from_dict with orchestrationId."""
        data = {
            "message": "Test",
-            "correlationId": "corr-128",
            "orchestrationId": "orch-789",
+            "thread_id": "thread-orch-from-dict",
        }
        request = RunRequest.from_dict(data)

        assert request.message == "Test"
        assert request.orchestration_id == "orch-789"
+        assert request.thread_id == "thread-orch-from-dict"

    def test_round_trip_with_orchestration_id(self) -> None:
        """Test round-trip to_dict and from_dict with orchestration_id."""
        original = RunRequest(
            message="Test message",
+            thread_id="thread-123",
            role=Role.SYSTEM,
-            correlation_id="corr-129",
+            correlation_id="corr-123",
            orchestration_id="orch-123",
        )

@@ -282,6 +399,20 @@ class TestRunRequest:
        assert restored.role == original.role
        assert restored.correlation_id == original.correlation_id
        assert restored.orchestration_id == original.orchestration_id
+        assert restored.thread_id == original.thread_id
+
+
+class TestModelIntegration:
+    """Test suite for integration between models."""
+
+    def test_run_request_with_session_id(self) -> None:
+        """Test using RunRequest with AgentSessionId."""
+        session_id = AgentSessionId.with_random_key("AgentEntity")
+        request = RunRequest(message="Test message", thread_id=str(session_id))
+
+        assert request.thread_id is not None
+        assert request.thread_id == str(session_id)
+        assert request.thread_id.startswith("@AgentEntity@")


 if __name__ == "__main__":
@@ -6,11 +6,11 @@ from typing import Any
 from unittest.mock import Mock

 import pytest
-from agent_framework import AgentRunResponse, ChatMessage, Role
-from agent_framework_durabletask import DurableAIAgent
+from agent_framework import AgentRunResponse, AgentThread, ChatMessage
 from azure.durable_functions.models.Task import TaskBase, TaskState

-from agent_framework_azurefunctions import AgentFunctionApp
+from agent_framework_azurefunctions import AgentFunctionApp, DurableAIAgent
+from agent_framework_azurefunctions._models import AgentSessionId, DurableAgentThread
 from agent_framework_azurefunctions._orchestration import AgentTask


@@ -38,96 +38,46 @@ def _create_entity_task(task_id: int = 1) -> TaskBase:
    return _FakeTask(task_id)


-@pytest.fixture
-def mock_context():
-    """Create a mock orchestration context with UUID support."""
-    context = Mock()
-    context.instance_id = "test-instance"
-    context.current_utc_datetime = Mock()
-    return context
-
-
-@pytest.fixture
-def mock_context_with_uuid() -> tuple[Mock, str]:
-    """Create a mock context with a single UUID."""
-    from uuid import UUID
-
-    context = Mock()
-    context.instance_id = "test-instance"
-    context.current_utc_datetime = Mock()
-    test_uuid = UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
-    context.new_uuid = Mock(return_value=test_uuid)
-    return context, test_uuid.hex
-
-
-@pytest.fixture
-def mock_context_with_multiple_uuids() -> tuple[Mock, list[str]]:
-    """Create a mock context with multiple UUIDs via side_effect."""
-    from uuid import UUID
-
-    context = Mock()
-    context.instance_id = "test-instance"
-    context.current_utc_datetime = Mock()
-    uuids = [
-        UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"),
-        UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"),
-        UUID("cccccccc-cccc-cccc-cccc-cccccccccccc"),
-    ]
-    context.new_uuid = Mock(side_effect=uuids)
-    # Return the hex versions for assertion checking
-    hex_uuids = [uuid.hex for uuid in uuids]
-    return context, hex_uuids
-
-
-@pytest.fixture
-def executor_with_uuid() -> tuple[Any, Mock, str]:
-    """Create an executor with a mocked generate_unique_id method."""
-    from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
-
-    context = Mock()
-    context.instance_id = "test-instance"
-    context.current_utc_datetime = Mock()
-
-    executor = AzureFunctionsAgentExecutor(context)
-    test_uuid_hex = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
-    executor.generate_unique_id = Mock(return_value=test_uuid_hex)
-
-    return executor, context, test_uuid_hex
-
-
-@pytest.fixture
-def executor_with_multiple_uuids() -> tuple[Any, Mock, list[str]]:
-    """Create an executor with multiple mocked UUIDs."""
-    from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
-
-    context = Mock()
-    context.instance_id = "test-instance"
-    context.current_utc_datetime = Mock()
-
-    executor = AzureFunctionsAgentExecutor(context)
-    uuid_hexes = [
-        "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
-        "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
-        "cccccccc-cccc-cccc-cccc-cccccccccccc",
-        "dddddddd-dddd-dddd-dddd-dddddddddddd",
-        "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee",
-    ]
-    executor.generate_unique_id = Mock(side_effect=uuid_hexes)
-
-    return executor, context, uuid_hexes
-
-
-@pytest.fixture
-def executor_with_context(mock_context_with_uuid: tuple[Mock, str]) -> tuple[Any, Mock]:
-    """Create an executor with a mocked context."""
-    from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
-
-    context, _ = mock_context_with_uuid
-    return AzureFunctionsAgentExecutor(context), context
-
-
 class TestAgentResponseHelpers:
-    """Tests for response handling through public AgentTask API."""
+    """Tests for helper utilities that prepare AgentRunResponse values."""
+
+    @staticmethod
+    def _create_agent_task() -> AgentTask:
+        entity_task = _create_entity_task()
+        return AgentTask(entity_task, None, "correlation-id")
+
+    def test_load_agent_response_from_instance(self) -> None:
+        task = self._create_agent_task()
+        response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"foo": "bar"}')])
+
+        loaded = task._load_agent_response(response)
+
+        assert loaded is response
+        assert loaded.value is None
+
+    def test_load_agent_response_from_serialized(self) -> None:
+        task = self._create_agent_task()
+        serialized = AgentRunResponse(messages=[ChatMessage(role="assistant", text="structured")]).to_dict()
+        serialized["value"] = {"answer": 42}
+
+        loaded = task._load_agent_response(serialized)
+
+        assert loaded is not None
+        assert loaded.value == {"answer": 42}
+        loaded_dict = loaded.to_dict()
+        assert loaded_dict["type"] == "agent_run_response"
+
+    def test_load_agent_response_rejects_none(self) -> None:
+        task = self._create_agent_task()
+
+        with pytest.raises(ValueError):
+            task._load_agent_response(None)
+
+    def test_load_agent_response_rejects_unsupported_type(self) -> None:
+        task = self._create_agent_task()
+
+        with pytest.raises(TypeError, match="Unsupported type"):
+            task._load_agent_response(["invalid", "list"])  # type: ignore[arg-type]

    def test_try_set_value_success(self) -> None:
        """Test try_set_value correctly processes successful task completion."""
@@ -194,10 +144,336 @@ class TestAgentResponseHelpers:
        assert isinstance(task.result.value, TestSchema)
        assert task.result.value.answer == "42"

+    def test_ensure_response_format_parses_value(self) -> None:
+        """Test _ensure_response_format correctly parses response value."""
+        from pydantic import BaseModel
+
+        class SampleSchema(BaseModel):
+            name: str
+
+        task = self._create_agent_task()
+        response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"name": "test"}')])
+
+        # Value should be None initially
+        assert response.value is None
+
+        # Parse the value
+        task._ensure_response_format(SampleSchema, "test-correlation", response)
+
+        # Value should now be parsed
+        assert isinstance(response.value, SampleSchema)
+        assert response.value.name == "test"
+
+    def test_ensure_response_format_skips_if_already_parsed(self) -> None:
+        """Test _ensure_response_format does not re-parse if value already matches format."""
+        from pydantic import BaseModel
+
+        class SampleSchema(BaseModel):
+            name: str
+
+        task = self._create_agent_task()
+        existing_value = SampleSchema(name="existing")
+        response = AgentRunResponse(
+            messages=[ChatMessage(role="assistant", text='{"name": "new"}')],
+            value=existing_value,
+        )
+
+        # Call _ensure_response_format
+        task._ensure_response_format(SampleSchema, "test-correlation", response)
+
+        # Value should remain unchanged (not re-parsed)
+        assert response.value is existing_value
+        assert response.value.name == "existing"
+
+
+class TestDurableAIAgent:
+    """Test suite for DurableAIAgent wrapper."""
+
+    def test_init(self) -> None:
+        """Test DurableAIAgent initialization."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-123"
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        assert agent.context == mock_context
+        assert agent.agent_name == "TestAgent"
+
+    def test_implements_agent_protocol(self) -> None:
+        """Test that DurableAIAgent implements AgentProtocol."""
+        from agent_framework import AgentProtocol
+
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # Check that agent satisfies AgentProtocol
+        assert isinstance(agent, AgentProtocol)
+
+    def test_has_agent_protocol_properties(self) -> None:
+        """Test that DurableAIAgent has AgentProtocol properties."""
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # AgentProtocol properties
+        assert hasattr(agent, "id")
+        assert hasattr(agent, "name")
+        assert hasattr(agent, "description")
+        assert hasattr(agent, "display_name")
+
+        # Verify values
+        assert agent.name == "TestAgent"
+        assert agent.description == "Durable agent proxy for TestAgent"
+        assert agent.display_name == "TestAgent"
+        assert agent.id is not None  # Auto-generated UUID
+
+    def test_get_new_thread(self) -> None:
+        """Test creating a new agent thread."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-456"
+        mock_context.new_uuid = Mock(return_value="test-guid-456")
+
+        agent = DurableAIAgent(mock_context, "WriterAgent")
+        thread = agent.get_new_thread()
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.session_id is not None
+        session_id = thread.session_id
+        assert isinstance(session_id, AgentSessionId)
+        assert session_id.name == "WriterAgent"
+        assert session_id.key == "test-guid-456"
+        mock_context.new_uuid.assert_called_once()
+
+    def test_get_new_thread_deterministic(self) -> None:
+        """Test that get_new_thread creates deterministic session IDs."""
+
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-789"
+        mock_context.new_uuid = Mock(side_effect=["session-guid-1", "session-guid-2"])
+
+        agent = DurableAIAgent(mock_context, "EditorAgent")
+
+        # Create multiple threads - they should have unique session IDs
+        thread1 = agent.get_new_thread()
+        thread2 = agent.get_new_thread()
+
+        assert isinstance(thread1, DurableAgentThread)
+        assert isinstance(thread2, DurableAgentThread)
+
+        session_id1 = thread1.session_id
+        session_id2 = thread2.session_id
+        assert session_id1 is not None and session_id2 is not None
+        assert isinstance(session_id1, AgentSessionId)
+        assert isinstance(session_id2, AgentSessionId)
+        assert session_id1.name == "EditorAgent"
+        assert session_id2.name == "EditorAgent"
+        assert session_id1.key == "session-guid-1"
+        assert session_id2.key == "session-guid-2"
+        assert mock_context.new_uuid.call_count == 2
+
+    def test_run_creates_entity_call(self) -> None:
+        """Test that run() creates proper entity call and returns a Task."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-001"
+        mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
+
+        entity_task = _create_entity_task()
+        mock_context.call_entity = Mock(return_value=entity_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # Create thread
+        thread = agent.get_new_thread()
+
+        # Call run() - returns AgentTask directly
+        task = agent.run(messages="Test message", thread=thread, enable_tool_calls=True)
+
+        assert isinstance(task, AgentTask)
+        assert task.children[0] == entity_task
+
+        # Verify call_entity was called with correct parameters
+        assert mock_context.call_entity.called
+        call_args = mock_context.call_entity.call_args
+        entity_id, operation, request = call_args[0]
+
+        assert operation == "run"
+        assert request["message"] == "Test message"
+        assert request["enable_tool_calls"] is True
+        assert "correlationId" in request
+        assert request["correlationId"] == "correlation-guid"
+        assert "thread_id" in request
+        assert request["thread_id"] == "thread-guid"
+        # Verify orchestration ID is set from context.instance_id
+        assert "orchestrationId" in request
+        assert request["orchestrationId"] == "test-instance-001"
+
+    def test_run_sets_orchestration_id(self) -> None:
+        """Test that run() sets the orchestration_id from context.instance_id."""
+        mock_context = Mock()
+        mock_context.instance_id = "my-orchestration-123"
+        mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
+
+        entity_task = _create_entity_task()
+        mock_context.call_entity = Mock(return_value=entity_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+        thread = agent.get_new_thread()
+
+        agent.run(messages="Test", thread=thread)
+
+        call_args = mock_context.call_entity.call_args
+        request = call_args[0][2]
+
+        assert request["orchestrationId"] == "my-orchestration-123"
+
+    def test_run_without_thread(self) -> None:
+        """Test that run() works without explicit thread (creates unique session key)."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-002"
+        mock_context.new_uuid = Mock(side_effect=["auto-generated-guid", "correlation-guid"])
+
+        entity_task = _create_entity_task()
+        mock_context.call_entity = Mock(return_value=entity_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # Call without thread
+        task = agent.run(messages="Test message")
+
+        assert isinstance(task, AgentTask)
+        assert task.children[0] == entity_task
+
+        # Verify the entity ID uses the auto-generated GUID with dafx- prefix
+        call_args = mock_context.call_entity.call_args
+        entity_id = call_args[0][0]
+        assert entity_id.name == "dafx-TestAgent"
+        assert entity_id.key == "auto-generated-guid"
+        # Should be called twice: once for session_key, once for correlationId
+        assert mock_context.new_uuid.call_count == 2
+
+    def test_run_with_response_format(self) -> None:
+        """Test that run() passes response format correctly."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-003"
+
+        entity_task = _create_entity_task()
+        mock_context.call_entity = Mock(return_value=entity_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        from pydantic import BaseModel
+
+        class SampleSchema(BaseModel):
+            key: str
+
+        # Create thread and call
+        thread = agent.get_new_thread()
+
+        task = agent.run(messages="Test message", thread=thread, response_format=SampleSchema)
+
+        assert isinstance(task, AgentTask)
+        assert task.children[0] == entity_task
+
+        # Verify schema was passed in the call_entity arguments
+        call_args = mock_context.call_entity.call_args
+        input_data = call_args[0][2]  # Third argument is input_data
+        assert "response_format" in input_data
+        assert input_data["response_format"]["__response_schema_type__"] == "pydantic_model"
+        assert input_data["response_format"]["module"] == SampleSchema.__module__
+        assert input_data["response_format"]["qualname"] == SampleSchema.__qualname__
+
+    def test_messages_to_string(self) -> None:
+        """Test converting ChatMessage list to string."""
+        from agent_framework import ChatMessage
+
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        messages = [
+            ChatMessage(role="user", text="Hello"),
+            ChatMessage(role="assistant", text="Hi there"),
+            ChatMessage(role="user", text="How are you?"),
+        ]
+
+        result = agent._messages_to_string(messages)
+
+        assert result == "Hello\nHi there\nHow are you?"
+
+    def test_run_with_chat_message(self) -> None:
+        """Test that run() handles ChatMessage input."""
+        from agent_framework import ChatMessage
+
+        mock_context = Mock()
+        mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
+        entity_task = _create_entity_task()
+        mock_context.call_entity = Mock(return_value=entity_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+        thread = agent.get_new_thread()
+
+        # Call with ChatMessage
+        msg = ChatMessage(role="user", text="Hello")
+        task = agent.run(messages=msg, thread=thread)
+
+        assert isinstance(task, AgentTask)
+        assert task.children[0] == entity_task
+
+        # Verify message was converted to string
+        call_args = mock_context.call_entity.call_args
+        request = call_args[0][2]
+        assert request["message"] == "Hello"
+
+    def test_run_stream_raises_not_implemented(self) -> None:
+        """Test that run_stream() method raises NotImplementedError."""
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        with pytest.raises(NotImplementedError) as exc_info:
+            agent.run_stream("Test message")
+
+        error_msg = str(exc_info.value)
+        assert "Streaming is not supported" in error_msg
+
+    def test_entity_id_format(self) -> None:
+        """Test that EntityId is created with correct format (name, key)."""
+        from azure.durable_functions import EntityId
+
+        mock_context = Mock()
+        mock_context.new_uuid = Mock(return_value="test-guid-789")
+        mock_context.call_entity = Mock(return_value=_create_entity_task())
+
+        agent = DurableAIAgent(mock_context, "WriterAgent")
+        thread = agent.get_new_thread()
+
+        # Call run() to trigger entity ID creation
+        agent.run("Test", thread=thread)
+
+        # Verify call_entity was called with correct EntityId
+        call_args = mock_context.call_entity.call_args
+        entity_id = call_args[0][0]
+
+        # EntityId should be EntityId(name="dafx-WriterAgent", key="test-guid-789")
+        # Which formats as "@dafx-writeragent@test-guid-789"
+        assert isinstance(entity_id, EntityId)
+        assert entity_id.name == "dafx-WriterAgent"
+        assert entity_id.key == "test-guid-789"
+        assert str(entity_id) == "@dafx-writeragent@test-guid-789"
+

 class TestAgentFunctionAppGetAgent:
    """Test suite for AgentFunctionApp.get_agent."""

+    def test_get_agent_method(self) -> None:
+        """Test get_agent method creates DurableAIAgent for registered agent."""
+        app = _app_with_registered_agents("MyAgent")
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-100"
+
+        agent = app.get_agent(mock_context, "MyAgent")
+
+        assert isinstance(agent, DurableAIAgent)
+        assert agent.agent_name == "MyAgent"
+        assert agent.context == mock_context
+
    def test_get_agent_raises_for_unregistered_agent(self) -> None:
        """Test get_agent raises ValueError when agent is not registered."""
        app = _app_with_registered_agents("KnownAgent")
@@ -206,87 +482,18 @@ class TestAgentFunctionAppGetAgent:
            app.get_agent(Mock(), "MissingAgent")


-class TestAzureFunctionsFireAndForget:
-    """Test fire-and-forget mode for AzureFunctionsAgentExecutor."""
-
-    def test_fire_and_forget_calls_signal_entity(self, executor_with_uuid: tuple[Any, Mock, str]) -> None:
-        """Verify wait_for_response=False calls signal_entity instead of call_entity."""
-        executor, context, _ = executor_with_uuid
-        context.signal_entity = Mock()
-        context.call_entity = Mock(return_value=_create_entity_task())
-
-        agent = DurableAIAgent(executor, "TestAgent")
-        thread = agent.get_new_thread()
-
-        # Run with wait_for_response=False
-        result = agent.run("Test message", thread=thread, wait_for_response=False)
-
-        # Verify signal_entity was called and call_entity was not
-        assert context.signal_entity.call_count == 1
-        assert context.call_entity.call_count == 0
-
-        # Should still return an AgentTask
-        assert isinstance(result, AgentTask)
-
-    def test_fire_and_forget_returns_completed_task(self, executor_with_uuid: tuple[Any, Mock, str]) -> None:
-        """Verify wait_for_response=False returns pre-completed AgentTask."""
-        executor, context, _ = executor_with_uuid
-        context.signal_entity = Mock()
-
-        agent = DurableAIAgent(executor, "TestAgent")
-        thread = agent.get_new_thread()
-
-        result = agent.run("Test message", thread=thread, wait_for_response=False)
-
-        # Task should be immediately complete
-        assert isinstance(result, AgentTask)
-        assert result.is_completed
-
-    def test_fire_and_forget_returns_acceptance_response(self, executor_with_uuid: tuple[Any, Mock, str]) -> None:
-        """Verify wait_for_response=False returns acceptance response."""
-        executor, context, _ = executor_with_uuid
-        context.signal_entity = Mock()
-
-        agent = DurableAIAgent(executor, "TestAgent")
-        thread = agent.get_new_thread()
-
-        result = agent.run("Test message", thread=thread, wait_for_response=False)
-
-        # Get the result
-        response = result.result
-        assert isinstance(response, AgentRunResponse)
-        assert len(response.messages) == 1
-        assert response.messages[0].role == Role.SYSTEM
-        # Check message contains key information
-        message_text = response.messages[0].text
-        assert "accepted" in message_text.lower()
-        assert "background" in message_text.lower()
-
-    def test_blocking_mode_still_works(self, executor_with_uuid: tuple[Any, Mock, str]) -> None:
-        """Verify wait_for_response=True uses call_entity as before."""
-        executor, context, _ = executor_with_uuid
-        context.signal_entity = Mock()
-        context.call_entity = Mock(return_value=_create_entity_task())
-
-        agent = DurableAIAgent(executor, "TestAgent")
-        thread = agent.get_new_thread()
-
-        result = agent.run("Test message", thread=thread, wait_for_response=True)
-
-        # Verify call_entity was called and signal_entity was not
-        assert context.call_entity.call_count == 1
-        assert context.signal_entity.call_count == 0
-
-        # Should return an AgentTask
-        assert isinstance(result, AgentTask)
-
-
 class TestOrchestrationIntegration:
    """Integration tests for orchestration scenarios."""

-    def test_sequential_agent_calls_simulation(self, executor_with_multiple_uuids: tuple[Any, Mock, list[str]]) -> None:
+    def test_sequential_agent_calls_simulation(self) -> None:
        """Simulate sequential agent calls in an orchestration."""
-        executor, context, uuid_hexes = executor_with_multiple_uuids
+        mock_context = Mock()
+        mock_context.instance_id = "test-orchestration-001"
+        # new_uuid will be called 3 times:
+        # 1. thread creation
+        # 2. correlationId for first call
+        # 3. correlationId for second call
+        mock_context.new_uuid = Mock(side_effect=["deterministic-guid-001", "corr-1", "corr-2"])

        # Track entity calls
        entity_calls: list[dict[str, Any]] = []
@@ -295,10 +502,10 @@ class TestOrchestrationIntegration:
            entity_calls.append({"entity_id": str(entity_id), "operation": operation, "input": input_data})
            return _create_entity_task()

-        context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
+        mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)

-        # Create agent directly with executor (not via app.get_agent)
-        agent = DurableAIAgent(executor, "WriterAgent")
+        app = _app_with_registered_agents("WriterAgent")
+        agent = app.get_agent(mock_context, "WriterAgent")

        # Create thread
        thread = agent.get_new_thread()
@@ -314,15 +521,18 @@ class TestOrchestrationIntegration:
        # Verify both calls used the same entity (same session key)
        assert len(entity_calls) == 2
        assert entity_calls[0]["entity_id"] == entity_calls[1]["entity_id"]
-        # EntityId format is @dafx-writeragent@<uuid_hex>
-        expected_entity_id = f"@dafx-writeragent@{uuid_hexes[0]}"
-        assert entity_calls[0]["entity_id"] == expected_entity_id
-        # generate_unique_id called 3 times: thread + 2 correlation IDs
-        assert executor.generate_unique_id.call_count == 3
+        # EntityId format is @dafx-writeragent@deterministic-guid-001
+        assert entity_calls[0]["entity_id"] == "@dafx-writeragent@deterministic-guid-001"
+        # new_uuid called 3 times: thread + 2 correlation IDs
+        assert mock_context.new_uuid.call_count == 3

-    def test_multiple_agents_in_orchestration(self, executor_with_multiple_uuids: tuple[Any, Mock, list[str]]) -> None:
+    def test_multiple_agents_in_orchestration(self) -> None:
        """Test using multiple different agents in one orchestration."""
-        executor, context, uuid_hexes = executor_with_multiple_uuids
+        mock_context = Mock()
+        mock_context.instance_id = "test-orchestration-002"
+        # Mock new_uuid to return different GUIDs for each call
+        # Order: writer thread, editor thread, writer correlation, editor correlation
+        mock_context.new_uuid = Mock(side_effect=["writer-guid-001", "editor-guid-002", "writer-corr", "editor-corr"])

        entity_calls: list[str] = []

@@ -330,11 +540,11 @@ class TestOrchestrationIntegration:
            entity_calls.append(str(entity_id))
            return _create_entity_task()

-        context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
+        mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)

-        # Create agents directly with executor (not via app.get_agent)
-        writer = DurableAIAgent(executor, "WriterAgent")
-        editor = DurableAIAgent(executor, "EditorAgent")
+        app = _app_with_registered_agents("WriterAgent", "EditorAgent")
+        writer = app.get_agent(mock_context, "WriterAgent")
+        editor = app.get_agent(mock_context, "EditorAgent")

        writer_thread = writer.get_new_thread()
        editor_thread = editor.get_new_thread()
@@ -348,11 +558,62 @@ class TestOrchestrationIntegration:

        # Verify different entity IDs were used
        assert len(entity_calls) == 2
-        # EntityId format is @dafx-agentname@uuid_hex (lowercased agent name with dafx- prefix)
-        expected_writer_id = f"@dafx-writeragent@{uuid_hexes[0]}"
-        expected_editor_id = f"@dafx-editoragent@{uuid_hexes[1]}"
-        assert entity_calls[0] == expected_writer_id
-        assert entity_calls[1] == expected_editor_id
+        # EntityId format is @dafx-agentname@guid (lowercased agent name with dafx- prefix)
+        assert entity_calls[0] == "@dafx-writeragent@writer-guid-001"
+        assert entity_calls[1] == "@dafx-editoragent@editor-guid-002"
+
+
+class TestAgentThreadSerialization:
+    """Test that AgentThread can be serialized for orchestration state."""
+
+    async def test_agent_thread_serialize(self) -> None:
+        """Test that AgentThread can be serialized."""
+        thread = AgentThread()
+
+        # Serialize
+        serialized = await thread.serialize()
+
+        assert isinstance(serialized, dict)
+        assert "service_thread_id" in serialized
+
+    async def test_agent_thread_deserialize(self) -> None:
+        """Test that AgentThread can be deserialized."""
+        thread = AgentThread()
+        serialized = await thread.serialize()
+
+        # Deserialize
+        restored = await AgentThread.deserialize(serialized)
+
+        assert isinstance(restored, AgentThread)
+        assert restored.service_thread_id == thread.service_thread_id
+
+    async def test_durable_agent_thread_serialization(self) -> None:
+        """Test that DurableAgentThread persists session metadata during serialization."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-999"
+        mock_context.new_uuid = Mock(return_value="test-guid-999")
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+        thread = agent.get_new_thread()
+
+        assert isinstance(thread, DurableAgentThread)
+        # Verify custom attribute and property exist
+        assert thread.session_id is not None
+        session_id = thread.session_id
+        assert isinstance(session_id, AgentSessionId)
+        assert session_id.name == "TestAgent"
+        assert session_id.key == "test-guid-999"
+
+        # Standard serialization should still work
+        serialized = await thread.serialize()
+        assert isinstance(serialized, dict)
+        assert serialized.get("durable_session_id") == str(session_id)
+
+        # After deserialization, we'd need to restore the custom attribute
+        # This would be handled by the orchestration framework
+        restored = await DurableAgentThread.deserialize(serialized)
+        assert isinstance(restored, DurableAgentThread)
+        assert restored.session_id == session_id


 if __name__ == "__main__":
@@ -4,7 +4,7 @@ description = "Amazon Bedrock integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251120"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -4,7 +4,7 @@ description = "OpenAI ChatKit integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -4,7 +4,7 @@ description = "Copilot Studio integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -16,7 +16,6 @@ from typing import (
    Generic,
    Literal,
    Protocol,
-    TypedDict,
    TypeVar,
    cast,
    get_args,
@@ -74,7 +73,6 @@ __all__ = [
    "FunctionInvocationConfiguration",
    "HostedCodeInterpreterTool",
    "HostedFileSearchTool",
-    "HostedImageGenerationTool",
    "HostedMCPSpecificApproval",
    "HostedMCPTool",
    "HostedWebSearchTool",
@@ -326,41 +324,6 @@ class HostedWebSearchTool(BaseTool):
        super().__init__(**args)


-class HostedImageGenerationToolOptions(TypedDict, total=False):
-    """Options for HostedImageGenerationTool."""
-
-    count: int
-    image_size: str
-    media_type: str
-    model_id: str
-    response_format: Literal["uri", "data", "hosted"]
-    streaming_count: int
-
-
-class HostedImageGenerationTool(BaseTool):
-    """Represents a hosted tool that can be specified to an AI service to enable it to perform image generation."""
-
-    def __init__(
-        self,
-        *,
-        options: HostedImageGenerationToolOptions | None = None,
-        description: str | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        **kwargs: Any,
-    ):
-        """Initialize a HostedImageGenerationTool."""
-        if "name" in kwargs:
-            raise ValueError("The 'name' argument is reserved for the HostedImageGenerationTool and cannot be set.")
-
-        self.options = options
-        super().__init__(
-            name="image_generation",
-            description=description or "",
-            additional_properties=additional_properties,
-            **kwargs,
-        )
-
-
 class HostedMCPSpecificApproval(TypedDict, total=False):
    """Represents the specific mode for a hosted tool.

@@ -1456,11 +1419,14 @@ async def _auto_invoke_function(
    Raises:
        KeyError: If the requested function is not found in the tool map.
    """
+    from ._types import (
+        FunctionResultContent,
+    )
+
    # Note: The scenarios for approval_mode="always_require", declaration_only, and
    # terminate_on_unknown_calls are all handled in _try_execute_function_calls before
    # this function is called. This function only handles the actual execution of approved,
    # non-declaration-only functions.
-    from ._types import FunctionCallContent, FunctionResultContent

    tool: AIFunction[BaseModel, Any] | None = None
    if function_call_content.type == "function_call":
@@ -1478,14 +1444,11 @@ async def _auto_invoke_function(
    else:
        # Note: Unapproved tools (approved=False) are handled in _replace_approval_contents_with_results
        # and never reach this function, so we only handle approved=True cases here.
-        inner_call = function_call_content.function_call
-        if not isinstance(inner_call, FunctionCallContent):
-            return function_call_content
-        tool = tool_map.get(inner_call.name)
+        tool = tool_map.get(function_call_content.function_call.name)
        if tool is None:
            # we assume it is a hosted tool
            return function_call_content
-        function_call_content = inner_call
+        function_call_content = function_call_content.function_call

    parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})

@@ -40,8 +40,6 @@ __all__ = [
    "ChatResponse",
    "ChatResponseUpdate",
    "CitationAnnotation",
-    "CodeInterpreterToolCallContent",
-    "CodeInterpreterToolResultContent",
    "Contents",
    "DataContent",
    "ErrorContent",
@@ -52,10 +50,6 @@ __all__ = [
    "FunctionResultContent",
    "HostedFileContent",
    "HostedVectorStoreContent",
-    "ImageGenerationToolCallContent",
-    "ImageGenerationToolResultContent",
-    "MCPServerToolCallContent",
-    "MCPServerToolResultContent",
    "Role",
    "TextContent",
    "TextReasoningContent",
@@ -127,18 +121,6 @@ def _parse_content(content_data: MutableMapping[str, Any]) -> "Contents":
            return HostedFileContent.from_dict(content_data)
        case "hosted_vector_store":
            return HostedVectorStoreContent.from_dict(content_data)
-        case "code_interpreter_tool_call":
-            return CodeInterpreterToolCallContent.from_dict(content_data)
-        case "code_interpreter_tool_result":
-            return CodeInterpreterToolResultContent.from_dict(content_data)
-        case "image_generation_tool_call":
-            return ImageGenerationToolCallContent.from_dict(content_data)
-        case "image_generation_tool_result":
-            return ImageGenerationToolResultContent.from_dict(content_data)
-        case "mcp_server_tool_call":
-            return MCPServerToolCallContent.from_dict(content_data)
-        case "mcp_server_tool_result":
-            return MCPServerToolResultContent.from_dict(content_data)
        case "function_approval_request":
            return FunctionApprovalRequestContent.from_dict(content_data)
        case "function_approval_response":
@@ -1625,8 +1607,6 @@ class HostedFileContent(BaseContent):
        self,
        file_id: str,
        *,
-        media_type: str | None = None,
-        name: str | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -1635,8 +1615,6 @@ class HostedFileContent(BaseContent):

        Args:
            file_id: The identifier of the hosted file.
-            media_type: Optional media type of the hosted file.
-            name: Optional display name of the hosted file.

        Keyword Args:
            additional_properties: Optional additional properties associated with the content.
@@ -1649,14 +1627,8 @@ class HostedFileContent(BaseContent):
            **kwargs,
        )
        self.file_id = file_id
-        self.media_type = media_type
-        self.name = name
        self.type: Literal["hosted_file"] = "hosted_file"

-    def has_top_level_media_type(self, top_level_media_type: Literal["application", "audio", "image", "text"]) -> bool:
-        """Returns a boolean indicating if the media type has the specified top-level media type."""
-        return _has_top_level_media_type(self.media_type, top_level_media_type)
-

 class HostedVectorStoreContent(BaseContent):
    """Represents a hosted vector store content.
@@ -1704,234 +1676,6 @@ class HostedVectorStoreContent(BaseContent):
        self.type: Literal["hosted_vector_store"] = "hosted_vector_store"


-class CodeInterpreterToolCallContent(BaseContent):
-    """Represents a code interpreter tool call invocation by a hosted service."""
-
-    def __init__(
-        self,
-        *,
-        call_id: str | None = None,
-        inputs: Sequence["Contents | MutableMapping[str, Any]"] | None = None,
-        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(
-            annotations=annotations,
-            additional_properties=additional_properties,
-            raw_representation=raw_representation,
-            **kwargs,
-        )
-        self.call_id = call_id
-        self.inputs: list["Contents"] | None = None
-        if inputs:
-            normalized_inputs: Sequence["Contents | MutableMapping[str, Any]"] = (
-                inputs
-                if isinstance(inputs, Sequence) and not isinstance(inputs, (str, bytes, MutableMapping))
-                else [inputs]
-            )
-            self.inputs = _parse_content_list(list(normalized_inputs))
-        self.type: Literal["code_interpreter_tool_call"] = "code_interpreter_tool_call"
-
-
-class CodeInterpreterToolResultContent(BaseContent):
-    """Represents the result of a code interpreter tool invocation by a hosted service."""
-
-    def __init__(
-        self,
-        *,
-        call_id: str | None = None,
-        outputs: Sequence["Contents | MutableMapping[str, Any]"] | None = None,
-        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(
-            annotations=annotations,
-            additional_properties=additional_properties,
-            raw_representation=raw_representation,
-            **kwargs,
-        )
-        self.call_id = call_id
-        self.outputs: list["Contents"] | None = None
-        if outputs:
-            normalized_outputs: Sequence["Contents | MutableMapping[str, Any]"] = (
-                outputs
-                if isinstance(outputs, Sequence) and not isinstance(outputs, (str, bytes, MutableMapping))
-                else [outputs]
-            )
-            self.outputs = _parse_content_list(list(normalized_outputs))
-        self.type: Literal["code_interpreter_tool_result"] = "code_interpreter_tool_result"
-
-
-class ImageGenerationToolCallContent(BaseContent):
-    """Represents the invocation of an image generation tool call by a hosted service."""
-
-    def __init__(
-        self,
-        *,
-        image_id: str | None = None,
-        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Initializes an ImageGenerationToolCallContent instance.
-
-        Keyword Args:
-            image_id: The identifier of the image to be generated.
-            annotations: Optional annotations associated with the content.
-            additional_properties: Optional additional properties associated with the content.
-            raw_representation: Optional raw representation of the content.
-            **kwargs: Any additional keyword arguments.
-
-        """
-        super().__init__(
-            annotations=annotations,
-            additional_properties=additional_properties,
-            raw_representation=raw_representation,
-            **kwargs,
-        )
-        self.image_id = image_id
-        self.type: Literal["image_generation_tool_call"] = "image_generation_tool_call"
-
-
-class ImageGenerationToolResultContent(BaseContent):
-    """Represents the result of an image generation tool call invocation by a hosted service."""
-
-    def __init__(
-        self,
-        *,
-        image_id: str | None = None,
-        outputs: DataContent | UriContent | None = None,
-        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Initializes an ImageGenerationToolResultContent instance.
-
-        Keyword Args:
-            image_id: The identifier of the generated image.
-            outputs: The outputs of the image generation tool call.
-            annotations: Optional annotations associated with the content.
-            additional_properties: Optional additional properties associated with the content.
-            raw_representation: Optional raw representation of the content.
-            **kwargs: Any additional keyword arguments.
-
-        """
-        super().__init__(
-            annotations=annotations,
-            additional_properties=additional_properties,
-            raw_representation=raw_representation,
-            **kwargs,
-        )
-        self.image_id = image_id
-        self.outputs: DataContent | UriContent | None = outputs
-        self.type: Literal["image_generation_tool_result"] = "image_generation_tool_result"
-
-
-class MCPServerToolCallContent(BaseContent):
-    """Represents a tool call request to a MCP server."""
-
-    def __init__(
-        self,
-        call_id: str,
-        tool_name: str,
-        server_name: str | None = None,
-        *,
-        arguments: str | Mapping[str, Any] | None = None,
-        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Initializes a MCPServerToolCallContent instance.
-
-        Args:
-            call_id: The tool call identifier.
-            tool_name: The name of the tool requested.
-            server_name: The name of the MCP server where the tool is hosted.
-
-        Keyword Args:
-            arguments: The arguments requested to be provided to the tool,
-                can be a string to allow gradual completion of the args.
-            annotations: Optional annotations associated with the content.
-            additional_properties: Optional additional properties associated with the content.
-            raw_representation: Optional raw representation of the content.
-            **kwargs: Any additional keyword arguments.
-        """
-        if not call_id:
-            raise ValueError("call_id must be a non-empty string.")
-        if not tool_name:
-            raise ValueError("tool_name must be a non-empty string.")
-        super().__init__(
-            annotations=annotations,
-            additional_properties=additional_properties,
-            raw_representation=raw_representation,
-            **kwargs,
-        )
-        self.call_id = call_id
-        self.tool_name = tool_name
-        self.name = tool_name
-        self.server_name = server_name
-        self.arguments = arguments
-        self.type: Literal["mcp_server_tool_call"] = "mcp_server_tool_call"
-
-    def parse_arguments(self) -> dict[str, Any] | None:
-        """Returns the parsed arguments for the MCP server tool call, if any."""
-        if isinstance(self.arguments, str):
-            # If arguments are a string, try to parse it as JSON
-            try:
-                loaded = json.loads(self.arguments)
-                if isinstance(loaded, dict):
-                    return loaded  # type:ignore
-                return {"raw": loaded}
-            except (json.JSONDecodeError, TypeError):
-                return {"raw": self.arguments}
-        return cast(dict[str, Any] | None, self.arguments)
-
-
-class MCPServerToolResultContent(BaseContent):
-    """Represents the result of a MCP server tool call."""
-
-    def __init__(
-        self,
-        call_id: str,
-        *,
-        output: Any | None = None,
-        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Initializes a MCPServerToolResultContent instance.
-
-        Args:
-            call_id: The identifier of the tool call for which this is the result.
-
-        Keyword Args:
-            output: The output of the MCP server tool call.
-            annotations: Optional annotations associated with the content.
-            additional_properties: Optional additional properties associated with the content.
-            raw_representation: Optional raw representation of the content.
-            **kwargs: Any additional keyword arguments.
-        """
-        if not call_id:
-            raise ValueError("call_id must be a non-empty string.")
-        super().__init__(
-            annotations=annotations,
-            additional_properties=additional_properties,
-            raw_representation=raw_representation,
-            **kwargs,
-        )
-        self.call_id = call_id
-        self.output: Any | None = output
-        self.type: Literal["mcp_server_tool_result"] = "mcp_server_tool_result"
-
-
 class BaseUserInputRequest(BaseContent):
    """Base class for all user requests."""

@@ -1992,7 +1736,7 @@ class FunctionApprovalResponseContent(BaseContent):
        approved: bool,
        *,
        id: str,
-        function_call: FunctionCallContent | MCPServerToolCallContent | MutableMapping[str, Any],
+        function_call: FunctionCallContent | MutableMapping[str, Any],
        annotations: Sequence[Annotations | MutableMapping[str, Any]] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
@@ -2020,12 +1764,8 @@ class FunctionApprovalResponseContent(BaseContent):
        self.id = id
        self.approved = approved
        # Convert dict to FunctionCallContent if needed (for SerializationMixin support)
-        self.function_call: FunctionCallContent | MCPServerToolCallContent
        if isinstance(function_call, MutableMapping):
-            if function_call.get("type") == "mcp_server_tool_call":
-                self.function_call = MCPServerToolCallContent.from_dict(function_call)
-            else:
-                self.function_call = FunctionCallContent.from_dict(function_call)
+            self.function_call = FunctionCallContent.from_dict(function_call)
        else:
            self.function_call = function_call
        # Override the type for this specific subclass
@@ -2083,7 +1823,6 @@ class FunctionApprovalRequestContent(BaseContent):
            **kwargs,
        )
        self.id = id
-        self.function_call: FunctionCallContent
        # Convert dict to FunctionCallContent if needed (for SerializationMixin support)
        if isinstance(function_call, MutableMapping):
            self.function_call = FunctionCallContent.from_dict(function_call)
@@ -2115,12 +1854,6 @@ Contents = (
    | UsageContent
    | HostedFileContent
    | HostedVectorStoreContent
-    | CodeInterpreterToolCallContent
-    | CodeInterpreterToolResultContent
-    | ImageGenerationToolCallContent
-    | ImageGenerationToolResultContent
-    | MCPServerToolCallContent
-    | MCPServerToolResultContent
    | FunctionApprovalRequestContent
    | FunctionApprovalResponseContent
 )
@@ -2136,9 +1869,6 @@ def _prepare_function_call_results_as_dumpable(content: Contents | Any | list[Co
        return content.model_dump()
    if hasattr(content, "to_dict"):
        return content.to_dict(exclude={"raw_representation", "additional_properties"})
-    # Handle objects with text attribute (e.g., MCP TextContent)
-    if hasattr(content, "text") and isinstance(content.text, str):
-        return content.text
    return content


@@ -309,9 +309,6 @@ class WorkflowAgent(BaseAgent):
                if isinstance(executor, AgentExecutor) and not executor.output_response:
                    return None
                if update:
-                    # Enrich with executor identity if author_name is not already set
-                    if not update.author_name:
-                        update.author_name = executor_id
                    return update
                return None

@@ -3,7 +3,7 @@
 import json
 import sys
 from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, MutableSequence
-from typing import Any, cast
+from typing import Any

 from openai import AsyncOpenAI
 from openai.types.beta.threads import (
@@ -28,11 +28,9 @@ from .._types import (
    ChatOptions,
    ChatResponse,
    ChatResponseUpdate,
-    CodeInterpreterToolCallContent,
    Contents,
    FunctionCallContent,
    FunctionResultContent,
-    MCPServerToolCallContent,
    Role,
    TextContent,
    ToolMode,
@@ -379,37 +377,10 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):

        if event_data.required_action is not None:
            for tool_call in event_data.required_action.submit_tool_outputs.tool_calls:
-                tool_call_any = cast(Any, tool_call)
                call_id = json.dumps([response_id, tool_call.id])
-                tool_type = getattr(tool_call, "type", None)
-                if tool_type == "code_interpreter" and getattr(tool_call_any, "code_interpreter", None):
-                    code_input = getattr(tool_call_any.code_interpreter, "input", None)
-                    inputs = (
-                        [TextContent(text=code_input, raw_representation=tool_call)] if code_input is not None else None
-                    )
-                    contents.append(
-                        CodeInterpreterToolCallContent(
-                            call_id=call_id,
-                            inputs=inputs,
-                            raw_representation=tool_call,
-                        )
-                    )
-                elif tool_type == "mcp":
-                    contents.append(
-                        MCPServerToolCallContent(
-                            call_id=call_id,
-                            tool_name=getattr(tool_call, "name", "") or "",
-                            server_name=getattr(tool_call, "server_label", None),
-                            arguments=getattr(tool_call, "args", None),
-                            raw_representation=tool_call,
-                        )
-                    )
-                else:
-                    function_name = tool_call.function.name
-                    function_arguments = json.loads(tool_call.function.arguments)
-                    contents.append(
-                        FunctionCallContent(call_id=call_id, name=function_name, arguments=function_arguments)
-                    )
+                function_name = tool_call.function.name
+                function_arguments = json.loads(tool_call.function.arguments)
+                contents.append(FunctionCallContent(call_id=call_id, name=function_name, arguments=function_arguments))

        return contents

@@ -1,14 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.

-from collections.abc import (
-    AsyncIterable,
-    Awaitable,
-    Callable,
-    Mapping,
-    MutableMapping,
-    MutableSequence,
-    Sequence,
-)
+from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, MutableSequence, Sequence
 from datetime import datetime, timezone
 from itertools import chain
 from typing import Any, TypeVar, cast
@@ -20,9 +12,7 @@ from openai.types.responses.parsed_response import (
    ParsedResponse,
 )
 from openai.types.responses.response import Response as OpenAIResponse
-from openai.types.responses.response_stream_event import (
-    ResponseStreamEvent as OpenAIResponseStreamEvent,
-)
+from openai.types.responses.response_stream_event import ResponseStreamEvent as OpenAIResponseStreamEvent
 from openai.types.responses.response_usage import ResponseUsage
 from openai.types.responses.tool_param import (
    CodeInterpreter,
@@ -30,9 +20,7 @@ from openai.types.responses.tool_param import (
    Mcp,
    ToolParam,
 )
-from openai.types.responses.web_search_tool_param import (
-    UserLocation as WebSearchUserLocation,
-)
+from openai.types.responses.web_search_tool_param import UserLocation as WebSearchUserLocation
 from openai.types.responses.web_search_tool_param import WebSearchToolParam
 from pydantic import BaseModel, ValidationError

@@ -43,7 +31,6 @@ from .._tools import (
    AIFunction,
    HostedCodeInterpreterTool,
    HostedFileSearchTool,
-    HostedImageGenerationTool,
    HostedMCPTool,
    HostedWebSearchTool,
    ToolProtocol,
@@ -55,8 +42,6 @@ from .._types import (
    ChatResponse,
    ChatResponseUpdate,
    CitationAnnotation,
-    CodeInterpreterToolCallContent,
-    CodeInterpreterToolResultContent,
    Contents,
    DataContent,
    FunctionApprovalRequestContent,
@@ -65,10 +50,6 @@ from .._types import (
    FunctionResultContent,
    HostedFileContent,
    HostedVectorStoreContent,
-    ImageGenerationToolCallContent,
-    ImageGenerationToolResultContent,
-    MCPServerToolCallContent,
-    MCPServerToolResultContent,
    Role,
    TextContent,
    TextReasoningContent,
@@ -76,7 +57,6 @@ from .._types import (
    UriContent,
    UsageContent,
    UsageDetails,
-    _parse_content,
    prepare_function_call_results,
 )
 from ..exceptions import (
@@ -151,17 +131,13 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
            if "text_format" not in run_options:
                async for chunk in await client.responses.create(stream=True, **run_options):
                    yield self._parse_chunk_from_openai(
-                        chunk,
-                        chat_options=chat_options,
-                        function_call_ids=function_call_ids,
+                        chunk, chat_options=chat_options, function_call_ids=function_call_ids
                    )
                return
            async with client.responses.stream(**run_options) as response:
                async for chunk in response:
                    yield self._parse_chunk_from_openai(
-                        chunk,
-                        chat_options=chat_options,
-                        function_call_ids=function_call_ids,
+                        chunk, chat_options=chat_options, function_call_ids=function_call_ids
                    )
        except BadRequestError as ex:
            if ex.code == "content_filter":
@@ -338,28 +314,39 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                                else None,
                            )
                        )
-                    case HostedImageGenerationTool():
-                        mapped_tool: dict[str, Any] = {"type": "image_generation"}
-                        if tool.options:
-                            option_mapping = {
-                                "image_size": "size",
-                                "media_type": "output_format",
-                                "model_id": "model",
-                                "streaming_count": "partial_images",
-                            }
-                            # count and response_format are not supported by Responses API
-                            for key, value in tool.options.items():
-                                mapped_key = option_mapping.get(key, key)
-                                mapped_tool[mapped_key] = value
-                        if tool.additional_properties:
-                            mapped_tool.update(tool.additional_properties)
-                        response_tools.append(mapped_tool)
                    case _:
                        logger.debug("Unsupported tool passed (type: %s)", type(tool))
            else:
                # Handle raw dictionary tools
                tool_dict = tool if isinstance(tool, dict) else dict(tool)
-                response_tools.append(tool_dict)
+
+                # Special handling for image_generation tools
+                if tool_dict.get("type") == "image_generation":
+                    # Create a copy to avoid modifying the original
+                    mapped_tool = tool_dict.copy()
+
+                    # Map user-friendly parameter names to OpenAI API parameter names
+                    parameter_mapping = {
+                        "format": "output_format",
+                        "compression": "output_compression",
+                    }
+
+                    for user_param, api_param in parameter_mapping.items():
+                        if user_param in mapped_tool:
+                            # Map the parameter name and remove the old one
+                            mapped_tool[api_param] = mapped_tool.pop(user_param)
+
+                    # Validate partial_images parameter for streaming image generation
+                    # OpenAI API requires partial_images to be between 0-3 (inclusive) for image_generation tool
+                    # Reference: https://platform.openai.com/docs/api-reference/responses/create#responses_create-tools-image_generation_tool-partial_images
+                    if "partial_images" in mapped_tool:
+                        partial_images = mapped_tool["partial_images"]
+                        if not isinstance(partial_images, int) or partial_images < 0 or partial_images > 3:
+                            raise ValueError("partial_images must be an integer between 0 and 3 (inclusive).")
+
+                    response_tools.append(mapped_tool)
+                else:
+                    response_tools.append(tool_dict)
        return response_tools

    @staticmethod
@@ -452,23 +439,16 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
        if (tool_choice := run_options.get("tool_choice")) and isinstance(tool_choice, dict) and "mode" in tool_choice:
            run_options["tool_choice"] = tool_choice["mode"]

-        # additional properties (excluding response_format which is handled separately)
+        # additional properties
        additional_options = {
-            key: value
-            for key, value in chat_options.additional_properties.items()
-            if value is not None and key != "response_format"
+            key: value for key, value in chat_options.additional_properties.items() if value is not None
        }
        if additional_options:
            run_options.update(additional_options)

        # response format and text config (after additional_properties so user can pass text via additional_properties)
-        # Check both chat_options.response_format and additional_properties for response_format
-        response_format: Any = (
-            chat_options.response_format
-            if chat_options.response_format is not None
-            else chat_options.additional_properties.get("response_format")
-        )
-        text_config: Any = run_options.pop("text", None)
+        response_format = chat_options.response_format
+        text_config = run_options.pop("text", None)
        response_format, text_config = self._prepare_response_and_text_format(
            response_format=response_format, text_config=text_config
        )
@@ -571,10 +551,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                    if status := props.get("status"):
                        ret["status"] = status
                    if reasoning_text := props.get("reasoning_text"):
-                        ret["content"] = {
-                            "type": "reasoning_text",
-                            "text": reasoning_text,
-                        }
+                        ret["content"] = {"type": "reasoning_text", "text": reasoning_text}
                    if encrypted_content := props.get("encrypted_content"):
                        ret["encrypted_content"] = encrypted_content
                return ret
@@ -620,17 +597,9 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                    return file_obj
                return {}
            case FunctionCallContent():
-                if not content.call_id:
-                    logger.warning(f"FunctionCallContent missing call_id for function '{content.name}'")
-                    return {}
-                # Use fc_id from additional_properties if available, otherwise fallback to call_id
-                fc_id = call_id_to_id.get(content.call_id, content.call_id)
-                # OpenAI Responses API requires IDs to start with `fc_`
-                if not fc_id.startswith("fc_"):
-                    fc_id = f"fc_{fc_id}"
                return {
                    "call_id": content.call_id,
-                    "id": fc_id,
+                    "id": call_id_to_id[content.call_id],
                    "type": "function_call",
                    "name": content.name,
                    "arguments": content.arguments,
@@ -766,17 +735,11 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                                                    )
                                                )
                                            case _:
-                                                logger.debug(
-                                                    "Unparsed annotation type: %s",
-                                                    annotation.type,
-                                                )
+                                                logger.debug("Unparsed annotation type: %s", annotation.type)
                                contents.append(text_content)
                            case "refusal":
                                contents.append(
-                                    TextContent(
-                                        text=message_content.refusal,
-                                        raw_representation=message_content,
-                                    )
+                                    TextContent(text=message_content.refusal, raw_representation=message_content)
                                )
                case "reasoning":  # ResponseOutputReasoning
                    if hasattr(item, "content") and item.content:
@@ -797,40 +760,22 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                                TextReasoningContent(text=summary.text, raw_representation=summary)  # type: ignore[arg-type]
                            )
                case "code_interpreter_call":  # ResponseOutputCodeInterpreterCall
-                    call_id = getattr(item, "call_id", None) or getattr(item, "id", None)
-                    outputs: list["Contents"] = []
-                    if item_outputs := getattr(item, "outputs", None):
-                        for code_output in item_outputs:
-                            if getattr(code_output, "type", None) == "logs":
-                                outputs.append(
-                                    TextContent(
-                                        text=code_output.logs,
-                                        raw_representation=code_output,
-                                    )
-                                )
-                            elif getattr(code_output, "type", None) == "image":
-                                outputs.append(
+                    if hasattr(item, "outputs") and item.outputs:
+                        for code_output in item.outputs:
+                            if code_output.type == "logs":
+                                contents.append(TextContent(text=code_output.logs, raw_representation=item))
+                            if code_output.type == "image":
+                                contents.append(
                                    UriContent(
                                        uri=code_output.url,
-                                        raw_representation=code_output,
+                                        raw_representation=item,
+                                        # no more specific media type then this can be inferred
                                        media_type="image",
                                    )
                                )
-                    if code := getattr(item, "code", None):
-                        contents.append(
-                            CodeInterpreterToolCallContent(
-                                call_id=call_id,
-                                inputs=[TextContent(text=code, raw_representation=item)],
-                                raw_representation=item,
-                            )
-                        )
-                    contents.append(
-                        CodeInterpreterToolResultContent(
-                            call_id=call_id,
-                            outputs=outputs,
-                            raw_representation=item,
-                        )
-                    )
+                    elif hasattr(item, "code") and item.code:
+                        # fallback if no output was returned is the code:
+                        contents.append(TextContent(text=item.code, raw_representation=item))
                case "function_call":  # ResponseOutputFunctionCall
                    contents.append(
                        FunctionCallContent(
@@ -854,49 +799,31 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                            ),
                        )
                    )
-                case "mcp_call":
-                    call_id = item.id
-                    contents.append(
-                        MCPServerToolCallContent(
-                            call_id=call_id,
-                            tool_name=item.name,
-                            server_name=item.server_label,
-                            arguments=item.arguments,
-                            raw_representation=item,
-                        )
-                    )
-                    if item.output is not None:
+                case "image_generation_call":  # ResponseOutputImageGenerationCall
+                    if item.result:
+                        # Handle the result as either a proper data URI or raw base64 string
+                        uri = item.result
+                        media_type = None
+                        if not uri.startswith("data:"):
+                            # Raw base64 string - convert to proper data URI format using helper
+                            uri, media_type = DataContent.create_data_uri_from_base64(uri)
+                        else:
+                            # Parse media type from existing data URI
+                            try:
+                                # Extract media type from data URI (e.g., "data:image/png;base64,...")
+                                if ";" in uri and uri.startswith("data:"):
+                                    media_type = uri.split(";")[0].split(":", 1)[1]
+                            except Exception:
+                                # Fallback if parsing fails
+                                media_type = "image"
                        contents.append(
-                            MCPServerToolResultContent(
-                                call_id=call_id,
-                                output=[TextContent(text=item.output)],
+                            DataContent(
+                                uri=uri,
+                                media_type=media_type,
                                raw_representation=item,
                            )
                        )
-                case "image_generation_call":  # ResponseOutputImageGenerationCall
-                    image_output: DataContent | None = None
-                    if item.result:
-                        base64_data = item.result
-                        image_format = DataContent.detect_image_format_from_base64(base64_data)
-                        image_output = DataContent(
-                            data=base64_data,
-                            media_type=f"image/{image_format}" if image_format else "image/png",
-                            raw_representation=item.result,
-                        )
-                    image_id = item.id
-                    contents.append(
-                        ImageGenerationToolCallContent(
-                            image_id=image_id,
-                            raw_representation=item,
-                        )
-                    )
-                    contents.append(
-                        ImageGenerationToolResultContent(
-                            image_id=image_id,
-                            outputs=image_output,
-                            raw_representation=item,
-                        )
-                    )
+                # TODO(peterychang): Add support for other content types
                case _:
                    logger.debug("Unparsed output of type: %s: %s", item.type, item)
        response_message = ChatMessage(role="assistant", contents=contents)
@@ -1046,10 +973,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                    # McpApprovalRequest,
                    # ResponseCustomToolCall,
                    case "function_call":
-                        function_call_ids[event.output_index] = (
-                            event_item.call_id,
-                            event_item.name,
-                        )
+                        function_call_ids[event.output_index] = (event_item.call_id, event_item.name)
                    case "mcp_approval_request":
                        contents.append(
                            FunctionApprovalRequestContent(
@@ -1063,78 +987,23 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                                ),
                            )
                        )
-                    case "mcp_call":
-                        call_id = getattr(event_item, "id", None) or getattr(event_item, "call_id", None) or ""
-                        contents.append(
-                            MCPServerToolCallContent(
-                                call_id=call_id,
-                                tool_name=getattr(event_item, "name", "") or "",
-                                server_name=getattr(event_item, "server_label", None),
-                                arguments=getattr(event_item, "arguments", None),
-                                raw_representation=event_item,
-                            )
-                        )
-                        result_output = (
-                            getattr(event_item, "result", None)
-                            or getattr(event_item, "output", None)
-                            or getattr(event_item, "outputs", None)
-                        )
-                        parsed_output: list[Contents] | None = None
-                        if result_output:
-                            normalized = (
-                                result_output
-                                if isinstance(result_output, Sequence)
-                                and not isinstance(result_output, (str, bytes, MutableMapping))
-                                else [result_output]
-                            )
-                            parsed_output = [_parse_content(output_item) for output_item in normalized]
-                        contents.append(
-                            MCPServerToolResultContent(
-                                call_id=call_id,
-                                output=parsed_output,
-                                raw_representation=event_item,
-                            )
-                        )
                    case "code_interpreter_call":  # ResponseOutputCodeInterpreterCall
-                        call_id = getattr(event_item, "call_id", None) or getattr(event_item, "id", None)
-                        outputs: list[Contents] = []
                        if hasattr(event_item, "outputs") and event_item.outputs:
                            for code_output in event_item.outputs:
-                                if getattr(code_output, "type", None) == "logs":
-                                    outputs.append(
-                                        TextContent(
-                                            text=cast(Any, code_output).logs,
-                                            raw_representation=code_output,
-                                        )
-                                    )
-                                elif getattr(code_output, "type", None) == "image":
-                                    outputs.append(
+                                if code_output.type == "logs":
+                                    contents.append(TextContent(text=code_output.logs, raw_representation=event_item))
+                                if code_output.type == "image":
+                                    contents.append(
                                        UriContent(
-                                            uri=cast(Any, code_output).url,
-                                            raw_representation=code_output,
+                                            uri=code_output.url,
+                                            raw_representation=event_item,
+                                            # no more specific media type then this can be inferred
                                            media_type="image",
                                        )
                                    )
-                        if hasattr(event_item, "code") and event_item.code:
-                            contents.append(
-                                CodeInterpreterToolCallContent(
-                                    call_id=call_id,
-                                    inputs=[
-                                        TextContent(
-                                            text=event_item.code,
-                                            raw_representation=event_item,
-                                        )
-                                    ],
-                                    raw_representation=event_item,
-                                )
-                            )
-                        contents.append(
-                            CodeInterpreterToolResultContent(
-                                call_id=call_id,
-                                outputs=outputs,
-                                raw_representation=event_item,
-                            )
-                        )
+                        elif hasattr(event_item, "code") and event_item.code:
+                            # fallback if no output was returned is the code:
+                            contents.append(TextContent(text=event_item.code, raw_representation=event_item))
                    case "reasoning":  # ResponseOutputReasoning
                        if hasattr(event_item, "content") and event_item.content:
                            for index, reasoning_content in enumerate(event_item.content):
@@ -1162,10 +1031,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                            call_id=call_id,
                            name=name,
                            arguments=event.delta,
-                            additional_properties={
-                                "output_index": event.output_index,
-                                "fc_id": event.item_id,
-                            },
+                            additional_properties={"output_index": event.output_index, "fc_id": event.item_id},
                            raw_representation=event,
                        )
                    )
@@ -1177,27 +1043,14 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                # Use helper function to create data URI from base64
                uri, media_type = DataContent.create_data_uri_from_base64(image_base64)

-                image_output = DataContent(
-                    uri=uri,
-                    media_type=media_type,
-                    additional_properties={
-                        "partial_image_index": partial_index,
-                        "is_partial_image": True,
-                    },
-                    raw_representation=event,
-                )
-
-                image_id = getattr(event, "item_id", None)
                contents.append(
-                    ImageGenerationToolCallContent(
-                        image_id=image_id,
-                        raw_representation=event,
-                    )
-                )
-                contents.append(
-                    ImageGenerationToolResultContent(
-                        image_id=image_id,
-                        outputs=image_output,
+                    DataContent(
+                        uri=uri,
+                        media_type=media_type,
+                        additional_properties={
+                            "partial_image_index": partial_index,
+                            "is_partial_image": True,
+                        },
                        raw_representation=event,
                    )
                )
@@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260107"
+version = "1.0.0b251223"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -50,7 +50,6 @@ all = [
    "agent-framework-copilotstudio",
    "agent-framework-declarative",
    "agent-framework-devui",
-    "agent-framework-durabletask",
    "agent-framework-lab",
    "agent-framework-mem0",
    "agent-framework-ollama",
@@ -552,24 +552,26 @@ async def test_azure_responses_client_agent_chat_options_agent_level() -> None:
 async def test_azure_responses_client_agent_hosted_mcp_tool() -> None:
    """Integration test for HostedMCPTool with Azure Response Agent using Microsoft Learn MCP."""

+    mcp_tool = HostedMCPTool(
+        name="Microsoft Learn MCP",
+        url="https://learn.microsoft.com/api/mcp",
+        description="A Microsoft Learn MCP server for documentation questions",
+        approval_mode="never_require",
+    )
+
    async with ChatAgent(
        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
        instructions="You are a helpful assistant that can help with microsoft documentation questions.",
-        tools=HostedMCPTool(
-            name="Microsoft Learn MCP",
-            url="https://learn.microsoft.com/api/mcp",
-            description="A Microsoft Learn MCP server for documentation questions",
-            approval_mode="never_require",
-        ),
+        tools=[mcp_tool],
    ) as agent:
        response = await agent.run(
            "How to create an Azure storage account using az cli?",
-            # this needs to be high enough to handle the full MCP tool response.
-            max_tokens=5000,
+            max_tokens=200,
        )

        assert isinstance(response, AgentRunResponse)
-        assert response.text
+        assert response.text is not None
+        assert len(response.text) > 0
        # Should contain Azure-related content since it's asking about Azure CLI
        assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])

@@ -10,7 +10,6 @@ from pydantic import BaseModel
 from agent_framework import (
    AIFunction,
    HostedCodeInterpreterTool,
-    HostedImageGenerationTool,
    HostedMCPTool,
    ToolProtocol,
    ai_function,
@@ -819,30 +818,6 @@ def test_hosted_code_interpreter_tool_with_unknown_input():
        HostedCodeInterpreterTool(inputs={"hosted_file": "file-single"})


-def test_hosted_image_generation_tool_defaults():
-    """HostedImageGenerationTool should default name and empty description."""
-    tool = HostedImageGenerationTool()
-
-    assert tool.name == "image_generation"
-    assert tool.description == ""
-    assert tool.options is None
-    assert str(tool) == "HostedImageGenerationTool(name=image_generation)"
-
-
-def test_hosted_image_generation_tool_with_options():
-    """HostedImageGenerationTool should store options."""
-    tool = HostedImageGenerationTool(
-        description="Generate images",
-        options={"format": "png", "size": "1024x1024"},
-        additional_properties={"quality": "high"},
-    )
-
-    assert tool.name == "image_generation"
-    assert tool.description == "Generate images"
-    assert tool.options == {"format": "png", "size": "1024x1024"}
-    assert tool.additional_properties == {"quality": "high"}
-
-
 # region HostedMCPTool tests


@@ -18,8 +18,6 @@ from agent_framework import (
    ChatResponse,
    ChatResponseUpdate,
    CitationAnnotation,
-    CodeInterpreterToolCallContent,
-    CodeInterpreterToolResultContent,
    DataContent,
    ErrorContent,
    FinishReason,
@@ -29,10 +27,6 @@ from agent_framework import (
    FunctionResultContent,
    HostedFileContent,
    HostedVectorStoreContent,
-    ImageGenerationToolCallContent,
-    ImageGenerationToolResultContent,
-    MCPServerToolCallContent,
-    MCPServerToolResultContent,
    Role,
    TextContent,
    TextReasoningContent,
@@ -275,78 +269,6 @@ def test_hosted_file_content_minimal():
    assert isinstance(content, BaseContent)


-def test_hosted_file_content_optional_fields():
-    """HostedFileContent should capture optional media type and name."""
-    content = HostedFileContent(file_id="file-789", media_type="image/png", name="plot.png")
-
-    assert content.media_type == "image/png"
-    assert content.name == "plot.png"
-    assert content.has_top_level_media_type("image")
-    assert content.has_top_level_media_type("application") is False
-
-
-# region: CodeInterpreter content
-
-
-def test_code_interpreter_tool_call_content_parses_inputs():
-    call = CodeInterpreterToolCallContent(
-        call_id="call-1",
-        inputs=[{"type": "text", "text": "print('hi')"}],
-    )
-
-    assert call.type == "code_interpreter_tool_call"
-    assert call.call_id == "call-1"
-    assert call.inputs and isinstance(call.inputs[0], TextContent)
-    assert call.inputs[0].text == "print('hi')"
-
-
-def test_code_interpreter_tool_result_content_outputs():
-    result = CodeInterpreterToolResultContent(
-        call_id="call-2",
-        outputs=[
-            {"type": "text", "text": "log output"},
-            {"type": "uri", "uri": "https://example.com/file.png", "media_type": "image/png"},
-        ],
-    )
-
-    assert result.type == "code_interpreter_tool_result"
-    assert result.call_id == "call-2"
-    assert result.outputs is not None
-    assert isinstance(result.outputs[0], TextContent)
-    assert isinstance(result.outputs[1], UriContent)
-
-
-# region: Image generation content
-
-
-def test_image_generation_tool_contents():
-    call = ImageGenerationToolCallContent(image_id="img-1")
-    outputs = [DataContent(data=b"1234", media_type="image/png")]
-    result = ImageGenerationToolResultContent(image_id="img-1", outputs=outputs)
-
-    assert call.type == "image_generation_tool_call"
-    assert call.image_id == "img-1"
-    assert result.type == "image_generation_tool_result"
-    assert result.image_id == "img-1"
-    assert result.outputs and isinstance(result.outputs[0], DataContent)
-
-
-# region: MCP server tool content
-
-
-def test_mcp_server_tool_call_and_result():
-    call = MCPServerToolCallContent(call_id="c-1", tool_name="tool", server_name="server", arguments={"x": 1})
-    assert call.type == "mcp_server_tool_call"
-    assert call.arguments == {"x": 1}
-
-    result = MCPServerToolResultContent(call_id="c-1", output=[{"type": "text", "text": "done"}])
-    assert result.type == "mcp_server_tool_result"
-    assert result.output
-
-    with raises(ValueError):
-        MCPServerToolCallContent(call_id="", tool_name="tool")
-
-
 # region: HostedVectorStoreContent


@@ -547,15 +469,6 @@ def test_function_approval_serialization_roundtrip():
    # The Contents union will need to be handled differently when we fully migrate


-def test_function_approval_accepts_mcp_call():
-    """Ensure FunctionApprovalRequestContent supports MCP server tool calls."""
-    mcp_call = MCPServerToolCallContent(call_id="c-mcp", tool_name="tool", server_name="srv", arguments={"x": 1})
-    req = FunctionApprovalRequestContent(id="req-mcp", function_call=mcp_call)
-
-    assert isinstance(req.function_call, MCPServerToolCallContent)
-    assert req.function_call.call_id == "c-mcp"
-
-
 # region BaseContent Serialization


@@ -2220,55 +2133,3 @@ def test_prepare_function_call_results_nested_pydantic_model():
    assert "Seattle" in json_result
    assert "rainy" in json_result
    assert "18.0" in json_result or "18" in json_result
-
-
-# region prepare_function_call_results with MCP TextContent-like objects
-
-
-def test_prepare_function_call_results_text_content_single():
-    """Test that objects with text attribute (like MCP TextContent) are properly handled."""
-    from dataclasses import dataclass
-
-    @dataclass
-    class MockTextContent:
-        text: str
-
-    result = [MockTextContent("Hello from MCP tool!")]
-    json_result = prepare_function_call_results(result)
-
-    # Should extract text and serialize as JSON array of strings
-    assert isinstance(json_result, str)
-    assert json_result == '["Hello from MCP tool!"]'
-
-
-def test_prepare_function_call_results_text_content_multiple():
-    """Test that multiple TextContent-like objects are serialized correctly."""
-    from dataclasses import dataclass
-
-    @dataclass
-    class MockTextContent:
-        text: str
-
-    result = [MockTextContent("First result"), MockTextContent("Second result")]
-    json_result = prepare_function_call_results(result)
-
-    # Should extract text from each and serialize as JSON array
-    assert isinstance(json_result, str)
-    assert json_result == '["First result", "Second result"]'
-
-
-def test_prepare_function_call_results_text_content_with_non_string_text():
-    """Test that objects with non-string text attribute are not treated as TextContent."""
-
-    class BadTextContent:
-        def __init__(self):
-            self.text = 12345  # Not a string!
-
-    result = [BadTextContent()]
-    json_result = prepare_function_call_results(result)
-
-    # Should not extract text since it's not a string, will serialize the object
-    assert isinstance(json_result, str)
-
-
-# endregion
@@ -26,8 +26,6 @@ from agent_framework import (
    ChatMessage,
    ChatResponse,
    ChatResponseUpdate,
-    CodeInterpreterToolCallContent,
-    CodeInterpreterToolResultContent,
    DataContent,
    FunctionApprovalRequestContent,
    FunctionApprovalResponseContent,
@@ -36,12 +34,9 @@ from agent_framework import (
    HostedCodeInterpreterTool,
    HostedFileContent,
    HostedFileSearchTool,
-    HostedImageGenerationTool,
    HostedMCPTool,
    HostedVectorStoreContent,
    HostedWebSearchTool,
-    ImageGenerationToolCallContent,
-    ImageGenerationToolResultContent,
    MCPStreamableHTTPTool,
    Role,
    TextContent,
@@ -617,14 +612,11 @@ def test_response_content_creation_with_code_interpreter() -> None:
    response = client._parse_response_from_openai(mock_response, chat_options=ChatOptions())  # type: ignore

    assert len(response.messages[0].contents) == 2
-    call_content, result_content = response.messages[0].contents
-    assert isinstance(call_content, CodeInterpreterToolCallContent)
-    assert call_content.inputs is not None
-    assert isinstance(call_content.inputs[0], TextContent)
-    assert isinstance(result_content, CodeInterpreterToolResultContent)
-    assert result_content.outputs is not None
-    assert any(isinstance(out, TextContent) for out in result_content.outputs)
-    assert any(isinstance(out, UriContent) for out in result_content.outputs)
+    assert isinstance(response.messages[0].contents[0], TextContent)
+    assert response.messages[0].contents[0].text == "Code execution log"
+    assert isinstance(response.messages[0].contents[1], UriContent)
+    assert response.messages[0].contents[1].uri == "https://example.com/image.png"
+    assert response.messages[0].contents[1].media_type == "image"


 def test_response_content_creation_with_function_call() -> None:
@@ -769,13 +761,14 @@ def test_prepare_tools_for_openai_with_raw_image_generation() -> None:
    """Test that raw image_generation tool dict is handled correctly with parameter mapping."""
    client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")

-    # Test with raw tool dict using OpenAI parameters directly
+    # Test with raw tool dict using user-friendly parameter names
    tool = {
        "type": "image_generation",
        "size": "1536x1024",
        "quality": "high",
-        "output_format": "webp",
-        "output_quality": 75,
+        "format": "webp",  # Will be mapped to output_format
+        "compression": 75,  # Will be mapped to output_compression
+        "background": "transparent",
    }

    resp_tools = client._prepare_tools_for_openai([tool])
@@ -787,8 +780,10 @@ def test_prepare_tools_for_openai_with_raw_image_generation() -> None:
    assert image_tool["type"] == "image_generation"
    assert image_tool["size"] == "1536x1024"
    assert image_tool["quality"] == "high"
+    assert image_tool["background"] == "transparent"
+    # Check parameter name mapping
    assert image_tool["output_format"] == "webp"
-    assert image_tool["output_quality"] == 75
+    assert image_tool["output_compression"] == 75


 def test_prepare_tools_for_openai_with_raw_image_generation_openai_responses_params() -> None:
@@ -802,7 +797,7 @@ def test_prepare_tools_for_openai_with_raw_image_generation_openai_responses_par
        "model": "gpt-image-1",
        "input_fidelity": "high",
        "moderation": "strict",
-        "output_format": "png",
+        "partial_images": 2,  # Should be integer 0-3
    }

    resp_tools = client._prepare_tools_for_openai([tool])
@@ -820,7 +815,7 @@ def test_prepare_tools_for_openai_with_raw_image_generation_openai_responses_par
    assert tool_dict["model"] == "gpt-image-1"
    assert tool_dict["input_fidelity"] == "high"
    assert tool_dict["moderation"] == "strict"
-    assert tool_dict["output_format"] == "png"
+    assert tool_dict["partial_images"] == 2


 def test_prepare_tools_for_openai_with_raw_image_generation_minimal() -> None:
@@ -841,24 +836,6 @@ def test_prepare_tools_for_openai_with_raw_image_generation_minimal() -> None:
    assert len(image_tool) == 1


-def test_prepare_tools_for_openai_with_hosted_image_generation() -> None:
-    """Test HostedImageGenerationTool conversion."""
-    client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
-    tool = HostedImageGenerationTool(
-        description="Generate images",
-        options={"output_format": "png", "size": "512x512"},
-        additional_properties={"quality": "high"},
-    )
-
-    resp_tools = client._prepare_tools_for_openai([tool])
-    assert len(resp_tools) == 1
-    image_tool = resp_tools[0]
-    assert image_tool["type"] == "image_generation"
-    assert image_tool["output_format"] == "png"
-    assert image_tool["size"] == "512x512"
-    assert image_tool["quality"] == "high"
-
-
 def test_parse_chunk_from_openai_with_mcp_approval_request() -> None:
    """Test that a streaming mcp_approval_request event is parsed into FunctionApprovalRequestContent."""
    client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
@@ -1301,11 +1278,9 @@ def test_parse_chunk_from_openai_code_interpreter() -> None:

    result = client._parse_chunk_from_openai(mock_event_image, chat_options, function_call_ids)  # type: ignore
    assert len(result.contents) == 1
-    assert isinstance(result.contents[0], CodeInterpreterToolResultContent)
-    assert result.contents[0].outputs
-    assert any(
-        isinstance(out, UriContent) and out.uri == "https://example.com/plot.png" for out in result.contents[0].outputs
-    )
+    assert isinstance(result.contents[0], UriContent)
+    assert result.contents[0].uri == "https://example.com/plot.png"
+    assert result.contents[0].media_type == "image"


 def test_parse_chunk_from_openai_reasoning() -> None:
@@ -1520,16 +1495,12 @@ def test_parse_response_from_openai_image_generation_raw_base64():
    with patch.object(client, "_get_metadata_from_response", return_value={}):
        response = client._parse_response_from_openai(mock_response, chat_options=ChatOptions())  # type: ignore

-    # Verify the response contains call + result with DataContent output
-    assert len(response.messages[0].contents) == 2
-    call_content, result_content = response.messages[0].contents
-    assert isinstance(call_content, ImageGenerationToolCallContent)
-    assert isinstance(result_content, ImageGenerationToolResultContent)
-    assert result_content.outputs
-    data_out = result_content.outputs
-    assert isinstance(data_out, DataContent)
-    assert data_out.uri.startswith("data:image/png;base64,")
-    assert data_out.media_type == "image/png"
+    # Verify the response contains DataContent with proper URI and media_type
+    assert len(response.messages[0].contents) == 1
+    content = response.messages[0].contents[0]
+    assert isinstance(content, DataContent)
+    assert content.uri.startswith("data:image/png;base64,")
+    assert content.media_type == "image/png"


 def test_parse_response_from_openai_image_generation_existing_data_uri():
@@ -1550,23 +1521,19 @@ def test_parse_response_from_openai_image_generation_existing_data_uri():
    valid_webp_base64 = base64.b64encode(webp_signature + b"VP8 fake_data").decode()
    mock_item = MagicMock()
    mock_item.type = "image_generation_call"
-    mock_item.result = valid_webp_base64
+    mock_item.result = f"data:image/webp;base64,{valid_webp_base64}"

    mock_response.output = [mock_item]

    with patch.object(client, "_get_metadata_from_response", return_value={}):
        response = client._parse_response_from_openai(mock_response, chat_options=ChatOptions())  # type: ignore

-    # Verify the response contains call + result with DataContent output
-    assert len(response.messages[0].contents) == 2
-    call_content, result_content = response.messages[0].contents
-    assert isinstance(call_content, ImageGenerationToolCallContent)
-    assert isinstance(result_content, ImageGenerationToolResultContent)
-    assert result_content.outputs
-    data_out = result_content.outputs
-    assert isinstance(data_out, DataContent)
-    assert data_out.uri == f"data:image/webp;base64,{valid_webp_base64}"
-    assert data_out.media_type == "image/webp"
+    # Verify the response contains DataContent with proper media_type parsed from URI
+    assert len(response.messages[0].contents) == 1
+    content = response.messages[0].contents[0]
+    assert isinstance(content, DataContent)
+    assert content.uri == f"data:image/webp;base64,{valid_webp_base64}"
+    assert content.media_type == "image/webp"


 def test_parse_response_from_openai_image_generation_format_detection():
@@ -1592,12 +1559,10 @@ def test_parse_response_from_openai_image_generation_format_detection():

    with patch.object(client, "_get_metadata_from_response", return_value={}):
        response_jpeg = client._parse_response_from_openai(mock_response_jpeg, chat_options=ChatOptions())  # type: ignore
-    result_contents = response_jpeg.messages[0].contents
-    assert isinstance(result_contents[1], ImageGenerationToolResultContent)
-    outputs = result_contents[1].outputs
-    assert outputs and isinstance(outputs, DataContent)
-    assert outputs.media_type == "image/jpeg"
-    assert "data:image/jpeg;base64," in outputs.uri
+    content_jpeg = response_jpeg.messages[0].contents[0]
+    assert isinstance(content_jpeg, DataContent)
+    assert content_jpeg.media_type == "image/jpeg"
+    assert "data:image/jpeg;base64," in content_jpeg.uri

    # Test WEBP detection
    webp_signature = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP"
@@ -1618,10 +1583,10 @@ def test_parse_response_from_openai_image_generation_format_detection():

    with patch.object(client, "_get_metadata_from_response", return_value={}):
        response_webp = client._parse_response_from_openai(mock_response_webp, chat_options=ChatOptions())  # type: ignore
-    outputs_webp = response_webp.messages[0].contents[1].outputs
-    assert outputs_webp and isinstance(outputs_webp, DataContent)
-    assert outputs_webp.media_type == "image/webp"
-    assert "data:image/webp;base64," in outputs_webp.uri
+    content_webp = response_webp.messages[0].contents[0]
+    assert isinstance(content_webp, DataContent)
+    assert content_webp.media_type == "image/webp"
+    assert "data:image/webp;base64," in content_webp.uri


 def test_parse_response_from_openai_image_generation_fallback():
@@ -1650,11 +1615,9 @@ def test_parse_response_from_openai_image_generation_fallback():
        response = client._parse_response_from_openai(mock_response, chat_options=ChatOptions())  # type: ignore

    # Verify it falls back to PNG format for unrecognized binary data
-    assert len(response.messages[0].contents) == 2
-    result_content = response.messages[0].contents[1]
-    assert isinstance(result_content, ImageGenerationToolResultContent)
-    assert result_content.outputs
-    content = result_content.outputs
+    assert len(response.messages[0].contents) == 1
+    content = response.messages[0].contents[0]
+    assert isinstance(content, DataContent)
    assert content.media_type == "image/png"
    assert f"data:image/png;base64,{unrecognized_base64}" == content.uri

@@ -2190,30 +2153,38 @@ async def test_openai_responses_client_agent_hosted_code_interpreter_tool():

@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
-async def test_openai_responses_client_agent_image_generation_tool():
+async def test_openai_responses_client_agent_raw_image_generation_tool():
    """Test OpenAI Responses Client agent with raw image_generation tool through OpenAIResponsesClient."""
    async with ChatAgent(
        chat_client=OpenAIResponsesClient(),
        instructions="You are a helpful assistant that can generate images.",
-        tools=HostedImageGenerationTool(options={"image_size": "1024x1024", "media_type": "png"}),
+        tools=[{"type": "image_generation", "size": "1024x1024", "quality": "low", "format": "png"}],
    ) as agent:
        # Test image generation functionality
        response = await agent.run("Generate an image of a cute red panda sitting on a tree branch in a forest.")

        assert isinstance(response, AgentRunResponse)
-        assert response.messages

-        # Verify we got image content - look for ImageGenerationToolResultContent
+        # For image generation, we expect to get some response content
+        # This could be DataContent with image data, UriContent
+        assert response.messages is not None and len(response.messages) > 0
+
+        # Check that we have some kind of content in the response
+        total_contents = sum(len(message.contents) for message in response.messages)
+        assert total_contents > 0, f"Expected some content in response messages, got {total_contents} contents"
+
+        # Verify we got image content - look for DataContent with URI starting with "data:image"
        image_content_found = False
        for message in response.messages:
            for content in message.contents:
-                if content.type == "image_generation_tool_result" and content.outputs:
+                uri = getattr(content, "uri", None)
+                if uri and uri.startswith("data:image"):
                    image_content_found = True
                    break
            if image_content_found:
                break

-        # The test passes if we got image content
+        # The test passes if we got image content (which we did based on the visible base64 output)
        assert image_content_found, "Expected to find image content in response"


@@ -2335,24 +2306,26 @@ async def test_openai_responses_client_agent_chat_options_agent_level() -> None:
 async def test_openai_responses_client_agent_hosted_mcp_tool() -> None:
    """Integration test for HostedMCPTool with OpenAI Response Agent using Microsoft Learn MCP."""

+    mcp_tool = HostedMCPTool(
+        name="Microsoft Learn MCP",
+        url="https://learn.microsoft.com/api/mcp",
+        description="A Microsoft Learn MCP server for documentation questions",
+        approval_mode="never_require",
+    )
+
    async with ChatAgent(
        chat_client=OpenAIResponsesClient(),
        instructions="You are a helpful assistant that can help with microsoft documentation questions.",
-        tools=HostedMCPTool(
-            name="Microsoft Learn MCP",
-            url="https://learn.microsoft.com/api/mcp",
-            description="A Microsoft Learn MCP server for documentation questions",
-            approval_mode="never_require",
-        ),
+        tools=[mcp_tool],
    ) as agent:
        response = await agent.run(
            "How to create an Azure storage account using az cli?",
-            # this needs to be high enough to handle the full MCP tool response.
-            max_tokens=5000,
+            max_tokens=200,
        )

        assert isinstance(response, AgentRunResponse)
-        assert response.text
+        assert response.text is not None
+        assert len(response.text) > 0
        # Should contain Azure-related content since it's asking about Azure CLI
        assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])

@@ -2382,91 +2355,3 @@ async def test_openai_responses_client_agent_local_mcp_tool() -> None:
        assert len(response.text) > 0
        # Should contain Azure-related content since it's asking about Azure CLI
        assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
-
-
-class ReleaseBrief(BaseModel):
-    """Structured output model for release brief testing."""
-
-    title: str
-    summary: str
-    highlights: list[str]
-    model_config = {"extra": "forbid"}
-
-
-@pytest.mark.flaky
-@skip_if_openai_integration_tests_disabled
-async def test_openai_responses_client_agent_with_response_format_pydantic() -> None:
-    """Integration test for response_format with Pydantic model using OpenAI Responses Client."""
-    async with ChatAgent(
-        chat_client=OpenAIResponsesClient(),
-        instructions="You are a helpful assistant that returns structured JSON responses.",
-    ) as agent:
-        response = await agent.run(
-            "Summarize the following release notes into a ReleaseBrief:\n\n"
-            "Version 2.0 Release Notes:\n"
-            "- Added new streaming API for real-time responses\n"
-            "- Improved error handling with detailed messages\n"
-            "- Performance boost of 50% in batch processing\n"
-            "- Fixed memory leak in connection pooling",
-            response_format=ReleaseBrief,
-        )
-
-        # Validate response
-        assert isinstance(response, AgentRunResponse)
-        assert response.value is not None
-        assert isinstance(response.value, ReleaseBrief)
-
-        # Validate structured output fields
-        brief = response.value
-        assert len(brief.title) > 0
-        assert len(brief.summary) > 0
-        assert len(brief.highlights) > 0
-
-
-@pytest.mark.flaky
-@skip_if_openai_integration_tests_disabled
-async def test_openai_responses_client_agent_with_runtime_json_schema() -> None:
-    """Integration test for response_format with runtime JSON schema using OpenAI Responses Client."""
-    runtime_schema = {
-        "title": "WeatherDigest",
-        "type": "object",
-        "properties": {
-            "location": {"type": "string"},
-            "conditions": {"type": "string"},
-            "temperature_c": {"type": "number"},
-            "advisory": {"type": "string"},
-        },
-        "required": ["location", "conditions", "temperature_c", "advisory"],
-        "additionalProperties": False,
-    }
-
-    async with ChatAgent(
-        chat_client=OpenAIResponsesClient(),
-        instructions="Return only JSON that matches the provided schema. Do not add commentary.",
-    ) as agent:
-        response = await agent.run(
-            "Give a brief weather digest for Seattle.",
-            additional_chat_options={
-                "response_format": {
-                    "type": "json_schema",
-                    "json_schema": {
-                        "name": runtime_schema["title"],
-                        "strict": True,
-                        "schema": runtime_schema,
-                    },
-                },
-            },
-        )
-
-        # Validate response
-        assert isinstance(response, AgentRunResponse)
-        assert response.text is not None
-
-        # Parse JSON and validate structure
-        import json
-
-        parsed = json.loads(response.text)
-        assert "location" in parsed
-        assert "conditions" in parsed
-        assert "temperature_c" in parsed
-        assert "advisory" in parsed
@@ -702,84 +702,6 @@ class TestWorkflowAgent:
        assert unique_text_count == 1, f"Response should appear exactly once, but appeared {unique_text_count} times"


-class TestWorkflowAgentAuthorName:
-    """Test cases for author_name enrichment in WorkflowAgent (GitHub issue #1331)."""
-
-    async def test_agent_run_update_event_gets_executor_id_as_author_name(self):
-        """Test that AgentRunUpdateEvent gets executor_id as author_name when not already set.
-
-        This validates the fix for GitHub issue #1331: agent responses should include
-        identification of which agent produced them in multi-agent workflows.
-        """
-        # Create workflow with executor that emits AgentRunUpdateEvent without author_name
-        executor1 = SimpleExecutor(id="my_executor_id", response_text="Response", emit_streaming=False)
-        workflow = WorkflowBuilder().set_start_executor(executor1).build()
-        agent = WorkflowAgent(workflow=workflow, name="Test Agent")
-
-        # Collect streaming updates
-        updates: list[AgentRunResponseUpdate] = []
-        async for update in agent.run_stream("Hello"):
-            updates.append(update)
-
-        # Verify at least one update was received
-        assert len(updates) >= 1
-
-        # Verify author_name is set to executor_id
-        assert updates[0].author_name == "my_executor_id"
-
-    async def test_agent_run_update_event_preserves_existing_author_name(self):
-        """Test that existing author_name is preserved and not overwritten."""
-
-        class AuthorNameExecutor(Executor):
-            """Executor that sets author_name explicitly."""
-
-            @handler
-            async def handle_message(self, message: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:
-                # Emit update with explicit author_name
-                update = AgentRunResponseUpdate(
-                    contents=[TextContent(text="Response with author")],
-                    role=Role.ASSISTANT,
-                    author_name="custom_author_name",  # Explicitly set
-                    message_id=str(uuid.uuid4()),
-                )
-                await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=update))
-
-        executor = AuthorNameExecutor(id="executor_id")
-        workflow = WorkflowBuilder().set_start_executor(executor).build()
-        agent = WorkflowAgent(workflow=workflow, name="Test Agent")
-
-        # Collect streaming updates
-        updates: list[AgentRunResponseUpdate] = []
-        async for update in agent.run_stream("Hello"):
-            updates.append(update)
-
-        # Verify author_name is preserved (not overwritten with executor_id)
-        assert len(updates) >= 1
-        assert updates[0].author_name == "custom_author_name"
-
-    async def test_multiple_executors_have_distinct_author_names(self):
-        """Test that multiple executors in a workflow have their own author_name."""
-        # Create workflow with two executors
-        executor1 = SimpleExecutor(id="first_executor", response_text="First", emit_streaming=False)
-        executor2 = SimpleExecutor(id="second_executor", response_text="Second", emit_streaming=False)
-
-        workflow = WorkflowBuilder().set_start_executor(executor1).add_edge(executor1, executor2).build()
-        agent = WorkflowAgent(workflow=workflow, name="Multi-Executor Agent")
-
-        # Collect streaming updates
-        updates: list[AgentRunResponseUpdate] = []
-        async for update in agent.run_stream("Hello"):
-            updates.append(update)
-
-        # Should have updates from both executors
-        assert len(updates) >= 2
-
-        # Verify each update has the correct author_name matching its executor
-        author_names = [u.author_name for u in updates]
-        assert "first_executor" in author_names
-        assert "second_executor" in author_names
-
-
 class TestWorkflowAgentMergeUpdates:
    """Test cases specifically for the WorkflowAgent.merge_updates static method."""

--- a/Show More
+++ b/Show More