refactor: unify tool call stream tracking

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-09-13 13:21:59 +02:00 · 2026-04-13 15:33:27 +02:00
parent 7b9c4d140c
commit 4e34d2abfc
9 changed files with 438 additions and 236 deletions
@@ -326,8 +326,7 @@ public class AgentWorkflowTurnRunner : ITurnWorkflowRunner
                command,
                requestInfo,
                state.ActiveAgent,
-                state.ToolNamesByCallId,
-                state.ToolCallHasArgumentsById);
+                state.ToolCalls);

            if (activity is null)
            {
@@ -70,7 +70,7 @@ internal sealed class CopilotApprovalCoordinator
        WorkflowNodeDto agent,
        PermissionRequest request,
        PermissionInvocation invocation,
-        IReadOnlyDictionary<string, string> toolNamesByCallId,
+        ToolCallRegistry toolCalls,
        Func<ApprovalRequestedEventDto, Task> 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<string, string> toolNamesByCallId,
+        ToolCallRegistry toolCalls,
        Func<AgentActivityEventDto, Task>? onActivity,
        Func<ApprovalRequestedEventDto, Task> 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<string, string>? 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<string, string>? 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<string, string>? 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;
        }
@@ -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),
@@ -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<string, ProviderToolExecutionSnapshot> _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<string, object?>? 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<string, object?>? 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();
+    }
+}
@@ -15,7 +15,6 @@ internal class TurnExecutionState
    private readonly ConcurrentQueue<SidecarEventDto> _pendingEvents = new();
    private readonly ConcurrentQueue<McpOauthRequiredEventDto> _pendingMcpOauthRequests = new();
    private readonly ConcurrentDictionary<string, AgentIdentity> _observedAgentsByMessageId = new(StringComparer.Ordinal);
-    private readonly ConcurrentDictionary<string, ProviderToolExecutionSnapshot> _toolExecutionsByCallId = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, ProviderReasoningSnapshot> _reasoningById = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, string> _latestIntentByAgentId = new(StringComparer.Ordinal);
    private readonly StreamingTranscriptBuffer _transcriptBuffer = new();
@@ -29,9 +28,7 @@ internal class TurnExecutionState
        _agentSubworkflowIndex = AgentIdentityResolver.BuildAgentSubworkflowIndex(command.Workflow, _workflowLibrary);
    }

-    public ConcurrentDictionary<string, string> ToolNamesByCallId { get; } = new(StringComparer.Ordinal);
-
-    public ConcurrentDictionary<string, bool> 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<string, object?>? toolArguments)
    {
-        ToolNamesByCallId[toolCallId] = toolName;
-        ToolCallHasArgumentsById[toolCallId] = toolArguments is { Count: > 0 };
-        TrackToolExecutionStart(toolCallId, toolName, toolArguments);
-    }
-
-    private void TrackToolExecutionStart(
-        string toolCallId,
-        string toolName,
-        IReadOnlyDictionary<string, object?>? 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)
@@ -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<string, string> toolNamesByCallId,
-        ConcurrentDictionary<string, bool> 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<string, string> toolNamesByCallId,
-        ConcurrentDictionary<string, bool> 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<string, string> toolNamesByCallId,
-        ConcurrentDictionary<string, bool> 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)