Skip to content

Commit 2f5713f

Browse files
quifoxquifox
authored
.NET: Preserve A2A run errors when session save fails (microsoft#8877)
* fix(a2a): preserve agent errors when session save fails * test(a2a): cover session save exception precedence
1 parent 4b0e3d8 commit 2f5713f

2 files changed

Lines changed: 107 additions & 10 deletions

File tree

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

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue
144144

145145
var updates = this._hostAgent.RunStreamingAsync(chatMessages, session, options, cancellationToken);
146146

147+
bool executionFailed = false;
147148
try
148149
{
149150
if (returnTask)
@@ -169,9 +170,14 @@ private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue
169170
await StreamMessageUpdatesAsync(contextId, updates, eventQueue, cancellationToken).ConfigureAwait(false);
170171
}
171172
}
173+
catch (Exception)
174+
{
175+
executionFailed = true;
176+
throw;
177+
}
172178
finally
173179
{
174-
await this._hostAgent.SaveSessionAsync(contextId, session, CancellationToken.None).ConfigureAwait(false);
180+
await this.SaveSessionAsync(contextId, session, suppressFailure: executionFailed).ConfigureAwait(false);
175181
}
176182
}
177183

@@ -185,6 +191,7 @@ private async Task HandleTaskUpdateAsync(RequestContext context, AgentEventQueue
185191
var options = CreateRunOptions(context);
186192

187193
AgentResponse response;
194+
bool executionFailed = false;
188195
try
189196
{
190197
response = await this._hostAgent.RunAsync(
@@ -195,17 +202,19 @@ private async Task HandleTaskUpdateAsync(RequestContext context, AgentEventQueue
195202
}
196203
catch (OperationCanceledException)
197204
{
205+
executionFailed = true;
198206
throw;
199207
}
200208
catch (Exception)
201209
{
210+
executionFailed = true;
202211
var failUpdater = new TaskUpdater(eventQueue, context.TaskId, contextId);
203212
await failUpdater.FailAsync(message: null, CancellationToken.None).ConfigureAwait(false);
204213
throw;
205214
}
206215
finally
207216
{
208-
await this._hostAgent.SaveSessionAsync(contextId, session, CancellationToken.None).ConfigureAwait(false);
217+
await this.SaveSessionAsync(contextId, session, suppressFailure: executionFailed).ConfigureAwait(false);
209218
}
210219

211220
if (response.ContinuationToken is null)
@@ -228,6 +237,20 @@ private async Task HandleTaskUpdateAsync(RequestContext context, AgentEventQueue
228237
}
229238
}
230239

240+
private async Task SaveSessionAsync(string contextId, AgentSession session, bool suppressFailure)
241+
{
242+
try
243+
{
244+
await this._hostAgent.SaveSessionAsync(contextId, session, CancellationToken.None).ConfigureAwait(false);
245+
}
246+
#pragma warning disable CA1031, RCS1075 // Preserve the exception already propagating from agent execution.
247+
catch (Exception) when (suppressFailure)
248+
#pragma warning restore CA1031, RCS1075
249+
{
250+
// Best-effort persistence must not replace the agent run failure.
251+
}
252+
}
253+
231254
/// <summary>
232255
/// Creates the <see cref="AgentRunOptions"/> for a run, forwarding the caller supplied A2A
233256
/// <c>MessageSendParams.metadata</c> and <c>MessageSendParams.configuration</c> to the hosted agent.

‎dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AAgentHandlerTests.cs‎

Lines changed: 82 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -487,16 +487,31 @@ await Assert.ThrowsAsync<InvalidOperationException>(() =>
487487
public async Task ExecuteAsync_OnContinuation_WhenOperationCancelled_DoesNotEmitFailedAsync()
488488
{
489489
// Arrange
490+
var mockSessionStore = new Mock<AgentSessionStore> { CallBase = true };
491+
mockSessionStore
492+
.Setup(x => x.GetSessionAsync(
493+
It.IsAny<AIAgent>(),
494+
It.IsAny<AgentSessionStoreKey>(),
495+
It.IsAny<CancellationToken>()))
496+
.ReturnsAsync(new TestAgentSession());
497+
mockSessionStore
498+
.Setup(x => x.SaveSessionAsync(
499+
It.IsAny<AIAgent>(),
500+
It.IsAny<AgentSessionStoreKey>(),
501+
It.IsAny<AgentSession>(),
502+
It.IsAny<CancellationToken>()))
503+
.Returns(() => new ValueTask(Task.FromException(new TimeoutException("Session save failed"))));
504+
490505
int callCount = 0;
491506
Mock<AIAgent> agentMock = CreateAgentMockWithCallCount(ref callCount, _ =>
492507
throw new OperationCanceledException("Cancelled"));
493-
A2AAgentHandler handler = CreateHandler(agentMock);
508+
A2AAgentHandler handler = CreateHandler(agentMock, agentSessionStore: mockSessionStore.Object);
494509

495510
// Act & Assert
496511
var events = new EventCollector();
497512
var eventQueue = new AgentEventQueue();
498513
var readerTask = ReadEventsAsync(eventQueue, events);
499-
await Assert.ThrowsAsync<OperationCanceledException>(() =>
514+
var exception = await Assert.ThrowsAsync<OperationCanceledException>(() =>
500515
handler.ExecuteAsync(
501516
new RequestContext
502517
{
@@ -509,11 +524,19 @@ await Assert.ThrowsAsync<OperationCanceledException>(() =>
509524
},
510525
eventQueue,
511526
CancellationToken.None));
527+
Assert.Equal("Cancelled", exception.Message);
512528
eventQueue.Complete(null);
513529
await readerTask;
514530

515531
// Assert - should NOT have emitted any status (OperationCanceledException is re-thrown without marking Failed)
516532
Assert.Empty(events.StatusUpdates);
533+
mockSessionStore.Verify(
534+
x => x.SaveSessionAsync(
535+
It.IsAny<AIAgent>(),
536+
It.Is<AgentSessionStoreKey>(key => key.SessionId == "ctx-1"),
537+
It.IsAny<AgentSession>(),
538+
It.Is<CancellationToken>(ct => ct == CancellationToken.None)),
539+
Times.Once);
517540
}
518541

519542
/// <summary>
@@ -2161,6 +2184,54 @@ public async Task ExecuteAsync_OnContinuation_DoesNotInvokeDynamicModeCallbackAs
21612184
Assert.False(callbackInvoked);
21622185
}
21632186

2187+
/// <summary>
2188+
/// Verifies that a session save failure still propagates when agent execution succeeds.
2189+
/// </summary>
2190+
[Fact]
2191+
public async Task ExecuteAsync_NonStreaming_WhenRunSucceedsAndSaveFails_PropagatesSaveFailureAsync()
2192+
{
2193+
// Arrange
2194+
var mockSessionStore = new Mock<AgentSessionStore> { CallBase = true };
2195+
mockSessionStore
2196+
.Setup(x => x.GetSessionAsync(
2197+
It.IsAny<AIAgent>(),
2198+
It.IsAny<AgentSessionStoreKey>(),
2199+
It.IsAny<CancellationToken>()))
2200+
.ReturnsAsync(new TestAgentSession());
2201+
mockSessionStore
2202+
.Setup(x => x.SaveSessionAsync(
2203+
It.IsAny<AIAgent>(),
2204+
It.IsAny<AgentSessionStoreKey>(),
2205+
It.IsAny<AgentSession>(),
2206+
It.IsAny<CancellationToken>()))
2207+
.Returns(() => new ValueTask(Task.FromException(new TimeoutException("Session save failed"))));
2208+
2209+
AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Reply")]);
2210+
A2AAgentHandler handler = CreateHandler(CreateAgentMockWithResponse(response), agentSessionStore: mockSessionStore.Object);
2211+
2212+
// Act
2213+
var eventQueue = new AgentEventQueue();
2214+
var exception = await Assert.ThrowsAsync<TimeoutException>(() =>
2215+
handler.ExecuteAsync(
2216+
new RequestContext
2217+
{
2218+
TaskId = "", ContextId = "ctx-save-fail", StreamingResponse = false,
2219+
Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] }
2220+
},
2221+
eventQueue,
2222+
CancellationToken.None));
2223+
2224+
// Assert
2225+
Assert.Equal("Session save failed", exception.Message);
2226+
mockSessionStore.Verify(
2227+
x => x.SaveSessionAsync(
2228+
It.IsAny<AIAgent>(),
2229+
It.Is<AgentSessionStoreKey>(key => key.SessionId == "ctx-save-fail"),
2230+
It.IsAny<AgentSession>(),
2231+
It.Is<CancellationToken>(ct => ct == CancellationToken.None)),
2232+
Times.Once);
2233+
}
2234+
21642235
/// <summary>
21652236
/// Verifies that in the non-streaming endpoint path, SaveSessionAsync is called with
21662237
/// CancellationToken.None even when RunStreamingAsync throws an exception.
@@ -2182,7 +2253,7 @@ public async Task ExecuteAsync_NonStreaming_WhenRunStreamingAsyncThrows_SavesSes
21822253
It.IsAny<AgentSessionStoreKey>(),
21832254
It.IsAny<AgentSession>(),
21842255
It.IsAny<CancellationToken>()))
2185-
.Returns(ValueTask.CompletedTask);
2256+
.Returns(() => new ValueTask(Task.FromException(new TimeoutException("Session save failed"))));
21862257

