mirror of
https://github.com/davidkaya/aryx.git
synced 2026-08-28 05:43:57 +02:00
feat: add sidecar turn cancellation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
@@ -166,6 +166,11 @@ public sealed class RunTurnCommandDto : SidecarCommandEnvelope
|
|||||||
public RunTurnToolingConfigDto? Tooling { get; init; }
|
public RunTurnToolingConfigDto? Tooling { get; init; }
|
||||||
}
|
}
|
||||||
|
|
||||||
|
public sealed class CancelTurnCommandDto : SidecarCommandEnvelope
|
||||||
|
{
|
||||||
|
public string TargetRequestId { get; init; } = string.Empty;
|
||||||
|
}
|
||||||
|
|
||||||
public sealed class ResolveApprovalCommandDto : SidecarCommandEnvelope
|
public sealed class ResolveApprovalCommandDto : SidecarCommandEnvelope
|
||||||
{
|
{
|
||||||
public string ApprovalId { get; init; } = string.Empty;
|
public string ApprovalId { get; init; } = string.Empty;
|
||||||
@@ -232,6 +237,7 @@ public sealed class TurnCompleteEventDto : SidecarEventDto
|
|||||||
{
|
{
|
||||||
public string SessionId { get; init; } = string.Empty;
|
public string SessionId { get; init; } = string.Empty;
|
||||||
public IReadOnlyList<ChatMessageDto> Messages { get; init; } = [];
|
public IReadOnlyList<ChatMessageDto> Messages { get; init; } = [];
|
||||||
|
public bool Cancelled { get; init; }
|
||||||
}
|
}
|
||||||
|
|
||||||
public sealed class AgentActivityEventDto : SidecarEventDto
|
public sealed class AgentActivityEventDto : SidecarEventDto
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ public sealed class SidecarProtocolHost
|
|||||||
private const string DescribeCapabilitiesCommandType = "describe-capabilities";
|
private const string DescribeCapabilitiesCommandType = "describe-capabilities";
|
||||||
private const string ValidatePatternCommandType = "validate-pattern";
|
private const string ValidatePatternCommandType = "validate-pattern";
|
||||||
private const string RunTurnCommandType = "run-turn";
|
private const string RunTurnCommandType = "run-turn";
|
||||||
|
private const string CancelTurnCommandType = "cancel-turn";
|
||||||
private const string ResolveApprovalCommandType = "resolve-approval";
|
private const string ResolveApprovalCommandType = "resolve-approval";
|
||||||
|
|
||||||
private static readonly string[] AuthenticationErrorIndicators =
|
private static readonly string[] AuthenticationErrorIndicators =
|
||||||
@@ -34,6 +35,7 @@ public sealed class SidecarProtocolHost
|
|||||||
private readonly IReadOnlyDictionary<string, Func<CommandContext, Task>> _commandHandlers;
|
private readonly IReadOnlyDictionary<string, Func<CommandContext, Task>> _commandHandlers;
|
||||||
private readonly SemaphoreSlim _writeLock = new(1, 1);
|
private readonly SemaphoreSlim _writeLock = new(1, 1);
|
||||||
private readonly ConcurrentDictionary<string, Task> _inFlight = new(StringComparer.Ordinal);
|
private readonly ConcurrentDictionary<string, Task> _inFlight = new(StringComparer.Ordinal);
|
||||||
|
private readonly ConcurrentDictionary<string, CancellationTokenSource> _turnCancellations = new(StringComparer.Ordinal);
|
||||||
|
|
||||||
public SidecarProtocolHost()
|
public SidecarProtocolHost()
|
||||||
: this(new PatternValidator())
|
: this(new PatternValidator())
|
||||||
@@ -58,6 +60,7 @@ public sealed class SidecarProtocolHost
|
|||||||
[DescribeCapabilitiesCommandType] = HandleDescribeCapabilitiesAsync,
|
[DescribeCapabilitiesCommandType] = HandleDescribeCapabilitiesAsync,
|
||||||
[ValidatePatternCommandType] = HandleValidatePatternAsync,
|
[ValidatePatternCommandType] = HandleValidatePatternAsync,
|
||||||
[RunTurnCommandType] = HandleRunTurnAsync,
|
[RunTurnCommandType] = HandleRunTurnAsync,
|
||||||
|
[CancelTurnCommandType] = HandleCancelTurnAsync,
|
||||||
[ResolveApprovalCommandType] = HandleResolveApprovalAsync,
|
[ResolveApprovalCommandType] = HandleResolveApprovalAsync,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
@@ -160,21 +163,66 @@ public sealed class SidecarProtocolHost
|
|||||||
private async Task HandleRunTurnAsync(CommandContext context)
|
private async Task HandleRunTurnAsync(CommandContext context)
|
||||||
{
|
{
|
||||||
RunTurnCommandDto command = DeserializeCommand<RunTurnCommandDto>(context);
|
RunTurnCommandDto command = DeserializeCommand<RunTurnCommandDto>(context);
|
||||||
IReadOnlyList<ChatMessageDto> messages = await _workflowRunner.RunTurnAsync(
|
using CancellationTokenSource turnCancellation =
|
||||||
command,
|
CancellationTokenSource.CreateLinkedTokenSource(context.CancellationToken);
|
||||||
delta => WriteAsync(context.Output, delta, context.CancellationToken),
|
if (!_turnCancellations.TryAdd(context.Envelope.RequestId, turnCancellation))
|
||||||
activity => WriteAsync(context.Output, activity, context.CancellationToken),
|
|
||||||
approval => WriteAsync(context.Output, approval, context.CancellationToken),
|
|
||||||
context.CancellationToken)
|
|
||||||
.ConfigureAwait(false);
|
|
||||||
|
|
||||||
await WriteAsync(context.Output, new TurnCompleteEventDto
|
|
||||||
{
|
{
|
||||||
Type = "turn-complete",
|
throw new InvalidOperationException(
|
||||||
RequestId = context.Envelope.RequestId,
|
$"A turn with request ID '{context.Envelope.RequestId}' is already in progress.");
|
||||||
SessionId = command.SessionId,
|
}
|
||||||
Messages = messages,
|
|
||||||
}, context.CancellationToken).ConfigureAwait(false);
|
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)
|
private async Task HandleResolveApprovalAsync(CommandContext context)
|
||||||
@@ -200,6 +248,24 @@ public sealed class SidecarProtocolHost
|
|||||||
}, context.CancellationToken);
|
}, 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)
|
private Task WriteCommandErrorAsync(CommandContext context, string message)
|
||||||
{
|
{
|
||||||
return WriteAsync(context.Output, new CommandErrorEventDto
|
return WriteAsync(context.Output, new CommandErrorEventDto
|
||||||
|
|||||||
@@ -219,6 +219,7 @@ public sealed class SidecarProtocolHostTests
|
|||||||
{
|
{
|
||||||
Assert.Equal("turn-complete", completionEvent.GetProperty("type").GetString());
|
Assert.Equal("turn-complete", completionEvent.GetProperty("type").GetString());
|
||||||
Assert.Equal("session-1", completionEvent.GetProperty("sessionId").GetString());
|
Assert.Equal("session-1", completionEvent.GetProperty("sessionId").GetString());
|
||||||
|
Assert.False(completionEvent.GetProperty("cancelled").GetBoolean());
|
||||||
JsonElement[] messages = completionEvent.GetProperty("messages").EnumerateArray().ToArray();
|
JsonElement[] messages = completionEvent.GetProperty("messages").EnumerateArray().ToArray();
|
||||||
Assert.Single(messages);
|
Assert.Single(messages);
|
||||||
Assert.Equal("Hello world", messages[0].GetProperty("content").GetString());
|
Assert.Equal("Hello world", messages[0].GetProperty("content").GetString());
|
||||||
@@ -287,6 +288,7 @@ public sealed class SidecarProtocolHostTests
|
|||||||
completionEvent =>
|
completionEvent =>
|
||||||
{
|
{
|
||||||
Assert.Equal("turn-complete", completionEvent.GetProperty("type").GetString());
|
Assert.Equal("turn-complete", completionEvent.GetProperty("type").GetString());
|
||||||
|
Assert.False(completionEvent.GetProperty("cancelled").GetBoolean());
|
||||||
},
|
},
|
||||||
commandCompleteEvent =>
|
commandCompleteEvent =>
|
||||||
{
|
{
|
||||||
@@ -295,6 +297,75 @@ public sealed class SidecarProtocolHostTests
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
[Fact]
|
||||||
|
public async Task CancelTurnCommand_CancelsInProgressTurnAndCompletesBothCommands()
|
||||||
|
{
|
||||||
|
SidecarProtocolHost host = new(
|
||||||
|
new PatternValidator(),
|
||||||
|
new FakeWorkflowRunner(async (command, onDelta, onActivity, onApproval, cancellationToken) =>
|
||||||
|
{
|
||||||
|
await Task.Delay(Timeout.Infinite, cancellationToken);
|
||||||
|
return [];
|
||||||
|
}));
|
||||||
|
|
||||||
|
IReadOnlyList<JsonElement> events = await RunHostAsync(
|
||||||
|
[
|
||||||
|
CreateRunTurnCommand(requestId: "turn-cancel"),
|
||||||
|
new CancelTurnCommandDto
|
||||||
|
{
|
||||||
|
Type = "cancel-turn",
|
||||||
|
RequestId = "cancel-command-1",
|
||||||
|
TargetRequestId = "turn-cancel",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
host);
|
||||||
|
|
||||||
|
JsonElement turnCompleteEvent = AssertSingleEvent(events, "turn-complete", "turn-cancel");
|
||||||
|
Assert.Equal("session-1", turnCompleteEvent.GetProperty("sessionId").GetString());
|
||||||
|
Assert.True(turnCompleteEvent.GetProperty("cancelled").GetBoolean());
|
||||||
|
Assert.Empty(turnCompleteEvent.GetProperty("messages").EnumerateArray().ToArray());
|
||||||
|
|
||||||
|
AssertSingleEvent(events, "command-complete", "turn-cancel");
|
||||||
|
AssertSingleEvent(events, "command-complete", "cancel-command-1");
|
||||||
|
Assert.DoesNotContain(events, evt => evt.GetProperty("type").GetString() == "command-error");
|
||||||
|
}
|
||||||
|
|
||||||
|
[Fact]
|
||||||
|
public async Task CancelTurnCommand_UnknownTarget_CompletesWithoutError()
|
||||||
|
{
|
||||||
|
IReadOnlyList<JsonElement> events = await RunHostAsync(new CancelTurnCommandDto
|
||||||
|
{
|
||||||
|
Type = "cancel-turn",
|
||||||
|
RequestId = "cancel-command-unknown",
|
||||||
|
TargetRequestId = "missing-turn",
|
||||||
|
});
|
||||||
|
|
||||||
|
JsonElement completionEvent = Assert.Single(events);
|
||||||
|
Assert.Equal("command-complete", completionEvent.GetProperty("type").GetString());
|
||||||
|
Assert.Equal("cancel-command-unknown", completionEvent.GetProperty("requestId").GetString());
|
||||||
|
}
|
||||||
|
|
||||||
|
[Fact]
|
||||||
|
public async Task CancelTurnCommand_AfterTurnCompletion_IsNoOp()
|
||||||
|
{
|
||||||
|
SidecarProtocolHost host = new(
|
||||||
|
new PatternValidator(),
|
||||||
|
new FakeWorkflowRunner(async (command, onDelta, onActivity, onApproval, cancellationToken) => []));
|
||||||
|
|
||||||
|
await RunHostAsync(CreateRunTurnCommand(requestId: "turn-completed"), host);
|
||||||
|
|
||||||
|
IReadOnlyList<JsonElement> events = await RunHostAsync(new CancelTurnCommandDto
|
||||||
|
{
|
||||||
|
Type = "cancel-turn",
|
||||||
|
RequestId = "cancel-command-completed",
|
||||||
|
TargetRequestId = "turn-completed",
|
||||||
|
}, host);
|
||||||
|
|
||||||
|
JsonElement completionEvent = Assert.Single(events);
|
||||||
|
Assert.Equal("command-complete", completionEvent.GetProperty("type").GetString());
|
||||||
|
Assert.Equal("cancel-command-completed", completionEvent.GetProperty("requestId").GetString());
|
||||||
|
}
|
||||||
|
|
||||||
[Fact]
|
[Fact]
|
||||||
public async Task ResolveApprovalCommand_DelegatesToWorkflowRunnerAndCompletes()
|
public async Task ResolveApprovalCommand_DelegatesToWorkflowRunnerAndCompletes()
|
||||||
{
|
{
|
||||||
@@ -369,7 +440,17 @@ public sealed class SidecarProtocolHostTests
|
|||||||
object command,
|
object command,
|
||||||
SidecarProtocolHost? host = null)
|
SidecarProtocolHost? host = null)
|
||||||
{
|
{
|
||||||
string input = JsonSerializer.Serialize(command, JsonOptions) + Environment.NewLine;
|
return await RunHostAsync([command], host);
|
||||||
|
}
|
||||||
|
|
||||||
|
private static async Task<IReadOnlyList<JsonElement>> RunHostAsync(
|
||||||
|
IReadOnlyList<object> commands,
|
||||||
|
SidecarProtocolHost? host = null)
|
||||||
|
{
|
||||||
|
string input = string.Join(
|
||||||
|
Environment.NewLine,
|
||||||
|
commands.Select(command => JsonSerializer.Serialize(command, JsonOptions)))
|
||||||
|
+ Environment.NewLine;
|
||||||
|
|
||||||
using StringReader reader = new(input);
|
using StringReader reader = new(input);
|
||||||
using StringWriter writer = new();
|
using StringWriter writer = new();
|
||||||
@@ -378,6 +459,16 @@ public sealed class SidecarProtocolHostTests
|
|||||||
return ParseEvents(writer.ToString());
|
return ParseEvents(writer.ToString());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private static JsonElement AssertSingleEvent(
|
||||||
|
IEnumerable<JsonElement> events,
|
||||||
|
string eventType,
|
||||||
|
string requestId)
|
||||||
|
{
|
||||||
|
return Assert.Single(events.Where(evt =>
|
||||||
|
evt.GetProperty("type").GetString() == eventType
|
||||||
|
&& evt.GetProperty("requestId").GetString() == requestId));
|
||||||
|
}
|
||||||
|
|
||||||
private static SidecarProtocolHost CreateHostForTests()
|
private static SidecarProtocolHost CreateHostForTests()
|
||||||
{
|
{
|
||||||
return new SidecarProtocolHost(
|
return new SidecarProtocolHost(
|
||||||
@@ -474,6 +565,41 @@ public sealed class SidecarProtocolHostTests
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private static RunTurnCommandDto CreateRunTurnCommand(
|
||||||
|
string requestId = "turn-1",
|
||||||
|
string sessionId = "session-1")
|
||||||
|
{
|
||||||
|
return new RunTurnCommandDto
|
||||||
|
{
|
||||||
|
Type = "run-turn",
|
||||||
|
RequestId = requestId,
|
||||||
|
SessionId = sessionId,
|
||||||
|
ProjectPath = "C:\\workspace\\project",
|
||||||
|
Pattern = new PatternDefinitionDto
|
||||||
|
{
|
||||||
|
Id = "pattern-1",
|
||||||
|
Name = "Single Agent",
|
||||||
|
Mode = "single",
|
||||||
|
Availability = "available",
|
||||||
|
Agents =
|
||||||
|
[
|
||||||
|
CreateAgent(name: "Primary"),
|
||||||
|
],
|
||||||
|
},
|
||||||
|
Messages =
|
||||||
|
[
|
||||||
|
new ChatMessageDto
|
||||||
|
{
|
||||||
|
Id = "user-1",
|
||||||
|
Role = "user",
|
||||||
|
AuthorName = "You",
|
||||||
|
Content = "Hello",
|
||||||
|
CreatedAt = "2026-01-01T00:00:00.0000000Z",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
private sealed class FakeWorkflowRunner : ITurnWorkflowRunner
|
private sealed class FakeWorkflowRunner : ITurnWorkflowRunner
|
||||||
{
|
{
|
||||||
private readonly Func<
|
private readonly Func<
|
||||||
|
|||||||
Reference in New Issue
Block a user