Add Rust BYOK get_bearer_token glue + e2e

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: SteveSandersonMS

rust/src/lib.rs

@@ -22,6 +22,9 @@ pub mod hooks;
 mod jsonrpc;
 /// Permission-policy helpers that produce a [`handler::PermissionHandler`].
 pub mod permission;
+/// BYOK bearer-token provider callbacks.
+pub mod provider_token;
+mod provider_token_dispatch;
 /// GitHub Copilot CLI binary resolution (env var, embedded, dev cache).
 pub(crate) mod resolve;
 mod router;
@@ -72,6 +75,7 @@ pub(crate) use jsonrpc::{
     JsonRpcClient, JsonRpcError, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, error_codes,
 };
 pub use mode::{BUILTIN_TOOLS_ISOLATED, ClientMode, ToolSet};
+pub use provider_token::{BearerTokenError, BearerTokenProvider, ProviderTokenArgs};
 
 /// Re-exported JSON-RPC internals for integration tests (requires `test-support` feature).
 #[cfg(feature = "test-support")]

rust/src/provider_token.rs

@@ -0,0 +1,105 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+//! BYOK bearer-token provider callbacks.
+//!
+//! <div class="warning">
+//!
+//! **Experimental.** These types are part of an experimental wire-protocol
+//! surface and may change or be removed in future SDK or CLI releases.
+//!
+//! </div>
+
+use std::future::Future;
+
+use async_trait::async_trait;
+
+/// Arguments passed to a BYOK bearer-token provider callback.
+///
+/// <div class="warning">
+///
+/// **Experimental.** This type is part of an experimental wire-protocol
+/// surface and may change or be removed in future SDK or CLI releases.
+///
+/// </div>
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ProviderTokenArgs {
+    /// Name of the BYOK provider needing a token.
+    ///
+    /// This is `"default"` for the singular whole-session provider, otherwise
+    /// the named provider's `name`.
+    pub provider_name: String,
+}
+
+/// Error returned by a [`BearerTokenProvider`].
+///
+/// <div class="warning">
+///
+/// **Experimental.** This type is part of an experimental wire-protocol
+/// surface and may change or be removed in future SDK or CLI releases.
+///
+/// </div>
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct BearerTokenError {
+    message: String,
+}
+
+impl BearerTokenError {
+    /// Construct a bearer-token error with a human-readable message.
+    pub fn message(message: impl Into<String>) -> Self {
+        Self {
+            message: message.into(),
+        }
+    }
+
+    /// Return the human-readable error message.
+    pub fn as_str(&self) -> &str {
+        &self.message
+    }
+}
+
+impl std::fmt::Display for BearerTokenError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(&self.message)
+    }
+}
+
+impl std::error::Error for BearerTokenError {}
+
+impl From<String> for BearerTokenError {
+    fn from(message: String) -> Self {
+        Self::message(message)
+    }
+}
+
+impl From<&str> for BearerTokenError {
+    fn from(message: &str) -> Self {
+        Self::message(message)
+    }
+}
+
+/// Provider-side callback used to acquire bearer tokens for BYOK providers.
+///
+/// <div class="warning">
+///
+/// **Experimental.** This trait is part of an experimental wire-protocol
+/// surface and may change or be removed in future SDK or CLI releases.
+///
+/// </div>
+#[async_trait]
+pub trait BearerTokenProvider: Send + Sync {
+    /// Acquire a bearer token without the `Bearer ` prefix.
+    async fn get_token(&self, args: ProviderTokenArgs) -> Result<String, BearerTokenError>;
+}
+
+#[async_trait]
+impl<F, Fut> BearerTokenProvider for F
+where
+    F: Fn(ProviderTokenArgs) -> Fut + Send + Sync,
+    Fut: Future<Output = Result<String, BearerTokenError>> + Send,
+{
+    async fn get_token(&self, args: ProviderTokenArgs) -> Result<String, BearerTokenError> {
+        (self)(args).await
+    }
+}

rust/src/provider_token_dispatch.rs

