@@ -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