refactor: finalize backend streaming ownership

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
David Kaya
2026-04-13 15:47:39 +02:00
co-authored by Copilot
parent 4e34d2abfc
commit 1fdb1c27d5
6 changed files with 365 additions and 84 deletions
@@ -100,7 +100,7 @@ public class AgentWorkflowTurnRunner : ITurnWorkflowRunner
onDelta,
onEvent)
.ConfigureAwait(false);
await EmitPendingEventsAsync(state, onEvent).ConfigureAwait(false);
await EmitPendingEventsAsync(state, onDelta, onEvent).ConfigureAwait(false);
await EmitPendingMcpOauthRequestsAsync(state, onMcpOAuthRequired).ConfigureAwait(false);
if (shouldEndTurn)
{
@@ -108,13 +108,13 @@ public class AgentWorkflowTurnRunner : ITurnWorkflowRunner
}
}
await EmitPendingEventsAsync(state, onEvent).ConfigureAwait(false);
await EmitPendingEventsAsync(state, onDelta, onEvent).ConfigureAwait(false);
await EmitPendingMcpOauthRequestsAsync(state, onMcpOAuthRequired).ConfigureAwait(false);
return state.FinalizeCompletedMessages(transcriptProjector);
}
catch (OperationCanceledException) when (runCancellation.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
{
await EmitPendingEventsAsync(state, onEvent).ConfigureAwait(false);
await EmitPendingEventsAsync(state, onDelta, onEvent).ConfigureAwait(false);
await EmitPendingMcpOauthRequestsAsync(state, onMcpOAuthRequired).ConfigureAwait(false);
ExitPlanModeRequestedEventDto? exitPlanModeEvent =
_providerTurnSupport.ConsumePendingExitPlanModeRequest(command.RequestId);
@@ -230,10 +230,17 @@ public class AgentWorkflowTurnRunner : ITurnWorkflowRunner
private static async Task EmitPendingEventsAsync(
TurnExecutionState state,
Func<TurnDeltaEventDto, Task> onDelta,
Func<SidecarEventDto, Task> onEvent)
{
foreach (SidecarEventDto pendingEvent in state.DrainPendingEvents())
{
if (pendingEvent is TurnDeltaEventDto delta)
{
await onDelta(delta).ConfigureAwait(false);
continue;
}
await onEvent(pendingEvent).ConfigureAwait(false);
}
}
@@ -523,10 +530,14 @@ public class AgentWorkflowTurnRunner : ITurnWorkflowRunner
}
string messageId = state.CreateMessageId(update.Update.MessageId);
(string _, string currentAuthorName, string currentContent) = state.AppendDelta(
if (!state.TryAppendDelta(
messageId,
authorName,
update.Update.Text);
update.Update.Text,
out TranscriptSegment currentSegment))
{
return;
}
await onDelta(new TurnDeltaEventDto
{
@@ -534,9 +545,9 @@ public class AgentWorkflowTurnRunner : ITurnWorkflowRunner
RequestId = command.RequestId,
SessionId = command.SessionId,
MessageId = messageId,
AuthorName = currentAuthorName,
AuthorName = currentSegment.AuthorName,
ContentDelta = update.Update.Text,
Content = currentContent,
Content = currentSegment.Content,
}).ConfigureAwait(false);
}
@@ -1,16 +1,9 @@
using System.Text;
using Aryx.AgentHost.Contracts;
using GitHub.Copilot.SDK;
using Microsoft.Extensions.AI;
namespace Aryx.AgentHost.Services;
internal readonly record struct TranscriptSegment(string MessageId, string AuthorName, string Content)
{
public static TranscriptSegment FromTuple((string MessageId, string AuthorName, string Content) segment)
=> new(segment.MessageId, segment.AuthorName, segment.Content);
}
internal static class WorkflowTranscriptProjector
{
public static ChatMessage ToChatMessage(ChatMessageDto message)
@@ -146,10 +139,16 @@ internal static class WorkflowTranscriptProjector
ChatMessage message,
TranscriptSegment? matchedSegment)
{
if (matchedSegment is { IsFinalized: true } finalizedSegment
&& !string.IsNullOrWhiteSpace(finalizedSegment.Content))
{
return finalizedSegment.Content;
}
return FirstNonBlank(
message.Text,
matchedSegment?.Content,
TryGetAssistantMessageContent(message))
TryGetAssistantMessageContent(message),
matchedSegment?.Content)
?? string.Empty;
}
@@ -227,6 +226,16 @@ internal static class WorkflowTranscriptProjector
return null;
}
string? messageId = FirstNonBlank(message.MessageId);
if (messageId is not null
&& TryFindSegment(
remainingSegments,
segment => string.Equals(segment.MessageId, messageId, StringComparison.Ordinal),
out TranscriptSegment messageIdMatchedSegment))
{
return messageIdMatchedSegment;
}
string? messageText = string.IsNullOrWhiteSpace(message.Text) ? null : message.Text;
if (messageText is not null)
{
@@ -445,70 +454,3 @@ internal static class WorkflowTranscriptProjector
return null;
}
}
internal sealed class StreamingTranscriptBuffer
{
private readonly List<BufferedTranscriptSegment> _segments = [];
public int Count => _segments.Count;
public TranscriptSegment AppendDelta(
string messageId,
string authorName,
string delta)
{
BufferedTranscriptSegment segment = GetOrCreateSegment(messageId, authorName);
segment.SetContent(StreamingTextMerger.Merge(segment.Content.ToString(), delta));
segment.SetAuthorName(authorName);
return segment.ToSnapshot();
}
public IReadOnlyList<TranscriptSegment> Snapshot()
{
return _segments.Select(segment => segment.ToSnapshot()).ToList();
}
private BufferedTranscriptSegment GetOrCreateSegment(string messageId, string authorName)
{
BufferedTranscriptSegment? existing = _segments.LastOrDefault(segment => segment.MessageId == messageId);
if (existing is not null)
{
return existing;
}
BufferedTranscriptSegment created = new(messageId, authorName);
_segments.Add(created);
return created;
}
private sealed class BufferedTranscriptSegment
{
public BufferedTranscriptSegment(string messageId, string authorName)
{
MessageId = messageId;
AuthorName = authorName;
}
public string MessageId { get; }
public string AuthorName { get; private set; }
public StringBuilder Content { get; } = new();
public void SetContent(string value)
{
Content.Clear();
Content.Append(value);
}
public void SetAuthorName(string value)
{
AuthorName = value;
}
public TranscriptSegment ToSnapshot()
{
return new TranscriptSegment(MessageId, AuthorName, Content.ToString());
}
}
}
@@ -0,0 +1,142 @@
using System.Text;
namespace Aryx.AgentHost.Services;
internal readonly record struct TranscriptSegment(
string MessageId,
string AuthorName,
string Content,
bool IsFinalized = false)
{
public static TranscriptSegment FromTuple((string MessageId, string AuthorName, string Content) segment)
=> new(segment.MessageId, segment.AuthorName, segment.Content);
public void Deconstruct(out string messageId, out string authorName, out string content)
{
messageId = MessageId;
authorName = AuthorName;
content = Content;
}
}
internal sealed class StreamingTranscriptBuffer
{
private readonly List<BufferedTranscriptSegment> _segments = [];
public int Count => _segments.Count;
public TranscriptSegment AppendDelta(
string messageId,
string authorName,
string delta)
{
_ = TryAppendDelta(messageId, authorName, delta, out TranscriptSegment segment);
return segment;
}
public bool TryAppendDelta(
string messageId,
string authorName,
string delta,
out TranscriptSegment segment)
{
BufferedTranscriptSegment bufferedSegment = GetOrCreateSegment(messageId, authorName);
bool contentChanged = bufferedSegment.TryAppendDelta(authorName, delta);
segment = bufferedSegment.ToSnapshot();
return contentChanged;
}
public bool TryApplySnapshot(
string messageId,
string authorName,
string content,
out TranscriptSegment segment)
{
BufferedTranscriptSegment bufferedSegment = GetOrCreateSegment(messageId, authorName);
bool visibleContentChanged = bufferedSegment.TryApplySnapshot(authorName, content);
segment = bufferedSegment.ToSnapshot();
return visibleContentChanged;
}
public IReadOnlyList<TranscriptSegment> Snapshot()
{
return _segments.Select(segment => segment.ToSnapshot()).ToList();
}
private BufferedTranscriptSegment GetOrCreateSegment(string messageId, string authorName)
{
BufferedTranscriptSegment? existing = _segments.LastOrDefault(segment => segment.MessageId == messageId);
if (existing is not null)
{
return existing;
}
BufferedTranscriptSegment created = new(messageId, authorName);
_segments.Add(created);
return created;
}
private sealed class BufferedTranscriptSegment
{
public BufferedTranscriptSegment(string messageId, string authorName)
{
MessageId = messageId;
AuthorName = authorName;
}
public string MessageId { get; }
public string AuthorName { get; private set; }
public bool IsFinalized { get; private set; }
public StringBuilder Content { get; } = new();
public bool TryAppendDelta(string authorName, string delta)
{
SetAuthorName(authorName);
if (IsFinalized || string.IsNullOrEmpty(delta))
{
return false;
}
string currentContent = Content.ToString();
string mergedContent = StreamingTextMerger.Merge(currentContent, delta);
if (string.Equals(currentContent, mergedContent, StringComparison.Ordinal))
{
return false;
}
SetContent(mergedContent);
return true;
}
public bool TryApplySnapshot(string authorName, string content)
{
SetAuthorName(authorName);
string normalizedContent = content ?? string.Empty;
string currentContent = Content.ToString();
bool visibleContentChanged = !string.Equals(currentContent, normalizedContent, StringComparison.Ordinal);
SetContent(normalizedContent);
IsFinalized = true;
return visibleContentChanged;
}
private void SetContent(string value)
{
Content.Clear();
Content.Append(value);
}
private void SetAuthorName(string value)
{
AuthorName = value;
}
public TranscriptSegment ToSnapshot()
{
return new TranscriptSegment(MessageId, AuthorName, Content.ToString(), IsFinalized);
}
}
}
@@ -155,6 +155,20 @@ internal class TurnExecutionState
case ProviderAssistantMessageEvent assistantMessage:
RecordObservedAgentForMessage(agent, assistantMessage.MessageId);
QueueThinkingIfNeeded(agent);
if (!string.IsNullOrWhiteSpace(assistantMessage.Content)
&& TryFinalizeTranscriptMessage(
assistantMessage.MessageId,
agent.AgentName,
assistantMessage.Content,
out TranscriptSegment finalizedSegment))
{
_pendingEvents.Enqueue(CreateTurnDeltaEvent(
finalizedSegment.MessageId,
finalizedSegment.AuthorName,
string.Empty,
finalizedSegment.Content));
}
if (assistantMessage.HasToolRequests)
{
QueueMessageReclassifiedIfNeeded(assistantMessage.MessageId);
@@ -375,6 +389,24 @@ internal class TurnExecutionState
return _transcriptBuffer.AppendDelta(messageId, authorName, delta);
}
public bool TryAppendDelta(
string messageId,
string authorName,
string delta,
out TranscriptSegment segment)
{
return _transcriptBuffer.TryAppendDelta(messageId, authorName, delta, out segment);
}
public bool TryFinalizeTranscriptMessage(
string messageId,
string authorName,
string content,
out TranscriptSegment segment)
{
return _transcriptBuffer.TryApplySnapshot(messageId, authorName, content, out segment);
}
public void ClearActiveAgentIfMatching(AgentIdentity completedAgent)
{
if (ActiveAgent.HasValue
@@ -725,6 +757,24 @@ internal class TurnExecutionState
};
}
private TurnDeltaEventDto CreateTurnDeltaEvent(
string messageId,
string authorName,
string contentDelta,
string? content)
{
return new TurnDeltaEventDto
{
Type = "turn-delta",
RequestId = _command.RequestId,
SessionId = _command.SessionId,
MessageId = messageId,
AuthorName = authorName,
ContentDelta = contentDelta,
Content = content,
};
}
private SkillInvokedEventDto CreateSkillInvokedEvent(
AgentIdentity agent,
ProviderSkillInvokedEvent data)
@@ -847,6 +847,53 @@ public sealed class CopilotTurnExecutionStateTests
Assert.Equal("agent-1", evt.AgentId);
}
[Fact]
public void ObserveSessionEvent_AssistantMessage_FinalizesTranscriptAndSuppressesLateDeltas()
{
RunTurnCommandDto command = CreateCommand();
CopilotTurnExecutionState state = new(command);
Assert.True(state.TryAppendDelta("msg-final", "Primary", "Draft response", out TranscriptSegment streamed));
Assert.False(streamed.IsFinalized);
state.ObserveSessionEvent(
command.Workflow.GetAgentNodes()[0],
SessionEvent.FromJson(
"""
{
"type": "assistant.message",
"data": {
"messageId": "msg-final",
"content": "Provider final response"
},
"id": "12121212-1212-1212-1212-121212121212",
"timestamp": "2026-03-27T00:00:00Z"
}
"""));
TurnDeltaEventDto correction = Assert.Single(state.DrainPendingEvents().OfType<TurnDeltaEventDto>());
Assert.Equal("msg-final", correction.MessageId);
Assert.Equal("Primary", correction.AuthorName);
Assert.Equal(string.Empty, correction.ContentDelta);
Assert.Equal("Provider final response", correction.Content);
Assert.False(state.TryAppendDelta("msg-final", "Primary", " late delta", out TranscriptSegment unchanged));
Assert.True(unchanged.IsFinalized);
Assert.Equal("Provider final response", unchanged.Content);
ChatMessage output = new(ChatRole.Assistant, "Workflow wording")
{
MessageId = "msg-final",
AuthorName = "assistant",
};
state.UpdateCompletedMessages([output], []);
ChatMessageDto completed = Assert.Single(state.FinalizeCompletedMessages());
Assert.Equal("msg-final", completed.Id);
Assert.Equal("Primary", completed.AuthorName);
Assert.Equal("Provider final response", completed.Content);
}
private static SessionEvent CreateHookStartEvent()
{
return SessionEvent.FromJson(
@@ -451,6 +451,50 @@ public sealed class CopilotWorkflowRunnerTests
});
}
[Fact]
public void ProjectCompletedMessages_UsesExactMessageIdMatchBeforeFallbackHeuristics()
{
RunTurnCommandDto command = CreateCommand(
"group-chat",
CreateAgent(id: "agent-group-writer", name: "Writer"),
CreateAgent(id: "agent-group-reviewer", name: "Reviewer"));
IReadOnlyList<ChatMessageDto> messages = WorkflowTranscriptProjector.ProjectCompletedMessagesFromSegments(
command,
[
new ChatMessage(ChatRole.Assistant, "Revised draft with cleaner wording.")
{
AuthorName = "assistant",
MessageId = "msg-writer-2",
},
new ChatMessage(ChatRole.Assistant, "Review feedback with cleaner wording.")
{
AuthorName = "assistant",
MessageId = "msg-reviewer-1",
},
],
[
new TranscriptSegment("msg-writer-1", "Writer", "Initial draft."),
new TranscriptSegment("msg-reviewer-1", "Reviewer", "Review feedback with draft wording."),
new TranscriptSegment("msg-writer-2", "Writer", "Revised draft with draft wording."),
]);
Assert.Collection(
messages,
writer =>
{
Assert.Equal("msg-writer-2", writer.Id);
Assert.Equal("Writer", writer.AuthorName);
Assert.Equal("Revised draft with cleaner wording.", writer.Content);
},
reviewer =>
{
Assert.Equal("msg-reviewer-1", reviewer.Id);
Assert.Equal("Reviewer", reviewer.AuthorName);
Assert.Equal("Review feedback with cleaner wording.", reviewer.Content);
});
}
[Fact]
public void ProjectCompletedMessages_UsesFallbackAgentForGenericAssistantOutput()
{
@@ -531,6 +575,36 @@ public sealed class CopilotWorkflowRunnerTests
Assert.Equal("Done — GoogleClient.cs now contains the requested class with cleaned formatting.", message.Content);
}
[Fact]
public void ProjectCompletedMessages_PrefersFinalizedSegmentContentOverWorkflowOutput()
{
RunTurnCommandDto command = CreateCommand(
"single",
CreateAgent(id: "agent-single-primary", name: "Primary Agent"));
IReadOnlyList<ChatMessageDto> messages = WorkflowTranscriptProjector.ProjectCompletedMessagesFromSegments(
command,
[
new ChatMessage(ChatRole.Assistant, "Workflow wording that arrived later.")
{
AuthorName = "assistant",
MessageId = "msg-1",
},
],
[
new TranscriptSegment(
"msg-1",
"Primary Agent",
"Provider final wording that should win.",
IsFinalized: true),
]);
ChatMessageDto message = Assert.Single(messages);
Assert.Equal("msg-1", message.Id);
Assert.Equal("Primary Agent", message.AuthorName);
Assert.Equal("Provider final wording that should win.", message.Content);
}
[Fact]
public void ProjectCompletedMessages_DropsBlankAssistantOutputMessages()
{
@@ -602,7 +676,22 @@ public sealed class CopilotWorkflowRunnerTests
Assert.Equal("msg-2", second.MessageId);
Assert.Equal("Implementer", second.AuthorName);
Assert.Equal("B", second.Content);
});
});
}
[Fact]
public void StreamingTranscriptBuffer_IgnoresLateDeltasAfterFinalSnapshot()
{
StreamingTranscriptBuffer buffer = new();
buffer.AppendDelta("msg-1", "Architect", "Draft");
Assert.True(buffer.TryApplySnapshot("msg-1", "Architect", "Final", out TranscriptSegment finalized));
Assert.True(finalized.IsFinalized);
Assert.Equal("Final", finalized.Content);
Assert.False(buffer.TryAppendDelta("msg-1", "Architect", " late delta", out TranscriptSegment unchanged));
Assert.True(unchanged.IsFinalized);
Assert.Equal("Final", unchanged.Content);
}
[Fact]