From f4bf0819c2ce640d62e4d045160cdee8b154db35 Mon Sep 17 00:00:00 2001 From: Shyju Krishnankutty Date: Mon, 15 Jun 2026 14:55:07 -0700 Subject: [PATCH] scope MCP threadId to the current agent --- .../DurableAIAgentProxy.cs | 9 ++ .../BuiltInFunctions.cs | 5 +- .../AgentSessionIdTests.cs | 12 +++ .../DurableAIAgentProxyTests.cs | 87 +++++++++++++++++++ 4 files changed, 111 insertions(+), 2 deletions(-) create mode 100644 dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAIAgentProxyTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs index 36a9336c36..da8911a569 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs @@ -79,6 +79,15 @@ internal class DurableAIAgentProxy(string name, IDurableAgentClient agentClient) RunRequest request = new([.. messages], responseFormat, enableToolCalls, enableToolNames); AgentSessionId sessionId = durableSession.SessionId; + // The session must belong to this agent. + if (!string.Equals(sessionId.Name, this.Name, StringComparison.OrdinalIgnoreCase)) + { + throw new ArgumentException( + $"The provided session belongs to agent '{sessionId.Name}' but was passed to agent '{this.Name}'. " + + "Sessions cannot be reused across agents.", + paramName: nameof(session)); + } + AgentRunHandle agentRunHandle = await this._agentClient.RunAgentAsync(sessionId, request, cancellationToken); if (isFireAndForget) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.AzureFunctions/BuiltInFunctions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.AzureFunctions/BuiltInFunctions.cs index 376f2fa2ca..41c46a7652 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.AzureFunctions/BuiltInFunctions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.AzureFunctions/BuiltInFunctions.cs @@ -363,9 +363,10 @@ internal static class BuiltInFunctions string agentName = context.Name; - // Derive session id: try to parse provided threadId, otherwise create a new one. + // Bind the caller-supplied threadId as a session key under the current agent name, + // mirroring the behavior of RunAgentHttpAsync. AgentSessionId sessionId = context.Arguments.TryGetValue("threadId", out object? threadObj) && threadObj is string threadId && !string.IsNullOrWhiteSpace(threadId) - ? AgentSessionId.Parse(threadId) + ? new AgentSessionId(agentName, threadId) : new AgentSessionId(agentName, functionContext.InvocationId); AIAgent agentProxy = client.AsDurableAgentProxy(functionContext, agentName); diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentSessionIdTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentSessionIdTests.cs index 03d171b7b3..bcc6df48a2 100644 --- a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentSessionIdTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentSessionIdTests.cs @@ -53,4 +53,16 @@ public sealed class AgentSessionIdTests AgentSessionId sessionId = entityId; }); } + + // Ensures the 2-arg constructor treats the key as opaque and never re-interprets + // it as a serialized session id, so the resulting Name always comes from the first + // argument regardless of the key's shape. + [Fact] + public void ConstructorTreatsKeyAsOpaqueValue() + { + AgentSessionId sessionId = new("agentA", "@dafx-agentB@some-key"); + + Assert.Equal("agentA", sessionId.Name); + Assert.Equal("@dafx-agentB@some-key", sessionId.Key); + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAIAgentProxyTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAIAgentProxyTests.cs new file mode 100644 index 0000000000..28e833719b --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAIAgentProxyTests.cs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.DurableTask.UnitTests; + +public sealed class DurableAIAgentProxyTests +{ + // Verifies the proxy rejects a session whose agent name differs from its own, + // and that the durable client is never called when this happens. + [Fact] + public async Task RunAsync_ThrowsWhenSessionBelongsToDifferentAgent() + { + StubDurableAgentClient client = new(); + DurableAIAgentProxy proxy = new("agentA", client); + DurableAgentSession session = new(new AgentSessionId("agentB", "shared-key")); + + ArgumentException ex = await Assert.ThrowsAsync(() => + proxy.RunAsync(new ChatMessage(ChatRole.User, "hello"), session)); + + Assert.Equal("session", ex.ParamName); + Assert.Contains("agentB", ex.Message, StringComparison.Ordinal); + Assert.Contains("agentA", ex.Message, StringComparison.Ordinal); + Assert.Equal(0, client.CallCount); + } + + // Control test: when the session's agent name matches the proxy's name, + // the request is forwarded to the durable client. + [Fact] + public async Task RunAsync_AllowsSessionWhenAgentNameMatches() + { + AgentSessionId sessionId = new("agentA", "shared-key"); + InvalidOperationException sentinel = new("reached the client"); + StubDurableAgentClient client = new() { Throw = sentinel }; + DurableAIAgentProxy proxy = new("agentA", client); + DurableAgentSession session = new(sessionId); + + // Reaching the durable client (and therefore propagating the sentinel) proves the + // name-matching guard accepted this session. + InvalidOperationException ex = await Assert.ThrowsAsync(() => + proxy.RunAsync(new ChatMessage(ChatRole.User, "hello"), session)); + + Assert.Same(sentinel, ex); + Assert.Equal(1, client.CallCount); + Assert.Equal(sessionId, client.LastSessionId); + } + + // Ensures the agent-name comparison is case-insensitive, so casing differences + // are neither a false-positive rejection nor a bypass. + [Fact] + public async Task RunAsync_AgentNameComparisonIsCaseInsensitive() + { + AgentSessionId sessionId = new("AGENTA", "shared-key"); + InvalidOperationException sentinel = new("reached the client"); + StubDurableAgentClient client = new() { Throw = sentinel }; + DurableAIAgentProxy proxy = new("agentA", client); + DurableAgentSession session = new(sessionId); + + InvalidOperationException ex = await Assert.ThrowsAsync(() => + proxy.RunAsync(new ChatMessage(ChatRole.User, "hello"), session)); + + Assert.Same(sentinel, ex); + Assert.Equal(1, client.CallCount); + } + + private sealed class StubDurableAgentClient : IDurableAgentClient + { + public int CallCount { get; private set; } + public AgentSessionId LastSessionId { get; private set; } + public Exception? Throw { get; set; } + + public Task RunAgentAsync( + AgentSessionId sessionId, + RunRequest request, + CancellationToken cancellationToken) + { + this.CallCount++; + this.LastSessionId = sessionId; + if (this.Throw is not null) + { + return Task.FromException(this.Throw); + } + + throw new InvalidOperationException("Test did not configure a response."); + } + } +}