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
@@ -839,15 +871,13 @@
 
         try
         {
-#pragma warning disable GHCP001
             await session.Rpc.Options.UpdateAsync(
                 skipCustomInstructions: skipCustomInstructions,
                 customAgentsLocalOnly: customAgentsLocalOnly,
                 coauthorEnabled: coauthorEnabled,
                 manageScheduleEnabled: manageScheduleEnabled,
                 installedPlugins: installedPlugins,
                 cancellationToken: cancellationToken).ConfigureAwait(false);
-#pragma warning restore GHCP001
         }
         catch
         {
@@ -2436,7 +2466,6 @@
         IList<string>? PluginDirectories = null,
         LargeToolOutputConfig? LargeOutput = null,
         MemoryConfiguration? Memory = null,
-#pragma warning disable GHCP001
         IList<CanvasDeclaration>? Canvases = null,
         bool? RequestCanvasRenderer = null,
         bool? RequestExtensions = null,
@@ -2445,7 +2474,6 @@
         IList<NamedProviderConfig>? Providers = null,
         IList<ProviderModelConfig>? Models = null,
         OptionsUpdateToolFilterPrecedence? ToolFilterPrecedence = null);
-#pragma warning restore GHCP001
 
     internal record ToolDefinition(
         string Name,
@@ -2471,9 +2499,7 @@
         string SessionId,
         string? WorkspacePath,
         SessionCapabilities? Capabilities = null,
-#pragma warning disable GHCP001
         IList<OpenCanvasInstance>? OpenCanvases = null);
-#pragma warning restore GHCP001
 
     internal record ResumeSessionRequest(
         string SessionId,
@@ -2530,7 +2556,6 @@
         IList<string>? PluginDirectories = null,
         LargeToolOutputConfig? LargeOutput = null,
         MemoryConfiguration? Memory = null,
-#pragma warning disable GHCP001
         IList<CanvasDeclaration>? Canvases = null,
         bool? RequestCanvasRenderer = null,
         bool? RequestExtensions = null,
@@ -2540,15 +2565,12 @@
         IList<NamedProviderConfig>? Providers = null,
         IList<ProviderModelConfig>? Models = null,
         OptionsUpdateToolFilterPrecedence? ToolFilterPrecedence = null);
-#pragma warning restore GHCP001
 
     internal record ResumeSessionResponse(
         string SessionId,
         string? WorkspacePath,
         SessionCapabilities? Capabilities = null,
-#pragma warning disable GHCP001
         IList<OpenCanvasInstance>? OpenCanvases = null);
-#pragma warning restore GHCP001
 
     internal record CommandWireDefinition(
         string Name,

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 }
@@ -962,7 +1003,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
     {