21872258
Mock<AIAgent> agentMock = new() { CallBase = true };
21882259
agentMock.SetupGet(x => x.Name).Returns("TestAgent");
@@ -2202,7 +2273,7 @@ public async Task ExecuteAsync_NonStreaming_WhenRunStreamingAsyncThrows_SavesSes
22022273

22032274
// Act
22042275
var eventQueue = new AgentEventQueue();
2205-
await Assert.ThrowsAsync<InvalidOperationException>(() =>
2276+
var exception = await Assert.ThrowsAsync<InvalidOperationException>(() =>
22062277
handler.ExecuteAsync(
22072278
new RequestContext
22082279
{
@@ -2211,6 +2282,7 @@ await Assert.ThrowsAsync<InvalidOperationException>(() =>
22112282
},
22122283
eventQueue,
22132284
cts.Token));
2285+
Assert.Equal("Agent failed", exception.Message);
22142286

22152287
// Assert - SaveSessionAsync was called with CancellationToken.None despite the exception
22162288
mockSessionStore.Verify(
@@ -2243,7 +2315,7 @@ public async Task ExecuteAsync_Streaming_WhenRunStreamingAsyncThrows_SavesSessio
22432315
It.IsAny<AgentSessionStoreKey>(),
22442316
It.IsAny<AgentSession>(),
22452317
It.IsAny<CancellationToken>()))
2246-
.Returns(ValueTask.CompletedTask);
2318+
.Returns(() => new ValueTask(Task.FromException(new TimeoutException("Session save failed"))));
22472319

22482320
Mock<AIAgent> agentMock = new() { CallBase = true };
22492321
agentMock.SetupGet(x => x.Name).Returns("TestAgent");
@@ -2263,7 +2335,7 @@ public async Task ExecuteAsync_Streaming_WhenRunStreamingAsyncThrows_SavesSessio
22632335

22642336
// Act
22652337
var eventQueue = new AgentEventQueue();
2266-
await Assert.ThrowsAsync<InvalidOperationException>(() =>
2338+
var exception = await Assert.ThrowsAsync<InvalidOperationException>(() =>
22672339
handler.ExecuteAsync(
22682340
new RequestContext
22692341
{
@@ -2272,6 +2344,7 @@ await Assert.ThrowsAsync<InvalidOperationException>(() =>
22722344
},
22732345
eventQueue,
22742346
cts.Token));
2347+
Assert.Equal("Stream failed", exception.Message);
22752348

22762349
// Assert - SaveSessionAsync was called with CancellationToken.None despite the exception
22772350
mockSessionStore.Verify(
@@ -2304,7 +2377,7 @@ public async Task ExecuteAsync_OnContinuation_WhenRunAsyncThrows_SavesSessionWit
23042377
It.IsAny<AgentSessionStoreKey>(),
23052378
It.IsAny<AgentSession>(),
23062379
It.IsAny<CancellationToken>()))
2307-
.Returns(ValueTask.CompletedTask);
2380+
.Returns(() => new ValueTask(Task.FromException(new TimeoutException("Session save failed"))));
23082381

23092382
Mock<AIAgent> agentMock = new() { CallBase = true };
23102383
agentMock.SetupGet(x => x.Name).Returns("TestAgent");
@@ -2326,7 +2399,7 @@ public async Task ExecuteAsync_OnContinuation_WhenRunAsyncThrows_SavesSessionWit
23262399
var eventQueue = new AgentEventQueue();
23272400
var events = new EventCollector();
23282401
var readerTask = ReadEventsAsync(eventQueue, events);
2329-
await Assert.ThrowsAsync<InvalidOperationException>(() =>
2402+
var exception = await Assert.ThrowsAsync<InvalidOperationException>(() =>
23302403
handler.ExecuteAsync(
23312404
new RequestContext
23322405
{
@@ -2337,6 +2410,7 @@ await Assert.ThrowsAsync<InvalidOperationException>(() =>
23372410
},
23382411
eventQueue,
23392412
cts.Token));
2413+
Assert.Equal("Agent failed", exception.Message);
23402414
eventQueue.Complete(null);
23412415
await readerTask;
23422416

0 commit comments

Comments
 (0)