@@ -0,0 +1,157 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+//! Inbound `providerToken.*` JSON-RPC request dispatch helpers.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use serde::Serialize;
+use serde_json::Value;
+use tracing::warn;
+
+use crate::generated::api_types::{
+    ProviderTokenAcquireRequest, ProviderTokenAcquireResult, rpc_methods,
+};
+use crate::provider_token::{BearerTokenError, BearerTokenProvider, ProviderTokenArgs};
+use crate::{Client, JsonRpcRequest, JsonRpcResponse, error_codes};
+
+async fn respond<T: Serialize>(client: &Client, request_id: u64, result: T) {
+    let value = match serde_json::to_value(&result) {
+        Ok(value) => value,
+        Err(error) => {
+            warn!(error = %error, "failed to serialize provider token response");
+            send_error(
+                client,
+                request_id,
+                error_codes::INTERNAL_ERROR,
+                "serialization failure",
+            )
+            .await;
+            return;
+        }
+    };
+
+    let _ = client
+        .send_response(&JsonRpcResponse {
+            jsonrpc: "2.0".to_string(),
+            id: request_id,
+            result: Some(value),
+            error: None,
+        })
+        .await;
+}
+
+async fn send_error(client: &Client, request_id: u64, code: i32, message: &str) {
+    let _ = client
+        .send_response(&JsonRpcResponse {
+            jsonrpc: "2.0".to_string(),
+            id: request_id,
+            result: None,
+            error: Some(crate::JsonRpcError {
+                code,
+                message: message.to_string(),
+                data: None,
+            }),
+        })
+        .await;
+}
+
+async fn parse_params<T: serde::de::DeserializeOwned>(
+    client: &Client,
+    request: &JsonRpcRequest,
+) -> Option<T> {
+    let params = request
+        .params
+        .as_ref()
+        .cloned()
+        .unwrap_or(Value::Object(serde_json::Map::new()));
+    match serde_json::from_value(params) {
+        Ok(params) => Some(params),
+        Err(error) => {
+            send_error(
+                client,
+                request.id,
+                error_codes::INVALID_PARAMS,
+                &format!("invalid params: {error}"),
+            )
+            .await;
+            None
+        }
+    }
+}
+
+fn token_provider_or_err(
+    providers: &HashMap<String, Arc<dyn BearerTokenProvider>>,
+    provider_name: &str,
+) -> Result<Arc<dyn BearerTokenProvider>, BearerTokenError> {
+    providers.get(provider_name).cloned().ok_or_else(|| {
+        BearerTokenError::message(format!(
+            "No bearer-token provider installed for BYOK provider {provider_name:?}"
+        ))
+    })
+}
+
+async fn get_token(
+    client: &Client,
+    providers: &HashMap<String, Arc<dyn BearerTokenProvider>>,
+    request: JsonRpcRequest,
+) {
+    let Some(params) = parse_params::<ProviderTokenAcquireRequest>(client, &request).await else {
+        return;
+    };
+
+    let token_provider = match token_provider_or_err(providers, &params.provider_name) {
+        Ok(provider) => provider,
+        Err(error) => {
+            send_error(
+                client,
+                request.id,
+                error_codes::INTERNAL_ERROR,
+                &error.to_string(),
+            )
+            .await;
+            return;
+        }
+    };
+
+    match token_provider
+        .get_token(ProviderTokenArgs {
+            provider_name: params.provider_name,
+        })
+        .await
+    {
+        Ok(token) => respond(client, request.id, ProviderTokenAcquireResult { token }).await,
+        Err(error) => {
+            send_error(
+                client,
+                request.id,
+                error_codes::INTERNAL_ERROR,
+                &format!("Bearer-token provider failed: {error}"),
+            )
+            .await;
+        }
+    }
+}
+
+pub(crate) async fn dispatch(
+    client: &Client,
+    providers: &HashMap<String, Arc<dyn BearerTokenProvider>>,
+    request: JsonRpcRequest,
+) {
+    let method = request.method.as_str();
+    match method {
+        rpc_methods::PROVIDERTOKEN_GETTOKEN => get_token(client, providers, request).await,
+        _ => {
+            warn!(method = %method, "unknown providerToken.* method");
+            send_error(
+                client,
+                request.id,
+                error_codes::METHOD_NOT_FOUND,
+                &format!("unknown method: {method}"),
+            )
+            .await;
+        }
+    }
+}

rust/src/session.rs

@@ -21,6 +21,7 @@ use crate::handler::{
     PermissionHandler, PermissionResult, UserInputHandler, UserInputResponse,
 };
 use crate::hooks::SessionHooks;
+use crate::provider_token::BearerTokenProvider;
 use crate::session_fs::SessionFsProvider;
 use crate::trace_context::inject_trace_context;
 use crate::transforms::SystemMessageTransform;
@@ -893,6 +894,7 @@ impl Client {
         let command_handlers = build_command_handler_map(runtime.commands.as_deref());
         let canvas_handler = runtime.canvas_handler.take();
         let session_fs_provider = runtime.session_fs_provider.take();
+        let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
         if self.inner.session_fs_configured && session_fs_provider.is_none() {
             return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
         }
@@ -1011,6 +1013,7 @@ impl Client {
             command_handlers,
             canvas_handler,
             session_fs_provider,
+            bearer_token_providers,
             channels,
             idle_waiter.clone(),
             capabilities.clone(),
@@ -1149,6 +1152,7 @@ impl Client {
         let command_handlers = build_command_handler_map(runtime.commands.as_deref());
         let canvas_handler = runtime.canvas_handler.take();
         let session_fs_provider = runtime.session_fs_provider.take();
+        let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
         if self.inner.session_fs_configured && session_fs_provider.is_none() {
             return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
         }
@@ -1183,6 +1187,7 @@ impl Client {
             command_handlers,
             canvas_handler,
             session_fs_provider,
+            bearer_token_providers,
             channels,
             idle_waiter.clone(),
             capabilities.clone(),
@@ -1391,6 +1396,7 @@ fn spawn_event_loop(
     command_handlers: Arc<CommandHandlerMap>,
     canvas_handler: Option<Arc<dyn CanvasHandler>>,
     session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
+    bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
     channels: crate::router::SessionChannels,
     idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
     capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
@@ -1432,6 +1438,7 @@ fn spawn_event_loop(
                             transforms: transforms.as_deref(),
                             canvas_handler: canvas_handler.as_ref(),
                             session_fs_provider: session_fs_provider.as_ref(),
+                            bearer_token_providers: &bearer_token_providers,
                         };
                         handle_request(&session_id, ctx, request).await;
                     }
@@ -2010,6 +2017,7 @@ struct RequestDispatchContext<'a> {
     transforms: Option<&'a dyn SystemMessageTransform>,
     canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
     session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
+    bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
 }
 
 /// Process a JSON-RPC request from the CLI.
@@ -2025,6 +2033,7 @@ async fn handle_request(
     let transforms = ctx.transforms;
     let canvas_handler = ctx.canvas_handler;
     let session_fs_provider = ctx.session_fs_provider;
+    let bearer_token_providers = ctx.bearer_token_providers;
 
     if request.method.starts_with("sessionFs.") {
         crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
@@ -2036,6 +2045,11 @@ async fn handle_request(
         return;
     }
 
+    if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
+        crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
+        return;
+    }
+
     match request.method.as_str() {
         "hooks.invoke" => {
             let params = request.params.as_ref();