.Net: Add cancellation token & add headers to request (#11900)

### Motivation, Context and Description

This PR:
1. Propagates the cancellation token from all kernel functions of the
`SessionsPythonPlugin` plugin down to the APIs they call.
2. Refactors the way the plugin adds headers to the request. Instead of
adding them as default headers of the HTTP client, they are now added as
request headers. The result is the same - the headers are sent each call
of any kernel function of the plugin. This change may save time later by
preventing bug troubleshooting if/when an HTTP client is injected into
the plugin and it sets default headers to the client it does not own.
3. Replaces the usage of HTTP client methods for getting content, such
as `ReadAsStringAsync`, with the SK wrappers
`ReadAsStringWithExceptionMappingAsync` to align the plugin's behavior
with that of other SK components.

Contributes to:
#10070

Author: SergeyMenshykh

dotnet/samples/Demos/CodeInterpreterPlugin/Program.cs

@@ -35,15 +35,15 @@
 /// <summary>
 /// Acquire a token for the Azure Container Apps service
 /// </summary>
-async Task<string> TokenProvider()
+async Task<string> TokenProvider(CancellationToken cancellationToken)
 {
     if (cachedToken is null)
     {
         string resource = "https://acasessions.io/.default";
         var credential = new InteractiveBrowserCredential();
 
         // Attempt to get the token
-        var accessToken = await credential.GetTokenAsync(new Azure.Core.TokenRequestContext([resource])).ConfigureAwait(false);
+        var accessToken = await credential.GetTokenAsync(new Azure.Core.TokenRequestContext([resource]), cancellationToken).ConfigureAwait(false);
         if (logger.IsEnabled(LogLevel.Information))
         {
             logger.LogInformation("Access token obtained successfully");

dotnet/src/IntegrationTests/Plugins/Core/SessionsPythonPluginTests.cs

@@ -1,19 +1,20 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
-using System.Threading.Tasks;
-using Xunit;
-using Microsoft.SemanticKernel.Plugins.Core.CodeInterpreter;
-using Microsoft.Extensions.Configuration;
-using SemanticKernel.IntegrationTests.TestSettings;
+using System.Collections.Generic;
 using System.Net.Http;
-using Azure.Identity;
+using System.Threading;
+using System.Threading.Tasks;
 using Azure.Core;
-using System.Collections.Generic;
-using Microsoft.SemanticKernel;
+using Azure.Identity;
+using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.SemanticKernel;
 using Microsoft.SemanticKernel.ChatCompletion;
 using Microsoft.SemanticKernel.Connectors.AzureOpenAI;
+using Microsoft.SemanticKernel.Plugins.Core.CodeInterpreter;
+using SemanticKernel.IntegrationTests.TestSettings;
+using Xunit;
 
 namespace SemanticKernel.IntegrationTests.Plugins.Core;
 
@@ -143,13 +144,13 @@ public async Task LlmShouldUploadFileAndAccessItFromCodeInterpreterAsync()
     /// <summary>
     /// Acquires authentication token for the Azure Container App Session pool.
     /// </summary>
-    private static async Task<string> GetAuthTokenAsync()
+    private static async Task<string> GetAuthTokenAsync(CancellationToken cancellationToken)
     {
         string resource = "https://acasessions.io/.default";
 
         var credential = new AzureCliCredential();
 
-        AccessToken token = await credential.GetTokenAsync(new Azure.Core.TokenRequestContext([resource])).ConfigureAwait(false);
+        AccessToken token = await credential.GetTokenAsync(new Azure.Core.TokenRequestContext([resource]), cancellationToken).ConfigureAwait(false);
 
         return token.Token;
     }

dotnet/src/Plugins/Plugins.Core/CodeInterpreter/SessionsPythonPlugin.cs

@@ -26,21 +26,21 @@ public sealed partial class SessionsPythonPlugin
     private const string ApiVersion = "2024-10-02-preview";
     private readonly Uri _poolManagementEndpoint;
     private readonly SessionsPythonSettings _settings;
-    private readonly Func<Task<string>>? _authTokenProvider;
+    private readonly Func<CancellationToken, Task<string>>? _authTokenProvider;
     private readonly IHttpClientFactory _httpClientFactory;
     private readonly ILogger _logger;
 
     /// <summary>
     /// Initializes a new instance of the SessionsPythonTool class.
     /// </summary>
-    /// <param name="settings">The settings for the Python tool plugin. </param>
-    /// <param name="httpClientFactory">The HTTP client factory. </param>
-    /// <param name="authTokenProvider"> Optional provider for auth token generation. </param>
-    /// <param name="loggerFactory">The logger factory. </param>
+    /// <param name="settings">The settings for the Python tool plugin.</param>
+    /// <param name="httpClientFactory">The HTTP client factory.</param>
+    /// <param name="authTokenProvider">Optional provider for auth token generation.</param>
+    /// <param name="loggerFactory">The logger factory.</param>
     public SessionsPythonPlugin(
         SessionsPythonSettings settings,
         IHttpClientFactory httpClientFactory,
-        Func<Task<string>>? authTokenProvider = null,
+        Func<CancellationToken, Task<string>>? authTokenProvider = null,
         ILoggerFactory? loggerFactory = null)
     {
         Verify.NotNull(settings, nameof(settings));
@@ -66,7 +66,8 @@ public SessionsPythonPlugin(
     /// Keep everything in a single line; the \n sequences will represent line breaks
     /// when the string is processed or displayed.
     /// </summary>
-    /// <param name="code"> The valid Python code to execute. </param>
+    /// <param name="code"> The valid Python code to execute.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns> The result of the Python code execution. </returns>
     /// <exception cref="ArgumentNullException"></exception>
     /// <exception cref="HttpRequestException"></exception>
@@ -79,7 +80,9 @@ Add spaces directly after \n sequences to replicate indentation.
         Keep everything in a single line; the \n sequences will represent line breaks
         when the string is processed or displayed.
         """)]
-    public async Task<string> ExecuteCodeAsync([Description("The valid Python code to execute.")] string code)
+    public async Task<string> ExecuteCodeAsync(
+        [Description("The valid Python code to execute.")] string code,
+        CancellationToken cancellationToken = default)
     {
         Verify.NotNullOrWhiteSpace(code, nameof(code));
 
@@ -91,15 +94,14 @@ public async Task<string> ExecuteCodeAsync([Description("The valid Python code t
         this._logger.LogTrace("Executing Python code: {Code}", code);
 
         using var httpClient = this._httpClientFactory.CreateClient();
-        await this.AddHeadersAsync(httpClient).ConfigureAwait(false);
 
         var requestBody = new SessionsPythonCodeExecutionProperties(this._settings, code);
 
         using var content = new StringContent(JsonSerializer.Serialize(requestBody), Encoding.UTF8, "application/json");
 
-        using var response = await this.SendAsync(httpClient, HttpMethod.Post, "executions", content).ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Post, "executions", cancellationToken, content).ConfigureAwait(false);
 
-        var responseContent = JsonSerializer.Deserialize<JsonElement>(await response.Content.ReadAsStringAsync().ConfigureAwait(false));
+        var responseContent = JsonSerializer.Deserialize<JsonElement>(await response.Content.ReadAsStringWithExceptionMappingAsync(cancellationToken).ConfigureAwait(false));
 
         var result = responseContent.GetProperty("result");
 
@@ -120,21 +122,22 @@ public async Task<string> ExecuteCodeAsync([Description("The valid Python code t
     /// </summary>
     /// <param name="remoteFileName">The name of the remote file, relative to `/mnt/data`.</param>
     /// <param name="localFilePath">The path to the file on the local machine.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The metadata of the uploaded file.</returns>
     /// <exception cref="ArgumentNullException"></exception>
     /// <exception cref="HttpRequestException"></exception>
     [KernelFunction, Description("Uploads a file to the `/mnt/data` directory of the current session.")]
     public async Task<SessionsRemoteFileMetadata> UploadFileAsync(
         [Description("The name of the remote file, relative to `/mnt/data`.")] string remoteFileName,
-        [Description("The path to the file on the local machine.")] string localFilePath)
+        [Description("The path to the file on the local machine.")] string localFilePath,
+        CancellationToken cancellationToken = default)
     {
         Verify.NotNullOrWhiteSpace(remoteFileName, nameof(remoteFileName));
         Verify.NotNullOrWhiteSpace(localFilePath, nameof(localFilePath));
 
         this._logger.LogInformation("Uploading file: {LocalFilePath} to {RemoteFileName}", localFilePath, remoteFileName);
 
         using var httpClient = this._httpClientFactory.CreateClient();
-        await this.AddHeadersAsync(httpClient).ConfigureAwait(false);
 
         using var fileContent = new ByteArrayContent(File.ReadAllBytes(localFilePath));
 
@@ -143,9 +146,9 @@ public async Task<SessionsRemoteFileMetadata> UploadFileAsync(
             { fileContent, "file", remoteFileName },
         };
 
-        using var response = await this.SendAsync(httpClient, HttpMethod.Post, "files", multipartFormDataContent).ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Post, "files", cancellationToken, multipartFormDataContent).ConfigureAwait(false);
 
-        var stringContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
+        var stringContent = await response.Content.ReadAsStringWithExceptionMappingAsync(cancellationToken).ConfigureAwait(false);
 
         return JsonSerializer.Deserialize<SessionsRemoteFileMetadata>(stringContent)!;
     }
@@ -155,22 +158,23 @@ public async Task<SessionsRemoteFileMetadata> UploadFileAsync(
     /// </summary>
     /// <param name="remoteFileName">The name of the remote file to download, relative to `/mnt/data`.</param>
     /// <param name="localFilePath">The path to save the downloaded file to. If not provided won't save it in the disk.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The data of the downloaded file as byte array.</returns>
     [KernelFunction, Description("Downloads a file from the `/mnt/data` directory of the current session.")]
     public async Task<byte[]> DownloadFileAsync(
         [Description("The name of the remote file to download, relative to `/mnt/data`.")] string remoteFileName,
-        [Description("The path to save the downloaded file to. If not provided won't save it in the disk.")] string? localFilePath = null)
+        [Description("The path to save the downloaded file to. If not provided won't save it in the disk.")] string? localFilePath = null,
+        CancellationToken cancellationToken = default)
     {
         Verify.NotNullOrWhiteSpace(remoteFileName, nameof(remoteFileName));
 
         this._logger.LogTrace("Downloading file: {RemoteFileName} to {LocalFileName}", remoteFileName, localFilePath);
 
         using var httpClient = this._httpClientFactory.CreateClient();
-        await this.AddHeadersAsync(httpClient).ConfigureAwait(false);
 
-        using var response = await this.SendAsync(httpClient, HttpMethod.Get, $"files/{Uri.EscapeDataString(remoteFileName)}/content").ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Get, $"files/{Uri.EscapeDataString(remoteFileName)}/content", cancellationToken).ConfigureAwait(false);
 
-        var fileContent = await response.Content.ReadAsByteArrayAsync().ConfigureAwait(false);
+        var fileContent = await response.Content.ReadAsByteArrayAndTranslateExceptionAsync(cancellationToken).ConfigureAwait(false);
 
         if (!string.IsNullOrWhiteSpace(localFilePath))
         {
@@ -190,18 +194,18 @@ public async Task<byte[]> DownloadFileAsync(
     /// <summary>
     /// Lists all entities: files or directories in the `/mnt/data` directory of the current session.
     /// </summary>
+    /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The list of files in the session.</returns>
     [KernelFunction, Description("Lists all entities: files or directories in the `/mnt/data` directory of the current session.")]
-    public async Task<IReadOnlyList<SessionsRemoteFileMetadata>> ListFilesAsync()
+    public async Task<IReadOnlyList<SessionsRemoteFileMetadata>> ListFilesAsync(CancellationToken cancellationToken = default)
     {
         this._logger.LogTrace("Listing files for Session ID: {SessionId}", this._settings.SessionId);
 
         using var httpClient = this._httpClientFactory.CreateClient();
-        await this.AddHeadersAsync(httpClient).ConfigureAwait(false);
 
-        using var response = await this.SendAsync(httpClient, HttpMethod.Get, "files").ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Get, "files", cancellationToken).ConfigureAwait(false);
 
-        var jsonElementResult = JsonSerializer.Deserialize<JsonElement>(await response.Content.ReadAsStringAsync().ConfigureAwait(false));
+        var jsonElementResult = JsonSerializer.Deserialize<JsonElement>(await response.Content.ReadAsStringWithExceptionMappingAsync(cancellationToken).ConfigureAwait(false));
 
         var files = jsonElementResult.GetProperty("value");
 
@@ -241,16 +245,17 @@ private static string SanitizeCodeInput(string code)
     }
 
     /// <summary>
-    /// Add headers to the HTTP client.
+    /// Add headers to the HTTP request.
     /// </summary>
-    /// <param name="httpClient">The HTTP client to add headers to.</param>
-    private async Task AddHeadersAsync(HttpClient httpClient)
+    /// <param name="request">The HTTP request to add headers to.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    private async Task AddHeadersAsync(HttpRequestMessage request, CancellationToken cancellationToken)
     {
-        httpClient.DefaultRequestHeaders.Add("User-Agent", $"{HttpHeaderConstant.Values.UserAgent}/{s_assemblyVersion} (Language=dotnet)");
+        request.Headers.Add("User-Agent", $"{HttpHeaderConstant.Values.UserAgent}/{s_assemblyVersion} (Language=dotnet)");
 
         if (this._authTokenProvider is not null)
         {
-            httpClient.DefaultRequestHeaders.Add("Authorization", $"Bearer {(await this._authTokenProvider().ConfigureAwait(false))}");
+            request.Headers.Add("Authorization", $"Bearer {(await this._authTokenProvider(cancellationToken).ConfigureAwait(false))}");
         }
     }
 
@@ -260,9 +265,10 @@ private async Task AddHeadersAsync(HttpClient httpClient)
     /// <param name="httpClient">The HTTP client to use.</param>
     /// <param name="method">The HTTP method to use.</param>
     /// <param name="path">The path to send the request to.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
     /// <param name="httpContent">The content to send with the request.</param>
     /// <returns>The HTTP response message.</returns>
-    private async Task<HttpResponseMessage> SendAsync(HttpClient httpClient, HttpMethod method, string path, HttpContent? httpContent = null)
+    private async Task<HttpResponseMessage> SendAsync(HttpClient httpClient, HttpMethod method, string path, CancellationToken cancellationToken, HttpContent? httpContent = null)
     {
         // The query string is the same for all operations
         var pathWithQueryString = $"{path}?identifier={this._settings.SessionId}&api-version={ApiVersion}";
@@ -281,7 +287,9 @@ private async Task<HttpResponseMessage> SendAsync(HttpClient httpClient, HttpMet
             Content = httpContent,
         };
 
-        return await httpClient.SendWithSuccessCheckAsync(request, CancellationToken.None).ConfigureAwait(false);
+        await this.AddHeadersAsync(request, cancellationToken).ConfigureAwait(false);
+
+        return await httpClient.SendWithSuccessCheckAsync(request, cancellationToken).ConfigureAwait(false);
     }
 
 #if NET

dotnet/src/Plugins/Plugins.UnitTests/Core/SessionsPythonPluginTests.cs

@@ -3,11 +3,14 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Linq;
 using System.Net;
 using System.Net.Http;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.SemanticKernel;
+using Microsoft.SemanticKernel.Http;
 using Microsoft.SemanticKernel.Plugins.Core.CodeInterpreter;
 using Moq;
 using Xunit;
@@ -22,6 +25,7 @@ public sealed class SessionsPythonPluginTests : IDisposable
     private const string ListFilesTestDataFilePath = "./TestData/sessions_python_plugin_file_list.json";
     private const string UpdaloadFileTestDataFilePath = "./TestData/sessions_python_plugin_file_upload.json";
     private const string FileTestDataFilePath = "./TestData/sessions_python_plugin_file.txt";
+    private readonly static string s_assemblyVersion = typeof(Kernel).Assembly.GetName().Version!.ToString();
 
     private readonly SessionsPythonSettings _defaultSettings = new(
         sessionId: Guid.NewGuid().ToString(),
@@ -97,7 +101,7 @@ public async Task ItShouldCallTokenProviderWhenProvidedAsync(string methodName)
         // Arrange
         var tokenProviderCalled = false;
 
-        Task<string> tokenProviderAsync()
+        Task<string> tokenProviderAsync(CancellationToken _)
         {
             tokenProviderCalled = true;
             return Task.FromResult("token");
@@ -332,6 +336,34 @@ public async Task ItShouldRespectAllowedDomainsAsync(string allowedDomain, strin
 #pragma warning restore CA1031 // Do not catch general exception types
     }
 
+    [Fact]
+    public async Task ItShouldAddHeadersAsync()
+    {
+        // Arrange
+        var responseContent = await File.ReadAllTextAsync(UpdaloadFileTestDataFilePath);
+
+        this._messageHandlerStub.ResponseToReturn = new HttpResponseMessage(HttpStatusCode.OK)
+        {
+            Content = new StringContent(responseContent),
+        };
+
+        var plugin = new SessionsPythonPlugin(this._defaultSettings, this._httpClientFactory, (_) => Task.FromResult("test-auth-token"));
+
+        // Act
+        var result = await plugin.UploadFileAsync("test-file.txt", FileTestDataFilePath);
+
+        // Assert
+        Assert.NotNull(this._messageHandlerStub.RequestHeaders);
+
+        var userAgentHeaderValues = this._messageHandlerStub.RequestHeaders.GetValues("User-Agent").ToArray();
+        Assert.Equal(2, userAgentHeaderValues.Length);
+        Assert.Equal($"{HttpHeaderConstant.Values.UserAgent}/{s_assemblyVersion}", userAgentHeaderValues[0]);
+        Assert.Equal("(Language=dotnet)", userAgentHeaderValues[1]);
+
+        var authorizationHeaderValues = this._messageHandlerStub.RequestHeaders.GetValues("Authorization");
+        Assert.Single(authorizationHeaderValues, value => value == "Bearer test-auth-token");
+    }
+
     public void Dispose()
     {
         this._httpClient.Dispose();

dotnet/src/Plugins/Plugins.UnitTests/Plugins.UnitTests.csproj

@@ -38,13 +38,17 @@
     <ProjectReference Include="..\Plugins.Web\Plugins.Web.csproj" />
     <ProjectReference Include="..\Plugins.Memory\Plugins.Memory.csproj" />
   </ItemGroup>
-  
+
+  <ItemGroup>
+    <Compile Include="$(RepoRoot)/dotnet/src/InternalUtilities/src/Http/HttpHeaderConstant.cs" Link="%(RecursiveDir)%(Filename)%(Extension)" />
+  </ItemGroup>
+
   <ItemGroup>
     <None Update="TestData\*">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </None>
   </ItemGroup>
-  
+
   <ItemGroup>
     <Folder Include="Web\Tavily\" />
   </ItemGroup>