Skip to content

Commit b53d25e

Browse files
atty57Copilot
andauthored
.NET: fix: preserve agent continuation token when wrapped with UseOpenTelemetry (#8265)
* .NET: fix: preserve agent continuation token through OpenTelemetryAgent OpenTelemetryAgent routes runs through OpenTelemetryChatClient, so every response makes an AgentResponse -> ChatResponse -> AgentResponse round trip. AsChatResponse/AsChatResponseUpdate return the inner agent's own ChatResponse when that is the response's raw representation, so the round trip dropped agent-level state: for a ChatClientAgent the caller received the provider's raw continuation token instead of the ChatClientAgentContinuationToken, and resuming with it threw. AgentId was lost the same way. Wrap the inner agent's response in a ChatResponse that keeps the response itself as the raw representation, so RunCoreAsync/RunCoreStreamingAsync hand back the original instance unchanged. Fixes #8243 * Add Model Id to OpenTelemetryAgent response Co-authored-by: Copilot Autofix powered by AI <[email protected]> * Carry ModelId through the streaming update wrapper The ChatResponseUpdate the agent's update is wrapped in for telemetry dropped ModelId, so gen_ai.response.model was missing on streaming runs. --------- Co-authored-by: Copilot Autofix powered by AI <[email protected]>
1 parent 726dcff commit b53d25e

2 files changed

Lines changed: 174 additions & 4 deletions

File tree

‎dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs‎

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ public async Task<ChatResponse> GetResponseAsync(
311311
var response = await parentAgent.InnerAgent.RunAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false);
312312

313313
// Wrap the response in a ChatResponse so we can pass it back through OpenTelemetryChatClient.
314-
return response.AsChatResponse();
314+
return ToChatResponse(response);
315315
}
316316

317317
public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
@@ -327,10 +327,62 @@ public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
327327
await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false))
328328
{
329329
// Wrap the response updates in ChatResponseUpdates so we can pass them back through OpenTelemetryChatClient.
330-
yield return update.AsChatResponseUpdate();
330+
yield return ToChatResponseUpdate(update);
331331
}
332332
}
333333

334+
/// <summary>
335+
/// Wraps an <see cref="AgentResponse"/> in a <see cref="ChatResponse"/> that keeps the original response as its
336+
/// <see cref="ChatResponse.RawRepresentation"/>, so <see cref="RunCoreAsync"/> can hand the
337+
/// inner agent's own response back to the caller unchanged.
338+
/// </summary>
339+
/// <remarks>
340+
/// <see cref="AgentResponseExtensions.AsChatResponse"/> isn't used here because it returns the inner agent's
341+
/// <see cref="ChatResponse"/> directly whenever that is the response's raw representation. Round-tripping
342+
/// through that <see cref="ChatResponse"/> discards agent-level state such as
343+
/// <see cref="AgentResponse.ContinuationToken"/> and <see cref="AgentResponse.AgentId"/>, which for a
344+
/// <see cref="ChatClientAgent"/> means the caller receives the underlying provider's continuation token rather
345+
/// than the agent's own, and resuming with it fails.
346+
/// </remarks>
347+
private static ChatResponse ToChatResponse(AgentResponse response) =>
348+
new()
349+
{
350+
AdditionalProperties = response.AdditionalProperties,
351+
ContinuationToken = response.ContinuationToken,
352+
ConversationId = (response.RawRepresentation as ChatResponse)?.ConversationId,
353+
ModelId = (response.RawRepresentation as ChatResponse)?.ModelId,
354+
CreatedAt = response.CreatedAt,
355+
FinishReason = response.FinishReason,
356+
Messages = response.Messages,
357+
RawRepresentation = response,
358+
ResponseId = response.ResponseId,
359+
Usage = response.Usage,
360+
};
361+
362+
/// <summary>
363+
/// Wraps an <see cref="AgentResponseUpdate"/> in a <see cref="ChatResponseUpdate"/> that keeps the original
364+
/// update as its <see cref="ChatResponseUpdate.RawRepresentation"/>, so
365+
/// <see cref="RunCoreStreamingAsync"/> can hand the inner agent's own update back to the
366+
/// caller unchanged. See <see cref="ToChatResponse"/> for why
367+
/// <see cref="AgentResponseExtensions.AsChatResponseUpdate"/> isn't used.
368+
/// </summary>
369+
private static ChatResponseUpdate ToChatResponseUpdate(AgentResponseUpdate update) =>
370+
new()
371+
{
372+
AdditionalProperties = update.AdditionalProperties,
373+
AuthorName = update.AuthorName,
374+
Contents = update.Contents,
375+
ContinuationToken = update.ContinuationToken,
376+
ConversationId = (update.RawRepresentation as ChatResponseUpdate)?.ConversationId,
377+
CreatedAt = update.CreatedAt,
378+
FinishReason = update.FinishReason,
379+
MessageId = update.MessageId,
380+
ModelId = (update.RawRepresentation as ChatResponseUpdate)?.ModelId,
381+
RawRepresentation = update,
382+
ResponseId = update.ResponseId,
383+
Role = update.Role,
384+
};
385+
334386
public object? GetService(Type serviceType, object? serviceKey = null) =>
335387
// Delegate any inquiries made by the OpenTelemetryChatClient back to the parent agent.
336388
parentAgent.GetService(serviceType, serviceKey);

