chore: import upstream snapshot with attribution

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs

@@ -0,0 +1,475 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
+using System.Net.Http;
+using System.Runtime.CompilerServices;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Threading;
+using System.Threading.Tasks;
+using AGUI.Abstractions;
+using AGUI.Client;
+using FluentAssertions;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.TestHost;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests;
+
+public sealed class BasicStreamingTests : IAsyncDisposable
+{
+    private WebApplication? _app;
+    private HttpClient? _client;
+
+    [Fact]
+    public async Task ClientReceivesStreamedAssistantMessageAsync()
+    {
+        // Arrange
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession? session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "hello");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        session.Should().NotBeNull();
+
+        updates.Should().NotBeEmpty();
+        updates.Should().AllSatisfy(u => u.Role.Should().Be(ChatRole.Assistant));
+
+        // Verify assistant response message
+        AgentResponse response = updates.ToAgentResponse();
+        response.Messages.Should().HaveCount(1);
+        response.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        response.Messages[0].Text.Should().Be("Hello from fake agent!");
+    }
+
+    [Fact]
+    public async Task ClientReceivesRunLifecycleEventsAsync()
+    {
+        // Arrange
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession? session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "test");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert - RunStarted should be the first update
+        updates.Should().NotBeEmpty();
+        updates[0].ResponseId.Should().NotBeNullOrEmpty();
+        ChatResponseUpdate firstUpdate = updates[0].AsChatResponseUpdate();
+        // The AG-UI thread id is surfaced on the RUN_STARTED event (the new AGUI.Client keeps the
+        // client stateless and never populates ChatResponseUpdate.ConversationId).
+        string? threadId = (firstUpdate.RawRepresentation as RunStartedEvent)?.ThreadId;
+        string? runId = updates[0].ResponseId;
+        threadId.Should().NotBeNullOrEmpty();
+        runId.Should().NotBeNullOrEmpty();
+
+        // Should have received text updates
+        updates.Should().Contain(u => !string.IsNullOrEmpty(u.Text));
+
+        // All text content updates should have the same message ID
+        List<AgentResponseUpdate> textUpdates = updates.Where(u => !string.IsNullOrEmpty(u.Text)).ToList();
+        textUpdates.Should().NotBeEmpty();
+        string? firstMessageId = textUpdates.FirstOrDefault()?.MessageId;
+        firstMessageId.Should().NotBeNullOrEmpty();
+        textUpdates.Should().AllSatisfy(u => u.MessageId.Should().Be(firstMessageId));
+
+        // RunFinished should be the last update
+        AgentResponseUpdate lastUpdate = updates[^1];
+        lastUpdate.ResponseId.Should().Be(runId);
+        ChatResponseUpdate lastChatUpdate = lastUpdate.AsChatResponseUpdate();
+        // The stateless client never populates ChatResponseUpdate.ConversationId; thread identity stays
+        // on the AG-UI wire events instead, so verify the RUN_FINISHED event carries the same ids.
+        lastChatUpdate.ConversationId.Should().BeNull();
+        RunFinishedEvent? runFinished = updates
+            .Select(u => u.AsChatResponseUpdate().RawRepresentation as RunFinishedEvent)
+            .FirstOrDefault(e => e is not null);
+        runFinished.Should().NotBeNull();
+        runFinished!.ThreadId.Should().Be(threadId);
+        runFinished.RunId.Should().Be(runId);
+    }
+
+    [Fact]
+    public async Task RunAsyncAggregatesStreamingUpdatesAsync()
+    {
+        // Arrange
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession? session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "hello");
+
+        // Act
+        AgentResponse response = await agent.RunAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None);
+
+        // Assert
+        response.Messages.Should().NotBeEmpty();
+        response.Messages.Should().Contain(m => m.Role == ChatRole.Assistant);
+        response.Messages.Should().Contain(m => m.Text == "Hello from fake agent!");
+    }
+
+    [Fact]
+    public async Task MultiTurnConversationPreservesAllMessagesInSessionAsync()
+    {
+        // Arrange
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession chatClientSession = (ChatClientAgentSession)await agent.CreateSessionAsync();
+        ChatMessage firstUserMessage = new(ChatRole.User, "First question");
+
+        // Act - First turn
+        List<AgentResponseUpdate> firstTurnUpdates = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([firstUserMessage], chatClientSession, new AgentRunOptions(), CancellationToken.None))
+        {
+            firstTurnUpdates.Add(update);
+        }
+
+        // Assert first turn completed
+        firstTurnUpdates.Should().Contain(u => !string.IsNullOrEmpty(u.Text));
+
+        // Act - Second turn with another message
+        ChatMessage secondUserMessage = new(ChatRole.User, "Second question");
+        List<AgentResponseUpdate> secondTurnUpdates = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([secondUserMessage], chatClientSession, new AgentRunOptions(), CancellationToken.None))
+        {
+            secondTurnUpdates.Add(update);
+        }
+
+        // Assert second turn completed
+        secondTurnUpdates.Should().Contain(u => !string.IsNullOrEmpty(u.Text));
+
+        // Verify first turn assistant response
+        AgentResponse firstResponse = firstTurnUpdates.ToAgentResponse();
+        firstResponse.Messages.Should().HaveCount(1);
+        firstResponse.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        firstResponse.Messages[0].Text.Should().Be("Hello from fake agent!");
+
+        // Verify second turn assistant response
+        AgentResponse secondResponse = secondTurnUpdates.ToAgentResponse();
+        secondResponse.Messages.Should().HaveCount(1);
+        secondResponse.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        secondResponse.Messages[0].Text.Should().Be("Hello from fake agent!");
+    }
+
+    [Fact]
+    public async Task AgentSendsMultipleMessagesInOneTurnAsync()
+    {
+        // Arrange
+        await this.SetupTestServerAsync(useMultiMessageAgent: true);
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession chatClientSession = (ChatClientAgentSession)await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Tell me a story");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], chatClientSession, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert - Should have received text updates with different message IDs
+        List<AgentResponseUpdate> textUpdates = updates.Where(u => !string.IsNullOrEmpty(u.Text)).ToList();
+        textUpdates.Should().NotBeEmpty();
+
+        // Extract unique message IDs
+        List<string> messageIds = textUpdates.Select(u => u.MessageId).Where(id => !string.IsNullOrEmpty(id)).Distinct().ToList()!;
+        messageIds.Should().HaveCountGreaterThan(1, "agent should send multiple messages");
+
+        // Verify assistant messages from updates
+        AgentResponse response = updates.ToAgentResponse();
+        response.Messages.Should().HaveCountGreaterThan(1);
+        response.Messages.Should().AllSatisfy(m => m.Role.Should().Be(ChatRole.Assistant));
+    }
+
+    [Fact]
+    public async Task UserSendsMultipleMessagesAtOnceAsync()
+    {
+        // Arrange
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession chatClientSession = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        // Multiple user messages sent in one turn
+        ChatMessage[] userMessages =
+        [
+            new ChatMessage(ChatRole.User, "First part of question"),
+            new ChatMessage(ChatRole.User, "Second part of question"),
+            new ChatMessage(ChatRole.User, "Third part of question")
+        ];
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync(userMessages, chatClientSession, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert - Should have received assistant response
+        updates.Should().Contain(u => !string.IsNullOrEmpty(u.Text));
+        updates.Should().Contain(u => u.Role == ChatRole.Assistant);
+
+        // Verify assistant response message
+        AgentResponse response = updates.ToAgentResponse();
+        response.Messages.Should().HaveCount(1);
+        response.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        response.Messages[0].Text.Should().Be("Hello from fake agent!");
+    }
+
+    [Fact]
+    public async Task PostMalformedOrEmptyBody_ReturnsBadRequestAsync()
+    {
+        // Arrange
+        await this.SetupTestServerAsync();
+
+        var endpoint = new Uri("http://localhost/agent");
+
+        // Act - malformed JSON body
+        using var malformed = new StringContent("{ not valid json", System.Text.Encoding.UTF8, "application/json");
+        using HttpResponseMessage malformedResponse = await this._client!.PostAsync(endpoint, malformed);
+
+        // Act - empty body
+        using var empty = new StringContent(string.Empty, System.Text.Encoding.UTF8, "application/json");
+        using HttpResponseMessage emptyResponse = await this._client!.PostAsync(endpoint, empty);
+
+        // Assert - the hosting glue rejects both with 400 rather than 5xx.
+        malformedResponse.StatusCode.Should().Be(System.Net.HttpStatusCode.BadRequest);
+        emptyResponse.StatusCode.Should().Be(System.Net.HttpStatusCode.BadRequest);
+    }
+
+    private async Task SetupTestServerAsync(bool useMultiMessageAgent = false)
+    {
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.WebHost.UseTestServer();
+
+        builder.Services.AddAGUIServer();
+
+        if (useMultiMessageAgent)
+        {
+            builder.Services.AddSingleton<FakeMultiMessageAgent>();
+        }
+        else
+        {
+            builder.Services.AddSingleton<FakeChatClientAgent>();
+        }
+
+        this._app = builder.Build();
+
+        AIAgent agent = useMultiMessageAgent
+            ? this._app.Services.GetRequiredService<FakeMultiMessageAgent>()
+            : this._app.Services.GetRequiredService<FakeChatClientAgent>();
+
+        this._app.MapAGUIServer("/agent", agent);
+
+        await this._app.StartAsync();
+
+        TestServer testServer = this._app.Services.GetRequiredService<IServer>() as TestServer
+            ?? throw new InvalidOperationException("TestServer not found");
+
+        this._client = testServer.CreateClient();
+        this._client.BaseAddress = new Uri("http://localhost/agent");
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        this._client?.Dispose();
+        if (this._app != null)
+        {
+            await this._app.DisposeAsync();
+        }
+    }
+}
+
+[SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")]
+internal sealed class FakeChatClientAgent : AIAgent
+{
+    protected override string? IdCore => "fake-agent";
+
+    public override string? Description => "A fake agent for testing";
+
+    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default) =>
+        new(new FakeAgentSession());
+
+    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+        new(serializedState.Deserialize<FakeAgentSession>(jsonSerializerOptions)!);
+
+    protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        => throw new NotImplementedException();
+
+    protected override async Task<AgentResponse> RunCoreAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+    {
+        List<AgentResponseUpdate> updates = [];
+        await foreach (AgentResponseUpdate update in this.RunStreamingAsync(messages, session, options, cancellationToken).ConfigureAwait(false))
+        {
+            updates.Add(update);
+        }
+
+        return updates.ToAgentResponse();
+    }
+
+    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        string messageId = Guid.NewGuid().ToString("N");
+
+        // Simulate streaming a deterministic response
+        foreach (string chunk in new[] { "Hello", " ", "from", " ", "fake", " ", "agent", "!" })
+        {
+            yield return new AgentResponseUpdate
+            {
+                MessageId = messageId,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent(chunk)]
+            };
+
+            await Task.Yield();
+        }
+    }
+
+    private sealed class FakeAgentSession : AgentSession
+    {
+        public FakeAgentSession()
+        {
+        }
+
+        [JsonConstructor]
+        public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+        {
+        }
+    }
+}
+
+[SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")]
+internal sealed class FakeMultiMessageAgent : AIAgent
+{
+    protected override string? IdCore => "fake-multi-message-agent";
+
+    public override string? Description => "A fake agent that sends multiple messages for testing";
+
+    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default) =>
+        new(new FakeAgentSession());
+
+    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+        new(serializedState.Deserialize<FakeAgentSession>(jsonSerializerOptions)!);
+
+    protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        if (session is not FakeAgentSession fakeSession)
+        {
+            throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(FakeAgentSession)}' can be serialized by this agent.");
+        }
+
+        return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions));
+    }
+
+    protected override async Task<AgentResponse> RunCoreAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+    {
+        List<AgentResponseUpdate> updates = [];
+        await foreach (AgentResponseUpdate update in this.RunStreamingAsync(messages, session, options, cancellationToken).ConfigureAwait(false))
+        {
+            updates.Add(update);
+        }
+
+        return updates.ToAgentResponse();
+    }
+
+    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        // Simulate sending first message
+        string messageId1 = Guid.NewGuid().ToString("N");
+        foreach (string chunk in new[] { "First", " ", "message" })
+        {
+            yield return new AgentResponseUpdate
+            {
+                MessageId = messageId1,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent(chunk)]
+            };
+
+            await Task.Yield();
+        }
+
+        // Simulate sending second message
+        string messageId2 = Guid.NewGuid().ToString("N");
+        foreach (string chunk in new[] { "Second", " ", "message" })
+        {
+            yield return new AgentResponseUpdate
+            {
+                MessageId = messageId2,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent(chunk)]
+            };
+
+            await Task.Yield();
+        }
+
+        // Simulate sending third message
+        string messageId3 = Guid.NewGuid().ToString("N");
+        foreach (string chunk in new[] { "Third", " ", "message" })
+        {
+            yield return new AgentResponseUpdate
+            {
+                MessageId = messageId3,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent(chunk)]
+            };
+
+            await Task.Yield();
+        }
+    }
+
+    private sealed class FakeAgentSession : AgentSession
+    {
+        public FakeAgentSession()
+        {
+        }
+
+        [JsonConstructor]
+        public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+        {
+        }
+    }
+
+    public override object? GetService(Type serviceType, object? serviceKey = null) => null;
+}

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs

