Add getBearerToken callback for BYOK providers (Managed Identity)

dotnet/src/BearerTokenProvider.cs

@@ -0,0 +1,32 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+using System.Diagnostics.CodeAnalysis;
+
+namespace GitHub.Copilot;
+
+/// <summary>
+/// Arguments passed to a bearer-token callback (the <c>GetBearerToken</c> property
+/// on <see cref="ProviderConfig"/> / <see cref="NamedProviderConfig"/>) when the
+/// runtime needs a fresh bearer token for a BYOK provider.
+/// </summary>
+/// <remarks>
+/// Part of the experimental managed-identity / bearer-token-provider surface and
+/// may change or be removed in future SDK or CLI releases.
+/// </remarks>
+[Experimental(Diagnostics.Experimental)]
+public sealed class ProviderTokenArgs
+{
+    /// <summary>
+    /// Name of the BYOK provider needing a token. For the singular, whole-session
+    /// <see cref="ProviderConfig"/> this is the implicit provider name
+    /// (<c>"default"</c>); for <see cref="NamedProviderConfig"/> entries it is
+    /// <see cref="NamedProviderConfig.Name"/>.
+    /// </summary>
+    /// <remarks>
+    /// The callback closes over its own token scope/audience; the runtime is
+    /// provider-agnostic and forwards only the provider name.
+    /// </remarks>
+    public required string ProviderName { get; init; }
+}

dotnet/src/Client.cs

@@ -652,6 +652,7 @@
         }
         ConfigureSessionFsHandlers(session, config.CreateSessionFsProvider);
         session.SetCanvasHandler(config.CanvasHandler);
+        session.RegisterBearerTokenProviders(BuildBearerTokenCallbacks(config));
         RegisterSession(session);
         session.StartProcessingEvents();
         LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
@@ -664,6 +665,37 @@
         return session;
     }
 
+    /// <summary>
+    /// Implicit provider name for the singular, whole-session <see cref="ProviderConfig"/>.
+    /// </summary>
+    private const string DefaultBearerTokenProviderName = "default";
+
+    /// <summary>
+    /// Collects the per-provider <c>GetBearerToken</c> callbacks keyed by
+    /// provider name for session-side registration. The singular, whole-session
+    /// <see cref="ProviderConfig"/> uses the implicit
+    /// <see cref="DefaultBearerTokenProviderName"/>.
+    /// </summary>
+    private static Dictionary<string, Func<ProviderTokenArgs, Task<string>>> BuildBearerTokenCallbacks(SessionConfigBase config)
+    {
+        var callbacks = new Dictionary<string, Func<ProviderTokenArgs, Task<string>>>(StringComparer.Ordinal);
+        if (config.Provider?.GetBearerToken is { } singular)
+        {
+            callbacks[DefaultBearerTokenProviderName] = singular;
+        }
+        if (config.Providers != null)
+        {
+            foreach (var provider in config.Providers)
+            {
+                if (provider.GetBearerToken is { } callback)
+                {
+                    callbacks[provider.Name] = callback;
+                }
+            }
+        }
+        return callbacks;
+    }
+
     /// <summary>
     /// Catches misuse of <see cref="SessionConfigBase.AvailableTools"/> /
     /// <see cref="SessionConfigBase.ExcludedTools"/> at the SDK boundary so

dotnet/src/Session.cs

@@ -58,6 +58,7 @@ public sealed partial class CopilotSession : IAsyncDisposable
 {
     private readonly Dictionary<string, AIFunction> _toolHandlers = [];
     private readonly Dictionary<string, Func<CommandContext, Task>> _commandHandlers = [];
+    private readonly Dictionary<string, Func<ProviderTokenArgs, Task<string>>> _bearerTokenProviders = new(StringComparer.Ordinal);
     private readonly ILogger _logger;
     private readonly CopilotClient _parentClient;
 
@@ -76,9 +77,7 @@ private sealed record EventSubscription(Type EventType, Action<SessionEvent> Han
     private Dictionary<string, Func<string, Task<string>>>? _transformCallbacks;
     private readonly SemaphoreSlim _transformCallbacksLock = new(1, 1);
 
-#pragma warning disable GHCP001
     private IReadOnlyList<OpenCanvasInstance> _openCanvases = Array.Empty<OpenCanvasInstance>();
-#pragma warning restore GHCP001
 
     private int _isDisposed;
 
@@ -126,7 +125,6 @@ public SessionCapabilities Capabilities
         private set;
     }
 
-#pragma warning disable GHCP001
     /// <summary>
     /// Canvas instances currently known to be open for this session.
     /// </summary>
@@ -136,7 +134,6 @@ public SessionCapabilities Capabilities
     /// </remarks>
     [Experimental(Diagnostics.Experimental)]
     public IReadOnlyList<OpenCanvasInstance> OpenCanvases => _openCanvases;
-#pragma warning restore GHCP001
 
     /// <summary>
     /// Gets the UI API for eliciting information from the user during this session.
@@ -873,6 +870,51 @@ internal void RegisterAutoModeSwitchHandler(Func<AutoModeSwitchRequest, AutoMode
         _autoModeSwitchHandler = handler;
     }
 
+    /// <summary>
+    /// Registers per-provider <c>GetBearerToken</c> callbacks for BYOK
+    /// providers configured with managed-identity / on-demand bearer-token auth.
+    /// </summary>
+    /// <remarks>
+    /// The runtime never receives the callback itself; the SDK strips it from the
+    /// provider config and instead sends <c>hasBearerTokenProvider: true</c>. When
+    /// the runtime needs a token it issues a session-scoped
+    /// <c>providerToken.getToken</c> request, which this handler routes to the
+    /// matching per-provider callback.
+    /// </remarks>
+    /// <param name="providers">Map of provider name to callback, or null/empty to clear.</param>
+    internal void RegisterBearerTokenProviders(IReadOnlyDictionary<string, Func<ProviderTokenArgs, Task<string>>>? providers)
+    {
+        _bearerTokenProviders.Clear();
+        if (providers is null || providers.Count == 0)
+        {
+            ClientSessionApis.ProviderToken = null;
+            return;
+        }
+        foreach (var (name, callback) in providers)
+        {
+            _bearerTokenProviders[name] = callback;
+        }
+        ClientSessionApis.ProviderToken = new BearerTokenProviderHandler(this);
+    }
+
+    /// <summary>
+    /// Routes runtime <c>providerToken.getToken</c> requests to the matching
+    /// per-provider <c>GetBearerToken</c> callback registered on the session.
+    /// </summary>
+    private sealed class BearerTokenProviderHandler(CopilotSession session) : IProviderTokenHandler
+    {
+        public async Task<ProviderTokenAcquireResult> GetTokenAsync(ProviderTokenAcquireRequest request, CancellationToken cancellationToken = default)
+        {
+            if (!session._bearerTokenProviders.TryGetValue(request.ProviderName, out var callback))
+            {
+                throw new InvalidOperationException(
+                    $"No bearer-token provider registered for provider \"{request.ProviderName}\"");
+            }
+            var token = await callback(new ProviderTokenArgs { ProviderName = request.ProviderName }).ConfigureAwait(false);
+            return new ProviderTokenAcquireResult { Token = token };
+        }
+    }
+
     /// <summary>
     /// Sets the capabilities reported by the host for this session.
     /// </summary>
@@ -882,7 +924,6 @@ internal void SetCapabilities(SessionCapabilities? capabilities)
         Capabilities = capabilities ?? new SessionCapabilities();
     }
 
