.NET: Preserve streamed annotations in Foundry hosted responses

.github/workflows/dotnet-build-and-test.yml

@@ -600,12 +600,9 @@ jobs:
         env:
           AZURE_AI_PROJECT_ENDPOINT: ${{ vars.IT_HOSTED_AGENT_PROJECT_ENDPOINT }}
           AZURE_AI_MODEL_DEPLOYMENT_NAME: ${{ vars.IT_HOSTED_AGENT_MODEL_DEPLOYMENT_NAME }}
-          # Azure AI Search (for the azure-search-rag scenario). Reuses the integration
-          # environment secrets shared with python-sample-validation.yml. The index is
-          # provisioned out of band; see dotnet/tests/Foundry.Hosting.IntegrationTests/README.md
-          # for the required schema and seed content.
+          # Use the dedicated hosted-agent Search index for both Search scenarios.
           AZURE_SEARCH_ENDPOINT: ${{ secrets.AZURE_SEARCH_ENDPOINT }}
-          AZURE_SEARCH_INDEX_NAME: ${{ secrets.AZURE_SEARCH_INDEX_NAME }}
+          AZURE_SEARCH_INDEX_NAME: ${{ vars.FOUNDRY_HOSTED_AGENT_SEARCH_INDEX_NAME }}
           # IT_HOSTED_AGENT_IMAGE was exported into $GITHUB_ENV by the previous step.
 
   # This final job is required to satisfy the merge queue. It must only run (or succeed) if no tests failed

dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/OutputConverter.cs

@@ -15,6 +15,14 @@
 using Microsoft.Agents.AI.Workflows;
 using Microsoft.Extensions.AI;
 using MeaiTextContent = Microsoft.Extensions.AI.TextContent;
+using OpenAIContainerFileCitationMessageAnnotation = OpenAI.Responses.ContainerFileCitationMessageAnnotation;
+using OpenAIFileCitationMessageAnnotation = OpenAI.Responses.FileCitationMessageAnnotation;
+using OpenAIFilePathMessageAnnotation = OpenAI.Responses.FilePathMessageAnnotation;
+using OpenAIMessageResponseItem = OpenAI.Responses.MessageResponseItem;
+using OpenAIStreamingResponseOutputItemAddedUpdate = OpenAI.Responses.StreamingResponseOutputItemAddedUpdate;
+using OpenAIStreamingResponseOutputItemDoneUpdate = OpenAI.Responses.StreamingResponseOutputItemDoneUpdate;
+using OpenAIStreamingResponseOutputTextAnnotationAddedUpdate = OpenAI.Responses.StreamingResponseOutputTextAnnotationAddedUpdate;
+using OpenAIStreamingResponseOutputTextDeltaUpdate = OpenAI.Responses.StreamingResponseOutputTextDeltaUpdate;
 
 namespace Microsoft.Agents.AI.Foundry.Hosting;
 
@@ -100,13 +108,14 @@ public static async IAsyncEnumerable<ResponseStreamEvent> ConvertUpdatesToEvents
                 continue;
             }
 