@@ -0,0 +1,369 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.IO;
+using System.Net.Http;
+using System.Net.ServerSentEvents;
+using System.Runtime.CompilerServices;
+using System.Text;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Threading;
+using System.Threading.Tasks;
+using AGUI.Abstractions;
+using AGUI.Server;
+using FluentAssertions;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.TestHost;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests;
+
+public sealed class ForwardedPropertiesTests : IAsyncDisposable
+{
+    private WebApplication? _app;
+    private HttpClient? _client;
+
+    [Fact]
+    public async Task ForwardedProps_AreParsedAndPassedToAgent_WhenProvidedInRequestAsync()
+    {
+        // Arrange
+        FakeForwardedPropsAgent fakeAgent = new();
+        await this.SetupTestServerAsync(fakeAgent);
+
+        // Create request JSON with forwardedProps (per AG-UI protocol spec)
+        const string RequestJson = """
+            {
+                "threadId": "session-123",
+                "runId": "run-456",
+                "messages": [{ "id": "msg-1", "role": "user", "content": "test forwarded props" }],
+                "forwardedProps": { "customProp": "customValue", "sessionId": "test-session-123" }
+            }
+            """;
+
+        using StringContent content = new(RequestJson, Encoding.UTF8, "application/json");
+
+        // Act
+        HttpResponseMessage response = await this._client!.PostAsync(new Uri("/agent", UriKind.Relative), content);
+
+        // Assert
+        response.IsSuccessStatusCode.Should().BeTrue();
+        fakeAgent.ReceivedForwardedProperties.ValueKind.Should().Be(JsonValueKind.Object);
+        fakeAgent.ReceivedForwardedProperties.GetProperty("customProp").GetString().Should().Be("customValue");
+        fakeAgent.ReceivedForwardedProperties.GetProperty("sessionId").GetString().Should().Be("test-session-123");
+    }
+
+    [Fact]
+    public async Task ForwardedProps_WithNestedObjects_AreCorrectlyParsedAsync()
+    {
+        // Arrange
+        FakeForwardedPropsAgent fakeAgent = new();
+        await this.SetupTestServerAsync(fakeAgent);
+
+        const string RequestJson = """
+            {
+                "threadId": "session-123",
+                "runId": "run-456",
+                "messages": [{ "id": "msg-1", "role": "user", "content": "test nested props" }],
+                "forwardedProps": {
+                    "user": { "id": "user-1", "name": "Test User" },
+                    "metadata": { "version": "1.0", "feature": "test" }
+                }
+            }
+            """;
+
+        using StringContent content = new(RequestJson, Encoding.UTF8, "application/json");
+
+        // Act
+        HttpResponseMessage response = await this._client!.PostAsync(new Uri("/agent", UriKind.Relative), content);
+
+        // Assert
+        response.IsSuccessStatusCode.Should().BeTrue();
+        fakeAgent.ReceivedForwardedProperties.ValueKind.Should().Be(JsonValueKind.Object);
+
+        JsonElement user = fakeAgent.ReceivedForwardedProperties.GetProperty("user");
+        user.GetProperty("id").GetString().Should().Be("user-1");
+        user.GetProperty("name").GetString().Should().Be("Test User");
+
+        JsonElement metadata = fakeAgent.ReceivedForwardedProperties.GetProperty("metadata");
+        metadata.GetProperty("version").GetString().Should().Be("1.0");
+        metadata.GetProperty("feature").GetString().Should().Be("test");
+    }
+
+    [Fact]
+    public async Task ForwardedProps_WithArrays_AreCorrectlyParsedAsync()
+    {
+        // Arrange
+        FakeForwardedPropsAgent fakeAgent = new();
+        await this.SetupTestServerAsync(fakeAgent);
+
+        const string RequestJson = """
+            {
+                "threadId": "session-123",
+                "runId": "run-456",
+                "messages": [{ "id": "msg-1", "role": "user", "content": "test array props" }],
+                "forwardedProps": {
+                    "tags": ["tag1", "tag2", "tag3"],
+                    "scores": [1, 2, 3, 4, 5]
+                }
+            }
+            """;
+
+        using StringContent content = new(RequestJson, Encoding.UTF8, "application/json");
+
+        // Act
+        HttpResponseMessage response = await this._client!.PostAsync(new Uri("/agent", UriKind.Relative), content);
+
+        // Assert
+        response.IsSuccessStatusCode.Should().BeTrue();
+        fakeAgent.ReceivedForwardedProperties.ValueKind.Should().Be(JsonValueKind.Object);
+
+        JsonElement tags = fakeAgent.ReceivedForwardedProperties.GetProperty("tags");
+        tags.GetArrayLength().Should().Be(3);
+        tags[0].GetString().Should().Be("tag1");
+
+        JsonElement scores = fakeAgent.ReceivedForwardedProperties.GetProperty("scores");
+        scores.GetArrayLength().Should().Be(5);
+        scores[2].GetInt32().Should().Be(3);
+    }
+
+    [Fact]
+    public async Task ForwardedProps_WhenEmpty_DoesNotCauseErrorsAsync()
+    {
+        // Arrange
+        FakeForwardedPropsAgent fakeAgent = new();
+        await this.SetupTestServerAsync(fakeAgent);
+
+        const string RequestJson = """
+            {
+                "threadId": "session-123",
+                "runId": "run-456",
+                "messages": [{ "id": "msg-1", "role": "user", "content": "test empty props" }],
+                "forwardedProps": {}
+            }
+            """;
+
+        using StringContent content = new(RequestJson, Encoding.UTF8, "application/json");
+
+        // Act
+        HttpResponseMessage response = await this._client!.PostAsync(new Uri("/agent", UriKind.Relative), content);
+
+        // Assert
+        response.IsSuccessStatusCode.Should().BeTrue();
+    }
+
+    [Fact]
+    public async Task ForwardedProps_WhenNotProvided_AgentStillWorksAsync()
+    {
+        // Arrange
+        FakeForwardedPropsAgent fakeAgent = new();
+        await this.SetupTestServerAsync(fakeAgent);
+
+        const string RequestJson = """
+            {
+                "threadId": "session-123",
+                "runId": "run-456",
+                "messages": [{ "id": "msg-1", "role": "user", "content": "test no props" }]
+            }
+            """;
+
+        using StringContent content = new(RequestJson, Encoding.UTF8, "application/json");
+
+        // Act
+        HttpResponseMessage response = await this._client!.PostAsync(new Uri("/agent", UriKind.Relative), content);
+
+        // Assert
+        response.IsSuccessStatusCode.Should().BeTrue();
+        fakeAgent.ReceivedForwardedProperties.ValueKind.Should().Be(JsonValueKind.Undefined);
+    }
+
+    [Fact]
+    public async Task ForwardedProps_ReturnsValidSSEResponse_WithTextDeltaEventsAsync()
+    {
+        // Arrange
+        FakeForwardedPropsAgent fakeAgent = new();
+        await this.SetupTestServerAsync(fakeAgent);
+
+        const string RequestJson = """
+            {
+                "threadId": "session-123",
+                "runId": "run-456",
+                "messages": [{ "id": "msg-1", "role": "user", "content": "test response" }],
+                "forwardedProps": { "customProp": "value" }
+            }
+            """;
+
+        using StringContent content = new(RequestJson, Encoding.UTF8, "application/json");
+
+        // Act
+        HttpResponseMessage response = await this._client!.PostAsync(new Uri("/agent", UriKind.Relative), content);
+        response.EnsureSuccessStatusCode();
+
+        Stream stream = await response.Content.ReadAsStreamAsync();
+        List<SseItem<string>> events = [];
+        await foreach (SseItem<string> item in SseParser.Create(stream).EnumerateAsync())
+        {
+            events.Add(item);
+        }
+
+        // Assert
+        events.Should().NotBeEmpty();
+
+        // SSE events have EventType = "message" and the actual type is in the JSON data
+        // Should have run_started event
+        events.Should().Contain(e => e.Data != null && e.Data.Contains("\"type\":\"RUN_STARTED\""));
+
+        // Should have text_message_start event
+        events.Should().Contain(e => e.Data != null && e.Data.Contains("\"type\":\"TEXT_MESSAGE_START\""));
+
+        // Should have text_message_content event with the response text
+        events.Should().Contain(e => e.Data != null && e.Data.Contains("\"type\":\"TEXT_MESSAGE_CONTENT\""));
+
+        // Should have run_finished event
+        events.Should().Contain(e => e.Data != null && e.Data.Contains("\"type\":\"RUN_FINISHED\""));
+    }
+
+    [Fact]
+    public async Task ForwardedProps_WithMixedTypes_AreCorrectlyParsedAsync()
+    {
+        // Arrange
+        FakeForwardedPropsAgent fakeAgent = new();
+        await this.SetupTestServerAsync(fakeAgent);
+
+        const string RequestJson = """
+            {
+                "threadId": "session-123",
+                "runId": "run-456",
+                "messages": [{ "id": "msg-1", "role": "user", "content": "test mixed types" }],
+                "forwardedProps": {
+                    "stringProp": "text",
+                    "numberProp": 42,
+                    "boolProp": true,
+                    "nullProp": null,
+                    "arrayProp": [1, "two", false],
+                    "objectProp": { "nested": "value" }
+                }
+            }
+            """;
+
+        using StringContent content = new(RequestJson, Encoding.UTF8, "application/json");
+
+        // Act
+        HttpResponseMessage response = await this._client!.PostAsync(new Uri("/agent", UriKind.Relative), content);
+
+        // Assert
+        response.IsSuccessStatusCode.Should().BeTrue();
+        fakeAgent.ReceivedForwardedProperties.ValueKind.Should().Be(JsonValueKind.Object);
+
+        fakeAgent.ReceivedForwardedProperties.GetProperty("stringProp").GetString().Should().Be("text");
+        fakeAgent.ReceivedForwardedProperties.GetProperty("numberProp").GetInt32().Should().Be(42);
+        fakeAgent.ReceivedForwardedProperties.GetProperty("boolProp").GetBoolean().Should().BeTrue();
+        fakeAgent.ReceivedForwardedProperties.GetProperty("nullProp").ValueKind.Should().Be(JsonValueKind.Null);
+        fakeAgent.ReceivedForwardedProperties.GetProperty("arrayProp").GetArrayLength().Should().Be(3);
+        fakeAgent.ReceivedForwardedProperties.GetProperty("objectProp").GetProperty("nested").GetString().Should().Be("value");
+    }
+
+    private async Task SetupTestServerAsync(FakeForwardedPropsAgent fakeAgent)
+    {
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.Services.AddAGUIServer();
+        builder.WebHost.UseTestServer();
+
+        this._app = builder.Build();
+
+        this._app.MapAGUIServer("/agent", fakeAgent);
+
+        await this._app.StartAsync();
+
+        TestServer testServer = this._app.Services.GetRequiredService<IServer>() as TestServer
+            ?? throw new InvalidOperationException("TestServer not found");
+
+        this._client = testServer.CreateClient();
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        this._client?.Dispose();
+        if (this._app != null)
+        {
+            await this._app.DisposeAsync();
+        }
+    }
+}
+
+[SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated in tests")]
+internal sealed class FakeForwardedPropsAgent : AIAgent
+{
+    public FakeForwardedPropsAgent()
+    {
+    }
+
+    public override string? Description => "Agent for forwarded properties testing";
+
+    public JsonElement ReceivedForwardedProperties { get; private set; }
+
+    protected override Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        return this.RunCoreStreamingAsync(messages, session, options, cancellationToken).ToAgentResponseAsync(cancellationToken);
+    }
+
+    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        // Recover the originating AG-UI input from the request options (set by the hosting layer).
+        if (options is ChatClientAgentRunOptions { ChatOptions: { } chatOptions } &&
+            chatOptions.TryGetRunAgentInput(out RunAgentInput? agentInput) &&
+            agentInput.ForwardedProperties is { ValueKind: not JsonValueKind.Undefined } forwardedProps)
+        {
+            this.ReceivedForwardedProperties = forwardedProps;
+        }
+
+        // Always return a text response
+        string messageId = Guid.NewGuid().ToString("N");
+        yield return new AgentResponseUpdate
+        {
+            MessageId = messageId,
+            Role = ChatRole.Assistant,
+            Contents = [new TextContent("Forwarded props processed")]
+        };
+
+        await Task.CompletedTask;
+    }
+
+    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default) =>
+        new(new FakeAgentSession());
+
+    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+        new(serializedState.Deserialize<FakeAgentSession>(jsonSerializerOptions)!);
+
+    protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        if (session is not FakeAgentSession fakeSession)
+        {
+            throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(FakeAgentSession)}' can be serialized by this agent.");
+        }
+
+        return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions));
+    }
+
+    private sealed class FakeAgentSession : AgentSession
+    {
+        public FakeAgentSession()
+        {
+        }
+
+        [JsonConstructor]
+        public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+        {
+        }
+    }
+
+    public override object? GetService(Type serviceType, object? serviceKey = null) => null;
+}

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests.csproj

