From 4e34d2abfcaa5896945ee406bfc0c2df26b37682 Mon Sep 17 00:00:00 2001 From: David Kaya Date: Mon, 13 Apr 2026 15:33:27 +0200 Subject: [PATCH] refactor: unify tool call stream tracking Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Services/AgentWorkflowTurnRunner.cs | 3 +- .../Copilot/CopilotApprovalCoordinator.cs | 24 +- .../Copilot/CopilotTurnRunnerSupport.cs | 2 +- .../Services/ToolCallRegistry.cs | 165 +++++++++++++ .../Services/TurnExecutionState.cs | 107 ++------- .../WorkflowRequestInfoInterpreter.cs | 35 +-- .../CopilotTurnExecutionStateTests.cs | 23 +- .../CopilotWorkflowRunnerTests.cs | 227 +++++++++++++++--- .../WorkflowRequestInfoInterpreterTests.cs | 88 ++++--- 9 files changed, 438 insertions(+), 236 deletions(-) create mode 100644 sidecar/src/Aryx.AgentHost/Services/ToolCallRegistry.cs diff --git a/sidecar/src/Aryx.AgentHost/Services/AgentWorkflowTurnRunner.cs b/sidecar/src/Aryx.AgentHost/Services/AgentWorkflowTurnRunner.cs index 3b0a025..ed82041 100644 --- a/sidecar/src/Aryx.AgentHost/Services/AgentWorkflowTurnRunner.cs +++ b/sidecar/src/Aryx.AgentHost/Services/AgentWorkflowTurnRunner.cs @@ -326,8 +326,7 @@ public class AgentWorkflowTurnRunner : ITurnWorkflowRunner command, requestInfo, state.ActiveAgent, - state.ToolNamesByCallId, - state.ToolCallHasArgumentsById); + state.ToolCalls); if (activity is null) { diff --git a/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotApprovalCoordinator.cs b/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotApprovalCoordinator.cs index 6b65fb1..293094d 100644 --- a/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotApprovalCoordinator.cs +++ b/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotApprovalCoordinator.cs @@ -70,7 +70,7 @@ internal sealed class CopilotApprovalCoordinator WorkflowNodeDto agent, PermissionRequest request, PermissionInvocation invocation, - IReadOnlyDictionary toolNamesByCallId, + ToolCallRegistry toolCalls, Func onApproval, CancellationToken cancellationToken) { @@ -79,7 +79,7 @@ internal sealed class CopilotApprovalCoordinator agent, request, invocation, - toolNamesByCallId, + toolCalls, onActivity: null, onApproval, cancellationToken) @@ -91,12 +91,12 @@ internal sealed class CopilotApprovalCoordinator WorkflowNodeDto agent, PermissionRequest request, PermissionInvocation invocation, - IReadOnlyDictionary toolNamesByCallId, + ToolCallRegistry toolCalls, Func? onActivity, Func onApproval, CancellationToken cancellationToken) { - string? toolName = ResolveApprovalToolName(request, toolNamesByCallId); + string? toolName = ResolveApprovalToolName(request, toolCalls); string? autoApprovedToolName = ResolveAutoApprovedToolName(request); string? mcpServerApprovalKey = ResolveMcpServerApprovalKey(request, command.Tooling?.McpServers); string? approvalCacheKey = ResolveApprovalCacheKey(toolName, autoApprovedToolName); @@ -339,15 +339,15 @@ internal sealed class CopilotApprovalCoordinator internal static bool TryGetApprovalToolName( PermissionRequest request, - IReadOnlyDictionary? toolNamesByCallId, + ToolCallRegistry? toolCalls, out string? toolName) { - toolName = ResolveApprovalToolName(request, toolNamesByCallId); + toolName = ResolveApprovalToolName(request, toolCalls); return toolName is not null; } internal static bool TryGetApprovalToolName(PermissionRequest request, out string? toolName) - => TryGetApprovalToolName(request, toolNamesByCallId: null, out toolName); + => TryGetApprovalToolName(request, toolCalls: null, out toolName); internal void ClearRequestApprovals(string requestId) { @@ -404,10 +404,10 @@ internal sealed class CopilotApprovalCoordinator private static string? ResolveApprovalToolName( PermissionRequest request, - IReadOnlyDictionary? toolNamesByCallId) + ToolCallRegistry? toolCalls) { return GetDirectToolName(request) - ?? ResolveToolNameFromLookup(request, toolNamesByCallId) + ?? ResolveToolNameFromLookup(request, toolCalls) ?? GetFallbackToolName(request); } @@ -480,16 +480,16 @@ internal sealed class CopilotApprovalCoordinator private static string? ResolveToolNameFromLookup( PermissionRequest request, - IReadOnlyDictionary? toolNamesByCallId) + ToolCallRegistry? toolCalls) { - if (toolNamesByCallId is null) + if (toolCalls is null) { return null; } string? toolCallId = GetToolCallId(request); if (toolCallId is null - || !toolNamesByCallId.TryGetValue(toolCallId, out string? resolvedToolName)) + || !toolCalls.TryGetToolName(toolCallId, out string? resolvedToolName)) { return null; } diff --git a/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotTurnRunnerSupport.cs b/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotTurnRunnerSupport.cs index 20e7931..1d6746b 100644 --- a/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotTurnRunnerSupport.cs +++ b/sidecar/src/Aryx.AgentHost/Services/Providers/Copilot/CopilotTurnRunnerSupport.cs @@ -29,7 +29,7 @@ internal sealed class CopilotTurnRunnerSupport : IProviderTurnSupport agent, request, invocation, - state.ToolNamesByCallId, + state.ToolCalls, activity => AgentWorkflowTurnRunner.EmitActivityAsync(command, state, activity, onEvent), onApproval, runCancellation.Token), diff --git a/sidecar/src/Aryx.AgentHost/Services/ToolCallRegistry.cs b/sidecar/src/Aryx.AgentHost/Services/ToolCallRegistry.cs new file mode 100644 index 0000000..3db0fa6 --- /dev/null +++ b/sidecar/src/Aryx.AgentHost/Services/ToolCallRegistry.cs @@ -0,0 +1,165 @@ +using System.Collections.Concurrent; +using System.Diagnostics.CodeAnalysis; +using Aryx.AgentHost.Contracts; + +namespace Aryx.AgentHost.Services; + +internal sealed class ToolCallRegistry +{ + private readonly ConcurrentDictionary _toolExecutionsByCallId = new(StringComparer.Ordinal); + + public bool TryGetExecution(string? toolCallId, [NotNullWhen(true)] out ProviderToolExecutionSnapshot? snapshot) + { + snapshot = null; + return !string.IsNullOrWhiteSpace(toolCallId) + && _toolExecutionsByCallId.TryGetValue(toolCallId, out snapshot); + } + + public bool TryGetToolName(string? toolCallId, [NotNullWhen(true)] out string? toolName) + { + toolName = null; + return TryGetExecution(toolCallId, out ProviderToolExecutionSnapshot? snapshot) + && !string.IsNullOrWhiteSpace(snapshot.ToolName) + && (toolName = snapshot.ToolName) is not null; + } + + public bool HasTrackedArguments(string? toolCallId) + { + return TryGetExecution(toolCallId, out ProviderToolExecutionSnapshot? snapshot) + && snapshot.ToolArguments is { Count: > 0 }; + } + + public void RecordToolStart( + string toolCallId, + string toolName, + IReadOnlyDictionary? toolArguments) + { + _toolExecutionsByCallId.AddOrUpdate( + toolCallId, + static (id, state) => new ProviderToolExecutionSnapshot + { + ToolCallId = id, + ToolName = state.ToolName, + ToolArguments = state.ToolArguments, + Status = ProviderToolExecutionStatus.Running, + }, + static (_, existing, state) => existing with + { + ToolName = state.ToolName, + ToolArguments = state.ToolArguments, + Status = ProviderToolExecutionStatus.Running, + }, + (ToolName: toolName, ToolArguments: toolArguments)); + } + + public bool TryRecordToolRequest( + string? toolCallId, + string toolName, + IReadOnlyDictionary? toolArguments) + { + bool hasToolArguments = toolArguments is { Count: > 0 }; + string? normalizedToolCallId = NormalizeOptionalString(toolCallId); + if (normalizedToolCallId is null) + { + return true; + } + + if (_toolExecutionsByCallId.TryGetValue(normalizedToolCallId, out ProviderToolExecutionSnapshot? existing)) + { + bool trackedHasArguments = existing.ToolArguments is { Count: > 0 }; + if (trackedHasArguments || !hasToolArguments) + { + return false; + } + } + + _toolExecutionsByCallId.AddOrUpdate( + normalizedToolCallId, + id => new ProviderToolExecutionSnapshot + { + ToolCallId = id, + ToolName = toolName, + ToolArguments = toolArguments, + Status = ProviderToolExecutionStatus.Running, + }, + (_, existing) => existing with + { + ToolName = toolName, + ToolArguments = toolArguments, + Status = existing.Status is ProviderToolExecutionStatus.Completed or ProviderToolExecutionStatus.Failed + ? existing.Status + : ProviderToolExecutionStatus.Running, + }); + return true; + } + + public void RecordProgress(string toolCallId, string? progressMessage) + { + string? normalizedProgress = NormalizeOptionalString(progressMessage); + _toolExecutionsByCallId.AddOrUpdate( + toolCallId, + id => new ProviderToolExecutionSnapshot + { + ToolCallId = id, + Status = ProviderToolExecutionStatus.Running, + LatestProgressMessage = normalizedProgress, + }, + (_, existing) => existing with + { + Status = existing.Status is ProviderToolExecutionStatus.Completed or ProviderToolExecutionStatus.Failed + ? existing.Status + : ProviderToolExecutionStatus.Running, + LatestProgressMessage = normalizedProgress ?? existing.LatestProgressMessage, + }); + } + + public void RecordPartialResult(string toolCallId, string? partialOutput) + { + if (string.IsNullOrEmpty(partialOutput)) + { + return; + } + + _toolExecutionsByCallId.AddOrUpdate( + toolCallId, + id => new ProviderToolExecutionSnapshot + { + ToolCallId = id, + Status = ProviderToolExecutionStatus.Running, + PartialOutput = partialOutput, + }, + (_, existing) => existing with + { + Status = existing.Status is ProviderToolExecutionStatus.Completed or ProviderToolExecutionStatus.Failed + ? existing.Status + : ProviderToolExecutionStatus.Running, + PartialOutput = string.Concat(existing.PartialOutput, partialOutput), + }); + } + + public void RecordCompletion(ProviderToolExecutionCompleteEvent toolExecution) + { + _toolExecutionsByCallId.AddOrUpdate( + toolExecution.ToolCallId, + id => new ProviderToolExecutionSnapshot + { + ToolCallId = id, + Status = toolExecution.Success ? ProviderToolExecutionStatus.Completed : ProviderToolExecutionStatus.Failed, + ResultContent = toolExecution.ResultContent, + DetailedResultContent = toolExecution.DetailedResultContent, + Error = toolExecution.Error, + }, + (_, existing) => existing with + { + Status = toolExecution.Success ? ProviderToolExecutionStatus.Completed : ProviderToolExecutionStatus.Failed, + ResultContent = toolExecution.ResultContent ?? existing.ResultContent, + DetailedResultContent = toolExecution.DetailedResultContent ?? existing.DetailedResultContent, + Error = toolExecution.Error ?? existing.Error, + }); + } + + private static string? NormalizeOptionalString(string? value) + { + return string.IsNullOrWhiteSpace(value) ? null : value.Trim(); + } +} diff --git a/sidecar/src/Aryx.AgentHost/Services/TurnExecutionState.cs b/sidecar/src/Aryx.AgentHost/Services/TurnExecutionState.cs index 605e358..9a2b425 100644 --- a/sidecar/src/Aryx.AgentHost/Services/TurnExecutionState.cs +++ b/sidecar/src/Aryx.AgentHost/Services/TurnExecutionState.cs @@ -15,7 +15,6 @@ internal class TurnExecutionState private readonly ConcurrentQueue _pendingEvents = new(); private readonly ConcurrentQueue _pendingMcpOauthRequests = new(); private readonly ConcurrentDictionary _observedAgentsByMessageId = new(StringComparer.Ordinal); - private readonly ConcurrentDictionary _toolExecutionsByCallId = new(StringComparer.Ordinal); private readonly ConcurrentDictionary _reasoningById = new(StringComparer.Ordinal); private readonly ConcurrentDictionary _latestIntentByAgentId = new(StringComparer.Ordinal); private readonly StreamingTranscriptBuffer _transcriptBuffer = new(); @@ -29,9 +28,7 @@ internal class TurnExecutionState _agentSubworkflowIndex = AgentIdentityResolver.BuildAgentSubworkflowIndex(command.Workflow, _workflowLibrary); } - public ConcurrentDictionary ToolNamesByCallId { get; } = new(StringComparer.Ordinal); - - public ConcurrentDictionary ToolCallHasArgumentsById { get; } = new(StringComparer.Ordinal); + public ToolCallRegistry ToolCalls { get; } = new(); public AgentIdentity? ActiveAgent { get; private set; } @@ -166,13 +163,16 @@ internal class TurnExecutionState case ProviderToolExecutionStartEvent toolExecutionStart: string toolCallId = toolExecutionStart.ToolCallId; string toolName = toolExecutionStart.ToolName; - TrackToolCall(toolCallId, toolName, toolExecutionStart.ToolArguments); + bool shouldQueueToolActivity = TrackToolCall(toolCallId, toolName, toolExecutionStart.ToolArguments); ActiveAgent = agent; - AgentActivityEventDto? toolActivity = CreateToolCallingActivity( - agent, toolName, toolCallId, toolExecutionStart.ToolArguments); - if (toolActivity is not null) + if (shouldQueueToolActivity) { - _pendingEvents.Enqueue(toolActivity); + AgentActivityEventDto? toolActivity = CreateToolCallingActivity( + agent, toolName, toolCallId, toolExecutionStart.ToolArguments); + if (toolActivity is not null) + { + _pendingEvents.Enqueue(toolActivity); + } } QueueMessageReclassifiedIfNeeded(_lastObservedMessageId); @@ -338,9 +338,7 @@ internal class TurnExecutionState public bool TryGetToolExecution(string? toolCallId, [NotNullWhen(true)] out ProviderToolExecutionSnapshot? snapshot) { - snapshot = null; - return !string.IsNullOrWhiteSpace(toolCallId) - && _toolExecutionsByCallId.TryGetValue(toolCallId, out snapshot); + return ToolCalls.TryGetExecution(toolCallId, out snapshot); } public bool TryGetReasoning(string? reasoningId, [NotNullWhen(true)] out ProviderReasoningSnapshot? snapshot) @@ -393,102 +391,27 @@ internal class TurnExecutionState _lastObservedMessageId = messageId; } - private void TrackToolCall( + private bool TrackToolCall( string toolCallId, string toolName, IReadOnlyDictionary? toolArguments) { - ToolNamesByCallId[toolCallId] = toolName; - ToolCallHasArgumentsById[toolCallId] = toolArguments is { Count: > 0 }; - TrackToolExecutionStart(toolCallId, toolName, toolArguments); - } - - private void TrackToolExecutionStart( - string toolCallId, - string toolName, - IReadOnlyDictionary? toolArguments) - { - _toolExecutionsByCallId.AddOrUpdate( - toolCallId, - static (id, state) => new ProviderToolExecutionSnapshot - { - ToolCallId = id, - ToolName = state.ToolName, - ToolArguments = state.ToolArguments, - Status = ProviderToolExecutionStatus.Running, - }, - static (_, existing, state) => existing with - { - ToolName = state.ToolName, - ToolArguments = state.ToolArguments, - Status = ProviderToolExecutionStatus.Running, - }, - (ToolName: toolName, ToolArguments: toolArguments)); + return ToolCalls.TryRecordToolRequest(toolCallId, toolName, toolArguments); } private void TrackToolExecutionProgress(string toolCallId, string? progressMessage) { - string? normalizedProgress = NormalizeOptionalString(progressMessage); - _toolExecutionsByCallId.AddOrUpdate( - toolCallId, - id => new ProviderToolExecutionSnapshot - { - ToolCallId = id, - Status = ProviderToolExecutionStatus.Running, - LatestProgressMessage = normalizedProgress, - }, - (_, existing) => existing with - { - Status = existing.Status is ProviderToolExecutionStatus.Completed or ProviderToolExecutionStatus.Failed - ? existing.Status - : ProviderToolExecutionStatus.Running, - LatestProgressMessage = normalizedProgress ?? existing.LatestProgressMessage, - }); + ToolCalls.RecordProgress(toolCallId, progressMessage); } private void TrackToolExecutionPartialResult(string toolCallId, string? partialOutput) { - if (string.IsNullOrEmpty(partialOutput)) - { - return; - } - - _toolExecutionsByCallId.AddOrUpdate( - toolCallId, - id => new ProviderToolExecutionSnapshot - { - ToolCallId = id, - Status = ProviderToolExecutionStatus.Running, - PartialOutput = partialOutput, - }, - (_, existing) => existing with - { - Status = existing.Status is ProviderToolExecutionStatus.Completed or ProviderToolExecutionStatus.Failed - ? existing.Status - : ProviderToolExecutionStatus.Running, - PartialOutput = string.Concat(existing.PartialOutput, partialOutput), - }); + ToolCalls.RecordPartialResult(toolCallId, partialOutput); } private void TrackToolExecutionComplete(ProviderToolExecutionCompleteEvent toolExecution) { - _toolExecutionsByCallId.AddOrUpdate( - toolExecution.ToolCallId, - id => new ProviderToolExecutionSnapshot - { - ToolCallId = id, - Status = toolExecution.Success ? ProviderToolExecutionStatus.Completed : ProviderToolExecutionStatus.Failed, - ResultContent = toolExecution.ResultContent, - DetailedResultContent = toolExecution.DetailedResultContent, - Error = toolExecution.Error, - }, - (_, existing) => existing with - { - Status = toolExecution.Success ? ProviderToolExecutionStatus.Completed : ProviderToolExecutionStatus.Failed, - ResultContent = toolExecution.ResultContent ?? existing.ResultContent, - DetailedResultContent = toolExecution.DetailedResultContent ?? existing.DetailedResultContent, - Error = toolExecution.Error ?? existing.Error, - }); + ToolCalls.RecordCompletion(toolExecution); } private void TrackLatestIntent(string agentId, string? intent) diff --git a/sidecar/src/Aryx.AgentHost/Services/WorkflowRequestInfoInterpreter.cs b/sidecar/src/Aryx.AgentHost/Services/WorkflowRequestInfoInterpreter.cs index a55ebd6..21b9152 100644 --- a/sidecar/src/Aryx.AgentHost/Services/WorkflowRequestInfoInterpreter.cs +++ b/sidecar/src/Aryx.AgentHost/Services/WorkflowRequestInfoInterpreter.cs @@ -1,4 +1,3 @@ -using System.Collections.Concurrent; using System.Text.Json; using Aryx.AgentHost.Contracts; using Microsoft.Agents.AI.Workflows; @@ -20,8 +19,7 @@ internal static class WorkflowRequestInfoInterpreter RunTurnCommandDto command, RequestInfoEvent requestInfo, AgentIdentity? activeAgent, - ConcurrentDictionary toolNamesByCallId, - ConcurrentDictionary toolCallHasArgumentsById) + ToolCallRegistry toolCalls) { RequestInterpretation interpretation = InterpretRequest(command, requestInfo); return interpretation switch @@ -29,7 +27,7 @@ internal static class WorkflowRequestInfoInterpreter HandoffRequestInterpretation handoff => CreateHandoffActivity(command, handoff.TargetAgent, activeAgent), ToolRequestInterpretation tool when activeAgent.HasValue => - CreateToolCallingActivity(command, activeAgent.Value, tool, toolNamesByCallId, toolCallHasArgumentsById), + CreateToolCallingActivity(command, activeAgent.Value, tool, toolCalls), _ => null, }; } @@ -66,22 +64,13 @@ internal static class WorkflowRequestInfoInterpreter RunTurnCommandDto command, AgentIdentity activeAgent, ToolRequestInterpretation tool, - ConcurrentDictionary toolNamesByCallId, - ConcurrentDictionary toolCallHasArgumentsById) + ToolCallRegistry toolCalls) { - bool hasToolArguments = tool.ToolArguments is { Count: > 0 }; - if (tool.ToolCallId is not null && toolNamesByCallId.ContainsKey(tool.ToolCallId)) + if (!toolCalls.TryRecordToolRequest(tool.ToolCallId, tool.ToolName, tool.ToolArguments)) { - bool trackedHasArguments = toolCallHasArgumentsById.TryGetValue(tool.ToolCallId, out bool hasTrackedArguments) - && hasTrackedArguments; - if (trackedHasArguments || !hasToolArguments) - { - return null; - } + return null; } - TrackToolCallId(toolNamesByCallId, toolCallHasArgumentsById, tool.ToolCallId, tool.ToolName, hasToolArguments); - return new AgentActivityEventDto { Type = "agent-activity", @@ -98,20 +87,6 @@ internal static class WorkflowRequestInfoInterpreter }; } - private static void TrackToolCallId( - ConcurrentDictionary toolNamesByCallId, - ConcurrentDictionary toolCallHasArgumentsById, - string? toolCallId, - string toolName, - bool hasToolArguments) - { - if (toolCallId is not null) - { - toolNamesByCallId[toolCallId] = toolName; - toolCallHasArgumentsById[toolCallId] = hasToolArguments; - } - } - private static RequestInterpretation InterpretRequest( RunTurnCommandDto command, RequestInfoEvent requestInfo) diff --git a/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs b/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs index be08f4b..a2a9430 100644 --- a/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs +++ b/sidecar/tests/Aryx.AgentHost.Tests/CopilotTurnExecutionStateTests.cs @@ -110,10 +110,9 @@ public sealed class CopilotTurnExecutionStateTests Assert.Equal("tool-call-1", toolActivity.ToolCallId); Assert.NotNull(toolActivity.ToolArguments); Assert.Equal("/src/main.ts", toolActivity.ToolArguments["path"]); - Assert.True(state.ToolNamesByCallId.TryGetValue("tool-call-1", out string? toolName)); + Assert.True(state.ToolCalls.TryGetToolName("tool-call-1", out string? toolName)); Assert.Equal("view", toolName); - Assert.True(state.ToolCallHasArgumentsById.TryGetValue("tool-call-1", out bool hasArguments)); - Assert.True(hasArguments); + Assert.True(state.ToolCalls.HasTrackedArguments("tool-call-1")); } [Fact] @@ -129,8 +128,7 @@ public sealed class CopilotTurnExecutionStateTests AgentActivityEventDto toolActivity = Assert.Single(state.DrainPendingEvents().OfType()); Assert.Null(toolActivity.ToolArguments); - Assert.True(state.ToolCallHasArgumentsById.TryGetValue("tool-call-1", out bool hasArguments)); - Assert.False(hasArguments); + Assert.False(state.ToolCalls.HasTrackedArguments("tool-call-1")); } [Fact] @@ -253,10 +251,9 @@ public sealed class CopilotTurnExecutionStateTests """{"type":"tool.execution_start","data":{"toolCallId":"tool-call-1","toolName":"handoff_to_specialist"},"id":"1ce9d1dc-68f1-4df5-9728-f97017233279","timestamp":"2026-03-27T00:00:00Z"}""")); Assert.Empty(state.DrainPendingEvents().OfType()); - Assert.True(state.ToolNamesByCallId.TryGetValue("tool-call-1", out string? toolName)); + Assert.True(state.ToolCalls.TryGetToolName("tool-call-1", out string? toolName)); Assert.Equal("handoff_to_specialist", toolName); - Assert.True(state.ToolCallHasArgumentsById.TryGetValue("tool-call-1", out bool hasArguments)); - Assert.False(hasArguments); + Assert.False(state.ToolCalls.HasTrackedArguments("tool-call-1")); } [Fact] @@ -375,14 +372,12 @@ public sealed class CopilotTurnExecutionStateTests Assert.Contains(toolActivities, activity => activity.ToolCallId == "tool-call-2" && activity.ToolName == "view"); MessageReclassifiedEventDto reclassified = Assert.Single(pending.OfType()); Assert.Equal("msg-3", reclassified.MessageId); - Assert.True(state.ToolNamesByCallId.TryGetValue("tool-call-1", out string? firstToolName)); + Assert.True(state.ToolCalls.TryGetToolName("tool-call-1", out string? firstToolName)); Assert.Equal("rg", firstToolName); - Assert.True(state.ToolNamesByCallId.TryGetValue("tool-call-2", out string? secondToolName)); + Assert.True(state.ToolCalls.TryGetToolName("tool-call-2", out string? secondToolName)); Assert.Equal("view", secondToolName); - Assert.True(state.ToolCallHasArgumentsById.TryGetValue("tool-call-1", out bool firstHasArguments)); - Assert.False(firstHasArguments); - Assert.True(state.ToolCallHasArgumentsById.TryGetValue("tool-call-2", out bool secondHasArguments)); - Assert.False(secondHasArguments); + Assert.False(state.ToolCalls.HasTrackedArguments("tool-call-1")); + Assert.False(state.ToolCalls.HasTrackedArguments("tool-call-2")); } [Fact] diff --git a/sidecar/tests/Aryx.AgentHost.Tests/CopilotWorkflowRunnerTests.cs b/sidecar/tests/Aryx.AgentHost.Tests/CopilotWorkflowRunnerTests.cs index 4bb83e9..450ea68 100644 --- a/sidecar/tests/Aryx.AgentHost.Tests/CopilotWorkflowRunnerTests.cs +++ b/sidecar/tests/Aryx.AgentHost.Tests/CopilotWorkflowRunnerTests.cs @@ -776,8 +776,165 @@ public sealed class CopilotWorkflowRunnerTests Assert.Equal("tool-call-1", activity.ToolCallId); Assert.NotNull(activity.ToolArguments); Assert.Equal(@"C:\workspace\README.md", activity.ToolArguments["path"]); - Assert.True(state.ToolCallHasArgumentsById.TryGetValue("tool-call-1", out bool hasArguments)); - Assert.True(hasArguments); + Assert.True(state.ToolCalls.HasTrackedArguments("tool-call-1")); + } + + [Fact] + public async Task ObserveSessionEvent_ToolExecutionStart_DoesNotDuplicateTrackedRequestInfoActivity() + { + RunTurnCommandDto command = CreateApprovalCommand(); + CopilotTurnExecutionState state = new(command); + WorkflowNodeDto agent = CreateAgent("agent-1", "Primary"); + state.ObserveSessionEvent( + agent, + SessionEvent.FromJson( + """ + { + "type": "assistant.message_delta", + "data": { + "messageId": "msg-1", + "deltaContent": "Inspecting" + }, + "id": "b61652d1-120e-4a9f-8f0e-1dbf04fb18da", + "timestamp": "2026-03-27T00:00:00Z" + } + """)); + _ = state.DrainPendingEvents(); + + RequestInfoEvent requestInfo = CreateRequestInfoEvent( + new FunctionCallContent("tool-call-1", "view", new Dictionary + { + ["path"] = @"C:\workspace\README.md", + })); + List requestActivities = []; + + MethodInfo handleWorkflowEvent = typeof(CopilotWorkflowRunner).GetMethod( + "HandleWorkflowEventAsync", + BindingFlags.NonPublic | BindingFlags.Static)!; + Task handleTask = (Task)handleWorkflowEvent.Invoke( + null, + [ + command, + requestInfo, + Array.Empty(), + state, + (Func)(_ => Task.CompletedTask), + (Func)(sidecarEvent => + { + requestActivities.Add(Assert.IsType(sidecarEvent)); + return Task.CompletedTask; + }), + ])!; + + bool shouldEndTurn = await handleTask; + + Assert.False(shouldEndTurn); + AgentActivityEventDto requestActivity = Assert.Single(requestActivities); + Assert.Equal("tool-calling", requestActivity.ActivityType); + Assert.True(state.ToolCalls.HasTrackedArguments("tool-call-1")); + + state.ObserveSessionEvent( + agent, + SessionEvent.FromJson( + """ + { + "type": "tool.execution_start", + "data": { + "toolCallId": "tool-call-1", + "toolName": "view", + "arguments": { + "path": "C:\\workspace\\README.md" + } + }, + "id": "c61652d1-120e-4a9f-8f0e-1dbf04fb18da", + "timestamp": "2026-03-27T00:00:01Z" + } + """)); + + IReadOnlyList pending = state.DrainPendingEvents(); + Assert.DoesNotContain( + pending.OfType(), + activity => activity.ActivityType == "tool-calling"); + + MessageReclassifiedEventDto reclassified = Assert.Single(pending.OfType()); + Assert.Equal("msg-1", reclassified.MessageId); + } + + [Fact] + public async Task ObserveSessionEvent_ToolExecutionStart_EmitsEnrichmentWhenRequestInfoWasMissingArguments() + { + RunTurnCommandDto command = CreateApprovalCommand(); + CopilotTurnExecutionState state = new(command); + WorkflowNodeDto agent = CreateAgent("agent-1", "Primary"); + state.ObserveSessionEvent( + agent, + SessionEvent.FromJson( + """ + { + "type": "assistant.message_delta", + "data": { + "messageId": "msg-2", + "deltaContent": "Inspecting" + }, + "id": "d61652d1-120e-4a9f-8f0e-1dbf04fb18da", + "timestamp": "2026-03-27T00:00:00Z" + } + """)); + _ = state.DrainPendingEvents(); + + RequestInfoEvent requestInfo = CreateRequestInfoEvent( + new FunctionCallContent("tool-call-1", "view", new Dictionary())); + List requestActivities = []; + + MethodInfo handleWorkflowEvent = typeof(CopilotWorkflowRunner).GetMethod( + "HandleWorkflowEventAsync", + BindingFlags.NonPublic | BindingFlags.Static)!; + Task handleTask = (Task)handleWorkflowEvent.Invoke( + null, + [ + command, + requestInfo, + Array.Empty(), + state, + (Func)(_ => Task.CompletedTask), + (Func)(sidecarEvent => + { + requestActivities.Add(Assert.IsType(sidecarEvent)); + return Task.CompletedTask; + }), + ])!; + + bool shouldEndTurn = await handleTask; + + Assert.False(shouldEndTurn); + AgentActivityEventDto requestActivity = Assert.Single(requestActivities); + Assert.Null(requestActivity.ToolArguments); + Assert.False(state.ToolCalls.HasTrackedArguments("tool-call-1")); + + state.ObserveSessionEvent( + agent, + SessionEvent.FromJson( + """ + { + "type": "tool.execution_start", + "data": { + "toolCallId": "tool-call-1", + "toolName": "view", + "arguments": { + "path": "C:\\workspace\\README.md" + } + }, + "id": "e61652d1-120e-4a9f-8f0e-1dbf04fb18da", + "timestamp": "2026-03-27T00:00:01Z" + } + """)); + + AgentActivityEventDto enrichment = Assert.Single( + state.DrainPendingEvents().OfType(), + activity => activity.ActivityType == "tool-calling"); + Assert.NotNull(enrichment.ToolArguments); + Assert.Equal(@"C:\workspace\README.md", enrichment.ToolArguments["path"]); + Assert.True(state.ToolCalls.HasTrackedArguments("tool-call-1")); } [Fact] @@ -1372,14 +1529,12 @@ public sealed class CopilotWorkflowRunnerTests [Fact] public void TryGetApprovalToolName_UsesToolCallLookupForPermissionCategoriesWithoutDirectToolNames() { - Dictionary toolNamesByCallId = new(StringComparer.Ordinal) - { - ["tool-call-url"] = "web_fetch", - ["tool-call-shell"] = "shell", - ["tool-call-read"] = "view", - ["tool-call-write"] = "write_file", - ["tool-call-memory"] = "store_memory", - }; + ToolCallRegistry toolCalls = CreateToolCallRegistry( + ("tool-call-url", "web_fetch"), + ("tool-call-shell", "shell"), + ("tool-call-read", "view"), + ("tool-call-write", "write_file"), + ("tool-call-memory", "store_memory")); Assert.True( CopilotApprovalCoordinator.TryGetApprovalToolName( @@ -1390,7 +1545,7 @@ public sealed class CopilotWorkflowRunnerTests Intention = "Fetch the requested page", Url = "https://example.com/docs", }, - toolNamesByCallId, + toolCalls, out string? urlToolName)); Assert.Equal("web_fetch", urlToolName); @@ -1408,7 +1563,7 @@ public sealed class CopilotWorkflowRunnerTests HasWriteFileRedirection = false, CanOfferSessionApproval = false, }, - toolNamesByCallId, + toolCalls, out string? shellToolName)); Assert.Equal("shell", shellToolName); @@ -1421,7 +1576,7 @@ public sealed class CopilotWorkflowRunnerTests Intention = "Inspect a file", Path = "README.md", }, - toolNamesByCallId, + toolCalls, out string? readToolName)); Assert.Equal("view", readToolName); @@ -1435,7 +1590,7 @@ public sealed class CopilotWorkflowRunnerTests FileName = "README.md", Diff = "@@ -1 +1 @@", }, - toolNamesByCallId, + toolCalls, out string? writeToolName)); Assert.Equal("write_file", writeToolName); @@ -1449,7 +1604,7 @@ public sealed class CopilotWorkflowRunnerTests Fact = "Use Bun for script execution.", Citations = "package.json", }, - toolNamesByCallId, + toolCalls, out string? memoryToolName)); Assert.Equal("store_memory", memoryToolName); } @@ -1960,7 +2115,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal), + CreateToolCallRegistry(), approval => { observedApproval = approval; @@ -2007,10 +2162,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal) - { - ["tool-call-write-1"] = "apply_patch", - }, + CreateToolCallRegistry(("tool-call-write-1", "apply_patch")), activity => { observedActivity = activity; @@ -2067,7 +2219,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal), + CreateToolCallRegistry(), approval => { sawApproval = true; @@ -2104,7 +2256,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal), + CreateToolCallRegistry(), approval => { sawApproval = true; @@ -2137,10 +2289,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal) - { - ["tool-call-read-1"] = "view", - }, + CreateToolCallRegistry(("tool-call-read-1", "view")), approval => { firstApproval = approval; @@ -2178,10 +2327,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal) - { - ["tool-call-read-2"] = "grep", - }, + CreateToolCallRegistry(("tool-call-read-2", "grep")), approval => { sawSecondApproval = true; @@ -2214,10 +2360,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal) - { - ["tool-call-read-1"] = "view", - }, + CreateToolCallRegistry(("tool-call-read-1", "view")), approval => { firstApproval = approval; @@ -2254,10 +2397,7 @@ public sealed class CopilotWorkflowRunnerTests { SessionId = "copilot-session-1", }, - new Dictionary(StringComparer.Ordinal) - { - ["tool-call-read-2"] = "grep", - }, + CreateToolCallRegistry(("tool-call-read-2", "grep")), approval => { secondApproval = approval; @@ -2548,6 +2688,17 @@ public sealed class CopilotWorkflowRunnerTests }; } + private static ToolCallRegistry CreateToolCallRegistry(params (string ToolCallId, string ToolName)[] toolCalls) + { + ToolCallRegistry registry = new(); + foreach ((string toolCallId, string toolName) in toolCalls) + { + registry.RecordToolStart(toolCallId, toolName, toolArguments: null); + } + + return registry; + } + private static RunTurnCommandDto CreateRequestPortCommand() { return new RunTurnCommandDto diff --git a/sidecar/tests/Aryx.AgentHost.Tests/WorkflowRequestInfoInterpreterTests.cs b/sidecar/tests/Aryx.AgentHost.Tests/WorkflowRequestInfoInterpreterTests.cs index 173eead..c2d886a 100644 --- a/sidecar/tests/Aryx.AgentHost.Tests/WorkflowRequestInfoInterpreterTests.cs +++ b/sidecar/tests/Aryx.AgentHost.Tests/WorkflowRequestInfoInterpreterTests.cs @@ -1,5 +1,4 @@ using System.Collections; -using System.Collections.Concurrent; using System.Runtime.CompilerServices; using System.Text.Json; using Aryx.AgentHost.Contracts; @@ -27,8 +26,7 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("tool-calling", activity.ActivityType); @@ -38,8 +36,9 @@ public sealed class WorkflowRequestInfoInterpreterTests Assert.NotNull(activity.ToolArguments); Assert.Equal(@"C:\workspace\file.txt", activity.ToolArguments["path"]); Assert.Equal([10, 25], Assert.IsAssignableFrom>(activity.ToolArguments["viewRange"])); - Assert.Equal("view", tracking.ToolNamesByCallId["call-1"]); - Assert.True(tracking.ToolCallHasArgumentsById["call-1"]); + Assert.True(tracking.TryGetToolName("call-1", out string? toolName)); + Assert.Equal("view", toolName); + Assert.True(tracking.HasTrackedArguments("call-1")); } [Fact] @@ -61,8 +60,7 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("tool-calling", activity.ActivityType); @@ -70,8 +68,9 @@ public sealed class WorkflowRequestInfoInterpreterTests Assert.NotNull(activity.ToolArguments); Assert.Equal(@"C:\workspace", activity.ToolArguments["path"]); Assert.Equal(true, activity.ToolArguments["includeIgnored"]); - Assert.Equal("git.status", tracking.ToolNamesByCallId["call-1"]); - Assert.True(tracking.ToolCallHasArgumentsById["call-1"]); + Assert.True(tracking.TryGetToolName("call-1", out string? toolName)); + Assert.Equal("git.status", toolName); + Assert.True(tracking.HasTrackedArguments("call-1")); } [Fact] @@ -85,8 +84,7 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("tool-calling", activity.ActivityType); @@ -95,8 +93,9 @@ public sealed class WorkflowRequestInfoInterpreterTests Assert.Equal( ["print('hello')"], Assert.IsAssignableFrom>(activity.ToolArguments["inputs"])); - Assert.Equal("code interpreter", tracking.ToolNamesByCallId["call-1"]); - Assert.True(tracking.ToolCallHasArgumentsById["call-1"]); + Assert.True(tracking.TryGetToolName("call-1", out string? toolName)); + Assert.Equal("code interpreter", toolName); + Assert.True(tracking.HasTrackedArguments("call-1")); } [Fact] @@ -109,15 +108,14 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("tool-calling", activity.ActivityType); Assert.Equal("image generation", activity.ToolName); Assert.Null(activity.ToolArguments); - Assert.Empty(tracking.ToolNamesByCallId); - Assert.Empty(tracking.ToolCallHasArgumentsById); + Assert.False(tracking.TryGetToolName("call-1", out _)); + Assert.False(tracking.HasTrackedArguments("call-1")); } [Fact] @@ -135,12 +133,11 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Null(activity.ToolArguments); - Assert.False(tracking.ToolCallHasArgumentsById["call-1"]); + Assert.False(tracking.HasTrackedArguments("call-1")); } [Fact] @@ -160,21 +157,25 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.NotNull(activity.ToolArguments); Assert.Equal("[truncated]", activity.ToolArguments["command"]); - Assert.True(tracking.ToolCallHasArgumentsById["call-1"]); + Assert.True(tracking.HasTrackedArguments("call-1")); } [Fact] public void TryCreateActivityFromRequest_SkipsDuplicateTrackedToolCallIdsThatAlreadyHaveArguments() { var tracking = CreateToolTracking(); - tracking.ToolNamesByCallId["call-1"] = "view"; - tracking.ToolCallHasArgumentsById["call-1"] = true; + tracking.RecordToolStart( + "call-1", + "view", + new Dictionary + { + ["path"] = @"C:\workspace\seed.txt", + }); RequestInfoEvent requestInfo = CreateRequestInfoEvent( new FunctionCallContent("call-1", "view", new Dictionary { @@ -185,20 +186,19 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.Null(activity); - Assert.Equal("view", tracking.ToolNamesByCallId["call-1"]); - Assert.True(tracking.ToolCallHasArgumentsById["call-1"]); + Assert.True(tracking.TryGetToolName("call-1", out string? toolName)); + Assert.Equal("view", toolName); + Assert.True(tracking.HasTrackedArguments("call-1")); } [Fact] public void TryCreateActivityFromRequest_EmitsEnrichmentWhenTrackedToolCallWasMissingArguments() { var tracking = CreateToolTracking(); - tracking.ToolNamesByCallId["call-1"] = "view"; - tracking.ToolCallHasArgumentsById["call-1"] = false; + tracking.RecordToolStart("call-1", "view", toolArguments: null); RequestInfoEvent requestInfo = CreateRequestInfoEvent( new FunctionCallContent("call-1", "view", new Dictionary { @@ -209,16 +209,16 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateSingleAgentCommand(), requestInfo, new AgentIdentity("agent-1", "Primary"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("tool-calling", activity.ActivityType); Assert.Equal("call-1", activity.ToolCallId); Assert.NotNull(activity.ToolArguments); Assert.Equal(@"C:\workspace\file.txt", activity.ToolArguments["path"]); - Assert.Equal("view", tracking.ToolNamesByCallId["call-1"]); - Assert.True(tracking.ToolCallHasArgumentsById["call-1"]); + Assert.True(tracking.TryGetToolName("call-1", out string? toolName)); + Assert.Equal("view", toolName); + Assert.True(tracking.HasTrackedArguments("call-1")); } [Fact] @@ -232,8 +232,7 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateHandoffCommand(), requestInfo, new AgentIdentity("agent-handoff-triage", "Triage"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("handoff", activity.ActivityType); @@ -242,8 +241,8 @@ public sealed class WorkflowRequestInfoInterpreterTests Assert.Equal("agent-handoff-triage", activity.SourceAgentId); Assert.Equal("Triage", activity.SourceAgentName); Assert.Null(activity.ToolName); - Assert.Empty(tracking.ToolNamesByCallId); - Assert.Empty(tracking.ToolCallHasArgumentsById); + Assert.False(tracking.TryGetToolName("call-1", out _)); + Assert.False(tracking.HasTrackedArguments("call-1")); } [Fact] @@ -263,8 +262,7 @@ public sealed class WorkflowRequestInfoInterpreterTests "agent-reviewer", "Reviewer", new SubworkflowContext("subworkflow-review", "Review Lane")), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("tool-calling", activity.ActivityType); @@ -283,8 +281,7 @@ public sealed class WorkflowRequestInfoInterpreterTests CreateHandoffCommandWithReferencedSubworkflow(), requestInfo, new AgentIdentity("agent-handoff-triage", "Triage"), - tracking.ToolNamesByCallId, - tracking.ToolCallHasArgumentsById); + tracking); Assert.NotNull(activity); Assert.Equal("handoff", activity.ActivityType); @@ -417,10 +414,7 @@ public sealed class WorkflowRequestInfoInterpreterTests workflowLibrary: [nestedWorkflow]); } - private static ( - ConcurrentDictionary ToolNamesByCallId, - ConcurrentDictionary ToolCallHasArgumentsById) CreateToolTracking() - => (new(StringComparer.Ordinal), new(StringComparer.Ordinal)); + private static ToolCallRegistry CreateToolTracking() => new(); private static RunTurnCommandDto CreateCommand( string orchestrationMode,