+            string? currentMessageId = ResolveMessageId(update);
             foreach (var content in update.Contents)
             {
                 switch (content)
                 {
                     case MeaiTextContent textContent:
                     {
-                        if (!IsSameMessage(update.MessageId, previousMessageId) && currentMessageBuilder is not null)
+                        if (!IsSameMessage(currentMessageId, previousMessageId) && currentMessageBuilder is not null)
                         {
                             foreach (var evt in CloseCurrentMessage(currentMessageBuilder, currentTextBuilder, accumulatedText, accumulatedAnnotations))
                             {
@@ -119,7 +128,10 @@ public static async IAsyncEnumerable<ResponseStreamEvent> ConvertUpdatesToEvents
                             accumulatedAnnotations = null;
                         }
 
-                        previousMessageId = update.MessageId;
+                        if (currentMessageId is { Length: > 0 })
+                        {
+                            previousMessageId = currentMessageId;
+                        }
 
                         if (currentMessageBuilder is null)
                         {
@@ -138,14 +150,6 @@ public static async IAsyncEnumerable<ResponseStreamEvent> ConvertUpdatesToEvents
                             yield return currentTextBuilder!.EmitDelta(textContent.Text);
                         }
 
-                        if (textContent.Annotations is { Count: > 0 })
-                        {
-                            foreach (var sdkAnnotation in ConvertToSdkAnnotations(textContent.Annotations))
-                            {
-                                (accumulatedAnnotations ??= []).Add(sdkAnnotation);
-                            }
-                        }
-
                         break;
                     }
 
@@ -336,6 +340,27 @@ public static async IAsyncEnumerable<ResponseStreamEvent> ConvertUpdatesToEvents
                     default:
                         break;
                 }
+
+                var isTextContent = content is MeaiTextContent;
+                var isAnnotationOnlyContentForCurrentMessage =
+                    content.GetType() == typeof(AIContent) &&
+                    IsSameAnnotationMessage(currentMessageId, previousMessageId);
+
+                // MEAI OpenAI sends streaming citations in a separate annotation-only AIContent after
+                // the text deltas. Accumulate them because AgentServer emits annotations only after
+                // output_text.done, and de-duplicate providers that report the same citation twice.
+                if ((isTextContent || isAnnotationOnlyContentForCurrentMessage)
+                    && content.Annotations is { Count: > 0 }
+                    && currentMessageBuilder is not null)
+                {
+                    foreach (var sdkAnnotation in ConvertToSdkAnnotations(content.Annotations))
+                    {
+                        if (accumulatedAnnotations?.Any(existing => AreEquivalentAnnotations(existing, sdkAnnotation)) is not true)
+                        {
+                            (accumulatedAnnotations ??= []).Add(sdkAnnotation);
+                        }
+                    }
+                }
             }
         }
 