@@ -0,0 +1,32 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFrameworks>$(TargetFrameworksCore)</TargetFrameworks>
+  </PropertyGroup>
+
+  <PropertyGroup>
+    <InjectSharedIntegrationTestCode>true</InjectSharedIntegrationTestCode>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Azure.Identity" />
+    <PackageReference Include="Azure.AI.OpenAI" />
+    <PackageReference Include="FluentAssertions" />
+    <PackageReference Include="Microsoft.AspNetCore.TestHost" />
+    <PackageReference Include="Microsoft.CodeAnalysis.CSharp" />
+    <PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
+    <PackageReference Include="Microsoft.Extensions.Configuration.UserSecrets" />
+    <PackageReference Include="Microsoft.Extensions.Configuration.EnvironmentVariables" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\src\Microsoft.Agents.AI\Microsoft.Agents.AI.csproj" />
+    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.Hosting\Microsoft.Agents.AI.Hosting.csproj" />
+    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <PackageReference Include="AGUI.Client" />
+    <PackageReference Include="AGUI.Abstractions" />
+    <PackageReference Include="AGUI.Server" />
+    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
+  </ItemGroup>
+
+</Project>

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SessionPersistenceTests.cs

@@ -0,0 +1,253 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
+using System.Net.Http;
+using System.Runtime.CompilerServices;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Threading;
+using System.Threading.Tasks;
+using AGUI.Abstractions;
+using AGUI.Client;
+using FluentAssertions;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.TestHost;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests;
+
+public sealed class SessionPersistenceTests : IAsyncDisposable
+{
+    private WebApplication? _app;
+    private HttpClient? _client;
+
+    [Fact]
+    public async Task MultiTurnWithSessionStore_PersistsSessionAcrossRequestsAsync()
+    {
+        // Arrange - use hosting DI pattern with InMemorySessionStore.
+        // FakeSessionAgent tracks turn count in session StateBag so we can verify
+        // that state survives the serialization round-trip through the session store.
+        await this.SetupTestServerWithSessionStoreAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        // Act - First turn
+        ChatMessage firstUserMessage = new(ChatRole.User, "First message");
+        List<AgentResponseUpdate> firstTurnUpdates = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([firstUserMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            firstTurnUpdates.Add(update);
+        }
+
+        // Act - Second turn (continue the same AG-UI run on the same thread).
+        // The new AGUI.Client is stateless and never round-trips ConversationId. AG-UI continuation is
+        // expressed on the wire with the same threadId plus a parentRunId equal to the previous turn's
+        // runId. We capture both from turn 1's RUN_STARTED event and supply them to the stateless client
+        // through RawRepresentationFactory, sending only the new message (the continuation subset).
+        RunStartedEvent? firstRunStarted = firstTurnUpdates
+            .Select(u => u.AsChatResponseUpdate().RawRepresentation as RunStartedEvent)
+            .FirstOrDefault(e => e is not null);
+        firstRunStarted.Should().NotBeNull();
+        string threadId = firstRunStarted!.ThreadId;
+        string previousRunId = firstRunStarted.RunId;
+        threadId.Should().NotBeNullOrEmpty();
+        previousRunId.Should().NotBeNullOrEmpty();
+
+        ChatMessage secondUserMessage = new(ChatRole.User, "Second message");
+        var continuationOptions = new ChatClientAgentRunOptions
+        {
+            ChatOptions = new ChatOptions
+            {
+                RawRepresentationFactory = _ => new RunAgentInput
+                {
+                    ThreadId = threadId,
+                    ParentRunId = previousRunId,
+                    Messages = new[] { secondUserMessage }.AsAGUIMessages().ToList(),
+                },
+            },
+        };
+        List<AgentResponseUpdate> secondTurnUpdates = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([secondUserMessage], session, continuationOptions, CancellationToken.None))
+        {
+            secondTurnUpdates.Add(update);
+        }
+
+        // Assert - Verify turn count proves session state was persisted.
+        // If session persistence were broken, both turns would return "Turn 1"
+        // because a fresh session (with turn count 0) would be created each time.
+        AgentResponse firstResponse = firstTurnUpdates.ToAgentResponse();
+        firstResponse.Messages.Should().HaveCount(1);
+        firstResponse.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        firstResponse.Messages[0].Text.Should().Contain("Turn 1:");
+
+        AgentResponse secondResponse = secondTurnUpdates.ToAgentResponse();
+        secondResponse.Messages.Should().HaveCount(1);
+        secondResponse.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        secondResponse.Messages[0].Text.Should().Contain("Turn 2:");
+    }
+
+    [Fact]
+    public async Task MapAGUIServer_WithAgentName_StreamsResponseCorrectlyAsync()
+    {
+        // Arrange - use the MapAGUIServer(agentName, pattern) overload via hosting DI
+        await this.SetupTestServerWithSessionStoreAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "hello");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        updates.Should().NotBeEmpty();
+        updates.Should().AllSatisfy(u => u.Role.Should().Be(ChatRole.Assistant));
+
+        AgentResponse response = updates.ToAgentResponse();
+        response.Messages.Should().HaveCount(1);
+        response.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        response.Messages[0].Text.Should().Be("Turn 1: Hello from session agent!");
+    }
+
+    private async Task SetupTestServerWithSessionStoreAsync()
+    {
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.WebHost.UseTestServer();
+
+        builder.Services.AddAGUIServer();
+
+        // Register agent using hosting DI pattern with InMemorySessionStore
+        builder.Services.AddAIAgent("session-test-agent", (_, name) => new FakeSessionAgent(name))
+            .WithInMemorySessionStore(withIsolation: false);
+
+        this._app = builder.Build();
+
+        // Use the agentName overload of MapAGUIServer
+        this._app.MapAGUIServer("session-test-agent", "/agent");
+
+        await this._app.StartAsync();
+
+        TestServer testServer = this._app.Services.GetRequiredService<IServer>() as TestServer
+            ?? throw new InvalidOperationException("TestServer not found");
+
+        this._client = testServer.CreateClient();
+        this._client.BaseAddress = new Uri("http://localhost/agent");
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        this._client?.Dispose();
+        if (this._app != null)
+        {
+            await this._app.DisposeAsync();
+        }
+    }
+}
+
+[SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")]
+internal sealed class FakeSessionAgent : AIAgent
+{
+    private readonly string _name;
+
+    public FakeSessionAgent(string name)
+    {
+        this._name = name;
+    }
+
+    protected override string? IdCore => this._name;
+
+    public override string? Name => this._name;
+
+    public override string? Description => "A fake agent with session support for testing";
+
+    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default) =>
+        new(new FakeSessionAgentSession());
+
+    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+        new(serializedState.Deserialize<FakeSessionAgentSession>(jsonSerializerOptions)!);
+
+    protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        if (session is not FakeSessionAgentSession fakeSession)
+        {
+            throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent.");
+        }
+
+        return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions));
+    }
+
+    protected override async Task<AgentResponse> RunCoreAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+    {
+        List<AgentResponseUpdate> updates = [];
+        await foreach (AgentResponseUpdate update in this.RunStreamingAsync(messages, session, options, cancellationToken).ConfigureAwait(false))
+        {
+            updates.Add(update);
+        }
+
+        return updates.ToAgentResponse();
+    }
+
+    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        // Track turn count in session state to enable persistence verification.
+        // If the session store works correctly, the turn count increments across requests.
+        int turnCount = 1;
+        if (session != null)
+        {
+            var counter = session.StateBag.GetValue<TurnCounter>("turnCounter");
+            turnCount = (counter?.Count ?? 0) + 1;
+            session.StateBag.SetValue("turnCounter", new TurnCounter { Count = turnCount });
+        }
+
+        string messageId = Guid.NewGuid().ToString("N");
+        string prefix = $"Turn {turnCount}: ";
+
+        foreach (string chunk in new[] { prefix, "Hello", " ", "from", " ", "session", " ", "agent", "!" })
+        {
+            yield return new AgentResponseUpdate
+            {
+                MessageId = messageId,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent(chunk)]
+            };
+
+            await Task.Yield();
+        }
+    }
+
+    internal sealed class TurnCounter
+    {
+        public int Count { get; set; }
+    }
+
+    private sealed class FakeSessionAgentSession : AgentSession
+    {
+        public FakeSessionAgentSession()
+        {
+        }
+
+        [JsonConstructor]
+        public FakeSessionAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+        {
+        }
+    }
+}

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs

