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)