Skip to content

.NET: Preserve streamed annotations in Foundry hosted responses #7984

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Roger Barreto (rogerbarreto) wants to merge 2 commits into
base: main
Choose a base branch
from
Open

.NET: Preserve streamed annotations in Foundry hosted responses #7984

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8690a70
.NET: Preserve streamed response annotations
rogerbarreto Sep 1, 2026
ea91a4e
.NET: Refine streamed annotation correlation
rogerbarreto Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 115 additions & 13 deletions dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/OutputConverter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 OpenAIStreamingResponseOutputTextDeltaUpdate = OpenAI.Responses.StreamingResponseOutputTextDeltaUpdate;
using OpenAIStreamingResponseTextAnnotationAddedUpdate = OpenAI.Responses.StreamingResponseTextAnnotationAddedUpdate;

namespace Microsoft.Agents.AI.Foundry.Hosting;

Expand Down Expand Up @@ -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))
{
Expand All @@ -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)
{
Expand All @@ -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;
}

Expand Down Expand Up @@ -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);
}
}
}
}
}

Expand Down Expand Up @@ -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
{
OpenAIStreamingResponseOutputTextDeltaUpdate textDelta => textDelta.ItemId,
OpenAIStreamingResponseTextAnnotationAddedUpdate annotationAdded => annotationAdded.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;
}
Expand Down
2 changes: 2 additions & 0 deletions dotnet/src/Shared/IntegrationTests/TestSettings.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down
68 changes: 61 additions & 7 deletions dotnet/tests/Foundry.Hosting.IntegrationTests.TestContainer/Program.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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),
Expand All @@ -66,14 +69,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();
Expand Down Expand Up @@ -196,6 +197,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) =>
Expand Down
Loading
Loading