@@ -0,0 +1,398 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
+using System.Net.Http;
+using System.Runtime.CompilerServices;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Threading;
+using System.Threading.Tasks;
+using AGUI.Abstractions;
+using AGUI.Client;
+using AGUI.Server;
+using FluentAssertions;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.TestHost;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests;
+
+public sealed class SharedStateTests : IAsyncDisposable
+{
+    private WebApplication? _app;
+    private HttpClient? _client;
+
+    [Fact]
+    public async Task StateSnapshot_IsSurfacedAsRawStateSnapshotEventAsync()
+    {
+        // Arrange
+        var fakeAgent = new FakeStateAgent();
+        await this.SetupTestServerAsync(fakeAgent);
+        AIAgent agent = this.CreateAgent();
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        // The AG-UI thread state travels on RunAgentInput.State, supplied to the stateless client via
+        // RawRepresentationFactory. The agent echoes it back as a STATE_SNAPSHOT event which the client
+        // surfaces as ChatResponseUpdate.RawRepresentation (issue #4869: no DataContent / ConversationId).
+        var initialState = JsonSerializer.SerializeToElement(new { counter = 42, status = "active" });
+        ChatMessage userMessage = new(ChatRole.User, "update state");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, StateRunOptions(initialState), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert - the state snapshot is surfaced as a StateSnapshotEvent raw representation.
+        updates.Should().NotBeEmpty();
+
+        StateSnapshotEvent? snapshot = FindStateSnapshot(updates);
+        snapshot.Should().NotBeNull("should receive a STATE_SNAPSHOT event");
+        snapshot!.Snapshot.GetProperty("counter").GetInt32().Should().Be(43, "state should be incremented");
+        snapshot.Snapshot.GetProperty("status").GetString().Should().Be("active");
+    }
+
+    [Fact]
+    public async Task StateSnapshot_UpdateHasAssistantRoleAndNoConversationIdAsync()
+    {
+        // Arrange
+        var fakeAgent = new FakeStateAgent();
+        await this.SetupTestServerAsync(fakeAgent);
+        AIAgent agent = this.CreateAgent();
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        var initialState = JsonSerializer.SerializeToElement(new { step = 1 });
+        ChatMessage userMessage = new(ChatRole.User, "process");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, StateRunOptions(initialState), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert - the state update carries the StateSnapshotEvent and the stateless client leaves
+        // ConversationId unset (state identity stays on the AG-UI wire events).
+        AgentResponseUpdate? stateUpdate = updates
+            .FirstOrDefault(u => u.AsChatResponseUpdate().RawRepresentation is StateSnapshotEvent);
+        stateUpdate.Should().NotBeNull();
+
+        ChatResponseUpdate chatUpdate = stateUpdate!.AsChatResponseUpdate();
+        chatUpdate.RawRepresentation.Should().BeOfType<StateSnapshotEvent>();
+        chatUpdate.ConversationId.Should().BeNull();
+        chatUpdate.Role.Should().Be(ChatRole.Assistant);
+    }
+
+    [Fact]
+    public async Task ComplexState_WithNestedObjectsAndArrays_RoundTripsCorrectlyAsync()
+    {
+        // Arrange
+        var fakeAgent = new FakeStateAgent();
+        await this.SetupTestServerAsync(fakeAgent);
+        AIAgent agent = this.CreateAgent();
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        JsonElement complexState = JsonDocument.Parse(
+            """{"sessionId":"test-123","nested":{"value":"test","count":10},"array":[1,2,3],"tags":["tag1","tag2"]}""").RootElement.Clone();
+        ChatMessage userMessage = new(ChatRole.User, "process complex state");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, StateRunOptions(complexState), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        StateSnapshotEvent? snapshot = FindStateSnapshot(updates);
+        snapshot.Should().NotBeNull();
+
+        JsonElement receivedState = snapshot!.Snapshot;
+        receivedState.GetProperty("sessionId").GetString().Should().Be("test-123");
+        receivedState.GetProperty("nested").GetProperty("count").GetInt32().Should().Be(10);
+        receivedState.GetProperty("array").GetArrayLength().Should().Be(3);
+        receivedState.GetProperty("tags").GetArrayLength().Should().Be(2);
+    }
+
+    [Fact]
+    public async Task StateSnapshot_CanBeUsedInSubsequentRequest_ForStateRoundTripAsync()
+    {
+        // Arrange
+        var fakeAgent = new FakeStateAgent();
+        await this.SetupTestServerAsync(fakeAgent);
+        AIAgent agent = this.CreateAgent();
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        var initialState = JsonSerializer.SerializeToElement(new { counter = 1, sessionId = "round-trip-test" });
+        ChatMessage userMessage = new(ChatRole.User, "increment");
+
+        List<AgentResponseUpdate> firstRoundUpdates = [];
+
+        // Act - First round
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, StateRunOptions(initialState), CancellationToken.None))
+        {
+            firstRoundUpdates.Add(update);
+        }
+
+        // Feed the returned state snapshot back into the second round.
+        StateSnapshotEvent? firstSnapshot = FindStateSnapshot(firstRoundUpdates);
+        firstSnapshot.Should().NotBeNull();
+        firstSnapshot!.Snapshot.GetProperty("counter").GetInt32().Should().Be(2);
+
+        ChatMessage secondUserMessage = new(ChatRole.User, "increment again");
+
+        List<AgentResponseUpdate> secondRoundUpdates = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([secondUserMessage], session, StateRunOptions(firstSnapshot.Snapshot), CancellationToken.None))
+        {
+            secondRoundUpdates.Add(update);
+        }
+
+        // Assert - Second round should have incremented counter again.
+        StateSnapshotEvent? secondSnapshot = FindStateSnapshot(secondRoundUpdates);
+        secondSnapshot.Should().NotBeNull();
+        secondSnapshot!.Snapshot.GetProperty("counter").GetInt32().Should().Be(3, "counter should be incremented twice: 1 -> 2 -> 3");
+    }
+
+    [Fact]
+    public async Task WithoutState_AgentBehavesNormally_NoStateSnapshotReturnedAsync()
+    {
+        // Arrange
+        var fakeAgent = new FakeStateAgent();
+        await this.SetupTestServerAsync(fakeAgent);
+        AIAgent agent = this.CreateAgent();
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        ChatMessage userMessage = new(ChatRole.User, "hello");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act - no RunAgentInput.State provided.
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        updates.Should().NotBeEmpty();
+        FindStateSnapshot(updates).Should().BeNull("should not return state snapshot when no state is provided");
+        updates.Should().Contain(u => u.Contents.Any(c => c is TextContent));
+    }
+
+    [Fact]
+    public async Task EmptyState_DoesNotTriggerStateHandlingAsync()
+    {
+        // Arrange
+        var fakeAgent = new FakeStateAgent();
+        await this.SetupTestServerAsync(fakeAgent);
+        AIAgent agent = this.CreateAgent();
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        var emptyState = JsonSerializer.SerializeToElement(new { });
+        ChatMessage userMessage = new(ChatRole.User, "hello");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, StateRunOptions(emptyState), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert - empty state {} should be treated as no state.
+        updates.Should().NotBeEmpty();
+        FindStateSnapshot(updates).Should().BeNull("empty state should be treated as no state");
+        updates.Should().Contain(u => u.Contents.Any(c => c is TextContent));
+    }
+
+    [Fact]
+    public async Task NonStreamingRunAsync_WithState_ReturnsTextResponseAsync()
+    {
+        // Arrange
+        var fakeAgent = new FakeStateAgent();
+        await this.SetupTestServerAsync(fakeAgent);
+        AIAgent agent = this.CreateAgent();
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        var initialState = JsonSerializer.SerializeToElement(new { counter = 5 });
+        ChatMessage userMessage = new(ChatRole.User, "process");
+
+        // Act - non-streaming run.
+        AgentResponse response = await agent.RunAsync([userMessage], session, StateRunOptions(initialState), CancellationToken.None);
+
+        // Assert - AG-UI state events are a streaming concern: the non-streaming aggregation drops the
+        // content-less STATE_SNAPSHOT update (Microsoft.Extensions.AI only materializes updates that carry
+        // content), so the non-streaming path surfaces the aggregated text response. The state round-trip
+        // itself is verified by the streaming tests above.
+        response.Should().NotBeNull();
+        response.Messages.Should().NotBeEmpty();
+        response.Text.Should().Contain("State processed");
+    }
+
+    private ChatClientAgent CreateAgent()
+    {
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        return chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+    }
+
+    private static ChatClientAgentRunOptions StateRunOptions(JsonElement state) =>
+        new()
+        {
+            ChatOptions = new ChatOptions
+            {
+                RawRepresentationFactory = _ => new RunAgentInput { State = state },
+            },
+        };
+
+    private static StateSnapshotEvent? FindStateSnapshot(IEnumerable<AgentResponseUpdate> updates) =>
+        updates
+            .Select(u => u.AsChatResponseUpdate().RawRepresentation as StateSnapshotEvent)
+            .FirstOrDefault(e => e is not null);
+
+    private async Task SetupTestServerAsync(FakeStateAgent fakeAgent)
+    {
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.Services.AddAGUIServer();
+        builder.WebHost.UseTestServer();
+
+        this._app = builder.Build();
+
+        this._app.MapAGUIServer("/agent", fakeAgent);
+
+        await this._app.StartAsync();
+
+        TestServer testServer = this._app.Services.GetRequiredService<IServer>() as TestServer
+            ?? throw new InvalidOperationException("TestServer not found");
+
+        this._client = testServer.CreateClient();
+        this._client.BaseAddress = new Uri("http://localhost/agent");
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        this._client?.Dispose();
+        if (this._app != null)
+        {
+            await this._app.DisposeAsync();
+        }
+    }
+}
+
+[SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated in tests")]
+internal sealed class FakeStateAgent : AIAgent
+{
+    public override string? Description => "Agent for state testing";
+
+    protected override Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        return this.RunCoreStreamingAsync(messages, session, options, cancellationToken).ToAgentResponseAsync(cancellationToken);
+    }
+
+    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        // Recover the originating AG-UI input from the request options (set by the hosting layer).
+        if (options is ChatClientAgentRunOptions { ChatOptions: { } chatOptions } &&
+            chatOptions.TryGetRunAgentInput(out RunAgentInput? agentInput) &&
+            agentInput.State is { ValueKind: JsonValueKind.Object } state &&
+            HasProperties(state))
+        {
+            Dictionary<string, object?> modifiedState = [];
+            foreach (JsonProperty prop in state.EnumerateObject())
+            {
+                if (prop.Name == "counter" && prop.Value.ValueKind == JsonValueKind.Number)
+                {
+                    modifiedState[prop.Name] = prop.Value.GetInt32() + 1;
+                }
+                else if (prop.Value.ValueKind == JsonValueKind.Number)
+                {
+                    modifiedState[prop.Name] = prop.Value.GetInt32();
+                }
+                else if (prop.Value.ValueKind == JsonValueKind.String)
+                {
+                    modifiedState[prop.Name] = prop.Value.GetString();
+                }
+                else if (prop.Value.ValueKind is JsonValueKind.Object or JsonValueKind.Array)
+                {
+                    modifiedState[prop.Name] = prop.Value;
+                }
+            }
+
+            // Emit the modified state as an AG-UI STATE_SNAPSHOT event. An AIAgent surfaces AG-UI raw
+            // events by wrapping a ChatResponseUpdate (whose RawRepresentation is the event) so the
+            // AgentResponseUpdate -> ChatResponseUpdate bridge forwards it to the server's event stream.
+            JsonElement snapshot = JsonSerializer.SerializeToElement(modifiedState);
+            yield return new AgentResponseUpdate
+            {
+                Role = ChatRole.Assistant,
+                RawRepresentation = new ChatResponseUpdate
+                {
+                    Role = ChatRole.Assistant,
+                    RawRepresentation = new StateSnapshotEvent { Snapshot = snapshot }
+                }
+            };
+        }
+
+        // Always return a text response.
+        string messageId = Guid.NewGuid().ToString("N");
+        yield return new AgentResponseUpdate
+        {
+            MessageId = messageId,
+            Role = ChatRole.Assistant,
+            Contents = [new TextContent("State processed")]
+        };
+
+        await Task.CompletedTask;
+    }
+
+    private static bool HasProperties(JsonElement element)
+    {
+        foreach (JsonProperty _ in element.EnumerateObject())
+        {
+            return true;
+        }
+
+        return false;
+    }
+
+    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default) =>
+        new(new FakeAgentSession());
+
+    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+        new(serializedState.Deserialize<FakeAgentSession>(jsonSerializerOptions)!);
+
+    protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        if (session is not FakeAgentSession fakeSession)
+        {
+            throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(FakeAgentSession)}' can be serialized by this agent.");
+        }
+
+        return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions));
+    }
+
+    private sealed class FakeAgentSession : AgentSession
+    {
+        public FakeAgentSession()
+        {
+        }
+
+        [JsonConstructor]
+        public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+        {
+        }
+    }
+
+    public override object? GetService(Type serviceType, object? serviceKey = null) => null;
+}

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ToolCallingTests.cs

