feat: add sidecar turn cancellation

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
David Kaya
2026-03-25 22:04:14 +01:00
co-authored by Copilot
parent c8bb9d6f59
commit d84b3021f2
3 changed files with 213 additions and 15 deletions
@@ -166,6 +166,11 @@ public sealed class RunTurnCommandDto : SidecarCommandEnvelope
public RunTurnToolingConfigDto? Tooling { get; init; }
}
public sealed class CancelTurnCommandDto : SidecarCommandEnvelope
{
public string TargetRequestId { get; init; } = string.Empty;
}
public sealed class ResolveApprovalCommandDto : SidecarCommandEnvelope
{
public string ApprovalId { get; init; } = string.Empty;
@@ -232,6 +237,7 @@ public sealed class TurnCompleteEventDto : SidecarEventDto
{
public string SessionId { get; init; } = string.Empty;
public IReadOnlyList<ChatMessageDto> Messages { get; init; } = [];
public bool Cancelled { get; init; }
}
public sealed class AgentActivityEventDto : SidecarEventDto
@@ -12,6 +12,7 @@ public sealed class SidecarProtocolHost
private const string DescribeCapabilitiesCommandType = "describe-capabilities";
private const string ValidatePatternCommandType = "validate-pattern";
private const string RunTurnCommandType = "run-turn";
private const string CancelTurnCommandType = "cancel-turn";
private const string ResolveApprovalCommandType = "resolve-approval";
private static readonly string[] AuthenticationErrorIndicators =
@@ -34,6 +35,7 @@ public sealed class SidecarProtocolHost
private readonly IReadOnlyDictionary<string, Func<CommandContext, Task>> _commandHandlers;
private readonly SemaphoreSlim _writeLock = new(1, 1);
private readonly ConcurrentDictionary<string, Task> _inFlight = new(StringComparer.Ordinal);
private readonly ConcurrentDictionary<string, CancellationTokenSource> _turnCancellations = new(StringComparer.Ordinal);
public SidecarProtocolHost()
: this(new PatternValidator())
@@ -58,6 +60,7 @@ public sealed class SidecarProtocolHost
[DescribeCapabilitiesCommandType] = HandleDescribeCapabilitiesAsync,
[ValidatePatternCommandType] = HandleValidatePatternAsync,
[RunTurnCommandType] = HandleRunTurnAsync,
[CancelTurnCommandType] = HandleCancelTurnAsync,
[ResolveApprovalCommandType] = HandleResolveApprovalAsync,
};
}
@@ -160,21 +163,66 @@ public sealed class SidecarProtocolHost
private async Task HandleRunTurnAsync(CommandContext context)
{
RunTurnCommandDto command = DeserializeCommand<RunTurnCommandDto>(context);
IReadOnlyList<ChatMessageDto> messages = await _workflowRunner.RunTurnAsync(
command,
delta => WriteAsync(context.Output, delta, context.CancellationToken),
activity => WriteAsync(context.Output, activity, context.CancellationToken),
approval => WriteAsync(context.Output, approval, context.CancellationToken),
context.CancellationToken)
.ConfigureAwait(false);
await WriteAsync(context.Output, new TurnCompleteEventDto
using CancellationTokenSource turnCancellation =
CancellationTokenSource.CreateLinkedTokenSource(context.CancellationToken);
if (!_turnCancellations.TryAdd(context.Envelope.RequestId, turnCancellation))
{
Type = "turn-complete",
RequestId = context.Envelope.RequestId,
SessionId = command.SessionId,
Messages = messages,
}, context.CancellationToken).ConfigureAwait(false);
throw new InvalidOperationException(
$"A turn with request ID '{context.Envelope.RequestId}' is already in progress.");
}
try
{
IReadOnlyList<ChatMessageDto> messages = await _workflowRunner.RunTurnAsync(
command,
delta => WriteAsync(context.Output, delta, turnCancellation.Token),
activity => WriteAsync(context.Output, activity, turnCancellation.Token),
approval => WriteAsync(context.Output, approval, turnCancellation.Token),
turnCancellation.Token)
.ConfigureAwait(false);
await WriteTurnCompleteAsync(
context.Output,
context.Envelope.RequestId,
command.SessionId,
messages,
cancelled: false,
context.CancellationToken)
.ConfigureAwait(false);
}
catch (OperationCanceledException) when (turnCancellation.IsCancellationRequested)
{
await WriteTurnCompleteAsync(
context.Output,
context.Envelope.RequestId,
command.SessionId,
[],
cancelled: true,
context.CancellationToken)
.ConfigureAwait(false);
}
finally
{
_turnCancellations.TryRemove(context.Envelope.RequestId, out _);
}
}
private Task HandleCancelTurnAsync(CommandContext context)
{
CancelTurnCommandDto command = DeserializeCommand<CancelTurnCommandDto>(context);
if (_turnCancellations.TryGetValue(command.TargetRequestId, out CancellationTokenSource? turnCancellation))
{
try
{
turnCancellation.Cancel();
}
catch (ObjectDisposedException)
{
// The turn completed between lookup and cancellation.
}
}
return Task.CompletedTask;
}
private async Task HandleResolveApprovalAsync(CommandContext context)
@@ -200,6 +248,24 @@ public sealed class SidecarProtocolHost
}, context.CancellationToken);
}
private Task WriteTurnCompleteAsync(
TextWriter output,
string requestId,
string sessionId,
IReadOnlyList<ChatMessageDto> messages,
bool cancelled,
CancellationToken cancellationToken)
{
return WriteAsync(output, new TurnCompleteEventDto
{
Type = "turn-complete",
RequestId = requestId,
SessionId = sessionId,
Messages = messages,
Cancelled = cancelled,
}, cancellationToken);
}
private Task WriteCommandErrorAsync(CommandContext context, string message)
{
return WriteAsync(context.Output, new CommandErrorEventDto