diff --git a/dotnet/samples/05-end-to-end/A2AClientServer/A2AServer/Program.cs b/dotnet/samples/05-end-to-end/A2AClientServer/A2AServer/Program.cs index cb8c57f377f..a76e779461f 100644 --- a/dotnet/samples/05-end-to-end/A2AClientServer/A2AServer/Program.cs +++ b/dotnet/samples/05-end-to-end/A2AClientServer/A2AServer/Program.cs @@ -101,8 +101,8 @@ You specialize in handling queries related to logistics. throw new ArgumentException("Either A2AServer:ApiKey or A2AServer:ConnectionString & agentName must be provided"); } -// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions by authenticated caller. -// Without this, contextId alone is the session key — any caller who knows a contextId can access that session. +// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions and tasks by authenticated caller. +// Without this, contextId/taskId alone are the lookup keys — any caller who knows them can access another caller's data. // Example using claims-based identity: // builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier }); diff --git a/dotnet/samples/05-end-to-end/AgentWebChat/AgentWebChat.AgentHost/Program.cs b/dotnet/samples/05-end-to-end/AgentWebChat/AgentWebChat.AgentHost/Program.cs index c54efda7a5e..8fd96ef3bdd 100644 --- a/dotnet/samples/05-end-to-end/AgentWebChat/AgentWebChat.AgentHost/Program.cs +++ b/dotnet/samples/05-end-to-end/AgentWebChat/AgentWebChat.AgentHost/Program.cs @@ -28,8 +28,8 @@ builder.AddOpenAIChatCompletions(); builder.AddOpenAIResponses(); -// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions by authenticated caller. -// Without this, contextId alone is the session key — any caller who knows a contextId can access that session. +// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions and tasks by authenticated caller. +// Without this, contextId/taskId alone are the lookup keys — any caller who knows them can access another caller's data. // Example using claims-based identity: // builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier }); @@ -157,8 +157,8 @@ Once the user has deduced what type (knight or knave) both Alice and Bob are, te pirateAgentBuilder.AddA2AServer(); knightsKnavesAgentBuilder.AddA2AServer(); -// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions by authenticated caller. -// Without this, contextId alone is the session key — any caller who knows a contextId can access that session. +// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions and tasks by authenticated caller. +// Without this, contextId/taskId alone are the lookup keys — any caller who knows them can access another caller's data. // Example using claims-based identity: // builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier }); diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs index ad07155387e..a11cd6ba86c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs @@ -30,19 +30,20 @@ public static class A2AServerServiceCollectionExtensions /// The for chaining. /// /// - /// Trust model. The A2A contextId arrives from the wire - /// and is treated as a chain-resume identifier — not as an authorization - /// token. The contract carries no principal/owner - /// dimension, so when a persistent store is registered any caller who knows or - /// guesses another caller's contextId can resume that other caller's - /// persisted thread. Hosts that serve more than one user must compose a principal - /// dimension into the lookup key — typically by calling - /// UseClaimsBasedSessionIsolation(...) from + /// Trust model. The A2A contextId and taskId arrive + /// from the wire and are treated as chain-resume identifiers — not as + /// authorization tokens. Both the and + /// contracts carry no principal/owner dimension by default, + /// so when a persistent store is registered any caller who knows or guesses another + /// caller's contextId or taskId can access that other caller's data. + /// Hosts that serve more than one user must compose a principal dimension into the + /// lookup key — typically by calling UseClaimsBasedSessionIsolation(...) from /// Microsoft.Agents.AI.Hosting.AspNetCore (or by registering a custom - /// ). When no isolation provider is - /// registered, behavior is unchanged — the bare contextId is used as the - /// conversation identifier, which is appropriate for first-run / single-user / - /// prototyping scenarios but unsafe for multi-user hosts. + /// ). When a + /// is registered, both the session store and the task store are automatically wrapped + /// with tenant-scoped isolation. When no isolation provider is registered, behavior + /// is unchanged — the bare identifiers are used directly, which is appropriate for + /// first-run / single-user / prototyping scenarios but unsafe for multi-user hosts. /// /// public static IHostedAgentBuilder AddA2AServer(this IHostedAgentBuilder agentBuilder, Action? configureOptions = null) @@ -65,10 +66,10 @@ public static IHostedAgentBuilder AddA2AServer(this IHostedAgentBuilder agentBui /// The for chaining. /// /// See the trust-model remarks on - /// for guidance on multi-user hosts (the wire contextId is a chain-resume - /// identifier, not an authorization token; multi-user hosts must compose a - /// principal dimension via UseClaimsBasedSessionIsolation(...) or a custom - /// ). + /// for guidance on multi-user hosts (the wire contextId and taskId + /// are chain-resume identifiers, not authorization tokens; multi-user hosts must + /// compose a principal dimension via UseClaimsBasedSessionIsolation(...) or + /// a custom ). /// public static IHostApplicationBuilder AddA2AServer(this IHostApplicationBuilder builder, string agentName, Action? configureOptions = null) { @@ -91,10 +92,10 @@ public static IHostApplicationBuilder AddA2AServer(this IHostApplicationBuilder /// The for chaining. /// /// See the trust-model remarks on - /// for guidance on multi-user hosts (the wire contextId is a chain-resume - /// identifier, not an authorization token; multi-user hosts must compose a - /// principal dimension via UseClaimsBasedSessionIsolation(...) or a custom - /// ). + /// for guidance on multi-user hosts (the wire contextId and taskId + /// are chain-resume identifiers, not authorization tokens; multi-user hosts must + /// compose a principal dimension via UseClaimsBasedSessionIsolation(...) or + /// a custom ). /// public static IHostApplicationBuilder AddA2AServer(this IHostApplicationBuilder builder, AIAgent agent, Action? configureOptions = null) { @@ -116,10 +117,10 @@ public static IHostApplicationBuilder AddA2AServer(this IHostApplicationBuilder /// The for chaining. /// /// See the trust-model remarks on - /// for guidance on multi-user hosts (the wire contextId is a chain-resume - /// identifier, not an authorization token; multi-user hosts must compose a - /// principal dimension via UseClaimsBasedSessionIsolation(...) or a custom - /// ). + /// for guidance on multi-user hosts (the wire contextId and taskId + /// are chain-resume identifiers, not authorization tokens; multi-user hosts must + /// compose a principal dimension via UseClaimsBasedSessionIsolation(...) or + /// a custom ). /// public static IServiceCollection AddA2AServer(this IServiceCollection services, string agentName, Action? configureOptions = null) { @@ -154,10 +155,10 @@ public static IServiceCollection AddA2AServer(this IServiceCollection services, /// The for chaining. /// /// See the trust-model remarks on - /// for guidance on multi-user hosts (the wire contextId is a chain-resume - /// identifier, not an authorization token; multi-user hosts must compose a - /// principal dimension via UseClaimsBasedSessionIsolation(...) or a custom - /// ). + /// for guidance on multi-user hosts (the wire contextId and taskId + /// are chain-resume identifiers, not authorization tokens; multi-user hosts must + /// compose a principal dimension via UseClaimsBasedSessionIsolation(...) or + /// a custom ). /// public static IServiceCollection AddA2AServer(this IServiceCollection services, AIAgent agent, Action? configureOptions = null) { @@ -179,6 +180,8 @@ public static IServiceCollection AddA2AServer(this IServiceCollection services, private static A2AServer CreateA2AServer(IServiceProvider serviceProvider, AIAgent agent, A2AServerRegistrationOptions? options) { + var isolationKeyProvider = serviceProvider.GetService(); + var agentHandler = serviceProvider.GetKeyedService(agent.Name); if (agentHandler is null) { @@ -186,7 +189,6 @@ private static A2AServer CreateA2AServer(IServiceProvider serviceProvider, AIAge var runMode = options?.AgentRunMode ?? AgentRunMode.DisallowBackground; // Ensure that we have an IsolationKeyScopedAgentSessionStore registered. - var isolationKeyProvider = serviceProvider.GetService(); if (agentSessionStore?.GetService() is null) { agentSessionStore ??= new NoopAgentSessionStore(); @@ -201,7 +203,13 @@ private static A2AServer CreateA2AServer(IServiceProvider serviceProvider, AIAge } var loggerFactory = serviceProvider.GetService() ?? NullLoggerFactory.Instance; - var taskStore = serviceProvider.GetKeyedService(agent.Name) ?? new InMemoryTaskStore(); + ITaskStore taskStore = serviceProvider.GetKeyedService(agent.Name) ?? new InMemoryTaskStore(); + + // Wrap the task store with isolation key scoping, same as the session store above. + if (taskStore is not IsolationKeyScopedTaskStore) + { + taskStore = new IsolationKeyScopedTaskStore(taskStore, isolationKeyProvider, strict: isolationKeyProvider != null); + } return new A2AServer( agentHandler, diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/IsolationKeyScopedTaskStore.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/IsolationKeyScopedTaskStore.cs new file mode 100644 index 00000000000..d70cfe0fd44 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/IsolationKeyScopedTaskStore.cs @@ -0,0 +1,217 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using A2A; + +namespace Microsoft.Agents.AI.Hosting.A2A; + +/// +/// A delegating that scopes task keys by an isolation key +/// provided by a , ensuring that tasks are isolated +/// per logical partition (e.g., user, tenant, or composite key). +/// +/// +/// +/// This class mirrors the isolation pattern of +/// but applies it to the A2A task store, preventing cross-tenant task access in multi-tenant deployments. +/// +/// +/// Both the store key and the persisted are scoped with the isolation +/// key. Scoping the persisted context is what allows list queries to be constrained to the calling +/// tenant, because list results are matched against the task body rather than the store key. Scoped +/// values are stripped again before being returned, so callers only ever observe bare identifiers. +/// +/// +public sealed class IsolationKeyScopedTaskStore : ITaskStore +{ + private readonly ITaskStore _innerStore; + private readonly SessionIsolationKeyProvider? _keyProvider; + private readonly bool _strict; + + /// + /// Initializes a new instance of the class. + /// + /// The underlying to delegate to. + /// + /// The used to retrieve the isolation key for the current context. + /// + /// + /// When , an is thrown if the isolation key + /// cannot be determined. When , the task ID is passed through unmodified. + /// + /// is . + public IsolationKeyScopedTaskStore( + ITaskStore innerStore, + SessionIsolationKeyProvider? keyProvider, + bool strict) + { + ArgumentNullException.ThrowIfNull(innerStore); + + this._innerStore = innerStore; + this._keyProvider = keyProvider; + this._strict = strict; + } + + /// + public async Task GetTaskAsync(string taskId, CancellationToken cancellationToken = default) + { + string? key = await this.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false); + + var task = await this._innerStore.GetTaskAsync(ScopeId(taskId, key), cancellationToken).ConfigureAwait(false); + + return task is null ? null : UnscopeTask(task, key); + } + + /// + public async Task SaveTaskAsync(string taskId, AgentTask task, CancellationToken cancellationToken = default) + { + string? key = await this.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false); + + await this._innerStore.SaveTaskAsync(ScopeId(taskId, key), ScopeTask(task, key), cancellationToken).ConfigureAwait(false); + } + + /// + public async Task DeleteTaskAsync(string taskId, CancellationToken cancellationToken = default) + { + string? key = await this.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false); + + await this._innerStore.DeleteTaskAsync(ScopeId(taskId, key), cancellationToken).ConfigureAwait(false); + } + + /// + /// + /// is reported by the inner store. When no + /// filter is supplied it therefore counts tasks across all + /// isolation keys, because a wrapper cannot narrow the count without enumerating the whole store. + /// The returned tasks themselves are always constrained to the current isolation key. + /// + public async Task ListTasksAsync(ListTasksRequest request, CancellationToken cancellationToken = default) + { + ArgumentNullException.ThrowIfNull(request); + + string? key = await this.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false); + + if (key is not null && !string.IsNullOrEmpty(request.ContextId)) + { + // Clone the request to avoid mutating the caller's object. + request = CloneRequestWithContextId(request, ScopeId(request.ContextId!, key)); + } + + var response = await this._innerStore.ListTasksAsync(request, cancellationToken).ConfigureAwait(false); + + if (key is null) + { + return response; + } + + // Tasks are persisted with a scoped ContextId, so any entry that does not carry the current + // isolation key belongs to another tenant and must not be returned. + var scopedTasks = new List(response.Tasks.Count); + foreach (var task in response.Tasks) + { + if (IsInScope(task, key)) + { + scopedTasks.Add(UnscopeTask(task, key)); + } + } + + response.Tasks = scopedTasks; + response.PageSize = scopedTasks.Count; + + return response; + } + + /// + /// Asynchronously retrieves the isolation key from the provider and validates it if in strict mode. + /// + private async ValueTask GetIsolationKeyAsync(CancellationToken cancellationToken) + { + string? key = this._keyProvider != null + ? await this._keyProvider.GetSessionIsolationKeyAsync(cancellationToken).ConfigureAwait(false) + : null; + + if (this._strict && key == null) + { + throw new InvalidOperationException("Session isolation key is required but was not provided by the configured SessionIsolationKeyProvider."); + } + + return key; + } + + /// + /// Escapes special characters in the isolation key to ensure unambiguous scoped identifiers. + /// + private static string EscapeIsolationKey(string key) => key.Replace("\\", "\\\\").Replace(":", "\\:"); + + /// + /// Prefixes a bare identifier with the escaped isolation key, or returns it unchanged when no key applies. + /// + private static string ScopeId(string id, string? key) + => key is null ? id : $"{EscapeIsolationKey(key)}::{id}"; + + /// + /// Strips the isolation key prefix from a scoped identifier, or returns it unchanged when the prefix is absent. + /// + private static string UnscopeId(string scopedId, string? key) + { + if (key is null) + { + return scopedId; + } + + string prefix = $"{EscapeIsolationKey(key)}::"; + + return scopedId.StartsWith(prefix, StringComparison.Ordinal) + ? scopedId.Substring(prefix.Length) + : scopedId; + } + + /// + /// Determines whether a persisted task carries the supplied isolation key. + /// + private static bool IsInScope(AgentTask task, string key) + => task.ContextId?.StartsWith($"{EscapeIsolationKey(key)}::", StringComparison.Ordinal) == true; + + /// + /// Creates a copy of the task whose is scoped by the isolation key. + /// + /// + /// The task instance is copied rather than mutated because the A2A server reuses it for live event + /// notification after persisting; mutating it would surface the scoped context on the wire. + /// + private static AgentTask ScopeTask(AgentTask task, string? key) + => key is null ? task : CloneTaskWithContextId(task, ScopeId(task.ContextId, key)); + + /// + /// Creates a copy of the task whose has the isolation key removed. + /// + private static AgentTask UnscopeTask(AgentTask task, string? key) + => key is null ? task : CloneTaskWithContextId(task, UnscopeId(task.ContextId, key)); + + private static AgentTask CloneTaskWithContextId(AgentTask task, string contextId) + => new() + { + Id = task.Id, + ContextId = contextId, + Status = task.Status, + History = task.History, + Artifacts = task.Artifacts, + Metadata = task.Metadata, + }; + + private static ListTasksRequest CloneRequestWithContextId(ListTasksRequest request, string contextId) + => new() + { + ContextId = contextId, + Tenant = request.Tenant, + Status = request.Status, + PageSize = request.PageSize, + PageToken = request.PageToken, + HistoryLength = request.HistoryLength, + StatusTimestampAfter = request.StatusTimestampAfter, + IncludeArtifacts = request.IncludeArtifacts, + }; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs index 11ad8574d32..21581e79f7c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs @@ -105,6 +105,85 @@ public async Task AddA2AServer_WithCustomTaskStore_ResolvesSuccessfullyAsync() Assert.NotNull(server); } + /// + /// Verifies that when a SessionIsolationKeyProvider is registered, task operations + /// use scoped identifiers (DI wiring test). + /// + [Fact] + public async Task AddA2AServer_WithIsolationKeyProvider_TaskStoreReceivesScopedIdsAsync() + { + // Arrange + const string AgentName = "isolation-wiring-agent"; + const string TaskId = "task-1"; + const string IsolationKey = "alice"; + + var services = new ServiceCollection(); + services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object); + + var mockTaskStore = new Mock(); + mockTaskStore + .Setup(s => s.GetTaskAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(new AgentTask { Id = TaskId, ContextId = $"{IsolationKey}::ctx-1", Status = new global::A2A.TaskStatus { State = TaskState.Completed } }); + services.AddKeyedSingleton(AgentName, mockTaskStore.Object); + + var mockKeyProvider = new Mock(); + mockKeyProvider + .Setup(p => p.GetSessionIsolationKeyAsync(It.IsAny())) + .ReturnsAsync(IsolationKey); + services.AddSingleton(mockKeyProvider.Object); + + services.AddA2AServer(AgentName); + await using var provider = services.BuildServiceProvider(); + var server = provider.GetRequiredKeyedService(AgentName); + + // Act + var result = await server.GetTaskAsync(new GetTaskRequest { Id = TaskId }); + + // Assert - the inner store received the scoped task ID + mockTaskStore.Verify(s => s.GetTaskAsync($"{IsolationKey}::{TaskId}", It.IsAny()), Times.Once); + } + + /// + /// Verifies that when an already-wrapped IsolationKeyScopedTaskStore is registered, + /// it is not double-wrapped (no double scoping of task IDs). + /// + [Fact] + public async Task AddA2AServer_WithAlreadyWrappedTaskStore_DoesNotDoubleWrapAsync() + { + // Arrange + const string AgentName = "no-double-wrap-agent"; + const string TaskId = "task-1"; + const string IsolationKey = "alice"; + + var services = new ServiceCollection(); + services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object); + + var mockInnerStore = new Mock(); + mockInnerStore + .Setup(s => s.GetTaskAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(new AgentTask { Id = TaskId, ContextId = $"{IsolationKey}::ctx-1", Status = new global::A2A.TaskStatus { State = TaskState.Completed } }); + + var mockKeyProvider = new Mock(); + mockKeyProvider + .Setup(p => p.GetSessionIsolationKeyAsync(It.IsAny())) + .ReturnsAsync(IsolationKey); + + // Pre-wrap the task store + var wrappedStore = new IsolationKeyScopedTaskStore(mockInnerStore.Object, mockKeyProvider.Object, strict: true); + services.AddKeyedSingleton(AgentName, wrappedStore); + services.AddSingleton(mockKeyProvider.Object); + + services.AddA2AServer(AgentName); + await using var provider = services.BuildServiceProvider(); + var server = provider.GetRequiredKeyedService(AgentName); + + // Act + var result = await server.GetTaskAsync(new GetTaskRequest { Id = TaskId }); + + // Assert - only single scoping occurred (not alice::alice::task-1) + mockInnerStore.Verify(s => s.GetTaskAsync($"{IsolationKey}::{TaskId}", It.IsAny()), Times.Once); + } + /// /// Verifies that when a custom AgentSessionStore is registered, AddA2AServer uses it /// instead of the default InMemoryAgentSessionStore. diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/IsolationKeyScopedTaskStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/IsolationKeyScopedTaskStoreTests.cs new file mode 100644 index 00000000000..ed7e3c530d8 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/IsolationKeyScopedTaskStoreTests.cs @@ -0,0 +1,335 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Threading; +using System.Threading.Tasks; +using A2A; +using Moq; + +namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests; + +/// +/// Unit tests for the class. +/// +public sealed class IsolationKeyScopedTaskStoreTests +{ + private const string AliceKey = "alice"; + private const string BobKey = "bob"; + private const string TaskId = "task-001"; + + /// + /// Verifies that GetTaskAsync scopes the task ID with the isolation key. + /// + [Fact] + public async Task GetTaskAsync_ScopesTaskIdWithIsolationKeyAsync() + { + // Arrange + var innerStore = new Mock(); + var keyProvider = CreateKeyProvider(AliceKey); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + + // Act + await store.GetTaskAsync(TaskId); + + // Assert + innerStore.Verify(s => s.GetTaskAsync($"{AliceKey}::{TaskId}", It.IsAny()), Times.Once); + } + + /// + /// Verifies that SaveTaskAsync scopes the task ID and the persisted ContextId with the isolation key, + /// without mutating the caller's task instance. + /// + [Fact] + public async Task SaveTaskAsync_ScopesTaskIdAndContextIdWithIsolationKeyAsync() + { + // Arrange + var innerStore = new Mock(); + var keyProvider = CreateKeyProvider(AliceKey); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + var task = new AgentTask { Id = TaskId, ContextId = "ctx-1" }; + + // Act + await store.SaveTaskAsync(TaskId, task); + + // Assert - both the store key and the persisted ContextId are scoped + innerStore.Verify(s => s.SaveTaskAsync( + $"{AliceKey}::{TaskId}", + It.Is(t => t.ContextId == $"{AliceKey}::ctx-1"), + It.IsAny()), Times.Once); + + // Assert - the caller's instance is untouched + Assert.Equal("ctx-1", task.ContextId); + } + + /// + /// Verifies that GetTaskAsync strips the isolation key from the returned task's ContextId. + /// + [Fact] + public async Task GetTaskAsync_UnscopesContextIdOnReadAsync() + { + // Arrange + var innerStore = new InMemoryTaskStore(); + var store = new IsolationKeyScopedTaskStore(innerStore, CreateKeyProvider(AliceKey), strict: true); + + await store.SaveTaskAsync(TaskId, new AgentTask { Id = TaskId, ContextId = "ctx-1" }); + + // Act + var result = await store.GetTaskAsync(TaskId); + + // Assert - the caller observes the bare ContextId + Assert.NotNull(result); + Assert.Equal("ctx-1", result.ContextId); + } + + /// + /// Verifies that DeleteTaskAsync scopes the task ID with the isolation key. + /// + [Fact] + public async Task DeleteTaskAsync_ScopesTaskIdWithIsolationKeyAsync() + { + // Arrange + var innerStore = new Mock(); + var keyProvider = CreateKeyProvider(AliceKey); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + + // Act + await store.DeleteTaskAsync(TaskId); + + // Assert + innerStore.Verify(s => s.DeleteTaskAsync($"{AliceKey}::{TaskId}", It.IsAny()), Times.Once); + } + + /// + /// Verifies that different isolation keys produce different scoped task IDs, + /// preventing cross-tenant task access. + /// + [Fact] + public async Task GetTaskAsync_DifferentTenantsGetDifferentScopedIdsAsync() + { + // Arrange + var innerStore = new InMemoryTaskStore(); + var aliceStore = new IsolationKeyScopedTaskStore(innerStore, CreateKeyProvider(AliceKey), strict: true); + var bobStore = new IsolationKeyScopedTaskStore(innerStore, CreateKeyProvider(BobKey), strict: true); + + var aliceTask = new AgentTask { Id = TaskId, ContextId = "ctx-1" }; + + // Act - Alice saves a task + await aliceStore.SaveTaskAsync(TaskId, aliceTask); + + // Assert - Alice can read it + var aliceResult = await aliceStore.GetTaskAsync(TaskId); + Assert.NotNull(aliceResult); + + // Assert - Bob cannot read it (different isolation key → different scoped ID) + var bobResult = await bobStore.GetTaskAsync(TaskId); + Assert.Null(bobResult); + } + + /// + /// Verifies that ListTasksAsync scopes the ContextId filter with the isolation key. + /// + [Fact] + public async Task ListTasksAsync_ScopesContextIdFilterAsync() + { + // Arrange + var innerStore = new Mock(); + innerStore + .Setup(s => s.ListTasksAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(new ListTasksResponse()); + + var keyProvider = CreateKeyProvider(AliceKey); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + + var request = new ListTasksRequest { ContextId = "ctx-1" }; + + // Act + await store.ListTasksAsync(request); + + // Assert - ContextId was scoped with the isolation key + innerStore.Verify(s => s.ListTasksAsync( + It.Is(r => r.ContextId == $"{AliceKey}::ctx-1"), + It.IsAny()), Times.Once); + + // Assert - original request was not mutated + Assert.Equal("ctx-1", request.ContextId); + } + + /// + /// Verifies that ListTasksAsync does not modify the ContextId filter when it is null. + /// + [Fact] + public async Task ListTasksAsync_NullContextId_DoesNotScopeFilterAsync() + { + // Arrange + var innerStore = new Mock(); + innerStore + .Setup(s => s.ListTasksAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(new ListTasksResponse()); + + var keyProvider = CreateKeyProvider(AliceKey); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + + var request = new ListTasksRequest { ContextId = null }; + + // Act + await store.ListTasksAsync(request); + + // Assert - ContextId was not modified + innerStore.Verify(s => s.ListTasksAsync( + It.Is(r => r.ContextId == null), + It.IsAny()), Times.Once); + } + + /// + /// Verifies that an unfiltered ListTasksAsync only returns tasks belonging to the calling tenant. + /// + [Fact] + public async Task ListTasksAsync_NoContextIdFilter_ExcludesOtherTenantsTasksAsync() + { + // Arrange + var innerStore = new InMemoryTaskStore(); + var aliceStore = new IsolationKeyScopedTaskStore(innerStore, CreateKeyProvider(AliceKey), strict: true); + var bobStore = new IsolationKeyScopedTaskStore(innerStore, CreateKeyProvider(BobKey), strict: true); + + await aliceStore.SaveTaskAsync("alice-task", new AgentTask { Id = "alice-task", ContextId = "ctx-1" }); + await bobStore.SaveTaskAsync("bob-task", new AgentTask { Id = "bob-task", ContextId = "ctx-2" }); + + // Act - Bob lists without any filter + var response = await bobStore.ListTasksAsync(new ListTasksRequest()); + + // Assert - only Bob's task is returned, with a bare ContextId + var task = Assert.Single(response.Tasks); + Assert.Equal("bob-task", task.Id); + Assert.Equal("ctx-2", task.ContextId); + Assert.Equal(1, response.PageSize); + } + + /// + /// Verifies that filtering by ContextId returns the caller's own tasks, since the persisted + /// ContextId is scoped by the same isolation key as the filter. + /// + [Fact] + public async Task ListTasksAsync_WithContextIdFilter_ReturnsOwnTasksAsync() + { + // Arrange + var innerStore = new InMemoryTaskStore(); + var aliceStore = new IsolationKeyScopedTaskStore(innerStore, CreateKeyProvider(AliceKey), strict: true); + var bobStore = new IsolationKeyScopedTaskStore(innerStore, CreateKeyProvider(BobKey), strict: true); + + await aliceStore.SaveTaskAsync(TaskId, new AgentTask { Id = TaskId, ContextId = "ctx-1" }); + + // Act + var aliceResponse = await aliceStore.ListTasksAsync(new ListTasksRequest { ContextId = "ctx-1" }); + var bobResponse = await bobStore.ListTasksAsync(new ListTasksRequest { ContextId = "ctx-1" }); + + // Assert - Alice sees her task; Bob sees nothing for the same bare context + var task = Assert.Single(aliceResponse.Tasks); + Assert.Equal("ctx-1", task.ContextId); + Assert.Empty(bobResponse.Tasks); + } + + /// + /// Verifies that strict mode throws when the isolation key provider returns null. + /// + [Fact] + public async Task GetTaskAsync_StrictMode_NullKey_ThrowsAsync() + { + // Arrange + var innerStore = new Mock(); + var keyProvider = CreateKeyProvider(null); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + + // Act & Assert + await Assert.ThrowsAsync(() => store.GetTaskAsync(TaskId)); + } + + /// + /// Verifies that non-strict mode passes through the bare task ID when the key is null. + /// + [Fact] + public async Task GetTaskAsync_NonStrictMode_NullKey_PassesThroughAsync() + { + // Arrange + var innerStore = new Mock(); + var keyProvider = CreateKeyProvider(null); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: false); + + // Act + await store.GetTaskAsync(TaskId); + + // Assert - bare task ID was used (no scoping) + innerStore.Verify(s => s.GetTaskAsync(TaskId, It.IsAny()), Times.Once); + } + + /// + /// Verifies that when no key provider is registered and strict is false, + /// the bare task ID is passed through. + /// + [Fact] + public async Task GetTaskAsync_NoKeyProvider_NonStrict_PassesThroughAsync() + { + // Arrange + var innerStore = new Mock(); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider: null, strict: false); + + // Act + await store.GetTaskAsync(TaskId); + + // Assert - bare task ID was used + innerStore.Verify(s => s.GetTaskAsync(TaskId, It.IsAny()), Times.Once); + } + + /// + /// Verifies that colons in the isolation key are escaped to prevent ID ambiguity. + /// + [Fact] + public async Task GetTaskAsync_EscapesColonsInIsolationKeyAsync() + { + // Arrange + var innerStore = new Mock(); + var keyProvider = CreateKeyProvider("tenant:sub"); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + + // Act + await store.GetTaskAsync(TaskId); + + // Assert - colons are escaped + innerStore.Verify(s => s.GetTaskAsync(@"tenant\:sub::" + TaskId, It.IsAny()), Times.Once); + } + + /// + /// Verifies that backslashes in the isolation key are escaped to prevent ID ambiguity. + /// + [Fact] + public async Task GetTaskAsync_EscapesBackslashesInIsolationKeyAsync() + { + // Arrange + var innerStore = new Mock(); + var keyProvider = CreateKeyProvider(@"domain\user"); + var store = new IsolationKeyScopedTaskStore(innerStore.Object, keyProvider, strict: true); + + // Act + await store.GetTaskAsync(TaskId); + + // Assert - backslashes are escaped + innerStore.Verify(s => s.GetTaskAsync(@"domain\\user::" + TaskId, It.IsAny()), Times.Once); + } + + /// + /// Verifies that the constructor throws when the inner store is null. + /// + [Fact] + public void Constructor_NullInnerStore_ThrowsArgumentNullException() + { + // Act & Assert + Assert.Throws(() => new IsolationKeyScopedTaskStore(null!, null, strict: false)); + } + + private static SessionIsolationKeyProvider CreateKeyProvider(string? key) + { + var mock = new Mock(); + mock.Setup(p => p.GetSessionIsolationKeyAsync(It.IsAny())) + .ReturnsAsync(key); + return mock.Object; + } +}