Updated to fix edge cases, and add more tests.

This commit is contained in:
Peter Ibekwe
2026-03-03 14:07:24 -08:00
Unverified
parent 7dbd68ea97
commit e143c3de15
4 changed files with 284 additions and 65 deletions
@@ -68,10 +68,17 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor
throw new InvalidOperationException($"No pending UserInputRequest found with id '{response.Id}'.");
}
List<ChatMessage> implicitTurnMessages = [new ChatMessage(ChatRole.User, [response])];
// Merge the external response with any already-buffered regular messages so mixed-content
// resumes can be processed in one invocation.
return this.ProcessTurnMessagesAsync(async (pendingMessages, ctx, ct) =>
{
pendingMessages.Add(new ChatMessage(ChatRole.User, [response]));
// ContinueTurnAsync owns failing to emit a TurnToken if this response does not clear up all remaining outstanding requests.
return this.ContinueTurnAsync(implicitTurnMessages, context, this._currentTurnEmitEvents ?? false, cancellationToken);
await this.ContinueTurnAsync(pendingMessages, ctx, this._currentTurnEmitEvents ?? false, ct).ConfigureAwait(false);
// Clear the buffered turn messages because they were consumed by ContinueTurnAsync.
return null;
}, context, cancellationToken);
}
private ValueTask HandleFunctionResultAsync(
@@ -84,8 +91,17 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor
throw new InvalidOperationException($"No pending FunctionCall found with id '{result.CallId}'.");
}
List<ChatMessage> implicitTurnMessages = [new ChatMessage(ChatRole.Tool, [result])];
return this.ContinueTurnAsync(implicitTurnMessages, context, this._currentTurnEmitEvents ?? false, cancellationToken);
// Merge the external response with any already-buffered regular messages so mixed-content
// resumes can be processed in one invocation.
return this.ProcessTurnMessagesAsync(async (pendingMessages, ctx, ct) =>
{
pendingMessages.Add(new ChatMessage(ChatRole.Tool, [result]));
await this.ContinueTurnAsync(pendingMessages, ctx, this._currentTurnEmitEvents ?? false, ct).ConfigureAwait(false);
// Clear the buffered turn messages because they were consumed by ContinueTurnAsync.
return null;
}, context, cancellationToken);
}
public bool ShouldEmitStreamingEvents(bool? emitEvents)
@@ -164,8 +180,8 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor
private async ValueTask<AgentResponse> InvokeAgentAsync(IEnumerable<ChatMessage> messages, IWorkflowContext context, bool emitEvents, CancellationToken cancellationToken = default)
{
#pragma warning disable MEAI001
Dictionary<string, UserInputRequestContent> userInputRequests = new();
Dictionary<string, FunctionCallContent> functionCalls = new();
Dictionary<string, UserInputRequestContent> userInputRequests = [];
Dictionary<string, FunctionCallContent> functionCalls = [];
AgentResponse response;
if (emitEvents)
@@ -198,7 +214,7 @@ internal sealed class AIAgentHostExecutor : ChatProtocolExecutor
ExtractUnservicedRequests(response.Messages.SelectMany(message => message.Contents));
}
if (this._options.EmitAgentResponseEvents == true)
if (this._options.EmitAgentResponseEvents)
{
await context.YieldOutputAsync(response, cancellationToken).ConfigureAwait(false);
}
@@ -58,7 +58,9 @@ internal sealed class AIContentExternalHandler<TRequestContent, TResponseContent
{
if (!this._pendingRequests.TryAdd(id, requestContent))
{
throw new InvalidOperationException($"A pending request with ID '{id}' already exists.");
// Request is already pending; treat as an idempotent re-emission.
// Do not repost to the sink because request IDs must remain unique while pending.
return default;
}
return this.IsIntercepted
@@ -145,7 +145,7 @@ internal sealed class WorkflowSession : AgentSession
return update;
}
private async ValueTask<StreamingRun> CreateOrResumeRunAsync(List<ChatMessage> messages, CancellationToken cancellationToken = default)
private async ValueTask<ResumeRunResult> CreateOrResumeRunAsync(List<ChatMessage> messages, CancellationToken cancellationToken = default)
{
// The workflow is validated to be a ChatProtocol workflow by the WorkflowHostAgent before creating the session,
// and does not need to be checked again here.
@@ -159,46 +159,38 @@ internal sealed class WorkflowSession : AgentSession
.ConfigureAwait(false);
// Process messages: convert response content to ExternalResponse, send regular messages as-is
await this.SendMessagesWithResponseConversionAsync(run, messages).ConfigureAwait(false);
return run;
bool hasMatchedExternalResponses = await this.SendMessagesWithResponseConversionAsync(run, messages).ConfigureAwait(false);
return new ResumeRunResult(run, hasMatchedExternalResponses: hasMatchedExternalResponses);
}
return await this._executionEnvironment
StreamingRun newRun = await this._executionEnvironment
.RunStreamingAsync(this._workflow,
messages,
this.SessionId,
cancellationToken)
.ConfigureAwait(false);
return new ResumeRunResult(newRun);
}
/// <summary>
/// Sends messages to the run, converting FunctionResultContent and UserInputResponseContent
/// to ExternalResponse when there's a matching pending request.
/// </summary>
private async ValueTask SendMessagesWithResponseConversionAsync(StreamingRun run, List<ChatMessage> messages)
/// <returns>
/// <see langword="true"/> if any external responses were sent; otherwise, <see langword="false"/>.
/// </returns>
private async ValueTask<bool> SendMessagesWithResponseConversionAsync(StreamingRun run, List<ChatMessage> messages)
{
List<ChatMessage> regularMessages = [];
// Responses are deferred until after regular messages are queued so response handlers
// can merge buffered regular content in the same continuation turn.
List<(ExternalResponse Response, string? ContentId)> externalResponses = [];
bool hasMatchedExternalResponses = false;
foreach (ChatMessage message in messages)
{
List<AIContent> regularContents = [];
foreach (AIContent content in message.Contents)
{
if (this.TryCreateExternalResponse(content) is ExternalResponse response)
{
await run.SendResponseAsync(response).ConfigureAwait(false);
if (GetResponseContentId(content) is string contentId)
{
this.RemovePendingRequest(contentId);
}
}
else
{
regularContents.Add(content);
}
}
PartitionMessageContents(message, regularContents);
if (regularContents.Count > 0)
{
@@ -208,11 +200,41 @@ internal sealed class WorkflowSession : AgentSession
}
}
// Send any remaining regular messages
// Send regular messages first so response handlers can merge them with responses.
if (regularMessages.Count > 0)
{
await run.TrySendMessageAsync(regularMessages).ConfigureAwait(false);
}
// Send external responses after regular messages.
foreach ((ExternalResponse response, string? contentId) in externalResponses)
{
await run.SendResponseAsync(response).ConfigureAwait(false);
hasMatchedExternalResponses = true;
if (contentId is string id)
{
this.RemovePendingRequest(id);
}
}
return hasMatchedExternalResponses;
void PartitionMessageContents(ChatMessage message, List<AIContent> regularContents)
{
foreach (AIContent content in message.Contents)
{
string? contentId = GetResponseContentId(content);
if (this.TryCreateExternalResponse(content) is ExternalResponse response)
{
externalResponses.Add((response, contentId));
}
else
{
regularContents.Add(content);
}
}
}
}
/// <summary>
@@ -233,21 +255,8 @@ internal sealed class WorkflowSession : AgentSession
return null;
}
// Create the response data based on content type
object? responseData = content switch
{
FunctionResultContent functionResultContent => functionResultContent,
UserInputResponseContent userInputResponseContent => userInputResponseContent,
_ => null
};
if (responseData == null)
{
return null;
}
// Create ExternalResponse using the pending request's port info
return new ExternalResponse(pendingRequest.PortInfo, pendingRequest.RequestId, new PortableValue(responseData));
// Create ExternalResponse via the pending request to ensure proper validation and wrapping
return pendingRequest.CreateResponse(content);
}
/// <summary>
@@ -287,12 +296,19 @@ internal sealed class WorkflowSession : AgentSession
this.LastResponseId = Guid.NewGuid().ToString("N");
List<ChatMessage> messages = this.ChatHistoryProvider.GetFromBookmark(this).ToList();
#pragma warning disable CA2007 // Analyzer misfiring and not seeing .ConfigureAwait(false) below.
await using StreamingRun run =
ResumeRunResult resumeResult =
await this.CreateOrResumeRunAsync(messages, cancellationToken).ConfigureAwait(false);
#pragma warning disable CA2007 // Analyzer misfiring.
await using StreamingRun run = resumeResult.Run;
#pragma warning restore CA2007
await run.TrySendMessageAsync(new TurnToken(emitEvents: true)).ConfigureAwait(false);
// Send a TurnToken only when no external responses were delivered.
// External response handlers already drive continuation turns and can merge
// buffered regular messages, so an extra TurnToken would cause a redundant turn.
if (!resumeResult.HasMatchedExternalResponses)
{
await run.TrySendMessageAsync(new TurnToken(emitEvents: true)).ConfigureAwait(false);
}
await foreach (WorkflowEvent evt in run.WatchStreamAsync(blockOnPendingRequest: false, cancellationToken)
.ConfigureAwait(false)
.WithCancellation(cancellationToken))
@@ -391,6 +407,25 @@ internal sealed class WorkflowSession : AgentSession
/// <inheritdoc/>
public WorkflowChatHistoryProvider ChatHistoryProvider { get; }
/// <summary>
/// Captures the outcome of creating or resuming a workflow run,
/// indicating what types of messages were sent during resume.
/// </summary>
private readonly struct ResumeRunResult
{
/// <summary>The streaming run that was created or resumed.</summary>
public StreamingRun Run { get; }
/// <summary>Whether any external responses (e.g., <see cref="FunctionResultContent"/>) were delivered.</summary>
public bool HasMatchedExternalResponses { get; }
public ResumeRunResult(StreamingRun run, bool hasMatchedExternalResponses = false)
{
this.Run = Throw.IfNull(run);
this.HasMatchedExternalResponses = hasMatchedExternalResponses;
}
}
internal sealed class SessionState(
string sessionId,
CheckpointInfo? lastCheckpoint,
@@ -35,10 +35,22 @@ public sealed class ExpectedException : Exception
internal sealed class RequestEmittingAgent : AIAgent
{
private readonly AIContent _requestContent;
private readonly bool _completeOnResponse;
public RequestEmittingAgent(AIContent requestContent)
/// <summary>
/// Creates a new <see cref="RequestEmittingAgent"/> that emits the given request content.
/// </summary>
/// <param name="requestContent">The content to emit on each turn.</param>
/// <param name="completeOnResponse">
/// When <see langword="true"/>, the agent emits a text completion instead of re-emitting
/// the request when the incoming messages contain a <see cref="FunctionResultContent"/>
/// or <see cref="UserInputResponseContent"/>. This models realistic agent behaviour
/// where the agent processes the tool result and produces a final answer.
/// </param>
public RequestEmittingAgent(AIContent requestContent, bool completeOnResponse = false)
{
this._requestContent = requestContent;
this._completeOnResponse = completeOnResponse;
}
private sealed class Session : AgentSession
@@ -60,8 +72,16 @@ internal sealed class RequestEmittingAgent : AIAgent
protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
// Emit the request content
yield return new AgentResponseUpdate(ChatRole.Assistant, [this._requestContent]);
if (this._completeOnResponse && messages.Any(m => m.Contents.Any(c =>
c is FunctionResultContent || c is UserInputResponseContent)))
{
yield return new AgentResponseUpdate(ChatRole.Assistant, [new TextContent("Request processed")]);
}
else
{
// Emit the request content
yield return new AgentResponseUpdate(ChatRole.Assistant, [this._requestContent]);
}
}
}
@@ -230,7 +250,7 @@ public class WorkflowHostSmokeTests
const string CallId = "roundtrip-call-id";
const string FunctionName = "testFunction";
FunctionCallContent requestContent = new(CallId, FunctionName);
RequestEmittingAgent requestAgent = new(requestContent);
RequestEmittingAgent requestAgent = new(requestContent, completeOnResponse: true);
ExecutorBinding agentBinding = requestAgent.BindAsExecutor(
new AIAgentHostOptions { InterceptUnterminatedFunctionCalls = false, EmitAgentUpdateEvents = true });
Workflow workflow = new WorkflowBuilder(agentBinding).Build();
@@ -256,12 +276,17 @@ public class WorkflowHostSmokeTests
FunctionResultContent responseContent = new(CallId, "test result");
ChatMessage responseMessage = new(ChatRole.Tool, [responseContent]);
// This should work without throwing - the response should be converted to ExternalResponse
// and processed by the workflow
Func<Task> sendResponse = () => agent.RunStreamingAsync(responseMessage, session).ToListAsync().AsTask();
// Act 2: Run the workflow with the response and capture the resulting updates
List<AgentResponseUpdate> secondCallUpdates = await agent.RunStreamingAsync(responseMessage, session).ToListAsync();
// Assert 2: The response should be accepted without error
await sendResponse.Should().NotThrowAsync("the response should be converted to ExternalResponse and processed");
// Assert 2: The response should be processed and the original request should no longer be pending.
// Concretely, the workflow should not re-emit a FunctionCallContent with the same CallId.
secondCallUpdates.Should().NotBeNull("processing the response should produce updates");
secondCallUpdates.Should().NotBeEmpty("processing the response should progress the workflow");
secondCallUpdates
.SelectMany(u => u.Contents.OfType<FunctionCallContent>())
.Should()
.NotContain(c => c.CallId == CallId, "the original FunctionCallContent request should be cleared after processing the response");
}
/// <summary>
@@ -275,7 +300,7 @@ public class WorkflowHostSmokeTests
const string RequestId = "roundtrip-request-id";
McpServerToolCallContent mcpCall = new("mcp-call-id", "testMcpTool", "http://localhost");
McpServerToolApprovalRequestContent requestContent = new(RequestId, mcpCall);
RequestEmittingAgent requestAgent = new(requestContent);
RequestEmittingAgent requestAgent = new(requestContent, completeOnResponse: true);
ExecutorBinding agentBinding = requestAgent.BindAsExecutor(
new AIAgentHostOptions { InterceptUserInputRequests = false, EmitAgentUpdateEvents = true });
Workflow workflow = new WorkflowBuilder(agentBinding).Build();
@@ -301,11 +326,152 @@ public class WorkflowHostSmokeTests
UserInputResponseContent responseContent = requestContent.CreateResponse(approved: true);
ChatMessage responseMessage = new(ChatRole.User, [responseContent]);
// This should work without throwing - the response should be converted to ExternalResponse
// and processed by the workflow
Func<Task> sendResponse = () => agent.RunStreamingAsync(responseMessage, session).ToListAsync().AsTask();
// Act 2: Run the workflow again with the response and capture the updates
List<AgentResponseUpdate> secondCallUpdates = await agent.RunStreamingAsync(responseMessage, session).ToListAsync();
// Assert 2: The response should be accepted without error
await sendResponse.Should().NotThrowAsync("the response should be converted to ExternalResponse and processed");
// Assert 2: The response should be applied so that the original request is no longer pending
secondCallUpdates.Should().NotBeEmpty("handling the user input response should produce follow-up updates");
bool requestStillPresent = secondCallUpdates.Any(u => u.Contents.OfType<UserInputRequestContent>().Any(r => r.Id == RequestId));
requestStillPresent.Should().BeFalse("the original UserInputRequestContent should not be re-emitted after its response is processed");
}
/// <summary>
/// Tests the mixed-message scenario: resume contains both an external response
/// (FunctionResultContent matching a pending request) and regular non-response content
/// in the same message.
/// Verifies that regular content is still processed and that no duplicate
/// pending-request errors, redundant FunctionCallContent re-emissions,
/// or workflow errors occur.
/// </summary>
[Fact]
public async Task Test_AsAgent_MixedResponseAndRegularMessage_BothProcessedAsync()
{
// Arrange: Create an agent that emits a FunctionCallContent request
const string CallId = "mixed-call-id";
const string FunctionName = "mixedTestFunction";
FunctionCallContent requestContent = new(CallId, FunctionName);
RequestEmittingAgent requestAgent = new(requestContent, completeOnResponse: true);
ExecutorBinding agentBinding = requestAgent.BindAsExecutor(
new AIAgentHostOptions { InterceptUnterminatedFunctionCalls = false, EmitAgentUpdateEvents = true });
Workflow workflow = new WorkflowBuilder(agentBinding).Build();
AIAgent agent = workflow.AsAIAgent("WorkflowAgent");
// Act 1: First call - should receive the FunctionCallContent request
AgentSession session = await agent.CreateSessionAsync();
List<AgentResponseUpdate> firstCallUpdates = await agent.RunStreamingAsync(
new ChatMessage(ChatRole.User, "Start"),
session).ToListAsync();
// Assert 1: We should have received a FunctionCallContent
firstCallUpdates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent),
"the first call should emit a FunctionCallContent request");
// Act 2: Send a mixed message containing both the function result AND regular non-response content
FunctionResultContent responseContent = new(CallId, "tool output");
ChatMessage mixedMessage = new(ChatRole.Tool, [responseContent, new TextContent("additional context")]);
List<AgentResponseUpdate> secondCallUpdates = await agent.RunStreamingAsync(mixedMessage, session).ToListAsync();
// Assert 2: The workflow should have processed both parts without errors
secondCallUpdates.Should().NotBeEmpty("the mixed message should produce follow-up updates");
secondCallUpdates
.SelectMany(u => u.Contents.OfType<FunctionCallContent>())
.Should()
.NotContain(c => c.CallId == CallId, "the original FunctionCallContent should be cleared after the response is processed");
secondCallUpdates
.SelectMany(u => u.Contents.OfType<ErrorContent>())
.Should()
.BeEmpty("no workflow errors should occur when processing a mixed response-and-regular message");
}
[Fact]
public async Task Test_AsAgent_ResponseThenRegularAcrossMessages_NoDuplicateFunctionCallAsync()
{
const string CallId = "mixed-separate-call-id";
const string FunctionName = "mixedSeparateTestFunction";
RequestEmittingAgent requestAgent = new(new FunctionCallContent(CallId, FunctionName), completeOnResponse: true);
ExecutorBinding agentBinding = requestAgent.BindAsExecutor(
new AIAgentHostOptions { InterceptUnterminatedFunctionCalls = false, EmitAgentUpdateEvents = true });
Workflow workflow = new WorkflowBuilder(agentBinding).Build();
AIAgent agent = workflow.AsAIAgent("WorkflowAgent");
AgentSession session = await agent.CreateSessionAsync();
List<AgentResponseUpdate> firstCallUpdates = await agent.RunStreamingAsync(new ChatMessage(ChatRole.User, "Start"), session).ToListAsync();
firstCallUpdates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent));
ChatMessage[] resumeMessages =
[
new(ChatRole.Tool, [new FunctionResultContent(CallId, "tool output")]),
new(ChatRole.Tool, [new TextContent("extra context in separate message")])
];
List<AgentResponseUpdate> secondCallUpdates = await agent.RunStreamingAsync(resumeMessages, session).ToListAsync();
secondCallUpdates.Should().NotBeEmpty();
secondCallUpdates
.SelectMany(u => u.Contents.OfType<FunctionCallContent>())
.Should()
.NotContain(c => c.CallId == CallId, "response+regular content split across messages should not re-emit the handled request");
secondCallUpdates
.SelectMany(u => u.Contents.OfType<ErrorContent>())
.Should()
.BeEmpty();
}
[Fact]
public async Task Test_AsAgent_MatchingResponse_DoesNotCauseExtraTurnAsync()
{
const string CallId = "matching-response-call-id";
const string FunctionName = "matchingResponseFunction";
RequestEmittingAgent requestAgent = new(new FunctionCallContent(CallId, FunctionName), completeOnResponse: false);
ExecutorBinding agentBinding = requestAgent.BindAsExecutor(
new AIAgentHostOptions { InterceptUnterminatedFunctionCalls = false, EmitAgentUpdateEvents = true });
Workflow workflow = new WorkflowBuilder(agentBinding).Build();
AIAgent agent = workflow.AsAIAgent("WorkflowAgent");
AgentSession session = await agent.CreateSessionAsync();
List<AgentResponseUpdate> firstCallUpdates = await agent.RunStreamingAsync(new ChatMessage(ChatRole.User, "Start"), session).ToListAsync();
firstCallUpdates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent));
List<AgentResponseUpdate> secondCallUpdates = await agent.RunStreamingAsync(
new ChatMessage(ChatRole.Tool, [new FunctionResultContent(CallId, "tool output")]),
session).ToListAsync();
int functionCallCount = secondCallUpdates
.Where(u => u.RawRepresentation?.GetType().Name == "RequestInfoEvent")
.SelectMany(u => u.Contents.OfType<FunctionCallContent>())
.Count(c => c.CallId == CallId);
functionCallCount.Should().Be(1, "a matching external response should not trigger an extra TurnToken-driven turn");
}
[Fact]
public async Task Test_AsAgent_UnmatchedResponse_TriggersTurnAndKeepsProgressingAsync()
{
const string CallId = "unmatched-response-call-id";
const string FunctionName = "unmatchedResponseFunction";
RequestEmittingAgent requestAgent = new(new FunctionCallContent(CallId, FunctionName), completeOnResponse: false);
ExecutorBinding agentBinding = requestAgent.BindAsExecutor(
new AIAgentHostOptions { InterceptUnterminatedFunctionCalls = false, EmitAgentUpdateEvents = true });
Workflow workflow = new WorkflowBuilder(agentBinding).Build();
AIAgent agent = workflow.AsAIAgent("WorkflowAgent");
AgentSession session = await agent.CreateSessionAsync();
List<AgentResponseUpdate> firstCallUpdates = await agent.RunStreamingAsync(new ChatMessage(ChatRole.User, "Start"), session).ToListAsync();
firstCallUpdates.Should().Contain(u => u.Contents.Any(c => c is FunctionCallContent));
List<AgentResponseUpdate> secondCallUpdates = await agent.RunStreamingAsync(
new ChatMessage(ChatRole.Tool, [new FunctionResultContent("different-call-id", "tool output")]),
session).ToListAsync();
int functionCallCount = secondCallUpdates
.SelectMany(u => u.Contents.OfType<FunctionCallContent>())
.Count(c => c.CallId == CallId);
functionCallCount.Should().Be(1, "an unmatched response should be treated as regular input and still drive a TurnToken continuation without workflow errors");
secondCallUpdates.SelectMany(u => u.Contents.OfType<ErrorContent>()).Should().BeEmpty();
}
}