[girishkvs]

Problem

McpTaskStoreTests.CancelTaskAsync_ViaStore_TransitionsToCancelled runs slow-tool as a task and then cancels it, but slow-tool finishes by itself after 200 ms. If the tool finishes before the tasks/cancel handler marks the task cancelled, the task stays completed, because InMemoryMcpTaskStore.SetCancelledAsync does not change a finished task. tasks/get then returns CompletedTaskResult and the test fails:

Assert.IsType() Failure: Value is not the exact type
Expected: typeof(ModelContextProtocol.Extensions.Tasks.CancelledTaskResult)
Actual: typeof(ModelContextProtocol.Extensions.Tasks.CompletedTaskResult)

This happened in build (windows-latest, Debug) on .NET 9 for #1856, which does not touch Tasks; attempt 3 of that job passed. In the test log the server starts the tasks/cancel handler, then logs "slow-tool" completed. IsError = False., then finishes the handler after 250 ms.

Change

• Add long-running-tool to TaskStoreTestTools. It waits on Task.Delay(Timeout.Infinite, cancellationToken), like the tool in TaskCancellationIntegrationTests.
• CancelTaskAsync_ViaStore_TransitionsToCancelled uses it, so the task can only end through tasks/cancel. The cancel handler marks the task cancelled before it cancels the tool's token, so the tool stops and the task stays cancelled. If tasks/cancel arrives before the tool starts, the tool gets an already-cancelled token and the task still ends cancelled.
• slow-tool is unchanged. Two other tests still use it, and CallToolAsync_WithTaskStore_PollsToCompletion needs it to finish.

Validation

Built tests/ModelContextProtocol.Tests for net10.0, net9.0, net8.0 (Debug) and net472 (Release) with 0 warnings, then ran the tests with vstest, for example:

vstest.console.exe artifacts\bin\ModelContextProtocol.Tests\Debug\net10.0\ModelContextProtocol.Tests.dll /TestCaseFilter:"FullyQualifiedName=ModelContextProtocol.Tests.Server.McpTaskStoreTests.CancelTaskAsync_ViaStore_TransitionsToCancelled"

• Before/after: with a local-only await Task.Delay(1000, ct); before CancelTaskAsync (not in this PR), the test failed 10 of 10 times on main with the CI message, and passed 10 of 10 times with this change (.NET 10).
• Without that delay: 20/20 passed on .NET 10, 20/20 on .NET 9, and 10/10 on .NET 10 with all 16 logical CPUs held at 100% by busy loops.
• FullyQualifiedName~McpTaskStoreTests|FullyQualifiedName~TaskCancellationIntegrationTests (58 tests, including InMemoryMcpTaskStoreTests): all passed on .NET 10, 9 and 8. On .NET Framework on Windows, McpTaskStoreTests and TaskCancellationIntegrationTests are skipped (Running ModelContextProtocol.Tests in net472 is causing VS test platform to hang #587).
• Full ModelContextProtocol.Tests with Execution!=Manual (the CI filter): 0 failed on .NET 10 and on .NET Framework 4.7.2.

[girishkvs]

@tarekgh This is another fix for flaky tests with race condition which failed some of the jobs on #1856. Can you review?

[tarekgh]

Thanks @girishkvs!