@@ -0,0 +1,696 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net.Http;
+using System.Runtime.CompilerServices;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Threading;
+using System.Threading.Tasks;
+using AGUI.Client;
+using FluentAssertions;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.TestHost;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests;
+
+public sealed class ToolCallingTests : IAsyncDisposable
+{
+    private WebApplication? _app;
+    private HttpClient? _client;
+    private readonly ITestOutputHelper _output;
+
+    public ToolCallingTests(ITestOutputHelper output)
+    {
+        this._output = output;
+    }
+
+    [Fact]
+    public async Task ServerTriggersSingleFunctionCallAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        AIFunction serverTool = AIFunctionFactory.Create(() =>
+        {
+            callCount++;
+            return "Server function result";
+        }, "ServerFunction", "A function on the server");
+
+        await this.SetupTestServerAsync(serverTools: [serverTool]);
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: []);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Call the server function");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        callCount.Should().Be(1, "server function should be called once");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent), "should contain function call");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionResultContent), "should contain function result");
+
+        var functionCallUpdates = updates.Where(u => u.Contents.Any(c => c is FunctionCallContent)).ToList();
+        functionCallUpdates.Should().HaveCount(1);
+
+        var functionResultUpdates = updates.Where(u => u.Contents.Any(c => c is FunctionResultContent)).ToList();
+        functionResultUpdates.Should().HaveCount(1);
+
+        var resultContent = functionResultUpdates[0].Contents.OfType<FunctionResultContent>().First();
+        resultContent.Result.Should().NotBeNull();
+    }
+
+    [Fact]
+    public async Task ServerTriggersMultipleFunctionCallsAsync()
+    {
+        // Arrange
+        int getWeatherCallCount = 0;
+        int getTimeCallCount = 0;
+
+        AIFunction getWeatherTool = AIFunctionFactory.Create(() =>
+        {
+            getWeatherCallCount++;
+            return "Sunny, 75°F";
+        }, "GetWeather", "Gets the current weather");
+
+        AIFunction getTimeTool = AIFunctionFactory.Create(() =>
+        {
+            getTimeCallCount++;
+            return "3:45 PM";
+        }, "GetTime", "Gets the current time");
+
+        await this.SetupTestServerAsync(serverTools: [getWeatherTool, getTimeTool]);
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: []);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "What's the weather and time?");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        getWeatherCallCount.Should().Be(1, "GetWeather should be called once");
+        getTimeCallCount.Should().Be(1, "GetTime should be called once");
+
+        var functionCallUpdates = updates.Where(u => u.Contents.Any(c => c is FunctionCallContent)).ToList();
+        functionCallUpdates.Should().NotBeEmpty("should contain function calls");
+
+        var functionCalls = updates.SelectMany(u => u.Contents.OfType<FunctionCallContent>()).ToList();
+        functionCalls.Should().HaveCount(2, "should have 2 function calls");
+        functionCalls.Should().Contain(fc => fc.Name == "GetWeather");
+        functionCalls.Should().Contain(fc => fc.Name == "GetTime");
+
+        var functionResults = updates.SelectMany(u => u.Contents.OfType<FunctionResultContent>()).ToList();
+        functionResults.Should().HaveCount(2, "should have 2 function results");
+    }
+
+    [Fact]
+    public async Task ClientTriggersSingleFunctionCallAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        AIFunction clientTool = AIFunctionFactory.Create(() =>
+        {
+            callCount++;
+            return "Client function result";
+        }, "ClientFunction", "A function on the client");
+
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: [clientTool]);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Call the client function");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        callCount.Should().Be(1, "client function should be called once");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent), "should contain function call");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionResultContent), "should contain function result");
+
+        var functionCallUpdates = updates.Where(u => u.Contents.Any(c => c is FunctionCallContent)).ToList();
+        functionCallUpdates.Should().HaveCount(1);
+
+        var functionResultUpdates = updates.Where(u => u.Contents.Any(c => c is FunctionResultContent)).ToList();
+        functionResultUpdates.Should().HaveCount(1);
+
+        var resultContent = functionResultUpdates[0].Contents.OfType<FunctionResultContent>().First();
+        resultContent.Result.Should().NotBeNull();
+    }
+
+    [Fact]
+    public async Task ClientTriggersMultipleFunctionCallsAsync()
+    {
+        // Arrange
+        int calculateCallCount = 0;
+        int formatCallCount = 0;
+
+        AIFunction calculateTool = AIFunctionFactory.Create((int a, int b) =>
+        {
+            calculateCallCount++;
+            return a + b;
+        }, "Calculate", "Calculates sum of two numbers");
+
+        AIFunction formatTool = AIFunctionFactory.Create((string text) =>
+        {
+            formatCallCount++;
+            return text.ToUpperInvariant();
+        }, "FormatText", "Formats text to uppercase");
+
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: [calculateTool, formatTool]);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Calculate 5 + 3 and format 'hello'");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        calculateCallCount.Should().Be(1, "Calculate should be called once");
+        formatCallCount.Should().Be(1, "FormatText should be called once");
+
+        var functionCallUpdates = updates.Where(u => u.Contents.Any(c => c is FunctionCallContent)).ToList();
+        functionCallUpdates.Should().NotBeEmpty("should contain function calls");
+
+        var functionCalls = updates.SelectMany(u => u.Contents.OfType<FunctionCallContent>()).ToList();
+        functionCalls.Should().HaveCount(2, "should have 2 function calls");
+        functionCalls.Should().Contain(fc => fc.Name == "Calculate");
+        functionCalls.Should().Contain(fc => fc.Name == "FormatText");
+
+        var functionResults = updates.SelectMany(u => u.Contents.OfType<FunctionResultContent>()).ToList();
+        functionResults.Should().HaveCount(2, "should have 2 function results");
+    }
+
+    [Fact]
+    public async Task ServerAndClientTriggerFunctionCallsSimultaneouslyAsync()
+    {
+        // Arrange
+        int serverCallCount = 0;
+        int clientCallCount = 0;
+
+        AIFunction serverTool = AIFunctionFactory.Create(() =>
+        {
+            System.Diagnostics.Debug.Assert(true, "Server function is being called!");
+            serverCallCount++;
+            return "Server data";
+        }, "GetServerData", "Gets data from the server");
+
+        AIFunction clientTool = AIFunctionFactory.Create(() =>
+        {
+            System.Diagnostics.Debug.Assert(true, "Client function is being called!");
+            clientCallCount++;
+            return "Client data";
+        }, "GetClientData", "Gets data from the client");
+
+        await this.SetupTestServerAsync(serverTools: [serverTool]);
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: [clientTool]);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Get both server and client data");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+            this._output.WriteLine($"Update: {update.Contents.Count} contents");
+            foreach (var content in update.Contents)
+            {
+                this._output.WriteLine($"  Content: {content.GetType().Name}");
+                if (content is FunctionCallContent fc)
+                {
+                    this._output.WriteLine($"    FunctionCall: {fc.Name}");
+                }
+                if (content is FunctionResultContent fr)
+                {
+                    this._output.WriteLine($"    FunctionResult: {fr.CallId} - {fr.Result}");
+                }
+            }
+        }
+
+        // Assert
+        this._output.WriteLine($"serverCallCount={serverCallCount}, clientCallCount={clientCallCount}");
+
+        // Verify both the server and client tools executed and both results round-tripped through
+        // the streaming pipeline. This is now correct behavior thanks to
+        // ConfigureForMixedInvocation in the AGUI.Hosting.AspNetCore package.
+
+        serverCallCount.Should().Be(1, "server function should execute on server");
+        clientCallCount.Should().Be(1, "client function should execute on client");
+
+        var functionCallUpdates = updates.Where(u => u.Contents.Any(c => c is FunctionCallContent)).ToList();
+        functionCallUpdates.Should().NotBeEmpty("should contain function calls");
+
+        var functionCalls = updates.SelectMany(u => u.Contents.OfType<FunctionCallContent>()).ToList();
+        functionCalls.Should().HaveCount(2, "should have 2 function calls");
+        functionCalls.Should().Contain(fc => fc.Name == "GetServerData");
+        functionCalls.Should().Contain(fc => fc.Name == "GetClientData");
+
+        var functionResults = updates.SelectMany(u => u.Contents.OfType<FunctionResultContent>()).ToList();
+        functionResults.Should().HaveCount(2, "both server and client function results should be present");
+
+        var serverResult = functionResults.FirstOrDefault(fr =>
+            functionCalls.Any(fc => fc.Name == "GetServerData" && fc.CallId == fr.CallId));
+        serverResult.Should().NotBeNull("server function call should have a result");
+        serverResult!.Result?.ToString().Should().Contain("Server data");
+
+        var clientResult = functionResults.FirstOrDefault(fr =>
+            functionCalls.Any(fc => fc.Name == "GetClientData" && fc.CallId == fr.CallId));
+        clientResult.Should().NotBeNull("client function call should have a result");
+        clientResult!.Result?.ToString().Should().Contain("Client data");
+    }
+
+    [Fact]
+    public async Task FunctionCallsPreserveCallIdAndNameAsync()
+    {
+        // Arrange
+        AIFunction testTool = AIFunctionFactory.Create(() => "Test result", "TestFunction", "A test function");
+
+        await this.SetupTestServerAsync(serverTools: [testTool]);
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: []);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Call the test function");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        var functionCallContent = updates.SelectMany(u => u.Contents.OfType<FunctionCallContent>()).FirstOrDefault();
+        functionCallContent.Should().NotBeNull();
+        functionCallContent!.CallId.Should().NotBeNullOrEmpty();
+        functionCallContent.Name.Should().Be("TestFunction");
+
+        var functionResultContent = updates.SelectMany(u => u.Contents.OfType<FunctionResultContent>()).FirstOrDefault();
+        functionResultContent.Should().NotBeNull();
+        functionResultContent!.CallId.Should().Be(functionCallContent.CallId, "result should have same call ID as the call");
+    }
+
+    [Fact]
+    public async Task ParallelFunctionCallsFromServerAreHandledCorrectlyAsync()
+    {
+        // Arrange
+        int func1CallCount = 0;
+        int func2CallCount = 0;
+
+        AIFunction func1 = AIFunctionFactory.Create(() =>
+        {
+            func1CallCount++;
+            return "Result 1";
+        }, "Function1", "First function");
+
+        AIFunction func2 = AIFunctionFactory.Create(() =>
+        {
+            func2CallCount++;
+            return "Result 2";
+        }, "Function2", "Second function");
+
+        await this.SetupTestServerAsync(serverTools: [func1, func2], triggerParallelCalls: true);
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: []);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Call both functions in parallel");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        func1CallCount.Should().Be(1, "Function1 should be called once");
+        func2CallCount.Should().Be(1, "Function2 should be called once");
+
+        var functionCalls = updates.SelectMany(u => u.Contents.OfType<FunctionCallContent>()).ToList();
+        functionCalls.Should().HaveCount(2);
+        functionCalls.Select(fc => fc.Name).Should().Contain(s_expectedFunctionNames);
+
+        var functionResults = updates.SelectMany(u => u.Contents.OfType<FunctionResultContent>()).ToList();
+        functionResults.Should().HaveCount(2);
+
+        // Each result should match its corresponding call ID
+        foreach (var call in functionCalls)
+        {
+            functionResults.Should().Contain(r => r.CallId == call.CallId);
+        }
+    }
+
+    private static readonly string[] s_expectedFunctionNames = ["Function1", "Function2"];
+
+    [Fact]
+    public async Task AGUIChatClientCombinesCustomJsonSerializerOptionsAsync()
+    {
+        // This test verifies that custom JSON contexts work correctly with AGUIChatClient by testing
+        // that a client-defined type can be serialized successfully using the combined options
+
+        // Arrange
+        await this.SetupTestServerAsync();
+
+        // Client uses custom JSON context
+        var clientJsonOptions = new JsonSerializerOptions();
+        clientJsonOptions.TypeInfoResolverChain.Add(ClientJsonContext.Default);
+
+        _ = new AGUIChatClient(new(this._client!, "") { JsonSerializerOptions = clientJsonOptions });
+
+        // Act - Verify that both AG-UI types and custom types can be serialized
+        // The AGUIChatClient should have combined AGUIJsonSerializerContext with ClientJsonContext
+
+        // Try to serialize a custom type using the ClientJsonContext
+        var testResponse = new ClientForecastResponse(75, 60, "Rainy");
+        var json = JsonSerializer.Serialize(testResponse, ClientJsonContext.Default.ClientForecastResponse);
+
+        // Assert
+        var jsonElement = JsonElement.Parse(json);
+        jsonElement.GetProperty("MaxTemp").GetInt32().Should().Be(75);
+        jsonElement.GetProperty("MinTemp").GetInt32().Should().Be(60);
+        jsonElement.GetProperty("Outlook").GetString().Should().Be("Rainy");
+
+        this._output.WriteLine("Successfully serialized custom type: " + json);
+
+        // The actual integration is tested by the ClientToolCallWithCustomArgumentsAsync test
+        // which verifies that AG-UI protocol works end-to-end with custom types
+    }
+
+    [Fact]
+    public async Task ServerToolCallWithCustomArgumentsAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        AIFunction serverTool = AIFunctionFactory.Create(
+            (ServerForecastRequest request) =>
+            {
+                callCount++;
+                return new ServerForecastResponse(
+                    Temperature: 72,
+                    Condition: request.Location == "Seattle" ? "Rainy" : "Sunny",
+                    Humidity: 65);
+            },
+            "GetServerForecast",
+            "Gets the weather forecast from server",
+            ServerJsonContext.Default.Options);
+
+        await this.SetupTestServerAsync(serverTools: [serverTool], jsonSerializerOptions: ServerJsonContext.Default.Options);
+        var chatClient = new AGUIChatClient(new(this._client!, "") { JsonSerializerOptions = ServerJsonContext.Default.Options });
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: []);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Get server forecast for Seattle for 5 days");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        callCount.Should().Be(1, "server function with custom arguments should be called once");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent), "should contain function call");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionResultContent), "should contain function result");
+
+        var functionCallContent = updates.SelectMany(u => u.Contents.OfType<FunctionCallContent>()).FirstOrDefault();
+        functionCallContent.Should().NotBeNull();
+        functionCallContent!.Name.Should().Be("GetServerForecast");
+
+        var functionResultContent = updates.SelectMany(u => u.Contents.OfType<FunctionResultContent>()).FirstOrDefault();
+        functionResultContent.Should().NotBeNull();
+        functionResultContent!.Result.Should().NotBeNull();
+    }
+
+    [Fact]
+    public async Task ClientToolCallWithCustomArgumentsAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        AIFunction clientTool = AIFunctionFactory.Create(
+            (ClientForecastRequest request) =>
+            {
+                callCount++;
+                return new ClientForecastResponse(
+                    MaxTemp: request.City == "Portland" ? 68 : 75,
+                    MinTemp: 55,
+                    Outlook: "Partly Cloudy");
+            },
+            "GetClientForecast",
+            "Gets the weather forecast from client",
+            ClientJsonContext.Default.Options);
+
+        await this.SetupTestServerAsync();
+        var chatClient = new AGUIChatClient(new(this._client!, "") { JsonSerializerOptions = ClientJsonContext.Default.Options });
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Test assistant", tools: [clientTool]);
+        AgentSession session = await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "Get client forecast for Portland with hourly data");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        callCount.Should().Be(1, "client function with custom arguments should be called once");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent), "should contain function call");
+        updates.Should().Contain(u => u.Contents.Any(c => c is FunctionResultContent), "should contain function result");
+
+        var functionCallContent = updates.SelectMany(u => u.Contents.OfType<FunctionCallContent>()).FirstOrDefault();
+        functionCallContent.Should().NotBeNull();
+        functionCallContent!.Name.Should().Be("GetClientForecast");
+
+        var functionResultContent = updates.SelectMany(u => u.Contents.OfType<FunctionResultContent>()).FirstOrDefault();
+        functionResultContent.Should().NotBeNull();
+        functionResultContent!.Result.Should().NotBeNull();
+    }
+
+    private async Task SetupTestServerAsync(
+        IList<AITool>? serverTools = null,
+        bool triggerParallelCalls = false,
+        JsonSerializerOptions? jsonSerializerOptions = null)
+    {
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.Services.AddAGUIServer();
+        builder.WebHost.UseTestServer();
+
+        // Configure HTTP JSON options if custom serializer options provided
+        if (jsonSerializerOptions?.TypeInfoResolver != null)
+        {
+            builder.Services.ConfigureHttpJsonOptions(options =>
+                options.SerializerOptions.TypeInfoResolverChain.Add(jsonSerializerOptions.TypeInfoResolver));
+        }
+
+        this._app = builder.Build();
+        // FakeChatClient will receive options.Tools containing both server and client tools (merged by framework)
+        var fakeChatClient = new FakeToolCallingChatClient(triggerParallelCalls, this._output, jsonSerializerOptions: jsonSerializerOptions);
+        AIAgent baseAgent = fakeChatClient.AsAIAgent(instructions: null, name: "base-agent", description: "A base agent for tool testing", tools: serverTools ?? []);
+        this._app.MapAGUIServer("/agent", baseAgent);
+
+        await this._app.StartAsync();
+
+        TestServer testServer = this._app.Services.GetRequiredService<IServer>() as TestServer
+            ?? throw new InvalidOperationException("TestServer not found");
+
+        this._client = testServer.CreateClient();
+        this._client.BaseAddress = new Uri("http://localhost/agent");
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        this._client?.Dispose();
+        if (this._app != null)
+        {
+            await this._app.DisposeAsync();
+        }
+    }
+}
+
+internal sealed class FakeToolCallingChatClient : IChatClient
+{
+    private readonly bool _triggerParallelCalls;
+    private readonly ITestOutputHelper? _output;
+    public FakeToolCallingChatClient(bool triggerParallelCalls = false, ITestOutputHelper? output = null, JsonSerializerOptions? jsonSerializerOptions = null)
+    {
+        this._triggerParallelCalls = triggerParallelCalls;
+        this._output = output;
+    }
+
+    public ChatClientMetadata Metadata => new("fake-tool-calling-chat-client");
+
+    public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
+        IEnumerable<ChatMessage> messages,
+        ChatOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        string messageId = Guid.NewGuid().ToString("N");
+
+        var messageList = messages.ToList();
+        this._output?.WriteLine($"[FakeChatClient] Received {messageList.Count} messages");
+
+        // Check if there are function results in the messages - if so, we've already done the function call loop
+        var hasFunctionResults = messageList.Any(m => m.Contents.Any(c => c is FunctionResultContent));
+
+        if (hasFunctionResults)
+        {
+            this._output?.WriteLine("[FakeChatClient] Function results present, returning final response");
+            // Function results are present, return a final response
+            yield return new ChatResponseUpdate
+            {
+                MessageId = messageId,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent("Function calls completed successfully")]
+            };
+            yield break;
+        }
+
+        // options?.Tools contains all tools (server + client merged by framework)
+        var allTools = (options?.Tools ?? []).ToList();
+        this._output?.WriteLine($"[FakeChatClient] Received {allTools.Count} tools to advertise");
+
+        if (allTools.Count == 0)
+        {
+            // No tools available, just return a simple message
+            yield return new ChatResponseUpdate
+            {
+                MessageId = messageId,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent("No tools available")]
+            };
+            yield break;
+        }
+
+        // Determine which tools to call based on the scenario
+        var toolsToCall = new List<AITool>();
+
+        // Check message content to determine what to call
+        var lastUserMessage = messageList.LastOrDefault(m => m.Role == ChatRole.User)?.Text ?? "";
+
+        if (this._triggerParallelCalls)
+        {
+            // Call all available tools in parallel
+            toolsToCall.AddRange(allTools);
+        }
+        else if (lastUserMessage.Contains("both", StringComparison.OrdinalIgnoreCase) ||
+                 lastUserMessage.Contains("all", StringComparison.OrdinalIgnoreCase))
+        {
+            // Call all available tools
+            toolsToCall.AddRange(allTools);
+        }
+        else
+        {
+            // Default: call all available tools
+            // The fake LLM doesn't distinguish between server and client tools - it just requests them all
+            // The FunctionInvokingChatClient layers will handle executing what they can
+            toolsToCall.AddRange(allTools);
+        }
+
+        // Assert: Should have tools to call
+        System.Diagnostics.Debug.Assert(toolsToCall.Count > 0, "Should have at least one tool to call");
+
+        // Generate function calls
+        // Server's FunctionInvokingChatClient will execute server tools
+        // Client tool calls will be sent back to client, and client's FunctionInvokingChatClient will execute them
+        this._output?.WriteLine($"[FakeChatClient] Generating {toolsToCall.Count} function calls");
+        foreach (var tool in toolsToCall)
+        {
+            string callId = $"call_{Guid.NewGuid():N}";
+            var functionName = tool.Name ?? "UnknownFunction";
+            this._output?.WriteLine($"[FakeChatClient]   Calling: {functionName} (type: {tool.GetType().Name})");
+
+            // Generate sample arguments based on the function signature
+            var arguments = GenerateArgumentsForTool(functionName);
+
+            yield return new ChatResponseUpdate
+            {
+                MessageId = messageId,
+                Role = ChatRole.Assistant,
+                Contents = [new FunctionCallContent(callId, functionName, arguments)]
+            };
+
+            await Task.Yield();
+        }
+    }
+
+    private static Dictionary<string, object?> GenerateArgumentsForTool(string functionName)
+    {
+        // Generate sample arguments based on the function name
+        return functionName switch
+        {
+            "GetWeather" => new Dictionary<string, object?> { ["location"] = "Seattle" },
+            "GetTime" => [], // No parameters
+            "Calculate" => new Dictionary<string, object?> { ["a"] = 5, ["b"] = 3 },
+            "FormatText" => new Dictionary<string, object?> { ["text"] = "hello" },
+            "GetServerData" => [], // No parameters
+            "GetClientData" => [], // No parameters
+            // For custom types, the parameter name is "request" and the value is an instance of the request type
+            "GetServerForecast" => new Dictionary<string, object?> { ["request"] = new ServerForecastRequest("Seattle", 5) },
+            "GetClientForecast" => new Dictionary<string, object?> { ["request"] = new ClientForecastRequest("Portland", true) },
+            _ => [] // Default: no parameters
+        };
+    }
+
+    public Task<ChatResponse> GetResponseAsync(
+        IEnumerable<ChatMessage> messages,
+        ChatOptions? options = null,
+        CancellationToken cancellationToken = default)
+    {
+        throw new NotImplementedException();
+    }
+
+    public void Dispose()
+    {
+    }
+
+    public object? GetService(Type serviceType, object? serviceKey = null) => null;
+}
+
+// Custom types and serialization contexts for testing cross-boundary serialization
+public record ServerForecastRequest(string Location, int Days);
+public record ServerForecastResponse(int Temperature, string Condition, int Humidity);
+
+public record ClientForecastRequest(string City, bool IncludeHourly);
+public record ClientForecastResponse(int MaxTemp, int MinTemp, string Outlook);
+
+[JsonSourceGenerationOptions(WriteIndented = false)]
+[JsonSerializable(typeof(ServerForecastRequest))]
+[JsonSerializable(typeof(ServerForecastResponse))]
+internal sealed partial class ServerJsonContext : JsonSerializerContext;
+
+[JsonSourceGenerationOptions(WriteIndented = false)]
+[JsonSerializable(typeof(ClientForecastRequest))]
+[JsonSerializable(typeof(ClientForecastResponse))]
+internal sealed partial class ClientJsonContext : JsonSerializerContext;