@@ -385,16 +410,93 @@ private static IEnumerable<ResponseStreamEvent> CloseCurrentMessage(
     private static bool IsSameMessage(string? currentId, string? previousId) =>
         currentId is not { Length: > 0 } || previousId is not { Length: > 0 } || currentId == previousId;
 
+    private static bool IsSameAnnotationMessage(string? currentId, string? previousId) =>
+        currentId is { Length: > 0 }
+            ? currentId == previousId
+            : previousId is not { Length: > 0 };
+
+    private static string? ResolveMessageId(AgentResponseUpdate update)
+    {
+        if (update.MessageId is { Length: > 0 })
+        {
+            return update.MessageId;
+        }
+
+        object? rawRepresentation = update.RawRepresentation is ChatResponseUpdate chatUpdate
+            ? chatUpdate.RawRepresentation
+            : update.RawRepresentation;
+
+        return rawRepresentation switch
+        {
+            OpenAIStreamingResponseOutputTextAnnotationAddedUpdate annotationAdded => annotationAdded.ItemId,
+            OpenAIStreamingResponseOutputTextDeltaUpdate textDelta => textDelta.ItemId,
+            OpenAIStreamingResponseOutputItemAddedUpdate { Item: OpenAIMessageResponseItem message } => message.Id,
+            OpenAIStreamingResponseOutputItemDoneUpdate { Item: OpenAIMessageResponseItem message } => message.Id,
+            _ => null,
+        };
+    }
+
+    private static bool AreEquivalentAnnotations(Annotation left, Annotation right) =>
+        (left, right) switch
+        {
+            (UrlCitationBody leftUrl, UrlCitationBody rightUrl) =>
+                leftUrl.Url == rightUrl.Url &&
+                leftUrl.StartIndex == rightUrl.StartIndex &&
+                leftUrl.EndIndex == rightUrl.EndIndex &&
+                leftUrl.Title == rightUrl.Title,
+            (FileCitationBody leftFile, FileCitationBody rightFile) =>
+                leftFile.FileId == rightFile.FileId &&
+                leftFile.Index == rightFile.Index &&
+                leftFile.Filename == rightFile.Filename,
+            (ContainerFileCitationBody leftContainerFile, ContainerFileCitationBody rightContainerFile) =>
+                leftContainerFile.ContainerId == rightContainerFile.ContainerId &&
+                leftContainerFile.FileId == rightContainerFile.FileId &&
+                leftContainerFile.StartIndex == rightContainerFile.StartIndex &&
+                leftContainerFile.EndIndex == rightContainerFile.EndIndex &&
+                leftContainerFile.Filename == rightContainerFile.Filename,
+            (FilePath leftFilePath, FilePath rightFilePath) =>
+                leftFilePath.FileId == rightFilePath.FileId &&
+                leftFilePath.Index == rightFilePath.Index,
+            _ => false,
+        };
+
     /// <summary>
     /// Converts MEAI <see cref="AIAnnotation"/> instances to Responses SDK <see cref="Annotation"/> objects.
-    /// Only <see cref="CitationAnnotation"/> with a URL and at least one <see cref="TextSpanAnnotatedRegion"/>
-    /// with explicit start/end indices is converted; all other shapes are skipped.
+    /// Only supported <see cref="CitationAnnotation"/> shapes are converted; all others are skipped.
     /// </summary>
     private static IEnumerable<Annotation> ConvertToSdkAnnotations(IList<AIAnnotation> annotations)
     {
         foreach (var ann in annotations)
         {
-            if (ann is not CitationAnnotation citation || citation.Url is null)
+            if (ann is not CitationAnnotation citation)
+            {
+                continue;
+            }
+
+            if (citation.RawRepresentation is OpenAIContainerFileCitationMessageAnnotation containerFileCitation)
+            {
+                yield return new ContainerFileCitationBody(
+                    containerFileCitation.ContainerId,
+                    containerFileCitation.FileId,
+                    containerFileCitation.StartIndex,
+                    containerFileCitation.EndIndex,
+                    containerFileCitation.Filename);
+                continue;
+            }
+
+            if (citation.RawRepresentation is OpenAIFileCitationMessageAnnotation fileCitation)
+            {
+                yield return new FileCitationBody(fileCitation.FileId, fileCitation.Index, fileCitation.Filename);
+                continue;
+            }
+
+            if (citation.RawRepresentation is OpenAIFilePathMessageAnnotation filePath)
+            {
+                yield return new FilePath(filePath.FileId, filePath.Index);
+                continue;
+            }
+
+            if (citation.Url is null)
             {
                 continue;
             }

dotnet/src/Shared/IntegrationTests/TestSettings.cs

@@ -22,6 +22,8 @@ internal static class TestSettings
     public const string AzureAIProjectEndpoint = "AZURE_AI_PROJECT_ENDPOINT";
 
     // Azure AI Search (Foundry.Hosting integration tests, RAG scenario)
+    public const string AzureSearchConnectionId = "AZURE_SEARCH_CONNECTION_ID";
+    public const string AzureSearchConnectionName = "AZURE_SEARCH_CONNECTION_NAME";
     public const string AzureSearchEndpoint = "AZURE_SEARCH_ENDPOINT";
     public const string AzureSearchIndexName = "AZURE_SEARCH_INDEX_NAME";
 

dotnet/tests/Foundry.Hosting.IntegrationTests.TestContainer/Program.cs

@@ -3,6 +3,7 @@
 using System.ComponentModel;
 using Azure;
 using Azure.AI.Projects;
+using Azure.AI.Projects.Agents;
 using Azure.Identity;
 using Azure.Search.Documents;
 using Azure.Search.Documents.Models;
@@ -43,6 +44,8 @@
     "custom-storage" => CreateCustomStorageAgent(projectClient, deployment),
     "memory" => await CreateMemoryAgentAsync(projectClient, deployment).ConfigureAwait(false),
     "azure-search-rag" => CreateAzureSearchRagAgent(projectClient, deployment),
+    "azure-search-tool-annotations" => CreateAzureSearchToolAnnotationsAgent(projectClient, deployment),
+    "web-search-annotations" => CreateWebSearchAnnotationsAgent(projectClient, deployment),
     "session-files" => CreateSessionFilesAgent(projectClient, deployment),
     "agent-skills" => CreateAgentSkillsAgent(projectClient, deployment),
     "user-identity" => CreateUserIdentityAgent(projectClient, deployment),
@@ -77,14 +80,12 @@
     options.SteerableConversations = scenario == "steerable-long-running";
 });
 
-// toolbox-oauth-consent scenario: pre-register a Foundry toolbox whose tool source is fronted by a
-// per-user OAuth connection. IT_TOOLBOX_NAME names that toolbox (the fixture sets it). With the
-// startup-deferral fix the container stays routable even though the toolbox cannot enumerate without
-// a consented user, and the first user request surfaces an oauth_consent_request.
-var consentToolboxName = Environment.GetEnvironmentVariable("IT_TOOLBOX_NAME");
-if (!string.IsNullOrEmpty(consentToolboxName))
+// Scenarios that consume a project toolbox set IT_TOOLBOX_NAME through their fixture.
+// The hosting bridge resolves the toolbox through its MCP endpoint and adds its tools to every request.
+var toolboxName = Environment.GetEnvironmentVariable("IT_TOOLBOX_NAME");
+if (!string.IsNullOrEmpty(toolboxName))
 {
-    builder.Services.AddFoundryToolboxes(credential, consentToolboxName);
+    builder.Services.AddFoundryToolboxes(credential, toolboxName);
 }
 
 var app = builder.Build();
@@ -207,6 +208,59 @@ static AIAgent CreateAzureSearchRagAgent(AIProjectClient client, string deployme
     });
 }
 
