Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,12 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA

if (cappedRequests is not null)
{
this.EnsureRequestIdsUnique(
cappedRequests,
state,
session,
includeSurfacedRequests: recordedAnyCappedRequest);

// The first update carrying requests supersedes any earlier batch; later updates in this
// same turn add to it, since every one of them reaches the caller.
if (recordedAnyCappedRequest)
Expand Down Expand Up @@ -348,6 +354,8 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
yield break;
}

this.EnsureRequestIdsUnique(streamedApprovalRequests, state, session, includeSurfacedRequests: false);

// 4. Classify the collected approval requests against standing rules and auto-approval rules.
List<ToolApprovalRequestContent> unapproved = [];
foreach (var tarc in streamedApprovalRequests)
Expand Down Expand Up @@ -418,6 +426,7 @@ private void RecordSurfacedApprovalRequestsFromMessages(

if (requests is not null)
{
this.EnsureRequestIdsUnique(requests, state, session, includeSurfacedRequests: false);
ResetSurfacedApprovalRequests(state, requests);
this._sessionState.SaveState(session, state);
}
Expand All @@ -431,12 +440,6 @@ private void RecordSurfacedApprovalRequestsFromMessages(
/// <remarks>
/// A new batch from the inner agent supersedes any previous one, so stale entries from an abandoned
/// approval cycle cannot later authorize a standing rule.
/// <para>
/// Request ids are unique by construction, so keying by id loses nothing: they derive from tool call ids, and
/// inference services correlate a tool call to its result by that id alone. A duplicate id would already have
/// broken that correlation upstream. Nothing a caller sends reaches this dictionary — every entry originates
/// from the inner agent's own response — so a caller cannot manufacture a collision here.
/// </para>
/// </remarks>
private static void ResetSurfacedApprovalRequests(ToolApprovalState state, IReadOnlyList<ToolApprovalRequestContent> requests)
{
Expand Down Expand Up @@ -477,6 +480,61 @@ private static ToolApprovalRequestContent SnapshotRequest(ToolApprovalRequestCon
return request;
}

/// <summary>
/// Rejects duplicate active approval request IDs.
/// </summary>
private void EnsureRequestIdsUnique(
List<ToolApprovalRequestContent> requests,
ToolApprovalState state,
AgentSession? session,
bool includeSurfacedRequests)
{
if (requests.Count == 0)
{
return;
}

var knownRequestIds = new HashSet<string>(StringComparer.Ordinal);
if (includeSurfacedRequests)
{
foreach (string requestId in state.SurfacedApprovalRequests.Keys)
{
knownRequestIds.Add(requestId);
}
}

HashSet<string>? duplicateRequestIds = null;

foreach (var request in requests)
{
if (knownRequestIds.Add(request.RequestId))
{
continue;
}

(duplicateRequestIds ??= new(StringComparer.Ordinal)).Add(request.RequestId);
}

if (duplicateRequestIds is null || duplicateRequestIds.Count == 0)
{
return;
}

foreach (string requestId in duplicateRequestIds)
{
state.SurfacedApprovalRequests.Remove(requestId);
state.QueuedApprovalRequests.RemoveAll(
queued => string.Equals(queued.RequestId, requestId, StringComparison.Ordinal));
state.CollectedApprovalResponses.RemoveAll(
response => string.Equals(response.RequestId, requestId, StringComparison.Ordinal));
}

this._sessionState.SaveState(session, state);

throw new InvalidOperationException(
$"Tool approval request IDs '{string.Join("', '", duplicateRequestIds.OrderBy(id => id, StringComparer.Ordinal))}' were used more than once.");
}

/// <summary>
/// Re-evaluates queued approval requests against current rules and auto-approval rules, and auto-approves any that now match.
/// </summary>
Expand Down Expand Up @@ -606,6 +664,12 @@ private async ValueTask<bool> ProcessAndQueueOutboundApprovalRequestsAsync(
AgentRunOptions? options,
IReadOnlyCollection<ChatMessage> requestMessages)
{
var requests = responseMessages
.SelectMany(message => message.Contents)
.OfType<ToolApprovalRequestContent>()
.ToList();
this.EnsureRequestIdsUnique(requests, state, session, includeSurfacedRequests: false);

// Pass 1: Scan all response messages and classify each approval request.
// Auto-approved requests (matching a standing rule or auto-approval rule) have their
// responses collected immediately, preserving the original request order, and are
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3285,6 +3285,205 @@ public async Task RunStreamingAsync_AutoApprovalRule_ApprovesMatchingToolAsync()

#endregion

#region Duplicate approval request IDs

/// <summary>
/// Verify that one request ID cannot refer to two different tool calls.
/// </summary>
[Fact]
public async Task RunAsync_DifferentCallsWithSameRequestId_ThrowsAsync()
{
// Arrange
var harmless = new ToolApprovalRequestContent(
"shared-id",
new FunctionCallContent("shared-call", "BenignLookup", new Dictionary<string, object?> { ["query"] = "weather" }));
var privileged = new ToolApprovalRequestContent(
"shared-id",
new FunctionCallContent("shared-call", "RunPrivilegedShell", new Dictionary<string, object?> { ["command"] = "sensitive-command" }));
var innerAgent = CreateMockAgent(new AgentResponse([new ChatMessage(ChatRole.Assistant, [harmless, privileged])]));
var agent = new ToolApprovalAgent(innerAgent.Object);

// Act
var exception = await Assert.ThrowsAsync<InvalidOperationException>(
() => agent.RunAsync([new ChatMessage(ChatRole.User, "Go")], new ChatClientAgentSession()));

// Assert
Assert.Contains("shared-id", exception.Message, StringComparison.Ordinal);
Assert.DoesNotContain("RunPrivilegedShell", exception.Message, StringComparison.Ordinal);
Assert.DoesNotContain("sensitive-command", exception.Message, StringComparison.Ordinal);
}

/// <summary>
/// Verify that repeating the exact same request under one ID is rejected.
/// </summary>
[Fact]
public async Task RunAsync_EquivalentCallsWithSameRequestId_ThrowsAsync()
{
// Arrange
var first = new ToolApprovalRequestContent(
"shared-id",
new FunctionCallContent("shared-call", "BenignLookup", new Dictionary<string, object?> { ["query"] = "weather" }));
var repeated = new ToolApprovalRequestContent(
"shared-id",
new FunctionCallContent("shared-call", "BenignLookup", new Dictionary<string, object?> { ["query"] = "weather" }));
var agent = new ToolApprovalAgent(
CreateMockAgent(new AgentResponse([new ChatMessage(ChatRole.Assistant, [first, repeated])])).Object);

// Act & Assert
await Assert.ThrowsAsync<InvalidOperationException>(
() => agent.RunAsync([new ChatMessage(ChatRole.User, "Go")], new ChatClientAgentSession()));
}

/// <summary>
/// Verify that streaming batches are checked before any conflicting request is surfaced.
/// </summary>
[Fact]
public async Task RunStreamingAsync_DifferentCallsWithSameRequestId_ThrowsBeforeYieldAsync()
{
// Arrange
var harmless = new ToolApprovalRequestContent("shared-id", new FunctionCallContent("shared-call", "BenignLookup"));
var privileged = new ToolApprovalRequestContent("shared-id", new FunctionCallContent("shared-call", "RunPrivilegedShell"));
var innerAgent = CreateMockStreamingAgent(
[
new AgentResponseUpdate(ChatRole.Assistant, new List<AIContent> { harmless }),
new AgentResponseUpdate(ChatRole.Assistant, new List<AIContent> { privileged }),
]);
var agent = new ToolApprovalAgent(innerAgent.Object);
var yielded = new List<AgentResponseUpdate>();

// Act
var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () =>
{
await foreach (var update in agent.RunStreamingAsync(
[new ChatMessage(ChatRole.User, "Go")],
new ChatClientAgentSession()))
{
yielded.Add(update);
}
});

// Assert
Assert.Empty(yielded);
Assert.Contains("shared-id", exception.Message, StringComparison.Ordinal);
}

/// <summary>
/// Verify that the final non-streaming turn at the auto-approval cap also rejects conflicting IDs.
/// </summary>
[Fact]
public async Task RunAsync_CappedResponseWithDifferentCallsUnderSameRequestId_ThrowsAsync()
{
// Arrange
var callCount = 0;
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(() => ++callCount == 1
? new AgentResponse([new ChatMessage(ChatRole.Assistant,
[new ToolApprovalRequestContent("auto", new FunctionCallContent("auto-call", "AutoTool"))])])
: new AgentResponse([new ChatMessage(ChatRole.Assistant,
[
new ToolApprovalRequestContent("shared-id", new FunctionCallContent("shared-call", "BenignLookup")),
new ToolApprovalRequestContent("shared-id", new FunctionCallContent("shared-call", "RunPrivilegedShell")),
])]));
var agent = new ToolApprovalAgent(innerAgent.Object, new ToolApprovalAgentOptions
{
AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule],
MaxAutoApprovalIterations = 1,
});

// Act & Assert
await Assert.ThrowsAsync<InvalidOperationException>(
() => agent.RunAsync([new ChatMessage(ChatRole.User, "Go")], new ChatClientAgentSession()));
}

/// <summary>
/// Verify that the final streaming turn at the auto-approval cap invalidates every ID repeated by a later update.
/// </summary>
[Fact]
public async Task RunStreamingAsync_CappedUpdatesWithMultipleDuplicateIds_InvalidatesAllAsync()
{
// Arrange
var callCount = 0;
var firstA = new ToolApprovalRequestContent("id-a", new FunctionCallContent("call-a", "ToolA"));
var firstB = new ToolApprovalRequestContent("id-b", new FunctionCallContent("call-b", "ToolB"));
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<IAsyncEnumerable<AgentResponseUpdate>>("RunCoreStreamingAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.Returns(() => ++callCount == 1
? ToAsyncEnumerableAsync(
[
new AgentResponseUpdate(ChatRole.Assistant, new List<AIContent>
{
new ToolApprovalRequestContent("auto", new FunctionCallContent("auto-call", "AutoTool")),
}),
])
: ToAsyncEnumerableAsync(
[
new AgentResponseUpdate(ChatRole.Assistant, new List<AIContent>
{
firstA,
firstB,
}),
new AgentResponseUpdate(ChatRole.Assistant, new List<AIContent>
{
new ToolApprovalRequestContent("id-a", new FunctionCallContent("repeated-call-a", "DifferentToolA")),
new ToolApprovalRequestContent("id-b", new FunctionCallContent("repeated-call-b", "DifferentToolB")),
}),
]));
var agent = new ToolApprovalAgent(innerAgent.Object, new ToolApprovalAgentOptions
{
AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule],
MaxAutoApprovalIterations = 1,
});
var session = new ChatClientAgentSession();

// Act
var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () =>
{
await foreach (var _ in agent.RunStreamingAsync(
[new ChatMessage(ChatRole.User, "Go")],
session))
{
}
});

var nextA = new ToolApprovalRequestContent("next-a", new FunctionCallContent("next-call-a", "ToolA"));
var nextB = new ToolApprovalRequestContent("next-b", new FunctionCallContent("next-call-b", "ToolB"));
var followUpAgent = new ToolApprovalAgent(
CreateMockAgent(new AgentResponse([new ChatMessage(ChatRole.Assistant, [nextA, nextB])])).Object);

var firstResponse = await followUpAgent.RunAsync(
[new ChatMessage(ChatRole.User,
[
firstA.CreateAlwaysApproveToolResponse(),
firstB.CreateAlwaysApproveToolResponse(),
])],
session);

var secondResponse = await followUpAgent.RunAsync(
[new ChatMessage(ChatRole.User, [nextA.CreateResponse(approved: true)])],
session);

// Assert
Assert.Contains("id-a", exception.Message, StringComparison.Ordinal);
Assert.Contains("id-b", exception.Message, StringComparison.Ordinal);
Assert.Equal("next-a", firstResponse.Messages.SelectMany(m => m.Contents).OfType<ToolApprovalRequestContent>().Single().RequestId);
Assert.Equal("next-b", secondResponse.Messages.SelectMany(m => m.Contents).OfType<ToolApprovalRequestContent>().Single().RequestId);
}

#endregion

#region Auto-approval iteration cap

/// <summary>
Expand Down
Loading