agent-framework/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs

// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Net.ServerSentEvents;
using System.Text;
using System.Text.Encodings.Web;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using A2A;
using Microsoft.Extensions.AI;

namespace Microsoft.Agents.AI.A2A.UnitTests;

/// <summary>
/// Unit tests for the <see cref="A2AAgent"/> class.
/// </summary>
public sealed class A2AAgentTests : IDisposable
{
    private readonly HttpClient _httpClient;
    private readonly A2AClientHttpMessageHandlerStub _handler;
    private readonly A2AClient _a2aClient;
    private readonly A2AAgent _agent;

    public A2AAgentTests()
    {
        this._handler = new A2AClientHttpMessageHandlerStub();
        this._httpClient = new HttpClient(this._handler, false);
        this._a2aClient = new A2AClient(new Uri("http://test-endpoint"), this._httpClient);
        this._agent = new A2AAgent(this._a2aClient);
    }

    [Fact]
    public void Constructor_WithAllParameters_InitializesPropertiesCorrectly()
    {
        // Arrange
        const string TestId = "test-id";
        const string TestName = "test-name";
        const string TestDescription = "test-description";

        // Act
        var agent = new A2AAgent(this._a2aClient, TestId, TestName, TestDescription);

        // Assert
        Assert.Equal(TestId, agent.Id);
        Assert.Equal(TestName, agent.Name);
        Assert.Equal(TestDescription, agent.Description);
    }

    [Fact]
    public void Constructor_WithNullA2AClient_ThrowsArgumentNullException() =>
        // Act & Assert
        Assert.Throws<ArgumentNullException>(() => new A2AAgent(null!));

    [Fact]
    public void Constructor_WithDefaultParameters_UsesBaseProperties()
    {
        // Act
        var agent = new A2AAgent(this._a2aClient);

        // Assert
        Assert.NotNull(agent.Id);
        Assert.NotEmpty(agent.Id);
        Assert.Null(agent.Name);
        Assert.Null(agent.Description);
    }