‎dotnet/tests/Microsoft.Agents.AI.UnitTests/OpenTelemetryAgentTests.cs‎

Lines changed: 120 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -956,8 +956,121 @@ public async Task AutoWireChatClient_PlainAgentRunOptions_PreservesContinuationT
956956
Assert.Same(inputOptions, observedOptions);
957957
Assert.Same(token, observedOptions.ContinuationToken);
958958
}
959+
960+
[Fact]
961+
public async Task RunAsync_WrappingChatClientAgent_PreservesAgentContinuationTokenAsync()
962+
{
963+
// Arrange
964+
// The inner ChatClientAgent wraps the chat client's token in a ChatClientAgentContinuationToken. That
965+
// wrapper must survive the AgentResponse -> ChatResponse -> AgentResponse round trip that this agent
966+
// performs through OpenTelemetryChatClient, otherwise the caller gets the provider's raw token back and
967+
// resuming with it fails.
968+
var innerToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
969+
var fakeChatClient = new AutoWireTestChatClient { ContinuationToken = innerToken };
970+
var innerAgent = new ChatClientAgent(fakeChatClient);
971+
using var agent = new OpenTelemetryAgent(innerAgent);
972+
var session = await agent.CreateSessionAsync();
973+
974+
// Act
975+
var response = await agent.RunAsync("hi", session, new AgentRunOptions { AllowBackgroundResponses = true });
976+
977+
// Assert
978+
var agentToken = Assert.IsType<ChatClientAgentContinuationToken>(response.ContinuationToken);
979+
Assert.Same(innerToken, agentToken.InnerToken);
980+
Assert.Equal(innerAgent.Id, response.AgentId);
981+
}
982+
983+
[Fact]
984+
public async Task RunAsync_WrappingChatClientAgent_ResumesFromSerializedContinuationTokenAsync()
985+
{
986+
// Arrange
987+
var innerToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
988+
ChatOptions? observedChatOptions = null;
989+
var fakeChatClient = new AutoWireTestChatClient
990+
{
991+
ContinuationToken = innerToken,
992+
OnGetResponseAsync = (_, opts) => observedChatOptions = opts,
993+
};
994+
using var agent = new OpenTelemetryAgent(new ChatClientAgent(fakeChatClient));
995+
var session = await agent.CreateSessionAsync();
996+
997+
var first = await agent.RunAsync("hi", session, new AgentRunOptions { AllowBackgroundResponses = true });
998+
999+
// Act
1000+
// Serialize and deserialize the token the way a caller resuming the run from another process would.
1001+
Assert.NotNull(first.ContinuationToken);
1002+
var restoredToken = ResponseContinuationToken.FromBytes(first.ContinuationToken.ToBytes());
1003+
_ = await agent.RunAsync(session, new AgentRunOptions { ContinuationToken = restoredToken });
1004+
1005+
// Assert
1006+
// The agent must have unwrapped the restored token and passed the chat client's own token back down.
1007+
Assert.NotNull(observedChatOptions?.ContinuationToken);
1008+
Assert.Equal(innerToken.ToBytes().ToArray(), observedChatOptions.ContinuationToken.ToBytes().ToArray());
1009+
}
1010+
1011+
[Fact]
1012+
public async Task RunStreamingAsync_WrappingChatClientAgent_PreservesAgentContinuationTokenAsync()
1013+
{
1014+
// Arrange
1015+
var innerToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
1016+
var fakeChatClient = new AutoWireTestChatClient { ContinuationToken = innerToken };
1017+
var innerAgent = new ChatClientAgent(fakeChatClient);
1018+
using var agent = new OpenTelemetryAgent(innerAgent);
1019+
var session = await agent.CreateSessionAsync();
1020+
1021+
// Act
1022+
var updates = new List<AgentResponseUpdate>();
1023+
await foreach (var update in agent.RunStreamingAsync("hi", session, new AgentRunOptions { AllowBackgroundResponses = true }))
1024+
{
1025+
updates.Add(update);
1026+
}
1027+
1028+
// Assert
1029+
var tokenUpdate = Assert.Single(updates, u => u.ContinuationToken is not null);
1030+
var agentToken = Assert.IsType<ChatClientAgentContinuationToken>(tokenUpdate.ContinuationToken);
1031+
Assert.Same(innerToken, agentToken.InnerToken);
1032+
Assert.Equal(innerAgent.Id, tokenUpdate.AgentId);
1033+
}
9591034
#pragma warning restore MEAI001
9601035