+static AIAgent CreateWebSearchAnnotationsAgent(AIProjectClient client, string deployment) =>
+    client.AsAIAgent(new ChatClientAgentOptions
+    {
+        Name = "web-search-annotations-agent",
+        Description = "Hosted web search annotation test agent.",
+        ChatOptions = new ChatOptions
+        {
+            ModelId = deployment,
+            Instructions = """
+                Answer with current information from the web search results.
+                Include citations for the sources used in the answer.
+                """,
+            Tools = [new HostedWebSearchTool()],
+            ToolMode = ChatToolMode.RequireAny,
+        },
+    });
+
+static AIAgent CreateAzureSearchToolAnnotationsAgent(AIProjectClient client, string deployment)
+{
+    var connectionId = Environment.GetEnvironmentVariable("AZURE_SEARCH_CONNECTION_ID")
+        ?? throw new InvalidOperationException(
+            "AZURE_SEARCH_CONNECTION_ID is not set for IT_SCENARIO=azure-search-tool-annotations.");
+    var indexName = Environment.GetEnvironmentVariable("AZURE_SEARCH_INDEX_NAME")
+        ?? throw new InvalidOperationException(
+            "AZURE_SEARCH_INDEX_NAME is not set for IT_SCENARIO=azure-search-tool-annotations.");
+    var searchTool = FoundryAITool.CreateAzureAISearchTool(new AzureAISearchToolOptions(
+    [
+        new AzureAISearchToolIndex
+        {
+            ProjectConnectionId = connectionId,
+            IndexName = indexName,
+            QueryType = AzureAISearchQueryType.Simple,
+            TopK = 3,
+        }
+    ]));
+
+    return client.AsAIAgent(new ChatClientAgentOptions
+    {
+        Name = "azure-search-tool-annotations-agent",
+        Description = "Azure AI Search hosted tool annotation test agent.",
+        ChatOptions = new ChatOptions
+        {
+            ModelId = deployment,
+            Instructions = """
+                Answer only from the Azure AI Search results.
+                Include citations for the sources used in the answer.
+                """,
+            Tools = [searchTool],
+            ToolMode = ChatToolMode.RequireAny,
+        },
+    });
+}
+
 static Func<string, CancellationToken, Task<IEnumerable<TextSearchProvider.TextSearchResult>>>
     CreateAzureSearchAdapter(SearchClient client, int top = 3) =>
     async (query, cancellationToken) =>