.NET: Add streaming support to A2A agent handler

dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs

@@ -42,11 +42,19 @@ public A2AAgentHandler(
     /// <inheritdoc/>
     public Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken)
     {
+        // Handle task updates
         if (context.IsContinuation)
         {
             return this.HandleTaskUpdateAsync(context, eventQueue, cancellationToken);
         }
 
+        // Handle messages received via streaming endpoint
+        if (context.StreamingResponse)
+        {
+            return this.HandleNewMessageStreamingAsync(context, eventQueue, cancellationToken);
+        }
+
+        // Handle new messages received via non-streaming endpoint
         return this.HandleNewMessageAsync(context, eventQueue, cancellationToken);
     }
 
@@ -80,13 +88,19 @@ private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue
             ? new AgentRunOptions { AllowBackgroundResponses = allowBackgroundResponses }
             : new AgentRunOptions { AllowBackgroundResponses = allowBackgroundResponses, AdditionalProperties = context.Metadata.ToAdditionalProperties() };
 
-        var response = await this._hostAgent.RunAsync(
-            chatMessages,
-            session: session,
-            options: options,
-            cancellationToken: cancellationToken).ConfigureAwait(false);
-
-        await this._hostAgent.SaveSessionAsync(contextId, session, cancellationToken).ConfigureAwait(false);
+        AgentResponse response;
+        try
+        {
+            response = await this._hostAgent.RunAsync(
+                chatMessages,
+                session: session,
+                options: options,
+                cancellationToken: cancellationToken).ConfigureAwait(false);
+        }
+        finally
+        {
+            await this._hostAgent.SaveSessionAsync(contextId, session, CancellationToken.None).ConfigureAwait(false);
+        }
 
         if (response.ContinuationToken is null)
         {
@@ -108,6 +122,39 @@ private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue
         }
     }
 
+    private async Task HandleNewMessageStreamingAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken)
+    {
+        var contextId = context.ContextId ?? Guid.NewGuid().ToString("N");
+        var session = await this._hostAgent.GetOrCreateSessionAsync(contextId, cancellationToken).ConfigureAwait(false);
+
+        // AIAgent does not support resuming from arbitrary prior tasks.
+        // Throw explicitly so the client gets a clear error rather than a response
+        // that silently ignores the referenced task context.
+        if (context.Message?.ReferenceTaskIds is { Count: > 0 })
+        {
+            throw new NotSupportedException("ReferenceTaskIds is not supported. AIAgent cannot resume from arbitrary prior task context.");
+        }
+
+        List<ChatMessage> chatMessages = context.Message is not null ? [context.Message.ToChatMessage()] : [];
+
+        var options = context.Metadata is { Count: > 0 }
+            ? new AgentRunOptions { AdditionalProperties = context.Metadata.ToAdditionalProperties() }
+            : null;
+
+        try
+        {
+            await foreach (var update in this._hostAgent.RunStreamingAsync(chatMessages, session, options, cancellationToken).ConfigureAwait(false))
+            {
+                var message = CreateMessageFromUpdate(contextId, update);
+                await eventQueue.EnqueueMessageAsync(message, cancellationToken).ConfigureAwait(false);
+            }
+        }
+        finally
+        {
+            await this._hostAgent.SaveSessionAsync(contextId, session, CancellationToken.None).ConfigureAwait(false);
+        }
+    }
+
     private async Task HandleTaskUpdateAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken)
     {
         var contextId = context.ContextId ?? Guid.NewGuid().ToString("N");
@@ -141,8 +188,10 @@ private async Task HandleTaskUpdateAsync(RequestContext context, AgentEventQueue
             await failUpdater.FailAsync(message: null, CancellationToken.None).ConfigureAwait(false);
             throw;
         }
-
-        await this._hostAgent.SaveSessionAsync(contextId, session, cancellationToken).ConfigureAwait(false);
+        finally
+        {
+            await this._hostAgent.SaveSessionAsync(contextId, session, CancellationToken.None).ConfigureAwait(false);
+        }
 
         if (response.ContinuationToken is null)
         {
@@ -174,6 +223,16 @@ private static Message CreateMessageFromResponse(string contextId, AgentResponse
             Metadata = response.AdditionalProperties?.ToA2AMetadata()
         };
 
+    private static Message CreateMessageFromUpdate(string contextId, AgentResponseUpdate update) =>
+        new()
+        {
+            MessageId = update.ResponseId ?? Guid.NewGuid().ToString("N"),
+            ContextId = contextId,
+            Role = Role.Agent,
+            Parts = update.ToParts(),
+            Metadata = update.AdditionalProperties?.ToA2AMetadata()
+        };
+
     private static List<ChatMessage> ExtractChatMessagesFromTaskHistory(AgentTask? agentTask)
     {
         if (agentTask?.History is not { Count: > 0 })

dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Converters/MessageConverter.cs

@@ -8,6 +8,26 @@ namespace Microsoft.Agents.AI.Hosting.A2A.Converters;
 
 internal static class MessageConverter
 {
+    public static List<Part> ToParts(this AgentResponseUpdate update)
+    {
+        if (update is null || update.Contents is not { Count: > 0 })
+        {
+            return [];
+        }
+
+        var parts = new List<Part>();
+        foreach (var content in update.Contents)
+        {
+            var part = content.ToPart();
+            if (part is not null)
+            {
+                parts.Add(part);
+            }
+        }
+
+        return parts;
+    }
+
     public static List<Part> ToParts(this IList<ChatMessage> chatMessages)
     {
         if (chatMessages is null || chatMessages.Count == 0)