1036+
[Fact]
1037+
public async Task RunStreamingAsync_ReportsResponseModelFromInnerUpdateAsync()
1038+
{
1039+
// Arrange
1040+
// The telemetry chat client reads the response model off the ChatResponseUpdate the agent's update is
1041+
// wrapped in, so ModelId has to be carried over from the inner update's raw representation.
1042+
var sourceName = Guid.NewGuid().ToString();
1043+
var activities = new ConcurrentActivityList();
1044+
using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
1045+
.AddSource(sourceName)
1046+
.AddInMemoryExporter(activities)
1047+
.Build();
1048+
1049+
var innerAgent = new TestAIAgent { RunStreamingAsyncFunc = CallbackAsync };
1050+
1051+
async static IAsyncEnumerable<AgentResponseUpdate> CallbackAsync(
1052+
IEnumerable<ChatMessage> messages, AgentSession? session, AgentRunOptions? options, [EnumeratorCancellation] CancellationToken cancellationToken)
1053+
{
1054+
await Task.Yield();
1055+
yield return new AgentResponseUpdate(ChatRole.Assistant, "ok")
1056+
{
1057+
RawRepresentation = new ChatResponseUpdate { ModelId = "innermodel" },
1058+
};
1059+
}
1060+
1061+
using var agent = new OpenTelemetryAgent(innerAgent, sourceName);
1062+
1063+
// Act
1064+
await foreach (var update in agent.RunStreamingAsync("hi"))
1065+
{
1066+
await Task.Yield();
1067+
}
1068+
1069+
// Assert
1070+
var activity = Assert.Single(activities);
1071+
Assert.Equal("innermodel", activity.GetTagItem("gen_ai.response.model"));
1072+
}
1073+
9611074
[Fact]
9621075
public async Task AutoWireChatClient_ChatClientAgentRunOptions_NoUserFactory_PreservesChatOptions_Async()
9631076
{
@@ -1209,18 +1322,23 @@ private sealed class AutoWireTestChatClient : IChatClient
12091322
{
12101323
public Action<IEnumerable<ChatMessage>, ChatOptions?>? OnGetResponseAsync { get; set; }
12111324

1325+
#pragma warning disable MEAI001 // ResponseContinuationToken is experimental.
1326+
/// <summary>Gets or sets the continuation token to report on the response, simulating a background response.</summary>
1327+
public ResponseContinuationToken? ContinuationToken { get; set; }
1328+
12121329
public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
12131330
{
12141331
this.OnGetResponseAsync?.Invoke(messages, options);
1215-
return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")));
1332+
return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")) { ContinuationToken = this.ContinuationToken });
12161333
}
12171334

12181335
public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
12191336
{
12201337
this.OnGetResponseAsync?.Invoke(messages, options);
12211338
await Task.Yield();
1222-
yield return new ChatResponseUpdate(ChatRole.Assistant, "ok");
1339+
yield return new ChatResponseUpdate(ChatRole.Assistant, "ok") { ContinuationToken = this.ContinuationToken };
12231340
}
1341+
#pragma warning restore MEAI001
12241342

12251343
public object? GetService(Type serviceType, object? serviceKey = null) =>
12261344
serviceType?.IsInstanceOfType(this) == true ? this : null;

0 commit comments

Comments
 (0)