    [Fact]
    public async Task RunAsync_AllowsNonUserRoleMessagesAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.System, "I am a system message"),
            new(ChatRole.Assistant, "I am an assistant message"),
            new(ChatRole.User, "Valid user message")
        };

        // Act & Assert
        await this._agent.RunAsync(inputMessages);
    }

    [Fact]
    public async Task RunAsync_WithValidUserMessage_RunsSuccessfullyAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts =
            [
                new TextPart { Text = "Hello! How can I help you today?" }
            ]
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Hello, world!")
        };

        // Act
        var result = await this._agent.RunAsync(inputMessages);

        // Assert input message sent to A2AClient
        var inputMessage = this._handler.CapturedMessageSendParams?.Message;
        Assert.NotNull(inputMessage);
        Assert.Single(inputMessage.Parts);
        Assert.Equal(MessageRole.User, inputMessage.Role);
        Assert.Equal("Hello, world!", ((TextPart)inputMessage.Parts[0]).Text);

        // Assert response from A2AClient is converted correctly
        Assert.NotNull(result);
        Assert.Equal(this._agent.Id, result.AgentId);
        Assert.Equal("response-123", result.ResponseId);

        Assert.NotNull(result.RawRepresentation);
        Assert.IsType<AgentMessage>(result.RawRepresentation);
        Assert.Equal("response-123", ((AgentMessage)result.RawRepresentation).MessageId);

        Assert.Single(result.Messages);
        Assert.Equal(ChatRole.Assistant, result.Messages[0].Role);
        Assert.Equal("Hello! How can I help you today?", result.Messages[0].Text);
        Assert.Equal(ChatFinishReason.Stop, result.FinishReason);
    }

    [Fact]
    public async Task RunAsync_WithNewSession_UpdatesSessionConversationIdAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts =
            [
                new TextPart { Text = "Response" }
            ],
            ContextId = "new-context-id"
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        await this._agent.RunAsync(inputMessages, session);

        // Assert
        Assert.IsType<A2AAgentSession>(session);
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal("new-context-id", a2aSession.ContextId);
    }

    [Fact]
    public async Task RunAsync_WithExistingSession_SetConversationIdToMessageAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var session = await this._agent.CreateSessionAsync();
        var a2aSession = (A2AAgentSession)session;
        a2aSession.ContextId = "existing-context-id";

        // Act
        await this._agent.RunAsync(inputMessages, session);

        // Assert
        var message = this._handler.CapturedMessageSendParams?.Message;
        Assert.NotNull(message);
        Assert.Equal("existing-context-id", message.ContextId);
    }

    [Fact]
    public async Task RunAsync_WithSessionHavingDifferentContextId_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        this._handler.ResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts =
            [
                new TextPart { Text = "Response" }
            ],
            ContextId = "different-context"
        };

        var session = await this._agent.CreateSessionAsync();
        var a2aSession = (A2AAgentSession)session;
        a2aSession.ContextId = "existing-context-id";

        // Act & Assert
        await Assert.ThrowsAsync<InvalidOperationException>(() => this._agent.RunAsync(inputMessages, session));
    }

    [Fact]
    public async Task RunStreamingAsync_WithValidUserMessage_YieldsAgentResponseUpdatesAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Hello, streaming!")
        };

        this._handler.StreamingResponseToReturn = new AgentMessage()
        {
            MessageId = "stream-1",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Hello" }],
            ContextId = "stream-context"
        };

        // Act
        var updates = new List<AgentResponseUpdate>();
        await foreach (var update in this._agent.RunStreamingAsync(inputMessages))
        {
            updates.Add(update);
        }

        // Assert
        Assert.Single(updates);

        // Assert input message sent to A2AClient
        var inputMessage = this._handler.CapturedMessageSendParams?.Message;
        Assert.NotNull(inputMessage);
        Assert.Single(inputMessage.Parts);
        Assert.Equal(MessageRole.User, inputMessage.Role);
        Assert.Equal("Hello, streaming!", ((TextPart)inputMessage.Parts[0]).Text);

        // Assert response from A2AClient is converted correctly
        Assert.Equal(ChatRole.Assistant, updates[0].Role);
        Assert.Equal("Hello", updates[0].Text);
        Assert.Equal("stream-1", updates[0].MessageId);
        Assert.Equal(this._agent.Id, updates[0].AgentId);
        Assert.Equal("stream-1", updates[0].ResponseId);
        Assert.Equal(ChatFinishReason.Stop, updates[0].FinishReason);
        Assert.IsType<AgentMessage>(updates[0].RawRepresentation);
        Assert.Equal("stream-1", ((AgentMessage)updates[0].RawRepresentation!).MessageId);
    }

    [Fact]
    public async Task RunStreamingAsync_WithSession_UpdatesSessionConversationIdAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test streaming")
        };

        this._handler.StreamingResponseToReturn = new AgentMessage()
        {
            MessageId = "stream-1",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response" }],
            ContextId = "new-stream-context"
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, session))
        {
            // Just iterate through to trigger the logic
        }

        // Assert
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal("new-stream-context", a2aSession.ContextId);
    }

    [Fact]
    public async Task RunStreamingAsync_WithExistingSession_SetConversationIdToMessageAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test streaming")
        };

        this._handler.StreamingResponseToReturn = new AgentMessage();

        var session = await this._agent.CreateSessionAsync();
        var a2aSession = (A2AAgentSession)session;
        a2aSession.ContextId = "existing-context-id";

        // Act
        await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, session))
        {
            // Just iterate through to trigger the logic
        }

        // Assert
        var message = this._handler.CapturedMessageSendParams?.Message;
        Assert.NotNull(message);
        Assert.Equal("existing-context-id", message.ContextId);
    }

    [Fact]
    public async Task RunStreamingAsync_WithSessionHavingDifferentContextId_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        var session = await this._agent.CreateSessionAsync();
        var a2aSession = (A2AAgentSession)session;
        a2aSession.ContextId = "existing-context-id";

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test streaming")
        };

        this._handler.StreamingResponseToReturn = new AgentMessage()
        {
            MessageId = "stream-1",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response" }],
            ContextId = "different-context"
        };

        // Act
        await Assert.ThrowsAsync<InvalidOperationException>(async () =>
        {
            await foreach (var update in this._agent.RunStreamingAsync(inputMessages, session))
            {
            }
        });
    }

    [Fact]
    public async Task RunStreamingAsync_AllowsNonUserRoleMessagesAsync()
    {
        // Arrange
        this._handler.StreamingResponseToReturn = new AgentMessage()
        {
            MessageId = "stream-1",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response" }],
            ContextId = "new-stream-context"
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.System, "I am a system message"),
            new(ChatRole.Assistant, "I am an assistant message"),
            new(ChatRole.User, "Valid user message")
        };

        // Act & Assert
        await foreach (var _ in this._agent.RunStreamingAsync(inputMessages))
        {
            // Just iterate through to trigger the logic
        }
    }

    [Fact]
    public async Task RunAsync_WithHostedFileContent_ConvertsToFilePartAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User,
            [
                new TextContent("Check this file:"),
                new UriContent("https://example.com/file.pdf", "application/pdf")
            ])
        };

        // Act
        await this._agent.RunAsync(inputMessages);

        // Assert
        var message = this._handler.CapturedMessageSendParams?.Message;
        Assert.NotNull(message);
        Assert.Equal(2, message.Parts.Count);
        Assert.IsType<TextPart>(message.Parts[0]);
        Assert.Equal("Check this file:", ((TextPart)message.Parts[0]).Text);
        Assert.IsType<FilePart>(message.Parts[1]);
        Assert.Equal("https://example.com/file.pdf", ((FilePart)message.Parts[1]).File.Uri?.ToString());
    }

    [Fact]
    public async Task RunAsync_WithContinuationTokenAndMessages_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-123") };

        // Act & Assert
        await Assert.ThrowsAsync<InvalidOperationException>(() => this._agent.RunAsync(inputMessages, null, options));
    }

    [Fact]
    public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentTask
        {
            Id = "task-123",
            ContextId = "context-123"
        };

        var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-123") };

        // Act
        await this._agent.RunAsync([], options: options);

        // Assert
        Assert.Equal("tasks/get", this._handler.CapturedJsonRpcRequest?.Method);
        Assert.Equal("task-123", this._handler.CapturedTaskIdParams?.Id);
    }

    [Fact]
    public async Task RunAsync_WithTaskInSessionAndMessage_AddTaskAsReferencesToMessageAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response to task" }]
        };

        var session = (A2AAgentSession)await this._agent.CreateSessionAsync();
        session.TaskId = "task-123";

        var inputMessage = new ChatMessage(ChatRole.User, "Please make the background transparent");

        // Act
        await this._agent.RunAsync(inputMessage, session);

        // Assert
        var message = this._handler.CapturedMessageSendParams?.Message;
        Assert.Null(message?.TaskId);
        Assert.NotNull(message?.ReferenceTaskIds);
        Assert.Contains("task-123", message.ReferenceTaskIds);
    }

    [Fact]
    public async Task RunAsync_WithAgentTask_UpdatesSessionTaskIdAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentTask
        {
            Id = "task-456",
            ContextId = "context-789",
            Status = new() { State = TaskState.Submitted }
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        await this._agent.RunAsync("Start a task", session);

        // Assert
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal("task-456", a2aSession.TaskId);
    }

    [Fact]
    public async Task RunAsync_WithAgentTaskResponse_ReturnsTaskResponseCorrectlyAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentTask
        {
            Id = "task-789",
            ContextId = "context-456",
            Status = new() { State = TaskState.Submitted },
            Metadata = new Dictionary<string, JsonElement>
            {
                { "key1", JsonSerializer.SerializeToElement("value1") },
                { "count", JsonSerializer.SerializeToElement(42) }
            }
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        var result = await this._agent.RunAsync("Start a long-running task", session);

        // Assert - verify task is converted correctly
        Assert.NotNull(result);
        Assert.Equal(this._agent.Id, result.AgentId);
        Assert.Equal("task-789", result.ResponseId);
        Assert.Null(result.FinishReason);
        Assert.IsType<AgentTask>(result.RawRepresentation);
        Assert.Equal("task-789", ((AgentTask)result.RawRepresentation).Id);

        // Assert - verify continuation token is set for submitted task
        Assert.NotNull(result.ContinuationToken);
        Assert.IsType<A2AContinuationToken>(result.ContinuationToken);
        Assert.Equal("task-789", ((A2AContinuationToken)result.ContinuationToken).TaskId);

        // Assert - verify session is updated with context and task IDs
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal("context-456", a2aSession.ContextId);
        Assert.Equal("task-789", a2aSession.TaskId);

        // Assert - verify metadata is preserved
        Assert.NotNull(result.AdditionalProperties);
        Assert.NotNull(result.AdditionalProperties["key1"]);
        Assert.Equal("value1", ((JsonElement)result.AdditionalProperties["key1"]!).GetString());
        Assert.NotNull(result.AdditionalProperties["count"]);
        Assert.Equal(42, ((JsonElement)result.AdditionalProperties["count"]!).GetInt32());
    }

    [Theory]
    [InlineData(TaskState.Submitted)]
    [InlineData(TaskState.Working)]
    [InlineData(TaskState.Completed)]
    [InlineData(TaskState.Failed)]
    [InlineData(TaskState.Canceled)]
    public async Task RunAsync_WithVariousTaskStates_ReturnsCorrectTokenAsync(TaskState taskState)
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentTask
        {
            Id = "task-123",
            ContextId = "context-123",
            Status = new() { State = taskState }
        };

        // Act
        var result = await this._agent.RunAsync("Test message");

        // Assert
        if (taskState is TaskState.Submitted or TaskState.Working)
        {
            Assert.NotNull(result.ContinuationToken);
        }
        else
        {
            Assert.Null(result.ContinuationToken);
        }

        if (taskState is TaskState.Completed)
        {
            Assert.Equal(ChatFinishReason.Stop, result.FinishReason);
        }
        else
        {
            Assert.Null(result.FinishReason);
        }
    }

    [Fact]
    public async Task RunStreamingAsync_WithContinuationTokenAndMessages_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-123") };

        // Act & Assert
        await Assert.ThrowsAsync<InvalidOperationException>(async () =>
        {
            await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, null, options))
            {
                // Just iterate through to trigger the exception
            }
        });
    }

    [Fact]
    public async Task RunStreamingAsync_WithTaskInSessionAndMessage_AddTaskAsReferencesToMessageAsync()
    {
        // Arrange
        this._handler.StreamingResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response to task" }]
        };

        var session = (A2AAgentSession)await this._agent.CreateSessionAsync();
        session.TaskId = "task-123";

        // Act
        await foreach (var _ in this._agent.RunStreamingAsync("Please make the background transparent", session))
        {
            // Just iterate through to trigger the logic
        }

        // Assert
        var message = this._handler.CapturedMessageSendParams?.Message;
        Assert.Null(message?.TaskId);
        Assert.NotNull(message?.ReferenceTaskIds);
        Assert.Contains("task-123", message.ReferenceTaskIds);
    }

    [Fact]
    public async Task RunStreamingAsync_WithAgentTask_UpdatesSessionTaskIdAsync()
    {
        // Arrange
        this._handler.StreamingResponseToReturn = new AgentTask
        {
            Id = "task-456",
            ContextId = "context-789",
            Status = new() { State = TaskState.Submitted }
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        await foreach (var _ in this._agent.RunStreamingAsync("Start a task", session))
        {
            // Just iterate through to trigger the logic
        }

        // Assert
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal("task-456", a2aSession.TaskId);
    }

    [Fact]
    public async Task RunStreamingAsync_WithAgentMessage_YieldsResponseUpdateAsync()
    {
        // Arrange
        const string MessageId = "msg-123";
        const string ContextId = "ctx-456";
        const string MessageText = "Hello from agent!";

        this._handler.StreamingResponseToReturn = new AgentMessage
        {
            MessageId = MessageId,
            Role = MessageRole.Agent,
            ContextId = ContextId,
            Parts =
            [
                new TextPart { Text = MessageText }
            ]
        };

        // Act
        var updates = new List<AgentResponseUpdate>();
        await foreach (var update in this._agent.RunStreamingAsync("Test message"))
        {
            updates.Add(update);
        }

        // Assert - one update should be yielded
        Assert.Single(updates);

        var update0 = updates[0];
        Assert.Equal(ChatRole.Assistant, update0.Role);
        Assert.Equal(MessageId, update0.MessageId);
        Assert.Equal(MessageId, update0.ResponseId);
        Assert.Equal(this._agent.Id, update0.AgentId);
        Assert.Equal(MessageText, update0.Text);
        Assert.Equal(ChatFinishReason.Stop, update0.FinishReason);
        Assert.IsType<AgentMessage>(update0.RawRepresentation);
        Assert.Equal(MessageId, ((AgentMessage)update0.RawRepresentation!).MessageId);
    }

    [Fact]
    public async Task RunStreamingAsync_WithAgentTask_YieldsResponseUpdateAsync()
    {
        // Arrange
        const string TaskId = "task-789";
        const string ContextId = "ctx-012";

        this._handler.StreamingResponseToReturn = new AgentTask
        {
            Id = TaskId,
            ContextId = ContextId,
            Status = new() { State = TaskState.Submitted },
            Artifacts = [
                new()
                {
                    ArtifactId = "art-123",
                    Parts = [new TextPart { Text = "Task artifact content" }]
                }
            ]
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        var updates = new List<AgentResponseUpdate>();
        await foreach (var update in this._agent.RunStreamingAsync("Start long-running task", session))
        {
            updates.Add(update);
        }

        // Assert - one update should be yielded from artifact
        Assert.Single(updates);

        var update0 = updates[0];
        Assert.Equal(ChatRole.Assistant, update0.Role);
        Assert.Equal(TaskId, update0.ResponseId);
        Assert.Equal(this._agent.Id, update0.AgentId);
        Assert.Null(update0.FinishReason);
        Assert.IsType<AgentTask>(update0.RawRepresentation);
        Assert.Equal(TaskId, ((AgentTask)update0.RawRepresentation!).Id);

        // Assert - session should be updated with context and task IDs
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal(ContextId, a2aSession.ContextId);
        Assert.Equal(TaskId, a2aSession.TaskId);
    }

    [Fact]
    public async Task RunStreamingAsync_WithTaskStatusUpdateEvent_YieldsResponseUpdateAsync()
    {
        // Arrange
        const string TaskId = "task-status-123";
        const string ContextId = "ctx-status-456";

        this._handler.StreamingResponseToReturn = new TaskStatusUpdateEvent
        {
            TaskId = TaskId,
            ContextId = ContextId,
            Status = new() { State = TaskState.Working }
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        var updates = new List<AgentResponseUpdate>();
        await foreach (var update in this._agent.RunStreamingAsync("Check task status", session))
        {
            updates.Add(update);
        }

        // Assert - one update should be yielded
        Assert.Single(updates);

        var update0 = updates[0];
        Assert.Equal(ChatRole.Assistant, update0.Role);
        Assert.Equal(TaskId, update0.ResponseId);
        Assert.Equal(this._agent.Id, update0.AgentId);
        Assert.Null(update0.FinishReason);
        Assert.IsType<TaskStatusUpdateEvent>(update0.RawRepresentation);

        // Assert - session should be updated with context and task IDs
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal(ContextId, a2aSession.ContextId);
        Assert.Equal(TaskId, a2aSession.TaskId);
    }

    [Fact]
    public async Task RunStreamingAsync_WithTaskArtifactUpdateEvent_YieldsResponseUpdateAsync()
    {
        // Arrange
        const string TaskId = "task-artifact-123";
        const string ContextId = "ctx-artifact-456";
        const string ArtifactContent = "Task artifact data";

        this._handler.StreamingResponseToReturn = new TaskArtifactUpdateEvent
        {
            TaskId = TaskId,
            ContextId = ContextId,
            Artifact = new()
            {
                ArtifactId = "artifact-789",
                Parts = [new TextPart { Text = ArtifactContent }]
            }
        };

        var session = await this._agent.CreateSessionAsync();

        // Act
        var updates = new List<AgentResponseUpdate>();
        await foreach (var update in this._agent.RunStreamingAsync("Process artifact", session))
        {
            updates.Add(update);
        }

        // Assert - one update should be yielded
        Assert.Single(updates);

        var update0 = updates[0];
        Assert.Equal(ChatRole.Assistant, update0.Role);
        Assert.Equal(TaskId, update0.ResponseId);
        Assert.Equal(this._agent.Id, update0.AgentId);
        Assert.Null(update0.FinishReason);
        Assert.IsType<TaskArtifactUpdateEvent>(update0.RawRepresentation);

        // Assert - artifact content should be in the update
        Assert.NotEmpty(update0.Contents);
        Assert.Equal(ArtifactContent, update0.Text);

        // Assert - session should be updated with context and task IDs
        var a2aSession = (A2AAgentSession)session;
        Assert.Equal(ContextId, a2aSession.ContextId);
        Assert.Equal(TaskId, a2aSession.TaskId);
    }

    [Fact]
    public async Task RunAsync_WithAllowBackgroundResponsesAndNoSession_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var options = new AgentRunOptions { AllowBackgroundResponses = true };

        // Act & Assert
        await Assert.ThrowsAsync<InvalidOperationException>(() => this._agent.RunAsync(inputMessages, null, options));
    }

    [Fact]
    public async Task RunStreamingAsync_WithAllowBackgroundResponsesAndNoSession_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var options = new AgentRunOptions { AllowBackgroundResponses = true };

        // Act & Assert
        await Assert.ThrowsAsync<InvalidOperationException>(async () =>
        {
            await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, null, options))
            {
                // Just iterate through to trigger the exception
            }
        });
    }

    [Fact]
    public async Task RunAsync_WithAgentMessageResponseMetadata_ReturnsMetadataAsAdditionalPropertiesAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response with metadata" }],
            Metadata = new Dictionary<string, JsonElement>
            {
                { "responseKey1", JsonSerializer.SerializeToElement("responseValue1") },
                { "responseCount", JsonSerializer.SerializeToElement(99) }
            }
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        // Act
        var result = await this._agent.RunAsync(inputMessages);

        // Assert
        Assert.NotNull(result.AdditionalProperties);
        Assert.NotNull(result.AdditionalProperties["responseKey1"]);
        Assert.Equal("responseValue1", ((JsonElement)result.AdditionalProperties["responseKey1"]!).GetString());
        Assert.NotNull(result.AdditionalProperties["responseCount"]);
        Assert.Equal(99, ((JsonElement)result.AdditionalProperties["responseCount"]!).GetInt32());
    }

    [Fact]
    public async Task RunAsync_WithAdditionalProperties_PropagatesThemAsMetadataToMessageSendParamsAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response" }]
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var options = new AgentRunOptions
        {
            AdditionalProperties = new()
            {
                { "key1", "value1" },
                { "key2", 42 },
                { "key3", true }
            }
        };

        // Act
        await this._agent.RunAsync(inputMessages, null, options);

        // Assert
        Assert.NotNull(this._handler.CapturedMessageSendParams);
        Assert.NotNull(this._handler.CapturedMessageSendParams.Metadata);
        Assert.Equal("value1", this._handler.CapturedMessageSendParams.Metadata["key1"].GetString());
        Assert.Equal(42, this._handler.CapturedMessageSendParams.Metadata["key2"].GetInt32());
        Assert.True(this._handler.CapturedMessageSendParams.Metadata["key3"].GetBoolean());
    }

    [Fact]
    public async Task RunAsync_WithNullAdditionalProperties_DoesNotSetMetadataAsync()
    {
        // Arrange
        this._handler.ResponseToReturn = new AgentMessage
        {
            MessageId = "response-123",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Response" }]
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        var options = new AgentRunOptions
        {
            AdditionalProperties = null
        };

        // Act
        await this._agent.RunAsync(inputMessages, null, options);

        // Assert
        Assert.NotNull(this._handler.CapturedMessageSendParams);
        Assert.Null(this._handler.CapturedMessageSendParams.Metadata);
    }

    [Fact]
    public async Task RunStreamingAsync_WithAdditionalProperties_PropagatesThemAsMetadataToMessageSendParamsAsync()
    {
        // Arrange
        this._handler.StreamingResponseToReturn = new AgentMessage
        {
            MessageId = "stream-123",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Streaming response" }]
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test streaming message")
        };

        var options = new AgentRunOptions
        {
            AdditionalProperties = new()
            {
                { "streamKey1", "streamValue1" },
                { "streamKey2", 100 },
                { "streamKey3", false }
            }
        };

        // Act
        await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, null, options))
        {
        }

        // Assert
        Assert.NotNull(this._handler.CapturedMessageSendParams);
        Assert.NotNull(this._handler.CapturedMessageSendParams.Metadata);
        Assert.Equal("streamValue1", this._handler.CapturedMessageSendParams.Metadata["streamKey1"].GetString());
        Assert.Equal(100, this._handler.CapturedMessageSendParams.Metadata["streamKey2"].GetInt32());
        Assert.False(this._handler.CapturedMessageSendParams.Metadata["streamKey3"].GetBoolean());
    }

    [Fact]
    public async Task RunStreamingAsync_WithNullAdditionalProperties_DoesNotSetMetadataAsync()
    {
        // Arrange
        this._handler.StreamingResponseToReturn = new AgentMessage
        {
            MessageId = "stream-123",
            Role = MessageRole.Agent,
            Parts = [new TextPart { Text = "Streaming response" }]
        };

        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test streaming message")
        };

        var options = new AgentRunOptions
        {
            AdditionalProperties = null
        };

        // Act
        await foreach (var _ in this._agent.RunStreamingAsync(inputMessages, null, options))
        {
        }

        // Assert
        Assert.NotNull(this._handler.CapturedMessageSendParams);
        Assert.Null(this._handler.CapturedMessageSendParams.Metadata);
    }

    [Fact]
    public async Task RunAsync_WithInvalidSessionType_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        // Create a session from a different agent type
        var invalidSession = new CustomAgentSession();

        // Act & Assert
        await Assert.ThrowsAsync<InvalidOperationException>(() => this._agent.RunAsync(invalidSession));
    }

    [Fact]
    public async Task RunStreamingAsync_WithInvalidSessionType_ThrowsInvalidOperationExceptionAsync()
    {
        // Arrange
        var inputMessages = new List<ChatMessage>
        {
            new(ChatRole.User, "Test message")
        };

        // Create a session from a different agent type
        var invalidSession = new CustomAgentSession();

        // Act & Assert
        await Assert.ThrowsAsync<InvalidOperationException>(async () => await this._agent.RunStreamingAsync(inputMessages, invalidSession).ToListAsync());
    }

    #region GetService Method Tests

    /// <summary>
    /// Verify that GetService returns A2AClient when requested.
    /// </summary>
    [Fact]
    public void GetService_RequestingA2AClient_ReturnsA2AClient()
    {
        // Arrange & Act
        var result = this._agent.GetService(typeof(A2AClient));

        // Assert
        Assert.NotNull(result);
        Assert.Same(this._a2aClient, result);
    }

    /// <summary>
    /// Verify that GetService returns AIAgentMetadata when requested.
    /// </summary>
    [Fact]
    public void GetService_RequestingAIAgentMetadata_ReturnsMetadata()
    {
        // Arrange & Act
        var result = this._agent.GetService(typeof(AIAgentMetadata));

        // Assert
        Assert.NotNull(result);
        Assert.IsType<AIAgentMetadata>(result);
        var metadata = (AIAgentMetadata)result;
        Assert.Equal("a2a", metadata.ProviderName);
    }

    /// <summary>
    /// Verify that GetService returns null for unknown service types.
    /// </summary>
    [Fact]
    public void GetService_RequestingUnknownServiceType_ReturnsNull()
    {
        // Arrange & Act
        var result = this._agent.GetService(typeof(string));

        // Assert
        Assert.Null(result);
    }

    /// <summary>
    /// Verify that GetService with serviceKey parameter returns null for unknown service types.
    /// </summary>
    [Fact]
    public void GetService_WithServiceKey_ReturnsNull()
    {
        // Arrange & Act
        var result = this._agent.GetService(typeof(string), "test-key");

        // Assert
        Assert.Null(result);
    }

    /// <summary>
    /// Verify that GetService calls base.GetService() first and returns the agent itself when requesting A2AAgent type.
    /// </summary>
    [Fact]
    public void GetService_RequestingA2AAgentType_ReturnsBaseImplementation()
    {
        // Arrange & Act
        var result = this._agent.GetService(typeof(A2AAgent));

        // Assert
        Assert.NotNull(result);
        Assert.Same(this._agent, result);
    }

    /// <summary>
    /// Verify that GetService calls base.GetService() first and returns the agent itself when requesting AIAgent type.
    /// </summary>
    [Fact]
    public void GetService_RequestingAIAgentType_ReturnsBaseImplementation()
    {
        // Arrange & Act
        var result = this._agent.GetService(typeof(AIAgent));

        // Assert
        Assert.NotNull(result);
        Assert.Same(this._agent, result);
    }

    /// <summary>
    /// Verify that GetService calls base.GetService() first but continues to derived logic when base returns null.
    /// </summary>
    [Fact]
    public void GetService_RequestingA2AClientWithServiceKey_CallsBaseFirstThenDerivedLogic()
    {
        // Arrange & Act - Request A2AClient with a service key (base.GetService will return null due to serviceKey)
        var result = this._agent.GetService(typeof(A2AClient), "some-key");

        // Assert
        Assert.NotNull(result);
        Assert.Same(this._a2aClient, result);
    }

    /// <summary>
    /// Verify that GetService returns consistent AIAgentMetadata across multiple calls.
    /// </summary>
    [Fact]
    public void GetService_RequestingAIAgentMetadata_ReturnsConsistentMetadata()
    {
        // Arrange & Act
        var result1 = this._agent.GetService(typeof(AIAgentMetadata));
        var result2 = this._agent.GetService(typeof(AIAgentMetadata));

        // Assert
        Assert.NotNull(result1);
        Assert.NotNull(result2);
        Assert.Same(result1, result2); // Should return the same instance
        Assert.IsType<AIAgentMetadata>(result1);
        var metadata = (AIAgentMetadata)result1;
        Assert.Equal("a2a", metadata.ProviderName);
    }

    /// <summary>
    /// Verify that CreateSessionAsync with contextId creates a session with the correct context ID.
    /// </summary>
    [Fact]
    public async Task CreateSessionAsync_WithContextId_CreatesSessionWithContextIdAsync()
    {
        // Arrange
        const string ContextId = "test-context-123";

        // Act
        var session = await this._agent.CreateSessionAsync(ContextId);

        // Assert
        Assert.NotNull(session);
        Assert.IsType<A2AAgentSession>(session);
        var typedSession = (A2AAgentSession)session;
        Assert.Equal(ContextId, typedSession.ContextId);
        Assert.Null(typedSession.TaskId);
    }

    /// <summary>
    /// Verify that CreateSessionAsync with contextId and taskId creates a session with both IDs set correctly.
    /// </summary>
    [Fact]
    public async Task CreateSessionAsync_WithContextIdAndTaskId_CreatesSessionWithBothIdsAsync()
    {
        // Arrange
        const string ContextId = "test-context-456";
        const string TaskId = "test-task-789";

        // Act
        var session = await this._agent.CreateSessionAsync(ContextId, TaskId);

        // Assert
        Assert.NotNull(session);
        Assert.IsType<A2AAgentSession>(session);
        var typedSession = (A2AAgentSession)session;
        Assert.Equal(ContextId, typedSession.ContextId);
        Assert.Equal(TaskId, typedSession.TaskId);
    }

    /// <summary>
    /// Verify that CreateSessionAsync throws when contextId is null, empty, or whitespace.
    /// </summary>
    [Theory]
    [InlineData(null)]
    [InlineData("")]
    [InlineData(" ")]
    [InlineData("\t")]
    [InlineData("\r\n")]
    public async Task CreateSessionAsync_WithInvalidContextId_ThrowsArgumentExceptionAsync(string? contextId)
    {
        // Act & Assert
        await Assert.ThrowsAnyAsync<ArgumentException>(async () =>
            await this._agent.CreateSessionAsync(contextId!));
    }

    /// <summary>
    /// Verify that CreateSessionAsync with both parameters throws when contextId is null, empty, or whitespace.
    /// </summary>
    [Theory]
    [InlineData(null)]
    [InlineData("")]
    [InlineData(" ")]
    [InlineData("\t")]
    [InlineData("\r\n")]
    public async Task CreateSessionAsync_WithInvalidContextIdAndValidTaskId_ThrowsArgumentExceptionAsync(string? contextId)
    {
        // Arrange
        const string TaskId = "valid-task-id";

        // Act & Assert
        await Assert.ThrowsAnyAsync<ArgumentException>(async () =>
            await this._agent.CreateSessionAsync(contextId!, TaskId));
    }

    /// <summary>
    /// Verify that CreateSessionAsync with both parameters throws when taskId is null, empty, or whitespace.
    /// </summary>
    [Theory]
    [InlineData(null)]
    [InlineData("")]
    [InlineData(" ")]
    [InlineData("\t")]
    [InlineData("\r\n")]
    public async Task CreateSessionAsync_WithValidContextIdAndInvalidTaskId_ThrowsArgumentExceptionAsync(string? taskId)
    {
        // Arrange
        const string ContextId = "valid-context-id";

        // Act & Assert
        await Assert.ThrowsAnyAsync<ArgumentException>(async () =>
            await this._agent.CreateSessionAsync(ContextId, taskId!));
    }
    #endregion

    public void Dispose()
    {
        this._handler.Dispose();
        this._httpClient.Dispose();
    }

    /// <summary>
    /// Custom agent session class for testing invalid session type scenario.
    /// </summary>
    private sealed class CustomAgentSession : AgentSession;

    internal sealed class A2AClientHttpMessageHandlerStub : HttpMessageHandler
    {
        public JsonRpcRequest? CapturedJsonRpcRequest { get; set; }

        public MessageSendParams? CapturedMessageSendParams { get; set; }

        public TaskIdParams? CapturedTaskIdParams { get; set; }

        public A2AEvent? ResponseToReturn { get; set; }

        public A2AEvent? StreamingResponseToReturn { get; set; }

        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
        {
            // Capture the request content
#pragma warning disable CA2016 // Forward the 'CancellationToken' parameter to methods; overload doesn't exist downlevel
            var content = await request.Content!.ReadAsStringAsync();
#pragma warning restore CA2016

            this.CapturedJsonRpcRequest = JsonSerializer.Deserialize<JsonRpcRequest>(content);

            try
            {
                this.CapturedMessageSendParams = this.CapturedJsonRpcRequest?.Params?.Deserialize<MessageSendParams>();
            }
            catch { /* Ignore deserialization errors for non-MessageSendParams requests */ }

            try
            {
                this.CapturedTaskIdParams = this.CapturedJsonRpcRequest?.Params?.Deserialize<TaskIdParams>();
            }
            catch { /* Ignore deserialization errors for non-TaskIdParams requests */ }

            // Return the pre-configured non-streaming response
            if (this.ResponseToReturn is not null)
            {
                var jsonRpcResponse = JsonRpcResponse.CreateJsonRpcResponse("response-id", this.ResponseToReturn);

                return new HttpResponseMessage(HttpStatusCode.OK)
                {
                    Content = new StringContent(JsonSerializer.Serialize(jsonRpcResponse), Encoding.UTF8, "application/json")
                };
            }
            // Return the pre-configured streaming response
            else if (this.StreamingResponseToReturn is not null)
            {
                var stream = new MemoryStream();

                await SseFormatter.WriteAsync(
                    new SseItem<JsonRpcResponse>[]
                    {
                        new(JsonRpcResponse.CreateJsonRpcResponse("response-id", this.StreamingResponseToReturn!))
                    }.ToAsyncEnumerable(),
                    stream,
                    (item, writer) =>
                    {
                        using Utf8JsonWriter json = new(writer, new() { Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping });
                        JsonSerializer.Serialize(json, item.Data);
                    },
                    cancellationToken
                );

                stream.Position = 0;

                return new HttpResponseMessage(HttpStatusCode.OK)
                {
                    Content = new StreamContent(stream)
                    {
                        Headers = { { "Content-Type", "text/event-stream" } }
                    }
                };
            }
            else
            {
                var jsonRpcResponse = JsonRpcResponse.CreateJsonRpcResponse<A2AEvent>("response-id", new AgentMessage());

                return new HttpResponseMessage(HttpStatusCode.OK)
                {
                    Content = new StringContent(JsonSerializer.Serialize(jsonRpcResponse), Encoding.UTF8, "application/json")
                };
            }
        }
    }
}