-#pragma warning disable GHCP001
     internal void SetOpenCanvases(IList<OpenCanvasInstance>? canvases)
     {
         _openCanvases = canvases is { Count: > 0 }
@@ -959,7 +1000,6 @@ private static JsonElement SerializeActionResult(object? value)
         var element = CopilotClient.ToJsonElementForWire(value);
         return element ?? NullJsonElement;
     }
-#pragma warning restore GHCP001
 
     private sealed class CanvasHandlerAdapter(ICanvasHandler handler) : Rpc.ICanvasHandler
     {

dotnet/src/Types.cs

@@ -5,12 +5,14 @@
 using GitHub.Copilot.Rpc;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.Logging;
+using System;
 using System.ComponentModel;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json;
 using System.Text.Json.Nodes;
 using System.Text.Json.Serialization;
+using System.Threading.Tasks;
 
 namespace GitHub.Copilot;
 
@@ -2041,6 +2043,28 @@ public sealed class ProviderConfig
     [JsonPropertyName("bearerToken")]
     public string? BearerToken { get; set; }
 
+    /// <summary>
+    /// Wire-only flag, emitted automatically when <see cref="GetBearerToken"/> is set, that tells
+    /// the runtime to request a token over the session-scoped <c>providerToken.getToken</c> RPC
+    /// before each outbound request to this provider. Derived from <see cref="GetBearerToken"/>;
+    /// internal and never part of the public API.
+    /// </summary>
+    [JsonInclude]
+    [JsonPropertyName("hasBearerTokenProvider")]
+    internal bool? HasBearerTokenProvider => GetBearerToken is not null ? true : null;
+
+    /// <summary>
+    /// Per-request callback that resolves a bearer token on demand for this BYOK provider (for
+    /// example via Azure Managed Identity). The Copilot SDK takes no identity dependency: supply a
+    /// callback backed by your own identity library. Never serialized — setting it makes the SDK send
+    /// <c>hasBearerTokenProvider: true</c> on the wire and answer the runtime's
+    /// <c>providerToken.getToken</c> requests. Mutually exclusive with <see cref="ApiKey"/> and
+    /// <see cref="BearerToken"/>.
+    /// </summary>
+    [JsonIgnore]
+    [Experimental(Diagnostics.Experimental)]
+    public Func<ProviderTokenArgs, Task<string>>? GetBearerToken { get; set; }
+
     /// <summary>
     /// Azure-specific configuration options.
     /// </summary>
@@ -2173,6 +2197,28 @@ public sealed class NamedProviderConfig
     [JsonPropertyName("bearerToken")]
     public string? BearerToken { get; set; }
 
+    /// <summary>
+    /// Wire-only flag, emitted automatically when <see cref="GetBearerToken"/> is set, that tells
+    /// the runtime to request a token over the session-scoped <c>providerToken.getToken</c> RPC
+    /// before each outbound request to this provider. Derived from <see cref="GetBearerToken"/>;
+    /// internal and never part of the public API.
+    /// </summary>
+    [JsonInclude]
+    [JsonPropertyName("hasBearerTokenProvider")]
+    internal bool? HasBearerTokenProvider => GetBearerToken is not null ? true : null;
+
+    /// <summary>
+    /// Per-request callback that resolves a bearer token on demand for this BYOK provider (for
+    /// example via Azure Managed Identity). The Copilot SDK takes no identity dependency: supply a
+    /// callback backed by your own identity library. Never serialized — setting it makes the SDK send
+    /// <c>hasBearerTokenProvider: true</c> on the wire and answer the runtime's
+    /// <c>providerToken.getToken</c> requests. Mutually exclusive with <see cref="ApiKey"/> and
+    /// <see cref="BearerToken"/>.
+    /// </summary>
+    [JsonIgnore]
+    [Experimental(Diagnostics.Experimental)]
+    public Func<ProviderTokenArgs, Task<string>>? GetBearerToken { get; set; }
+
     /// <summary>
     /// Azure-specific configuration options.
     /// </summary>