From 4b9d7d989f73aee8143d9392a636d90bed111cfc Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Wed, 26 Aug 2026 16:38:43 +0200 Subject: [PATCH 1/4] Add session GitHub token providers Expose lifecycle-safe GitHub credential callbacks across all six SDKs, with idiomatic APIs, tagged results, tests, and documentation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: a447a8bc-8687-47ea-a545-4370979e8128 --- CHANGELOG.md | 6 + docs/auth/authenticate.md | 123 ++++++ docs/setup/multi-tenancy.md | 4 +- dotnet/README.md | 17 + dotnet/src/Client.cs | 168 ++++++-- dotnet/src/GitHubTokenProvider.cs | 77 ++++ dotnet/src/JsonRpc.cs | 26 +- dotnet/src/Session.cs | 15 + dotnet/src/Types.cs | 9 + .../test/Unit/ClientSessionLifetimeTests.cs | 236 ++++++++++- go/README.md | 17 + go/client.go | 266 ++++++++++--- go/github_token_provider.go | 82 ++++ go/github_token_provider_test.go | 344 ++++++++++++++++ go/session.go | 108 +++-- go/types.go | 8 + java/README.md | 14 + .../com/github/copilot/CopilotClient.java | 56 ++- .../com/github/copilot/CopilotSession.java | 10 + .../copilot/GitHubTokenProviderRegistry.java | 79 ++++ .../com/github/copilot/JsonRpcClient.java | 35 +- .../github/copilot/RpcHandlerDispatcher.java | 74 +++- .../copilot/rpc/CreateSessionRequest.java | 18 + .../copilot/rpc/GitHubTokenProvider.java | 29 ++ .../copilot/rpc/GitHubTokenProviderArgs.java | 22 + .../rpc/GitHubTokenProviderResult.java | 113 ++++++ .../copilot/rpc/ResumeSessionConfig.java | 30 ++ .../copilot/rpc/ResumeSessionRequest.java | 18 + .../com/github/copilot/rpc/SessionConfig.java | 39 +- .../GitHubTokenProviderRegistryTest.java | 74 ++++ .../com/github/copilot/JsonRpcClientTest.java | 16 + .../copilot/RpcHandlerDispatcherTest.java | 59 ++- .../copilot/SessionRequestBuilderTest.java | 22 + nodejs/README.md | 11 + nodejs/src/client.ts | 117 +++++- nodejs/src/index.ts | 5 + nodejs/src/session.ts | 18 +- nodejs/src/types.ts | 39 ++ nodejs/test/github-token-provider.test.ts | 289 ++++++++++++++ python/README.md | 13 + python/copilot/__init__.py | 16 + python/copilot/client.py | 197 ++++++++- python/copilot/session.py | 9 + python/test_github_token_provider.py | 269 +++++++++++++ rust/README.md | 23 ++ rust/src/errors.rs | 3 + rust/src/github_token.rs | 375 ++++++++++++++++++ rust/src/lib.rs | 36 ++ rust/src/router.rs | 5 + rust/src/session.rs | 30 ++ rust/src/types.rs | 71 +++- rust/src/wire.rs | 10 + rust/tests/session_test.rs | 208 ++++++++++ 53 files changed, 3793 insertions(+), 165 deletions(-) create mode 100644 dotnet/src/GitHubTokenProvider.cs create mode 100644 go/github_token_provider.go create mode 100644 go/github_token_provider_test.go create mode 100644 java/sdk/src/main/java/com/github/copilot/GitHubTokenProviderRegistry.java create mode 100644 java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java create mode 100644 java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderArgs.java create mode 100644 java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderResult.java create mode 100644 java/sdk/src/test/java/com/github/copilot/GitHubTokenProviderRegistryTest.java create mode 100644 nodejs/test/github-token-provider.test.ts create mode 100644 python/test_github_token_provider.py create mode 100644 rust/src/github_token.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a5974c67d7..bc3b41c71c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ See [GitHub Releases](https://github.com/github/copilot-sdk/releases) for the fu ## [Unreleased] +### Feature: rotating session-scoped GitHub credentials + +All six SDKs can now acquire short-lived GitHub credentials through a session-scoped callback. The SDK registers the callback before session create or resume, maps `initial` and `refresh` requests to the owning session, and removes registrations on rollback, replacement, session close, and client close. Static per-session `gitHubToken` credentials remain supported and are mutually exclusive with the callback. + +Token responses use the shared tagged token/cancelled shape and require `expiresIn`, expressed as the positive number of seconds remaining when the callback completes. See [github/copilot-agent-runtime#16381](https://github.com/github/copilot-agent-runtime/pull/16381) for the runtime credential-authority implementation. + ### Feature: extensions can request sensitive environment variables Copilot CLI extensions can now ask for named sensitive environment variables when they join a session. `joinSession()` accepts a `requestedEnvironmentVariables` option listing the variable names the extension needs. The CLI shows a permission prompt naming the extension and the exact variables requested. On approval, only those variables reach that extension and their values are written into the extension process's `process.env` before `joinSession()` resolves. On denial, `joinSession()` rejects, the extension does not load, and its tools never reach the model. diff --git a/docs/auth/authenticate.md b/docs/auth/authenticate.md index 2fc80dd0aa..1fe3246547 100644 --- a/docs/auth/authenticate.md +++ b/docs/auth/authenticate.md @@ -262,6 +262,129 @@ const client = new CopilotClient({ For more information, see [GitHub OAuth](../setup/github-oauth.md). +## Rotating session-scoped GitHub tokens + +For multi-user services and integrations, set a token provider on each session instead of storing one long-lived token. The runtime calls the provider for the effective GitHub host and identifies the request as `initial` or `refresh`. The session ID is absent only when a cloud session has not received its ID yet. + +Return a tagged token result or an explicit cancellation. Every token result must include `expiresIn`: the positive number of seconds remaining when the callback completes. Production GitHub tokens typically last eight hours, so `8 * 60 * 60` is a common value. Do not set both the static per-session token and the provider. + +
+TypeScript + + +```typescript +const session = await client.createSession({ + gitHubTokenProvider: async ({ host, sessionId, reason }) => { + const token = await acquireGitHubToken({ host, sessionId, reason }); + return { + kind: "token", + accessToken: token.value, + expiresIn: token.secondsRemaining, + }; + }, +}); +``` + +
+
+Python + + +```python +async def provide_github_token(args): + token = await acquire_github_token( + host=args["host"], + session_id=args["session_id"], + reason=args["reason"], + ) + return { + "kind": "token", + "accessToken": token.value, + "expiresIn": token.seconds_remaining, + } + + +session = await client.create_session(github_token_provider=provide_github_token) +``` + +
+
+Go + + +```go +session, err := client.CreateSession(ctx, &copilot.SessionConfig{ + GitHubTokenProvider: func(args copilot.GitHubTokenProviderArgs) (*copilot.GitHubTokenProviderResult, error) { + token, secondsRemaining, err := acquireGitHubToken(args.Host, args.SessionID, args.Reason) + if err != nil { + return nil, err + } + return copilot.GitHubTokenResult(&copilot.GitHubToken{ + AccessToken: token, + ExpiresIn: secondsRemaining, + }), nil + }, +}) +``` + +
+
+.NET + + +```csharp +await using var session = await client.CreateSessionAsync(new SessionConfig +{ + GitHubTokenProvider = async args => + { + var token = await AcquireGitHubTokenAsync(args.Host, args.SessionId, args.Reason); + return GitHubTokenProviderResult.FromToken(new GitHubToken + { + AccessToken = token.Value, + ExpiresIn = token.SecondsRemaining, + }); + }, +}); +``` + +
+
+Java + + +```java +var session = client.createSession(new SessionConfig() + .setGitHubTokenProvider(args -> + acquireGitHubToken(args.host(), args.sessionId(), args.reason()) + .thenApply(token -> GitHubTokenProviderResult.token( + token.value(), token.secondsRemaining()))) + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL) +).get(); +``` + +
+
+Rust + + +```rust +let provider = Arc::new(|args: GitHubTokenProviderArgs| async move { + let token = acquire_github_token(&args.host, args.session_id.as_ref(), args.reason).await?; + Ok(GitHubTokenProviderResult::Token(GitHubToken::new( + token.value, + token.seconds_remaining, + ))) +}); + +let session = client + .create_session(SessionConfig::default().with_github_token_provider(provider)) + .await?; +``` + +
+ +Return the language's cancelled result when the user or credential broker cancels acquisition. Provider exceptions follow the SDK's normal callback error behavior. The runtime refreshes during async preflight when one hour or less remains; it does not use background timers, rejection-driven replay, 401/403 challenge propagation, or upscope for this callback. + ## Environment variables For automation, CI/CD pipelines, and server-to-server scenarios, you can authenticate using environment variables. diff --git a/docs/setup/multi-tenancy.md b/docs/setup/multi-tenancy.md index 2f82dde0bf..d22f8fd3e1 100644 --- a/docs/setup/multi-tenancy.md +++ b/docs/setup/multi-tenancy.md @@ -24,7 +24,9 @@ This guide is a sister to [Scaling and multi-tenancy](./scaling.md). Use that gu | `baseDirectory` | Isolating `COPILOT_HOME` per runtime instance | Ignored when connecting to an existing runtime. | | `sessionFs` | Routing session filesystem storage off local disk | Pair with per-session filesystem providers. | | `RuntimeConnection.forUri(url)` | Sharing one already-running runtime | Language names vary; see samples below. | -| Per-session `gitHubToken` | Scoping auth to the requesting user | Prefer this over a single shared user token. | +| Per-session GitHub token or provider | Scoping auth to the requesting user | Prefer a rotating provider for short-lived credentials; use a static `gitHubToken` only when rotation is unnecessary. | + +For callback-backed credentials, see [Rotating session-scoped GitHub tokens](../auth/authenticate.md#rotating-session-scoped-github-tokens). Each session owns its provider registration, so concurrent sessions can use different GitHub hosts and accounts without sharing callback state. ### `mode: "empty"` diff --git a/dotnet/README.md b/dotnet/README.md index 6efd6e094c..ba935c190a 100644 --- a/dotnet/README.md +++ b/dotnet/README.md @@ -133,6 +133,7 @@ Create a new conversation session. - `InfiniteSessions` - Configure automatic context compaction (see below) - `WorkingDirectory` - Working directory for the session. When not set, the runtime uses its own process working directory. - `EnableSessionStore` - Enables the cross-session store for search and retrieval across sessions. When unset in `CopilotClientMode.CopilotCli`, the runtime default applies (enabled). In `CopilotClientMode.Empty`, defaults to disabled. +- `GitHubTokenProvider` - Acquires session-scoped GitHub tokens on demand. Return `GitHubTokenProviderResult.FromToken` with a positive `ExpiresIn` value (production GitHub tokens typically use `8 * 60 * 60` seconds), or `GitHubTokenProviderResult.Cancel()`. Cannot be combined with `GitHubToken`. - `OnPermissionRequest` - Optional handler called before each tool execution to approve or deny it. When omitted, permission requests are emitted as events and left pending for manual resolution. `PermissionHandler.ApproveAll` approves requests when managed settings are disabled and throws when `EnableManagedSettings` is true. Custom handlers can inspect `ManagedApprovalRequired` for human-facing confirmation logic. See [Permission Handling](#permission-handling) section. - `OnUserInputRequest` - Handler for user input requests from the agent (enables ask_user tool). See [User Input Requests](#user-input-requests) section. - `Hooks` - Hook handlers for session lifecycle events. See [Session Hooks](#session-hooks) section. @@ -144,6 +145,22 @@ Resume an existing session. Returns the session with `WorkspacePath` populated i **ResumeSessionConfig:** - `OnPermissionRequest` - Optional handler called before each tool execution to approve or deny it. See [Permission Handling](#permission-handling) section. +- `GitHubTokenProvider` - Replaces the session-scoped token provider when resuming. Cannot be combined with `GitHubToken`. + +```csharp +await using var session = await client.CreateSessionAsync(new SessionConfig +{ + GitHubTokenProvider = async args => + { + var token = await AcquireTokenAsync(args.Host); + return GitHubTokenProviderResult.FromToken(new GitHubToken + { + AccessToken = token, + ExpiresIn = 8 * 60 * 60 + }); + } +}); +``` ##### `PingAsync(string? message = null): Task` diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index 86178ba74b..b886a4a42b 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -70,6 +70,7 @@ public sealed partial class CopilotClient : IDisposable, IAsyncDisposable /// that has not been explicitly disposed or removed. /// internal readonly ConcurrentDictionary _sessions = new(); + private readonly ConcurrentDictionary>> _gitHubTokenProviders = new(); private readonly CopilotClientOptions _options; private readonly RuntimeConnection _connection; @@ -91,8 +92,7 @@ public sealed partial class CopilotClient : IDisposable, IAsyncDisposable /// /// Client-global RPC handlers (e.g. the LLM inference provider adapter), - /// built once at construction when the corresponding option is configured and - /// registered on every connection. Null when no client-global API is enabled. + /// built once at construction and registered on every connection. /// private readonly ClientGlobalApiHandlers? _clientGlobalApis; @@ -541,6 +541,7 @@ public async Task StopAsync() } _sessions.Clear(); + ClearGitHubTokenProviders(); await CleanupConnectionAsync(errors, gracefulRuntimeShutdown: true); @@ -572,6 +573,7 @@ public async Task StopAsync() public async Task ForceStopAsync() { _sessions.Clear(); + ClearGitHubTokenProviders(); var errors = new List(); await CleanupConnectionAsync(errors, gracefulRuntimeShutdown: false); @@ -1119,6 +1121,7 @@ await session.Rpc.Options.UpdateAsync( public async Task CreateSessionAsync(SessionConfig config, CancellationToken cancellationToken = default) { ArgumentNullException.ThrowIfNull(config); + ValidateGitHubTokenConfig(config); var connection = await EnsureConnectedAsync(cancellationToken); var totalTimestamp = Stopwatch.GetTimestamp(); @@ -1153,19 +1156,22 @@ public async Task CreateSessionAsync(SessionConfig config, Cance ? null : (string.IsNullOrEmpty(config.SessionId) ? Guid.NewGuid().ToString() : config.SessionId); + var registrationId = RegisterGitHubTokenProvider(config.GitHubTokenProvider); + var registrationTransferred = false; CopilotSession? session = null; - if (localSessionId != null) - { - session = InitializeSession( - localSessionId, - connection.Rpc, - config, - transformCallbacks, - hasHooks, - "CopilotClient.CreateSessionAsync"); - } try { + if (localSessionId != null) + { + session = InitializeSession( + localSessionId, + connection.Rpc, + config, + transformCallbacks, + hasHooks, + "CopilotClient.CreateSessionAsync"); + } + var (traceparent, tracestate) = TelemetryHelpers.GetTraceContext(); var request = new CreateSessionRequest( @@ -1222,6 +1228,7 @@ public async Task CreateSessionAsync(SessionConfig config, Cance Tracestate: tracestate, ModelCapabilities: config.ModelCapabilities, GitHubToken: config.GitHubToken, + GitHubTokenProviderRegistrationId: registrationId, RemoteSession: config.RemoteSession, Cloud: config.Cloud, InstructionDirectories: config.InstructionDirectories, @@ -1301,6 +1308,11 @@ public async Task CreateSessionAsync(SessionConfig config, Cance session.SetOpenCanvases(response.OpenCanvases); await UpdateSessionOptionsForModeAsync(session, config, cancellationToken).ConfigureAwait(false); + if (registrationId is not null) + { + session.SetGitHubTokenProviderRegistration(registrationId); + registrationTransferred = true; + } } catch (Exception ex) { @@ -1316,6 +1328,13 @@ public async Task CreateSessionAsync(SessionConfig config, Cance throw; } + finally + { + if (!registrationTransferred && registrationId is not null) + { + UnregisterGitHubTokenProvider(registrationId); + } + } LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null, "CopilotClient.CreateSessionAsync complete. Elapsed={Elapsed}, SessionId={SessionId}", @@ -1353,6 +1372,7 @@ public async Task ResumeSessionAsync(string sessionId, ResumeSes { ArgumentNullException.ThrowIfNull(sessionId); ArgumentNullException.ThrowIfNull(config); + ValidateGitHubTokenConfig(config); var connection = await EnsureConnectedAsync(cancellationToken); var totalTimestamp = Stopwatch.GetTimestamp(); @@ -1375,17 +1395,21 @@ public async Task ResumeSessionAsync(string sessionId, ResumeSes var (wireSystemMessage, transformCallbacks) = ExtractTransformCallbacks(config.SystemMessage); - // Create and register the session before issuing the RPC so that - // events emitted by the CLI (e.g. session.start) are not dropped. - var session = InitializeSession( - sessionId, - connection.Rpc, - config, - transformCallbacks, - hasHooks, - "CopilotClient.ResumeSessionAsync"); + var registrationId = RegisterGitHubTokenProvider(config.GitHubTokenProvider); + var registrationTransferred = false; + CopilotSession? session = null; try { + // Create and register the session before issuing the RPC so that + // events emitted by the CLI (e.g. session.start) are not dropped. + session = InitializeSession( + sessionId, + connection.Rpc, + config, + transformCallbacks, + hasHooks, + "CopilotClient.ResumeSessionAsync"); + var (traceparent, tracestate) = TelemetryHelpers.GetTraceContext(); var request = new ResumeSessionRequest( @@ -1443,6 +1467,7 @@ public async Task ResumeSessionAsync(string sessionId, ResumeSes Tracestate: tracestate, ModelCapabilities: config.ModelCapabilities, GitHubToken: config.GitHubToken, + GitHubTokenProviderRegistrationId: registrationId, RemoteSession: config.RemoteSession, ContinuePendingWork: config.ContinuePendingWork, InstructionDirectories: config.InstructionDirectories, @@ -1486,10 +1511,15 @@ public async Task ResumeSessionAsync(string sessionId, ResumeSes } await UpdateSessionOptionsForModeAsync(session, config, cancellationToken).ConfigureAwait(false); + if (registrationId is not null) + { + session.SetGitHubTokenProviderRegistration(registrationId); + registrationTransferred = true; + } } catch (Exception ex) { - session.RemoveFromClient(); + session?.RemoveFromClient(); if (ex is not OperationCanceledException) { LoggingHelpers.LogTiming(_logger, LogLevel.Warning, ex, @@ -1499,12 +1529,19 @@ public async Task ResumeSessionAsync(string sessionId, ResumeSes } throw; } + finally + { + if (!registrationTransferred && registrationId is not null) + { + UnregisterGitHubTokenProvider(registrationId); + } + } LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null, "CopilotClient.ResumeSessionAsync complete. Elapsed={Elapsed}, SessionId={SessionId}", totalTimestamp, sessionId); - return session; + return session!; } /// @@ -1946,25 +1983,94 @@ await Rpc.SessionFs.SetProviderAsync( /// /// Builds the client-global RPC handler bag at construction time. Registers /// the LLM inference provider adapter and/or the GitHub telemetry adapter - /// depending on which options are configured; returns null when no - /// client-global API is configured so the registration is skipped entirely. + /// depending on which options are configured. The GitHub token dispatcher is + /// always registered because providers are configured per session. /// private ClientGlobalApiHandlers? BuildClientGlobalApis() { var handler = _options.RequestHandler; var onGitHubTelemetry = _options.OnGitHubTelemetry; - if (handler is null && onGitHubTelemetry is null) - { - return null; - } - return new ClientGlobalApiHandlers { LlmInference = handler is null ? null : new LlmInferenceAdapter(handler, () => _serverRpc), GitHubTelemetry = onGitHubTelemetry is null ? null : new GitHubTelemetryAdapter(onGitHubTelemetry, _logger), + GitHubToken = new GitHubTokenAdapter(this), }; } + private static void ValidateGitHubTokenConfig(SessionConfigBase config) + { + if (!string.IsNullOrEmpty(config.GitHubToken) && config.GitHubTokenProvider is not null) + { + throw new ArgumentException( + $"{nameof(SessionConfigBase.GitHubToken)} and {nameof(SessionConfigBase.GitHubTokenProvider)} cannot be used together.", + nameof(config)); + } + } + + private string? RegisterGitHubTokenProvider( + Func>? provider) + { + if (provider is null) + { + return null; + } + + var registrationId = Guid.NewGuid().ToString(); + if (!_gitHubTokenProviders.TryAdd(registrationId, provider)) + { + throw new InvalidOperationException("Failed to register GitHub token provider."); + } + return registrationId; + } + + internal void UnregisterGitHubTokenProvider(string registrationId) + => _gitHubTokenProviders.TryRemove(registrationId, out _); + + private void ClearGitHubTokenProviders() => _gitHubTokenProviders.Clear(); + + private sealed class GitHubTokenAdapter(CopilotClient client) : IGitHubTokenHandler + { + public async Task GetTokenAsync( + GitHubTokenAcquireRequest request, + CancellationToken cancellationToken = default) + { + if (!client._gitHubTokenProviders.TryGetValue(request.RegistrationId, out var provider)) + { + throw new InvalidOperationException( + $"Unknown GitHub token provider registration ID '{request.RegistrationId}'."); + } + + var reason = request.Reason == GitHubTokenAcquireReason.Initial + ? GitHubTokenRequestReason.Initial + : request.Reason == GitHubTokenAcquireReason.Refresh + ? GitHubTokenRequestReason.Refresh + : throw new InvalidOperationException($"Unknown GitHub token request reason '{request.Reason}'."); + var result = await provider(new GitHubTokenProviderArgs + { + Host = request.Host, + SessionId = request.SessionId, + Reason = reason, + }).ConfigureAwait(false); + + if (result is { Cancelled: true }) + { + return new GitHubTokenAcquireResultCancelled(); + } + if (result?.Token is not { } token) + { + throw new InvalidOperationException( + "GitHub token provider returned neither a token nor cancellation."); + } + return new GitHubTokenAcquireResultToken + { + AccessToken = token.AccessToken, + TokenType = token.TokenType, + ExpiresIn = token.ExpiresIn, + }; + } + } + /// /// Tells the runtime to route its outbound model-layer requests through this /// client's LLM inference provider. No-op when interception is not configured. @@ -2802,6 +2908,7 @@ internal record CreateSessionRequest( string? Tracestate = null, ModelCapabilitiesOverride? ModelCapabilities = null, string? GitHubToken = null, + [property: JsonPropertyName("gitHubTokenProviderRegistrationId")] string? GitHubTokenProviderRegistrationId = null, RemoteSessionMode? RemoteSession = null, CloudSessionOptions? Cloud = null, IList? InstructionDirectories = null, @@ -2917,6 +3024,7 @@ internal record ResumeSessionRequest( string? Tracestate = null, ModelCapabilitiesOverride? ModelCapabilities = null, string? GitHubToken = null, + [property: JsonPropertyName("gitHubTokenProviderRegistrationId")] string? GitHubTokenProviderRegistrationId = null, RemoteSessionMode? RemoteSession = null, bool? ContinuePendingWork = null, IList? InstructionDirectories = null, diff --git a/dotnet/src/GitHubTokenProvider.cs b/dotnet/src/GitHubTokenProvider.cs new file mode 100644 index 0000000000..3a79d09662 --- /dev/null +++ b/dotnet/src/GitHubTokenProvider.cs @@ -0,0 +1,77 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +using System.Diagnostics.CodeAnalysis; + +namespace GitHub.Copilot; + +/// Why the runtime is requesting a GitHub token. +[Experimental(Diagnostics.Experimental)] +public enum GitHubTokenRequestReason +{ + /// The session needs its initial token. + Initial, + + /// The session needs a refreshed token. + Refresh, +} + +/// Arguments passed to a session-scoped GitHub token provider. +[Experimental(Diagnostics.Experimental)] +public sealed class GitHubTokenProviderArgs +{ + /// Gets the effective GitHub host for which a token is needed. + public required string Host { get; init; } + + /// + /// Gets the session receiving the token, or before a + /// cloud session has been assigned an identifier. + /// + public string? SessionId { get; init; } + + /// Gets whether the runtime needs an initial or refreshed token. + public required GitHubTokenRequestReason Reason { get; init; } +} + +/// A GitHub access token returned by a session-scoped provider. +[Experimental(Diagnostics.Experimental)] +public sealed class GitHubToken +{ + /// Gets or sets the GitHub access token. + public required string AccessToken { get; set; } + + /// Gets or sets the OAuth token type. The runtime defaults it to bearer. + public string? TokenType { get; set; } + + /// + /// Gets or sets the required positive number of seconds remaining when the + /// callback completes. Production GitHub tokens typically last eight hours. + /// + public required long ExpiresIn { get; set; } + + /// + public override string ToString() + => $"{nameof(GitHubToken)} {{ {nameof(TokenType)} = {TokenType}, {nameof(ExpiresIn)} = {ExpiresIn}, {nameof(AccessToken)} = }}"; +} + +/// The result returned by a session-scoped GitHub token provider. +[Experimental(Diagnostics.Experimental)] +public sealed class GitHubTokenProviderResult +{ + /// Gets whether token acquisition was cancelled. + public bool Cancelled { get; private init; } + + /// Gets the acquired token, if acquisition was not cancelled. + public GitHubToken? Token { get; private init; } + + /// Creates a successful token result. + public static GitHubTokenProviderResult FromToken(GitHubToken token) + { + ArgumentNullException.ThrowIfNull(token); + return new() { Token = token }; + } + + /// Creates a cancelled result. + public static GitHubTokenProviderResult Cancel() => new() { Cancelled = true }; +} diff --git a/dotnet/src/JsonRpc.cs b/dotnet/src/JsonRpc.cs index 36289d2e68..c5c444ea70 100644 --- a/dotnet/src/JsonRpc.cs +++ b/dotnet/src/JsonRpc.cs @@ -548,7 +548,11 @@ private async Task HandleIncomingMethodAsync(string methodName, JsonElement mess if (requestId.HasValue) { - await SendResultResponseAsync(requestId.Value, result, cancellationToken).ConfigureAwait(false); + await SendResultResponseAsync( + requestId.Value, + result, + registration.ResultType, + cancellationToken).ConfigureAwait(false); } } catch (Exception ex) when (ex is not OperationCanceledException) @@ -772,18 +776,20 @@ private static bool TryGetPropertyCaseInsensitive(JsonElement obj, string name, return doc.RootElement.Clone(); } - private async Task SendResultResponseAsync(JsonElement id, object? result, CancellationToken cancellationToken) + private async Task SendResultResponseAsync( + JsonElement id, + object? result, + Type? declaredResultType, + CancellationToken cancellationToken) { try { - // Convert the result to a JsonElement using the runtime type, looked up via - // the merged resolver. Source-gen serialization of an `object`-typed property - // would otherwise have no way to find metadata for the actual response type - // (e.g. SystemMessageTransformRpcResponse, SessionFsReadFileResult, ...). + // Prefer the handler's declared result type so polymorphic base types emit + // their discriminator. Fall back to the runtime type for untyped handlers. JsonElement? resultElement = null; if (result is not null) { - var typeInfo = _serializerOptions.GetTypeInfo(result.GetType()); + var typeInfo = _serializerOptions.GetTypeInfo(declaredResultType ?? result.GetType()); resultElement = JsonSerializer.SerializeToElement(result, typeInfo); } @@ -863,6 +869,11 @@ public MethodRegistration(Delegate handler, bool singleObjectParam) { ValueTaskAsTaskMethod = GetMethodFromGenericMethodDefinition(returnType, s_valueTaskAsTask); TaskResultGetter = GetMethodFromGenericMethodDefinition(ValueTaskAsTaskMethod.ReturnType, s_taskGetResult); + ResultType = returnType.GetGenericArguments()[0]; + } + else if (returnType != typeof(void) && returnType != typeof(Task) && returnType != typeof(ValueTask)) + { + ResultType = returnType; } } @@ -871,6 +882,7 @@ public MethodRegistration(Delegate handler, bool singleObjectParam) public ParameterInfo[] Parameters { get; } public MethodInfo? ValueTaskAsTaskMethod { get; } public MethodInfo? TaskResultGetter { get; } + public Type? ResultType { get; } } private static MethodInfo GetMethodFromGenericMethodDefinition(Type specializedType, MethodInfo genericMethodDefinition) diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index a6fad7dcfa..3176e7db23 100644 --- a/dotnet/src/Session.cs +++ b/dotnet/src/Session.cs @@ -82,6 +82,7 @@ private sealed record EventSubscription(Type EventType, Action Han private IReadOnlyList _openCanvases = Array.Empty(); private int _isDisposed; + private string? _gitHubTokenProviderRegistrationId; /// /// Channel that serializes event dispatch. enqueues; @@ -203,6 +204,19 @@ internal void RemoveFromClient() ((ICollection>)_parentClient._sessions).Remove(new(SessionId, this)); } + internal void SetGitHubTokenProviderRegistration(string registrationId) + { + _gitHubTokenProviderRegistrationId = registrationId; + } + + internal void ReleaseGitHubTokenProviderRegistration() + { + if (Interlocked.Exchange(ref _gitHubTokenProviderRegistrationId, null) is { } registrationId) + { + _parentClient.UnregisterGitHubTokenProvider(registrationId); + } + } + internal void StartProcessingEvents() { _ = ProcessEventsAsync(); @@ -1936,6 +1950,7 @@ await InvokeRpcAsync( } finally { + ReleaseGitHubTokenProviderRegistration(); RemoveFromClient(); GC.SuppressFinalize(this); } diff --git a/dotnet/src/Types.cs b/dotnet/src/Types.cs index 54595bfed7..d038721a5e 100644 --- a/dotnet/src/Types.cs +++ b/dotnet/src/Types.cs @@ -3206,6 +3206,7 @@ protected SessionConfigBase(SessionConfigBase? other) ContextTier = other.ContextTier; CreateSessionFsProvider = other.CreateSessionFsProvider; GitHubToken = other.GitHubToken; + GitHubTokenProvider = other.GitHubTokenProvider; RemoteSession = other.RemoteSession; ExpAssignments = other.ExpAssignments; EnableManagedSettings = other.EnableManagedSettings; @@ -3649,6 +3650,14 @@ protected SessionConfigBase(SessionConfigBase? other) /// public string? GitHubToken { get; set; } + /// + /// Gets or sets a callback that acquires session-scoped GitHub tokens on + /// demand. This cannot be combined with . + /// + [Experimental(Diagnostics.Experimental)] + [JsonIgnore] + public Func>? GitHubTokenProvider { get; set; } + /// /// Per-session remote behavior control: /// diff --git a/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index 47ee14bb69..356ef605a8 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -3,6 +3,7 @@ *--------------------------------------------------------------------------------------------*/ #if NET8_0_OR_GREATER +using System.Collections.Concurrent; using System.Net; using System.Net.Sockets; using System.Diagnostics; @@ -19,6 +20,170 @@ public sealed class ClientSessionLifetimeTests { private sealed record RpcRequestRecord(string Method, JsonElement Params); + [Fact] + public async Task GitHubTokenProvider_Is_Mutually_Exclusive_With_Static_Token() + { + await using var client = new CopilotClient(); + var config = new SessionConfig + { + GitHubToken = "static", + GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) + }; + + var error = await Assert.ThrowsAsync(() => client.CreateSessionAsync(config)); + + Assert.Contains("cannot be used together", error.Message); + } + + [Fact] + public async Task GitHubTokenProvider_Is_Serialized_And_Maps_Callbacks() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + GitHubTokenProviderArgs? callbackArgs = null; + var session = await client.CreateSessionAsync(new SessionConfig + { + GitHubTokenProvider = args => + { + callbackArgs = args; + return Task.FromResult(GitHubTokenProviderResult.FromToken(new GitHubToken + { + AccessToken = "secret-token", + TokenType = "bearer", + ExpiresIn = 8 * 60 * 60 + })); + } + }); + var request = Assert.Single(server.Requests, request => request.Method == "session.create"); + var registrationId = request.Params.GetProperty("gitHubTokenProviderRegistrationId").GetString(); + Assert.False(string.IsNullOrEmpty(registrationId)); + Assert.False(request.Params.TryGetProperty("gitHubToken", out _)); + + var result = await server.SendRequestAsync("gitHubToken.getToken", new Dictionary + { + ["registrationId"] = registrationId, + ["host"] = "github.example.com", + ["sessionId"] = session.SessionId, + ["reason"] = "refresh" + }); + + Assert.True(result.TryGetProperty("kind", out var kind), result.ToString()); + Assert.Equal("token", kind.GetString()); + Assert.Equal("secret-token", result.GetProperty("accessToken").GetString()); + Assert.Equal(8 * 60 * 60, result.GetProperty("expiresIn").GetInt64()); + Assert.NotNull(callbackArgs); + Assert.Equal("github.example.com", callbackArgs.Host); + Assert.Equal(session.SessionId, callbackArgs.SessionId); + Assert.Equal(GitHubTokenRequestReason.Refresh, callbackArgs.Reason); + Assert.DoesNotContain("secret-token", new GitHubToken + { + AccessToken = "secret-token", + ExpiresIn = 8 * 60 * 60 + }.ToString()); + + await session.DisposeAsync(); + var error = await Assert.ThrowsAsync(() => + server.SendRequestAsync("gitHubToken.getToken", new Dictionary + { + ["registrationId"] = registrationId, + ["host"] = "github.com", + ["reason"] = "initial" + })); + Assert.Contains("Unknown GitHub token provider registration ID", error.Message); + + server.ClearRequests(); + var resumed = await client.ResumeSessionAsync("resumed-session", new ResumeSessionConfig + { + GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) + }); + var resumeRequest = Assert.Single(server.Requests, request => request.Method == "session.resume"); + Assert.False(string.IsNullOrEmpty( + resumeRequest.Params.GetProperty("gitHubTokenProviderRegistrationId").GetString())); + await resumed.DisposeAsync(); + } + + [Fact] + public async Task GitHubTokenProvider_Handles_Cancellation_Errors_And_Rollback() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + var cancelledSession = await client.CreateSessionAsync(new SessionConfig + { + GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) + }); + var cancelledId = Assert.Single(server.Requests, request => request.Method == "session.create") + .Params.GetProperty("gitHubTokenProviderRegistrationId").GetString(); + var cancelled = await server.SendRequestAsync("gitHubToken.getToken", TokenRequest(cancelledId)); + Assert.True(cancelled.TryGetProperty("kind", out var cancelledKind), cancelled.ToString()); + Assert.Equal("cancelled", cancelledKind.GetString()); + await cancelledSession.DisposeAsync(); + + server.ClearRequests(); + var providerSession = await client.CreateSessionAsync(new SessionConfig + { + GitHubTokenProvider = _ => Task.FromException( + new InvalidOperationException("provider failed")) + }); + var providerId = Assert.Single(server.Requests, request => request.Method == "session.create") + .Params.GetProperty("gitHubTokenProviderRegistrationId").GetString(); + var callbackError = await Assert.ThrowsAsync(() => + server.SendRequestAsync("gitHubToken.getToken", TokenRequest(providerId))); + Assert.Contains("provider failed", callbackError.Message); + await providerSession.DisposeAsync(); + + server.ClearRequests(); + server.FailSessionCreate(); + await Assert.ThrowsAsync(() => client.CreateSessionAsync(new SessionConfig + { + GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) + })); + var rolledBackId = Assert.Single(server.Requests, request => request.Method == "session.create") + .Params.GetProperty("gitHubTokenProviderRegistrationId").GetString(); + var rollbackError = await Assert.ThrowsAsync(() => + server.SendRequestAsync("gitHubToken.getToken", TokenRequest(rolledBackId))); + Assert.Contains("Unknown GitHub token provider registration ID", rollbackError.Message); + } + + [Fact] + public async Task GitHubTokenProvider_Resume_Replaces_Ownership() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + var first = await client.CreateSessionAsync(new SessionConfig + { + SessionId = "replacement-session", + GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) + }); + var firstId = Assert.Single(server.Requests, request => request.Method == "session.create") + .Params.GetProperty("gitHubTokenProviderRegistrationId").GetString(); + + await first.DisposeAsync(); + await Assert.ThrowsAsync(() => + server.SendRequestAsync("gitHubToken.getToken", TokenRequest(firstId))); + + server.ClearRequests(); + var resumed = await client.ResumeSessionAsync("replacement-session", new ResumeSessionConfig + { + GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) + }); + var secondId = Assert.Single(server.Requests, request => request.Method == "session.resume") + .Params.GetProperty("gitHubTokenProviderRegistrationId").GetString(); + + var result = await server.SendRequestAsync("gitHubToken.getToken", TokenRequest(secondId)); + Assert.Equal("cancelled", result.GetProperty("kind").GetString()); + + await resumed.DisposeAsync(); + await Assert.ThrowsAsync(() => + server.SendRequestAsync("gitHubToken.getToken", TokenRequest(secondId))); + } + + private static Dictionary TokenRequest(string? registrationId) => new() + { + ["registrationId"] = registrationId, + ["host"] = "github.com", + ["reason"] = "initial" + }; + [Fact] public async Task StopAsync_Requests_Runtime_Shutdown_For_Owned_Process() { @@ -918,9 +1083,13 @@ private sealed class FakeCopilotServer : IAsyncDisposable private readonly Task _serverTask; private readonly List _requests = []; private readonly object _requestsLock = new(); + private readonly ConcurrentDictionary> _pendingRequests = new(); + private NetworkStream? _stream; + private int _nextRequestId; private string? _lastSessionId; private bool _delayDestroy; private bool _failRuntimeShutdown; + private bool _failSessionCreate; private FakeCopilotServer(TcpListener listener) { @@ -982,6 +1151,31 @@ public void FailRuntimeShutdown() _failRuntimeShutdown = true; } + public void FailSessionCreate() + { + _failSessionCreate = true; + } + + public async Task SendRequestAsync(string method, Dictionary parameters) + { + var stream = _stream ?? throw new InvalidOperationException("Client is not connected."); + var id = Interlocked.Increment(ref _nextRequestId); + var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + if (!_pendingRequests.TryAdd(id, completion)) + { + throw new InvalidOperationException("Failed to track callback request."); + } + + await WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["id"] = id, + ["method"] = method, + ["params"] = parameters + }, _cts.Token); + return await completion.Task.WaitAsync(_cts.Token); + } + public async ValueTask DisposeAsync() { _allowDestroy.TrySetResult(); @@ -1004,16 +1198,37 @@ private async Task RunAsync() { using var tcpClient = await _listener.AcceptTcpClientAsync(_cts.Token); using var stream = tcpClient.GetStream(); + _stream = stream; while (!_cts.Token.IsCancellationRequested) { - using var request = await ReadMessageAsync(stream, _cts.Token); - if (request is null) + using var message = await ReadMessageAsync(stream, _cts.Token); + if (message is null) { return; } - await HandleRequestAsync(stream, request.RootElement, _cts.Token); + var root = message.RootElement; + if (root.TryGetProperty("method", out _)) + { + await HandleRequestAsync(stream, root, _cts.Token); + continue; + } + + if (root.TryGetProperty("id", out var responseId) + && responseId.TryGetInt32(out var id) + && _pendingRequests.TryRemove(id, out var completion)) + { + if (root.TryGetProperty("error", out var error)) + { + completion.TrySetException(new InvalidOperationException( + error.GetProperty("message").GetString())); + } + else + { + completion.TrySetResult(root.GetProperty("result").Clone()); + } + } } } @@ -1049,6 +1264,21 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel { _requests.Add(new RpcRequestRecord(method!, paramsElement)); } + if (method == "session.create" && _failSessionCreate) + { + _failSessionCreate = false; + await WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["id"] = id, + ["error"] = new Dictionary + { + ["code"] = -32000, + ["message"] = "session create failed" + } + }, cancellationToken); + return; + } object? result = method switch { "connect" => new Dictionary diff --git a/go/README.md b/go/README.md index d8588699c4..80abbd19c8 100644 --- a/go/README.md +++ b/go/README.md @@ -222,6 +222,7 @@ Event types: `SessionLifecycleCreated`, `SessionLifecycleDeleted`, `SessionLifec - `InfiniteSessions` (\*InfiniteSessionConfig): Automatic context compaction configuration - `WorkingDirectory` (string): Working directory for the session (default: runtime process working directory) - `EnableSessionStore` (\*bool): Enables the cross-session store for search and retrieval across sessions. When unset in `ModeCopilotCli`, the runtime default applies (enabled). In `ModeEmpty`, defaults to disabled. +- `GitHubTokenProvider` (GitHubTokenProvider): Acquires session-scoped GitHub tokens on demand. Return `GitHubTokenResult` with a positive `ExpiresIn` value (production GitHub tokens typically use `8 * 60 * 60` seconds), or `GitHubTokenCancelled`. Cannot be combined with `GitHubToken`. - `OnPermissionRequest` (PermissionHandlerFunc): Optional handler called before each tool execution to approve or deny it. When nil, permission requests are emitted as events and left pending for manual resolution. `copilot.PermissionHandler.ApproveAll` approves requests when managed settings are disabled and returns an error when `EnableManagedSettings` is true. Custom handlers can inspect `RequiresManagedApproval()` for human-facing confirmation logic. See [Permission Handling](#permission-handling) section. - `OnUserInputRequest` (UserInputHandler): Handler for user input requests from the agent (enables ask_user tool). See [User Input Requests](#user-input-requests) section. - `Hooks` (\*SessionHooks): Hook handlers for session lifecycle events. See [Session Hooks](#session-hooks) section. @@ -237,6 +238,22 @@ Event types: `SessionLifecycleCreated`, `SessionLifecycleDeleted`, `SessionLifec - `Streaming` (*bool): Enable streaming delta events (nil = runtime default) - `Commands` ([]CommandDefinition): Slash-commands. See [Commands](#commands) section. - `OnElicitationRequest` (ElicitationHandler): Elicitation handler. See [Elicitation Requests](#elicitation-requests-serverclient) section. +- `GitHubTokenProvider` (GitHubTokenProvider): Replaces the session-scoped token provider when resuming. Cannot be combined with `GitHubToken`. + +```go +session, err := client.CreateSession(ctx, &copilot.SessionConfig{ + GitHubTokenProvider: func(args copilot.GitHubTokenProviderArgs) (*copilot.GitHubTokenProviderResult, error) { + token, err := acquireToken(args.Host) + if err != nil { + return nil, err + } + return copilot.GitHubTokenResult(&copilot.GitHubToken{ + AccessToken: token, + ExpiresIn: 8 * 60 * 60, + }), nil + }, +}) +``` ### Session diff --git a/go/client.go b/go/client.go index fb02897f91..bc7bd7be55 100644 --- a/go/client.go +++ b/go/client.go @@ -145,19 +145,23 @@ func validateEnvironmentOptions(connection RuntimeConnection, opts *ClientOption // } // defer client.Stop() type Client struct { - options ClientOptions - process *exec.Cmd - client *jsonrpc2.Client - actualPort int - actualHost string - state connectionState - sessions map[string]*Session - sessionsMux sync.Mutex - isExternalServer bool - conn net.Conn // stores net.Conn for external TCP connections - useStdio bool // resolved value from options - useInProcess bool // true for InProcessConnection (FFI transport) - ffiHost inProcessHost + options ClientOptions + process *exec.Cmd + client *jsonrpc2.Client + actualPort int + actualHost string + state connectionState + sessions map[string]*Session + sessionsMux sync.Mutex + gitHubTokenProviders map[string]GitHubTokenProvider + gitHubTokenProvidersMux sync.RWMutex + sessionOperations map[string]*sessionOperation + sessionOperationsMux sync.Mutex + isExternalServer bool + conn net.Conn // stores net.Conn for external TCP connections + useStdio bool // resolved value from options + useInProcess bool // true for InProcessConnection (FFI transport) + ffiHost inProcessHost // resolved process options for the spawned runtime (zero values for URIConnection) cliPath string cliArgs []string @@ -189,6 +193,11 @@ type Client struct { internalRPC *rpc.InternalServerRPC } +type sessionOperation struct { + mutex sync.Mutex + users int +} + // NewClient creates a new Copilot runtime client with the given options. // // If options is nil, default options are used (spawns the bundled runtime over @@ -215,12 +224,13 @@ func NewClient(options *ClientOptions) *Client { opts := ClientOptions{} client := &Client{ - options: opts, - state: stateDisconnected, - sessions: make(map[string]*Session), - actualHost: "localhost", - isExternalServer: false, - useStdio: true, + options: opts, + state: stateDisconnected, + sessions: make(map[string]*Session), + gitHubTokenProviders: make(map[string]GitHubTokenProvider), + actualHost: "localhost", + isExternalServer: false, + useStdio: true, } if options != nil { @@ -548,6 +558,7 @@ func (c *Client) Stop() error { c.sessionsMux.Lock() c.sessions = make(map[string]*Session) c.sessionsMux.Unlock() + c.clearGitHubTokenProviders() c.startStopMux.Lock() defer c.startStopMux.Unlock() @@ -663,6 +674,7 @@ func (c *Client) ForceStop() { c.sessionsMux.Lock() c.sessions = make(map[string]*Session) c.sessionsMux.Unlock() + c.clearGitHubTokenProviders() c.startStopMux.Lock() defer c.startStopMux.Unlock() @@ -780,6 +792,9 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses if config == nil { config = &SessionConfig{} } + if config.GitHubToken != "" && config.GitHubTokenProvider != nil { + return nil, fmt.Errorf("GitHubToken and GitHubTokenProvider cannot be used together") + } if err := c.ensureConnected(ctx); err != nil { return nil, err @@ -787,6 +802,14 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses c.applyConfigDefaultsForMode(config) + registrationID := c.registerGitHubTokenProvider(config.GitHubTokenProvider) + registrationTransferred := false + defer func() { + if !registrationTransferred { + c.unregisterGitHubTokenProvider(registrationID) + } + }() + req := createSessionRequest{} req.Model = config.Model req.ClientName = config.ClientName @@ -849,6 +872,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses req.ToolSearch = config.ToolSearch req.Memory = config.Memory req.GitHubToken = config.GitHubToken + req.GitHubTokenProviderRegistrationID = registrationID req.RemoteSession = config.RemoteSession req.Cloud = config.Cloud req.Canvases = config.Canvases @@ -1103,6 +1127,12 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses return nil, err } + if registrationID != "" { + session.setGitHubTokenProviderRegistrationRelease(func() { + c.unregisterGitHubTokenProvider(registrationID) + }) + registrationTransferred = true + } return session, nil } @@ -1130,9 +1160,15 @@ func (c *Client) ResumeSession(ctx context.Context, sessionID string, config *Re // Tools: []copilot.Tool{myNewTool}, // }) func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, config *ResumeSessionConfig) (*Session, error) { + unlockSession := c.lockSessionOperation(sessionID) + defer unlockSession() + if config == nil { config = &ResumeSessionConfig{} } + if config.GitHubToken != "" && config.GitHubTokenProvider != nil { + return nil, fmt.Errorf("GitHubToken and GitHubTokenProvider cannot be used together") + } if err := c.ensureConnected(ctx); err != nil { return nil, err @@ -1140,6 +1176,14 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, c.applyResumeDefaultsForMode(config) + registrationID := c.registerGitHubTokenProvider(config.GitHubTokenProvider) + registrationTransferred := false + defer func() { + if !registrationTransferred { + c.unregisterGitHubTokenProvider(registrationID) + } + }() + var req resumeSessionRequest req.SessionID = sessionID req.ClientName = config.ClientName @@ -1233,6 +1277,7 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, req.ToolSearch = config.ToolSearch req.Memory = config.Memory req.GitHubToken = config.GitHubToken + req.GitHubTokenProviderRegistrationID = registrationID req.RemoteSession = config.RemoteSession req.Canvases = config.Canvases req.OpenCanvases = config.OpenCanvases @@ -1318,22 +1363,31 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, } c.sessionsMux.Lock() + replacedSession := c.sessions[sessionID] c.sessions[sessionID] = session c.sessionsMux.Unlock() + restoreReplacedSession := func() { + c.sessionsMux.Lock() + if current := c.sessions[sessionID]; current == nil || current == session { + if replacedSession != nil { + c.sessions[sessionID] = replacedSession + } else { + delete(c.sessions, sessionID) + } + } + c.sessionsMux.Unlock() + } + if c.options.SessionFS != nil { if config.CreateSessionFSProvider == nil { - c.sessionsMux.Lock() - delete(c.sessions, sessionID) - c.sessionsMux.Unlock() + restoreReplacedSession() return nil, fmt.Errorf("CreateSessionFSProvider is required in session config when SessionFS is enabled in client options") } provider := config.CreateSessionFSProvider(session) if c.options.SessionFS.Capabilities != nil && c.options.SessionFS.Capabilities.Sqlite { if _, ok := provider.(SessionFSSqliteProvider); !ok { - c.sessionsMux.Lock() - delete(c.sessions, sessionID) - c.sessionsMux.Unlock() + restoreReplacedSession() return nil, fmt.Errorf("SessionFS capabilities declare SQLite support but the provider does not implement SessionFSSqliteProvider") } } @@ -1342,17 +1396,13 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, result, err := c.client.Request(ctx, "session.resume", req) if err != nil { - c.sessionsMux.Lock() - delete(c.sessions, sessionID) - c.sessionsMux.Unlock() + restoreReplacedSession() return nil, fmt.Errorf("failed to resume session: %w", err) } var response resumeSessionResponse if err := json.Unmarshal(result, &response); err != nil { - c.sessionsMux.Lock() - delete(c.sessions, sessionID) - c.sessionsMux.Unlock() + restoreReplacedSession() return nil, fmt.Errorf("failed to unmarshal response: %w", err) } @@ -1361,9 +1411,7 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, "sessionId": sessionID, "eventType": "mcp.oauth_required", }); err != nil { - c.sessionsMux.Lock() - delete(c.sessions, sessionID) - c.sessionsMux.Unlock() + restoreReplacedSession() return nil, err } } @@ -1378,9 +1426,19 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, CoauthorEnabled: config.CoauthorEnabled, ManageScheduleEnabled: config.ManageScheduleEnabled, }); err != nil { + restoreReplacedSession() return nil, err } + if registrationID != "" { + session.setGitHubTokenProviderRegistrationRelease(func() { + c.unregisterGitHubTokenProvider(registrationID) + }) + registrationTransferred = true + } + if replacedSession != nil && replacedSession != session { + replacedSession.releaseGitHubTokenProviderRegistration() + } return session, nil } @@ -1472,6 +1530,9 @@ func (c *Client) GetSessionMetadata(ctx context.Context, sessionID string) (*Ses // log.Fatal(err) // } func (c *Client) DeleteSession(ctx context.Context, sessionID string) error { + unlockSession := c.lockSessionOperation(sessionID) + defer unlockSession() + if err := c.ensureConnected(ctx); err != nil { return err } @@ -1496,8 +1557,12 @@ func (c *Client) DeleteSession(ctx context.Context, sessionID string) error { // Remove from local sessions map if present c.sessionsMux.Lock() + session := c.sessions[sessionID] delete(c.sessions, sessionID) c.sessionsMux.Unlock() + if session != nil { + session.releaseGitHubTokenProviderRegistration() + } return nil } @@ -2039,15 +2104,7 @@ func (c *Client) startCLIServer(ctx context.Context) error { // Create JSON-RPC client immediately c.client = jsonrpc2.NewClient(stdin, stdout) c.client.SetProcessDone(c.processDone, c.processErrorPtr) - c.client.SetOnClose(func() { - // Run in a goroutine to avoid deadlocking with Stop/ForceStop, - // which hold startStopMux while waiting for readLoop to finish. - go func() { - c.startStopMux.Lock() - defer c.startStopMux.Unlock() - c.state = stateDisconnected - }() - }) + c.client.SetOnClose(c.handleConnectionClose) c.RPC = rpc.NewServerRPC(c.client) c.internalRPC = rpc.NewInternalServerRPC(c.client) c.setupNotificationHandler() @@ -2166,15 +2223,7 @@ func (c *Client) startInProcess(ctx context.Context) error { } c.client = jsonrpc2.NewClient(host.Writer(), host.Reader()) - c.client.SetOnClose(func() { - // Run in a goroutine to avoid deadlocking with Stop/ForceStop, which hold - // startStopMux while waiting for readLoop to finish. - go func() { - c.startStopMux.Lock() - defer c.startStopMux.Unlock() - c.state = stateDisconnected - }() - }) + c.client.SetOnClose(c.handleConnectionClose) c.RPC = rpc.NewServerRPC(c.client) c.internalRPC = rpc.NewInternalServerRPC(c.client) c.setupNotificationHandler() @@ -2324,13 +2373,7 @@ func (c *Client) connectViaTCP(ctx context.Context) error { if c.processDone != nil { c.client.SetProcessDone(c.processDone, c.processErrorPtr) } - c.client.SetOnClose(func() { - go func() { - c.startStopMux.Lock() - defer c.startStopMux.Unlock() - c.state = stateDisconnected - }() - }) + c.client.SetOnClose(c.handleConnectionClose) c.RPC = rpc.NewServerRPC(c.client) c.internalRPC = rpc.NewInternalServerRPC(c.client) c.setupNotificationHandler() @@ -2361,8 +2404,10 @@ func (c *Client) setupNotificationHandler() { // payload's sessionId. Always register the global handlers so the generated // hooks.invoke handler is wired to our dispatcher. handlers := &rpc.ClientGlobalAPIHandlers{ - Hooks: &hooksAdapter{client: c}, + Hooks: &hooksAdapter{client: c}, + GitHubToken: &gitHubTokenAdapter{client: c}, } + if c.options.RequestHandler != nil { handlers.LlmInference = newCopilotRequestAdapter(c.options.RequestHandler, func() *rpc.ServerLlmInferenceAPI { if c.RPC == nil { @@ -2377,6 +2422,107 @@ func (c *Client) setupNotificationHandler() { rpc.RegisterClientGlobalAPIHandlers(c.client, handlers) } +func (c *Client) registerGitHubTokenProvider(provider GitHubTokenProvider) string { + if provider == nil { + return "" + } + registrationID := uuid.NewString() + c.gitHubTokenProvidersMux.Lock() + if c.gitHubTokenProviders == nil { + c.gitHubTokenProviders = make(map[string]GitHubTokenProvider) + } + c.gitHubTokenProviders[registrationID] = provider + c.gitHubTokenProvidersMux.Unlock() + return registrationID +} + +func (c *Client) unregisterGitHubTokenProvider(registrationID string) { + if registrationID == "" { + return + } + c.gitHubTokenProvidersMux.Lock() + delete(c.gitHubTokenProviders, registrationID) + c.gitHubTokenProvidersMux.Unlock() +} + +func (c *Client) clearGitHubTokenProviders() { + c.gitHubTokenProvidersMux.Lock() + c.gitHubTokenProviders = make(map[string]GitHubTokenProvider) + c.gitHubTokenProvidersMux.Unlock() +} + +func (c *Client) handleConnectionClose() { + c.clearGitHubTokenProviders() + // Avoid deadlocking with Stop/ForceStop, which hold startStopMux while + // waiting for the JSON-RPC read loop to finish. + go func() { + c.startStopMux.Lock() + defer c.startStopMux.Unlock() + c.state = stateDisconnected + }() +} + +func (c *Client) lockSessionOperation(sessionID string) func() { + c.sessionOperationsMux.Lock() + if c.sessionOperations == nil { + c.sessionOperations = make(map[string]*sessionOperation) + } + operation := c.sessionOperations[sessionID] + if operation == nil { + operation = &sessionOperation{} + c.sessionOperations[sessionID] = operation + } + operation.users++ + c.sessionOperationsMux.Unlock() + + operation.mutex.Lock() + return func() { + operation.mutex.Unlock() + c.sessionOperationsMux.Lock() + operation.users-- + if operation.users == 0 { + delete(c.sessionOperations, sessionID) + } + c.sessionOperationsMux.Unlock() + } +} + +type gitHubTokenAdapter struct { + client *Client +} + +func (a *gitHubTokenAdapter) GetToken(request *rpc.GitHubTokenAcquireRequest) (rpc.GitHubTokenAcquireResult, error) { + if request == nil { + return nil, fmt.Errorf("missing GitHub token acquire request") + } + a.client.gitHubTokenProvidersMux.RLock() + provider := a.client.gitHubTokenProviders[request.RegistrationID] + a.client.gitHubTokenProvidersMux.RUnlock() + if provider == nil { + return nil, fmt.Errorf("unknown GitHub token provider registration ID %q", request.RegistrationID) + } + + result, err := provider(GitHubTokenProviderArgs{ + Host: request.Host, + SessionID: request.SessionID, + Reason: request.Reason, + }) + if err != nil { + return nil, err + } + if result != nil && result.Cancelled { + return &rpc.GitHubTokenAcquireResultCancelled{}, nil + } + if result == nil || result.Token == nil { + return nil, fmt.Errorf("GitHub token provider returned neither a token nor cancellation") + } + return &rpc.GitHubTokenAcquireResultToken{ + AccessToken: result.Token.AccessToken, + TokenType: result.Token.TokenType, + ExpiresIn: result.Token.ExpiresIn, + }, nil +} + // gitHubTelemetryAdapter adapts the OnGitHubTelemetry option to the generated // rpc.GitHubTelemetryHandler interface. type gitHubTelemetryAdapter struct { diff --git a/go/github_token_provider.go b/go/github_token_provider.go new file mode 100644 index 0000000000..fa57e24ca0 --- /dev/null +++ b/go/github_token_provider.go @@ -0,0 +1,82 @@ +package copilot + +import ( + "fmt" + + "github.com/github/copilot-sdk/go/rpc" +) + +// GitHubTokenRequestReason describes why the runtime needs a GitHub token. +// +// Experimental: GitHubTokenRequestReason may change or be removed. +type GitHubTokenRequestReason = rpc.GitHubTokenAcquireReason + +const ( + // GitHubTokenRequestReasonInitial indicates the session needs its initial token. + GitHubTokenRequestReasonInitial = rpc.GitHubTokenAcquireReasonInitial + // GitHubTokenRequestReasonRefresh indicates the session needs a refreshed token. + GitHubTokenRequestReasonRefresh = rpc.GitHubTokenAcquireReasonRefresh +) + +// GitHubTokenProviderArgs contains the context for a GitHub token request. +// +// Experimental: GitHubTokenProviderArgs may change or be removed. +type GitHubTokenProviderArgs struct { + // Host is the effective GitHub host for which a token is needed. + Host string + // SessionID identifies the session receiving the token. It is nil before a + // cloud session has been assigned an ID. + SessionID *string + // Reason indicates whether this is the initial token or a refresh. + Reason GitHubTokenRequestReason +} + +// GitHubToken contains a GitHub access token returned by a provider. +// +// Experimental: GitHubToken may change or be removed. +type GitHubToken struct { + // AccessToken is the GitHub access token. + AccessToken string + // TokenType is the OAuth token type. The runtime defaults it to "bearer". + TokenType *string + // ExpiresIn is the required positive number of seconds remaining when the + // callback completes. Production GitHub tokens typically last eight hours. + ExpiresIn int64 +} + +// String returns a redacted description that never includes the access token. +func (t GitHubToken) String() string { + tokenType := "" + if t.TokenType != nil { + tokenType = *t.TokenType + } + return fmt.Sprintf("GitHubToken{TokenType:%q, ExpiresIn:%d, AccessToken:}", tokenType, t.ExpiresIn) +} + +// GoString returns a redacted Go-syntax description that never includes the access token. +func (t GitHubToken) GoString() string { + return t.String() +} + +// GitHubTokenProviderResult is the result of a GitHub token request. +// +// Experimental: GitHubTokenProviderResult may change or be removed. +type GitHubTokenProviderResult struct { + Cancelled bool + Token *GitHubToken +} + +// GitHubTokenResult returns a successful token-provider result. +func GitHubTokenResult(token *GitHubToken) *GitHubTokenProviderResult { + return &GitHubTokenProviderResult{Token: token} +} + +// GitHubTokenCancelled returns a result indicating that token acquisition was cancelled. +func GitHubTokenCancelled() *GitHubTokenProviderResult { + return &GitHubTokenProviderResult{Cancelled: true} +} + +// GitHubTokenProvider acquires session-scoped GitHub tokens on demand. +// +// Experimental: GitHubTokenProvider may change or be removed. +type GitHubTokenProvider func(args GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) diff --git a/go/github_token_provider_test.go b/go/github_token_provider_test.go new file mode 100644 index 0000000000..f00837379f --- /dev/null +++ b/go/github_token_provider_test.go @@ -0,0 +1,344 @@ +package copilot + +import ( + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "testing" + "time" + + "github.com/github/copilot-sdk/go/internal/jsonrpc2" + "github.com/github/copilot-sdk/go/rpc" +) + +func TestGitHubTokenProviderConfigValidation(t *testing.T) { + provider := func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + } + + if _, err := NewClient(nil).CreateSession(t.Context(), &SessionConfig{ + GitHubToken: "static", + GitHubTokenProvider: provider, + }); err == nil || !strings.Contains(err.Error(), "cannot be used together") { + t.Fatalf("CreateSession error = %v", err) + } + if _, err := NewClient(nil).ResumeSession(t.Context(), "session", &ResumeSessionConfig{ + GitHubToken: "static", + GitHubTokenProvider: provider, + }); err == nil || !strings.Contains(err.Error(), "cannot be used together") { + t.Fatalf("ResumeSession error = %v", err) + } +} + +func TestGitHubTokenProviderCreateRequestAndCallback(t *testing.T) { + rpcClient, server, _ := newRuntimeShutdownRpcPair(t) + t.Cleanup(server.Stop) + client := &Client{ + client: rpcClient, + RPC: rpc.NewServerRPC(rpcClient), + sessions: make(map[string]*Session), + } + client.setupNotificationHandler() + + var createParams json.RawMessage + server.SetRequestHandler("session.create", func(params json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + createParams = append(json.RawMessage(nil), params...) + sessionID := sessionIDFromParams(t, params) + return []byte(`{"sessionId":"` + sessionID + `","workspacePath":"/workspace"}`), nil + }) + server.SetRequestHandler("session.destroy", func(json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + return []byte(`{}`), nil + }) + + var gotArgs GitHubTokenProviderArgs + session, err := client.CreateSession(t.Context(), &SessionConfig{ + GitHubTokenProvider: func(args GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + gotArgs = args + return GitHubTokenResult(&GitHubToken{ + AccessToken: "secret-token", + TokenType: String("bearer"), + ExpiresIn: 8 * 60 * 60, + }), nil + }, + }) + if err != nil { + t.Fatalf("CreateSession failed: %v", err) + } + + var wire struct { + RegistrationID string `json:"gitHubTokenProviderRegistrationId"` + GitHubToken string `json:"gitHubToken"` + } + if err := json.Unmarshal(createParams, &wire); err != nil { + t.Fatal(err) + } + if wire.RegistrationID == "" { + t.Fatal("gitHubTokenProviderRegistrationId was not serialized") + } + if wire.GitHubToken != "" { + t.Fatal("static gitHubToken should not be serialized") + } + + sessionID := session.SessionID + raw, rpcErr := server.Request(t.Context(), "gitHubToken.getToken", &rpc.GitHubTokenAcquireRequest{ + RegistrationID: wire.RegistrationID, + Host: "github.example.com", + SessionID: &sessionID, + Reason: rpc.GitHubTokenAcquireReasonRefresh, + }) + if rpcErr != nil { + t.Fatalf("getToken failed: %v", rpcErr) + } + var tokenResult struct { + Kind string `json:"kind"` + AccessToken string `json:"accessToken"` + ExpiresIn int64 `json:"expiresIn"` + } + if err := json.Unmarshal(raw, &tokenResult); err != nil { + t.Fatal(err) + } + if tokenResult.Kind != "token" || tokenResult.AccessToken != "secret-token" || tokenResult.ExpiresIn != 8*60*60 { + t.Fatalf("unexpected token result: %+v", tokenResult) + } + if gotArgs.Host != "github.example.com" || gotArgs.SessionID == nil || + *gotArgs.SessionID != sessionID || gotArgs.Reason != GitHubTokenRequestReasonRefresh { + t.Fatalf("unexpected callback args: %+v", gotArgs) + } + + if err := session.Disconnect(); err != nil { + t.Fatal(err) + } + if len(client.gitHubTokenProviders) != 0 { + t.Fatal("provider registration was not removed on disconnect") + } + if _, rpcErr := server.Request(t.Context(), "gitHubToken.getToken", &rpc.GitHubTokenAcquireRequest{ + RegistrationID: wire.RegistrationID, + Host: "github.com", + Reason: rpc.GitHubTokenAcquireReasonInitial, + }); rpcErr == nil { + t.Fatal("unknown registration ID should return a handler error") + } + + var resumeParams json.RawMessage + server.SetRequestHandler("session.resume", func(params json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + resumeParams = append(json.RawMessage(nil), params...) + return []byte(`{"sessionId":"resumed-session","workspacePath":"/workspace"}`), nil + }) + resumed, err := client.ResumeSession(t.Context(), "resumed-session", &ResumeSessionConfig{ + GitHubTokenProvider: func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + }, + }) + if err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(resumeParams, &wire); err != nil { + t.Fatal(err) + } + if wire.RegistrationID == "" { + t.Fatal("resume did not serialize gitHubTokenProviderRegistrationId") + } + if err := resumed.Disconnect(); err != nil { + t.Fatal(err) + } +} + +func TestGitHubTokenProviderResultsErrorsAndRollback(t *testing.T) { + client := &Client{} + adapter := &gitHubTokenAdapter{client: client} + + cancelID := client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + }) + result, err := adapter.GetToken(&rpc.GitHubTokenAcquireRequest{RegistrationID: cancelID}) + if err != nil { + t.Fatal(err) + } + if _, ok := result.(*rpc.GitHubTokenAcquireResultCancelled); !ok { + t.Fatalf("result type = %T", result) + } + + sentinel := errors.New("provider failed") + errorID := client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return nil, sentinel + }) + if _, err := adapter.GetToken(&rpc.GitHubTokenAcquireRequest{RegistrationID: errorID}); !errors.Is(err, sentinel) { + t.Fatalf("provider error = %v", err) + } + + rpcClient, server, _ := newRuntimeShutdownRpcPair(t) + t.Cleanup(server.Stop) + rollbackClient := &Client{ + client: rpcClient, + RPC: rpc.NewServerRPC(rpcClient), + sessions: make(map[string]*Session), + } + server.SetRequestHandler("session.create", func(json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + return nil, &jsonrpc2.Error{Code: -32000, Message: "create failed"} + }) + if _, err := rollbackClient.CreateSession(t.Context(), &SessionConfig{ + GitHubTokenProvider: func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + }, + }); err == nil { + t.Fatal("expected create failure") + } + if len(rollbackClient.gitHubTokenProviders) != 0 { + t.Fatal("provider registration was not rolled back") + } +} + +func TestGitHubTokenStringRedactsAccessToken(t *testing.T) { + token := GitHubToken{AccessToken: "secret-token", ExpiresIn: 28_800} + + if got := fmt.Sprintf("%v %#v", token, token); strings.Contains(got, token.AccessToken) { + t.Fatalf("GitHubToken formatting exposed the access token: %s", got) + } +} + +func TestGitHubTokenProviderCleanupOnDisconnectError(t *testing.T) { + rpcClient, server, _ := newRuntimeShutdownRpcPair(t) + t.Cleanup(server.Stop) + server.SetRequestHandler("session.destroy", func(json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + return nil, &jsonrpc2.Error{Code: -32000, Message: "destroy failed"} + }) + client := &Client{} + registrationID := client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + }) + session := newSession("cleanup-session", rpcClient, "", false) + session.setGitHubTokenProviderRegistrationRelease(func() { + client.unregisterGitHubTokenProvider(registrationID) + }) + + if err := session.Disconnect(); err == nil || !strings.Contains(err.Error(), "destroy failed") { + t.Fatalf("Disconnect error = %v", err) + } + if len(client.gitHubTokenProviders) != 0 { + t.Fatal("provider registration was not removed after disconnect failed") + } +} + +func TestGitHubTokenProviderReleaseBeforeOwnershipTransfer(t *testing.T) { + client := &Client{} + registrationID := client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + }) + session := &Session{} + + session.releaseGitHubTokenProviderRegistration() + session.setGitHubTokenProviderRegistrationRelease(func() { + client.unregisterGitHubTokenProvider(registrationID) + }) + + if len(client.gitHubTokenProviders) != 0 { + t.Fatal("provider registration was not removed after a pending session had already been retired") + } +} + +func TestGitHubTokenProviderCleanupOnDelete(t *testing.T) { + rpcClient, server, _ := newRuntimeShutdownRpcPair(t) + t.Cleanup(server.Stop) + server.SetRequestHandler("session.delete", func(json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + return []byte(`{"success":true}`), nil + }) + client := &Client{ + client: rpcClient, + sessions: make(map[string]*Session), + state: stateConnected, + } + registrationID := client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + }) + session := newSession("delete-session", rpcClient, "", false) + session.setGitHubTokenProviderRegistrationRelease(func() { + client.unregisterGitHubTokenProvider(registrationID) + }) + client.sessions[session.SessionID] = session + + if err := client.DeleteSession(t.Context(), session.SessionID); err != nil { + t.Fatal(err) + } + if len(client.gitHubTokenProviders) != 0 { + t.Fatal("provider registration was not removed after session deletion") + } +} + +func TestSessionOperationsSerializeBySessionID(t *testing.T) { + client := &Client{} + unlockFirst := client.lockSessionOperation("same-session") + sameSessionAcquired := make(chan struct{}) + go func() { + unlock := client.lockSessionOperation("same-session") + close(sameSessionAcquired) + unlock() + }() + + select { + case <-sameSessionAcquired: + t.Fatal("same-session operation was not serialized") + case <-time.After(25 * time.Millisecond): + } + + otherSessionAcquired := make(chan struct{}) + go func() { + unlock := client.lockSessionOperation("other-session") + close(otherSessionAcquired) + unlock() + }() + select { + case <-otherSessionAcquired: + case <-time.After(time.Second): + t.Fatal("different-session operation was unnecessarily blocked") + } + + unlockFirst() + select { + case <-sameSessionAcquired: + case <-time.After(time.Second): + t.Fatal("same-session operation did not proceed after release") + } +} + +func TestGitHubTokenProvidersClearedOnConnectionClose(t *testing.T) { + client := &Client{state: stateConnected} + client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenCancelled(), nil + }) + + client.handleConnectionClose() + + if len(client.gitHubTokenProviders) != 0 { + t.Fatal("provider registrations were not cleared after connection closure") + } +} + +func TestGitHubTokenProviderConcurrentRegistrationsAreIsolated(t *testing.T) { + client := &Client{} + adapter := &gitHubTokenAdapter{client: client} + idA := client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenResult(&GitHubToken{AccessToken: "a", ExpiresIn: 1}), nil + }) + idB := client.registerGitHubTokenProvider(func(GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) { + return GitHubTokenResult(&GitHubToken{AccessToken: "b", ExpiresIn: 1}), nil + }) + + var wg sync.WaitGroup + for id, want := range map[string]string{idA: "a", idB: "b"} { + wg.Add(1) + go func() { + defer wg.Done() + result, err := adapter.GetToken(&rpc.GitHubTokenAcquireRequest{RegistrationID: id}) + if err != nil { + t.Error(err) + return + } + if got := result.(*rpc.GitHubTokenAcquireResultToken).AccessToken; got != want { + t.Errorf("token = %q, want %q", got, want) + } + }() + } + wg.Wait() +} diff --git a/go/session.go b/go/session.go index 600a4bbebc..3f6c4f605f 100644 --- a/go/session.go +++ b/go/session.go @@ -56,42 +56,45 @@ type sessionHandler struct { // }) type Session struct { // SessionID is the unique identifier for this session. - SessionID string - workspacePath string - client *jsonrpc2.Client - clientSessionAPIs *rpc.ClientSessionAPIHandlers - handlers []sessionHandler - nextHandlerID uint64 - handlerMutex sync.RWMutex - toolHandlers map[string]ToolHandler - toolHandlersM sync.RWMutex - permissionHandler PermissionHandlerFunc - permissionMux sync.RWMutex - managedSettings bool - mcpAuthHandler MCPAuthHandler - mcpAuthMu sync.RWMutex - userInputHandler UserInputHandler - userInputMux sync.RWMutex - exitPlanModeHandler ExitPlanModeRequestHandler - exitPlanModeMu sync.RWMutex - autoModeSwitchHandler AutoModeSwitchRequestHandler - autoModeSwitchMu sync.RWMutex - hooks *SessionHooks - hooksMux sync.RWMutex - transformCallbacks map[string]SectionTransformFn - transformMu sync.Mutex - commandHandlers map[string]CommandHandler - commandHandlersMu sync.RWMutex - elicitationHandler ElicitationHandler - elicitationMu sync.RWMutex - canvasHandler CanvasHandler - canvasMu sync.RWMutex - bearerTokenProviders map[string]BearerTokenProvider - bearerTokenMu sync.RWMutex - openCanvases []rpc.OpenCanvasInstance - openCanvasesMu sync.RWMutex - capabilities SessionCapabilities - capabilitiesMu sync.RWMutex + SessionID string + workspacePath string + client *jsonrpc2.Client + clientSessionAPIs *rpc.ClientSessionAPIHandlers + handlers []sessionHandler + nextHandlerID uint64 + handlerMutex sync.RWMutex + toolHandlers map[string]ToolHandler + toolHandlersM sync.RWMutex + permissionHandler PermissionHandlerFunc + permissionMux sync.RWMutex + managedSettings bool + mcpAuthHandler MCPAuthHandler + mcpAuthMu sync.RWMutex + userInputHandler UserInputHandler + userInputMux sync.RWMutex + exitPlanModeHandler ExitPlanModeRequestHandler + exitPlanModeMu sync.RWMutex + autoModeSwitchHandler AutoModeSwitchRequestHandler + autoModeSwitchMu sync.RWMutex + hooks *SessionHooks + hooksMux sync.RWMutex + transformCallbacks map[string]SectionTransformFn + transformMu sync.Mutex + commandHandlers map[string]CommandHandler + commandHandlersMu sync.RWMutex + elicitationHandler ElicitationHandler + elicitationMu sync.RWMutex + canvasHandler CanvasHandler + canvasMu sync.RWMutex + bearerTokenProviders map[string]BearerTokenProvider + bearerTokenMu sync.RWMutex + releaseGitHubTokenProvider func() + gitHubTokenProviderMu sync.Mutex + gitHubTokenProviderReleased bool + openCanvases []rpc.OpenCanvasInstance + openCanvasesMu sync.RWMutex + capabilities SessionCapabilities + capabilitiesMu sync.RWMutex // eventCh serializes user event handler dispatch. dispatchEvent enqueues; // a single goroutine (processEvents) dequeues and invokes handlers in FIFO order. @@ -1722,11 +1725,9 @@ func (s *Session) GetEvents(ctx context.Context) ([]SessionEvent, error) { // } func (s *Session) Disconnect() error { _, err := s.client.Request(context.Background(), "session.destroy", sessionDestroyRequest{SessionID: s.SessionID}) - if err != nil { - return fmt.Errorf("failed to disconnect session: %w", err) - } s.closeOnce.Do(func() { close(s.eventCh) }) + s.releaseGitHubTokenProviderRegistration() // Clear handlers s.handlerMutex.Lock() @@ -1749,9 +1750,38 @@ func (s *Session) Disconnect() error { s.elicitationHandler = nil s.elicitationMu.Unlock() + if err != nil { + return fmt.Errorf("failed to disconnect session: %w", err) + } return nil } +func (s *Session) releaseGitHubTokenProviderRegistration() { + s.gitHubTokenProviderMu.Lock() + if s.gitHubTokenProviderReleased { + s.gitHubTokenProviderMu.Unlock() + return + } + s.gitHubTokenProviderReleased = true + release := s.releaseGitHubTokenProvider + s.releaseGitHubTokenProvider = nil + s.gitHubTokenProviderMu.Unlock() + if release != nil { + release() + } +} + +func (s *Session) setGitHubTokenProviderRegistrationRelease(release func()) { + s.gitHubTokenProviderMu.Lock() + if !s.gitHubTokenProviderReleased { + s.releaseGitHubTokenProvider = release + s.gitHubTokenProviderMu.Unlock() + return + } + s.gitHubTokenProviderMu.Unlock() + release() +} + // Abort aborts the currently processing message in this session. // // Use this to cancel a long-running request. The session remains valid diff --git a/go/types.go b/go/types.go index ff77c9e0c7..f5fadba6b0 100644 --- a/go/types.go +++ b/go/types.go @@ -1322,6 +1322,9 @@ type SessionConfig struct { // When provided, the SDK can satisfy MCP server OAuth requests with host-provided // token data or cancellation. OnMCPAuthRequest MCPAuthHandler + // GitHubTokenProvider acquires session-scoped GitHub tokens on demand. It + // cannot be combined with GitHubToken. + GitHubTokenProvider GitHubTokenProvider // OnUserInputRequest is a handler for user input requests from the agent (enables ask_user tool) OnUserInputRequest UserInputHandler // Hooks configures hook handlers for session lifecycle events @@ -1798,6 +1801,9 @@ type ResumeSessionConfig struct { // ClientName identifies the application using the SDK. // Included in the User-Agent header for API requests. ClientName string + // GitHubTokenProvider acquires session-scoped GitHub tokens on demand. It + // cannot be combined with GitHubToken. + GitHubTokenProvider GitHubTokenProvider // Model to use for this session. Can change the model when resuming. Model string // Tools exposes caller-implemented tools to the CLI. A Tool with a nil Handler @@ -2522,6 +2528,7 @@ type createSessionRequest struct { RequestMCPApps *bool `json:"requestMcpApps,omitempty"` GitHubMCPToolConfig *GitHubMCPToolConfig `json:"githubMcpToolConfig,omitempty"` GitHubToken string `json:"gitHubToken,omitempty"` + GitHubTokenProviderRegistrationID string `json:"gitHubTokenProviderRegistrationId,omitempty"` RemoteSession rpc.RemoteSessionMode `json:"remoteSession,omitempty"` Cloud *CloudSessionOptions `json:"cloud,omitempty"` Canvases []CanvasDeclaration `json:"canvases,omitempty"` @@ -2620,6 +2627,7 @@ type resumeSessionRequest struct { RequestMCPApps *bool `json:"requestMcpApps,omitempty"` GitHubMCPToolConfig *GitHubMCPToolConfig `json:"githubMcpToolConfig,omitempty"` GitHubToken string `json:"gitHubToken,omitempty"` + GitHubTokenProviderRegistrationID string `json:"gitHubTokenProviderRegistrationId,omitempty"` RemoteSession rpc.RemoteSessionMode `json:"remoteSession,omitempty"` Canvases []CanvasDeclaration `json:"canvases,omitempty"` OpenCanvases []rpc.OpenCanvasInstance `json:"openCanvases,omitempty"` diff --git a/java/README.md b/java/README.md index 02ffcdefff..fd03aa5f6c 100644 --- a/java/README.md +++ b/java/README.md @@ -176,6 +176,20 @@ directly. `CopilotClientOptions.setCwd(...)` sets the runtime process working directory, which otherwise inherits the current process working directory. `SessionConfig.setWorkingDirectory(...)` sets the session working directory, which otherwise defaults to the runtime process working directory. +For rotating per-session GitHub credentials, use +`SessionConfig.setGitHubTokenProvider(...)` (or the equivalent +`ResumeSessionConfig` setter) instead of `setGitHubToken(...)`: + +```java +var config = new SessionConfig().setGitHubTokenProvider(args -> + acquireForHost(args.host()).thenApply(token -> + GitHubTokenProviderResult.token(token, 8 * 60 * 60))); +``` + +The remaining lifetime is required and must be positive when the callback +completes; production GitHub tokens typically last eight hours. A static token +and a provider are mutually exclusive. + ## Permission Handling `PermissionHandler.APPROVE_ALL` approves requests when managed settings are disabled. When `enableManagedSettings` is true, it completes exceptionally. Custom handlers can inspect `request.getManagedApprovalRequired()` for human-facing confirmation logic. diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java index fe9a3c3b80..639562224d 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java @@ -117,6 +117,7 @@ public final class CopilotClient implements AutoCloseable { private final CliServerManager serverManager; private final LifecycleEventManager lifecycleManager = new LifecycleEventManager(); private final Map sessions = new ConcurrentHashMap<>(); + private final GitHubTokenProviderRegistry gitHubTokenProviders = new GitHubTokenProviderRegistry(); private volatile CompletableFuture connectionFuture; private volatile boolean disposed = false; private final String optionsHost; @@ -558,7 +559,8 @@ private Connection startCoreBody() { inProcessTransport == null ? null : inProcessTransport.host()); // Register handlers for server-to-client calls - RpcHandlerDispatcher dispatcher = new RpcHandlerDispatcher(sessions, lifecycleManager::dispatch, executor); + RpcHandlerDispatcher dispatcher = new RpcHandlerDispatcher(sessions, lifecycleManager::dispatch, executor, + gitHubTokenProviders); dispatcher.registerHandlers(connectedRpc); // Register the LLM inference request handler when configured. @@ -727,6 +729,7 @@ public CompletableFuture stop() { closeFutures.add(future); } sessions.clear(); + gitHubTokenProviders.clear(); return CompletableFuture.allOf(closeFutures.toArray(new CompletableFuture[0])) .thenCompose(v -> cleanupConnection(true)); @@ -740,6 +743,7 @@ public CompletableFuture stop() { public CompletableFuture forceStop() { disposed = true; sessions.clear(); + gitHubTokenProviders.clear(); // Dispatch the blocking shutdownOwnedExecutor() on a dedicated thread: // cleanupConnection() is chained off async work running on the owned // executor, so a plain whenComplete(...) here could land the awaitTermination @@ -873,6 +877,10 @@ public CompletableFuture createSession(SessionConfig config) { + "For example, to allow all permissions, use: " + "new SessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL)")); } + if (config.getGitHubToken() != null && config.getGitHubTokenProvider() != null) { + return CompletableFuture.failedFuture( + new IllegalArgumentException("gitHubToken and gitHubTokenProvider are mutually exclusive")); + } return ensureConnected().thenCompose(connection -> { long totalNanos = System.nanoTime(); // For cloud sessions, let the CLI/server assign the session id @@ -978,6 +986,16 @@ public CompletableFuture createSession(SessionConfig config) { } } + GitHubTokenProviderRegistry.Registration tokenRegistration = config.getGitHubTokenProvider() == null + ? null + : gitHubTokenProviders.register(config.getGitHubTokenProvider()); + if (tokenRegistration != null) { + request.setGitHubTokenProviderRegistrationId(tokenRegistration.id()); + if (preRegisteredSessionHolder[0] != null) { + preRegisteredSessionHolder[0].setGitHubTokenProviderRegistration(tokenRegistration); + } + } + long rpcNanos = System.nanoTime(); return connection.rpc.invoke("session.create", request, CreateSessionResponse.class) .thenCompose(response -> { @@ -996,6 +1014,9 @@ public CompletableFuture createSession(SessionConfig config) { CopilotSession session = preRegisteredSessionHolder[0] != null ? preRegisteredSessionHolder[0] : initializeSession.apply(returnedId); + if (tokenRegistration != null) { + session.setGitHubTokenProviderRegistration(tokenRegistration); + } registeredIdHolder[0] = returnedId; CompletableFuture interest = config.getOnMcpAuthRequest() != null ? session.getRpc().eventLog.registerInterest( @@ -1012,6 +1033,11 @@ public CompletableFuture createSession(SessionConfig config) { config.getCoauthorEnabled().orElse(null), config.getManageScheduleEnabled().orElse(null)); }).thenApply(v -> { + if (tokenRegistration != null) { + tokenRegistration.claim(session.getSessionId()); + } else { + gitHubTokenProviders.retire(session.getSessionId()); + } LoggingHelpers.logTiming(LOG, Level.FINE, "CopilotClient.createSession complete. Elapsed={Elapsed}, SessionId=" + session.getSessionId(), @@ -1022,6 +1048,9 @@ public CompletableFuture createSession(SessionConfig config) { if (registeredIdHolder[0] != null) { sessions.remove(registeredIdHolder[0]); } + if (tokenRegistration != null) { + tokenRegistration.close(); + } LoggingHelpers.logTiming(LOG, Level.WARNING, ex, "CopilotClient.createSession failed. Elapsed={Elapsed}, SessionId=" + (registeredIdHolder[0] != null ? registeredIdHolder[0] : ""), @@ -1069,10 +1098,15 @@ public CompletableFuture resumeSession(String sessionId, ResumeS + "For example, to allow all permissions, use: " + "new ResumeSessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL)")); } + if (config.getGitHubToken() != null && config.getGitHubTokenProvider() != null) { + return CompletableFuture.failedFuture( + new IllegalArgumentException("gitHubToken and gitHubTokenProvider are mutually exclusive")); + } return ensureConnected().thenCompose(connection -> { long totalNanos = System.nanoTime(); // Register the session before the RPC call to avoid missing early events. long setupNanos = System.nanoTime(); + CopilotSession replacedSession = sessions.get(sessionId); var session = new CopilotSession(sessionId, connection.rpc); session.setExecutor(executor); SessionRequestBuilder.configureSession(session, config); @@ -1137,6 +1171,14 @@ public CompletableFuture resumeSession(String sessionId, ResumeS } } + GitHubTokenProviderRegistry.Registration tokenRegistration = config.getGitHubTokenProvider() == null + ? null + : gitHubTokenProviders.register(config.getGitHubTokenProvider()); + if (tokenRegistration != null) { + request.setGitHubTokenProviderRegistrationId(tokenRegistration.id()); + session.setGitHubTokenProviderRegistration(tokenRegistration); + } + long rpcNanos = System.nanoTime(); return connection.rpc.invoke("session.resume", request, ResumeSessionResponse.class) .thenCompose(response -> { @@ -1166,7 +1208,6 @@ public CompletableFuture resumeSession(String sessionId, ResumeS session.setActiveSessionId(returnedId); sessions.put(returnedId, session); } - return updateSessionOptionsForMode(session, config.getSkipCustomInstructions().orElse(null), config.getCustomAgentsLocalOnly().orElse(null), config.getCoauthorEnabled().orElse(null), @@ -1175,6 +1216,11 @@ public CompletableFuture resumeSession(String sessionId, ResumeS "CopilotClient.resumeSession complete. Elapsed={Elapsed}, SessionId=" + sessionId, totalNanos); + if (tokenRegistration != null) { + tokenRegistration.claim(session.getSessionId()); + } else { + gitHubTokenProviders.retire(session.getSessionId()); + } return session; }); }).exceptionally(ex -> { @@ -1184,6 +1230,12 @@ public CompletableFuture resumeSession(String sessionId, ResumeS if (!sessionId.equals(activeId)) { sessions.remove(activeId); } + if (replacedSession != null) { + sessions.putIfAbsent(sessionId, replacedSession); + } + if (tokenRegistration != null) { + tokenRegistration.close(); + } LoggingHelpers.logTiming(LOG, Level.WARNING, ex, "CopilotClient.resumeSession failed. Elapsed={Elapsed}, SessionId=" + sessionId, totalNanos); diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java index bccf914db7..bf1f17d413 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java @@ -198,6 +198,7 @@ public final class CopilotSession implements AutoCloseable { private volatile Map>> transformCallbacks; private final ScheduledExecutorService timeoutScheduler; private volatile Executor executor; + private volatile GitHubTokenProviderRegistry.Registration gitHubTokenProviderRegistration; /** Tracks whether this session instance has been terminated via close(). */ private volatile boolean isTerminated = false; @@ -252,6 +253,10 @@ void setExecutor(Executor executor) { this.executor = executor; } + void setGitHubTokenProviderRegistration(GitHubTokenProviderRegistry.Registration registration) { + this.gitHubTokenProviderRegistration = registration; + } + /** * Gets the unique identifier for this session. * @@ -2301,6 +2306,11 @@ public void close() { } timeoutScheduler.shutdownNow(); + GitHubTokenProviderRegistry.Registration tokenRegistration = gitHubTokenProviderRegistration; + gitHubTokenProviderRegistration = null; + if (tokenRegistration != null) { + tokenRegistration.close(); + } try { rpc.invoke("session.destroy", Map.of("sessionId", sessionId), Void.class).get(5, TimeUnit.SECONDS); diff --git a/java/sdk/src/main/java/com/github/copilot/GitHubTokenProviderRegistry.java b/java/sdk/src/main/java/com/github/copilot/GitHubTokenProviderRegistry.java new file mode 100644 index 0000000000..59b4f61469 --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/GitHubTokenProviderRegistry.java @@ -0,0 +1,79 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; + +import com.github.copilot.rpc.GitHubTokenProvider; + +final class GitHubTokenProviderRegistry { + + private final Map providers = new HashMap<>(); + private final Map sessionOwners = new HashMap<>(); + + synchronized Registration register(GitHubTokenProvider provider) { + String id = UUID.randomUUID().toString(); + providers.put(id, provider); + return new Registration(this, id); + } + + synchronized GitHubTokenProvider get(String registrationId) { + return providers.get(registrationId); + } + + private synchronized void claim(String registrationId, String sessionId) { + String previous = sessionOwners.put(sessionId, registrationId); + if (previous != null && !previous.equals(registrationId)) { + providers.remove(previous); + } + } + + private synchronized void unregister(String registrationId) { + providers.remove(registrationId); + sessionOwners.values().removeIf(registrationId::equals); + } + + synchronized void retire(String sessionId) { + String registrationId = sessionOwners.remove(sessionId); + if (registrationId != null) { + providers.remove(registrationId); + } + } + + synchronized void clear() { + providers.clear(); + sessionOwners.clear(); + } + + static final class Registration implements AutoCloseable { + + private final GitHubTokenProviderRegistry registry; + private final String id; + private boolean closed; + + private Registration(GitHubTokenProviderRegistry registry, String id) { + this.registry = registry; + this.id = id; + } + + String id() { + return id; + } + + void claim(String sessionId) { + registry.claim(id, sessionId); + } + + @Override + public synchronized void close() { + if (!closed) { + closed = true; + registry.unregister(id); + } + } + } +} diff --git a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java index 550bd4ca42..7eda069d23 100644 --- a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java +++ b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java @@ -28,6 +28,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializationFeature; +import com.fasterxml.jackson.databind.node.ObjectNode; import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; import com.github.copilot.rpc.JsonRpcError; import com.github.copilot.rpc.JsonRpcRequest; @@ -220,7 +221,39 @@ private synchronized void sendMessage(Object message) throws IOException { outputStream.write(content); outputStream.flush(); - LOG.fine("Sent: " + json); + if (LOG.isLoggable(Level.FINE)) { + LOG.fine("Sent: " + redactCredentialsForLogging(json)); + } + } + + static String redactCredentialsForLogging(String json) { + try { + JsonNode root = MAPPER.readTree(json); + redactCredentials(root); + return MAPPER.writeValueAsString(root); + } catch (JsonProcessingException error) { + return ""; + } + } + + private static void redactCredentials(JsonNode node) { + if (node.isObject()) { + var object = (ObjectNode) node; + object.properties().forEach(entry -> { + if (isCredentialField(entry.getKey())) { + object.put(entry.getKey(), ""); + } else { + redactCredentials(entry.getValue()); + } + }); + } else if (node.isArray()) { + node.forEach(JsonRpcClient::redactCredentials); + } + } + + private static boolean isCredentialField(String name) { + return name.equals("accessToken") || name.equals("gitHubToken") || name.equals("bearerToken") + || name.equals("apiKey"); } private void startReader() { diff --git a/java/sdk/src/main/java/com/github/copilot/RpcHandlerDispatcher.java b/java/sdk/src/main/java/com/github/copilot/RpcHandlerDispatcher.java index d2dff958dc..0dd8a5dc98 100644 --- a/java/sdk/src/main/java/com/github/copilot/RpcHandlerDispatcher.java +++ b/java/sdk/src/main/java/com/github/copilot/RpcHandlerDispatcher.java @@ -20,7 +20,12 @@ import com.github.copilot.rpc.AutoModeSwitchRequest; import com.github.copilot.rpc.ExitPlanModeRequest; import com.github.copilot.rpc.BearerTokenProvider; +import com.github.copilot.rpc.GitHubTokenProviderArgs; +import com.github.copilot.rpc.GitHubTokenProviderResult; import com.github.copilot.rpc.ProviderTokenArgs; +import com.github.copilot.generated.rpc.GitHubTokenAcquireRequest; +import com.github.copilot.generated.rpc.GitHubTokenAcquireResultCancelled; +import com.github.copilot.generated.rpc.GitHubTokenAcquireResultToken; import com.github.copilot.rpc.PermissionRequestResult; import com.github.copilot.rpc.PermissionRequestResultKind; import com.github.copilot.rpc.SessionLifecycleEvent; @@ -51,6 +56,7 @@ final class RpcHandlerDispatcher { private final Map sessions; private final LifecycleEventDispatcher lifecycleDispatcher; private final Executor executor; + private final GitHubTokenProviderRegistry gitHubTokenProviders; /** * Creates a dispatcher with session registry and lifecycle dispatcher. @@ -63,10 +69,11 @@ final class RpcHandlerDispatcher { * the executor for async dispatch, or {@code null} for default */ RpcHandlerDispatcher(Map sessions, LifecycleEventDispatcher lifecycleDispatcher, - Executor executor) { + Executor executor, GitHubTokenProviderRegistry gitHubTokenProviders) { this.sessions = sessions; this.lifecycleDispatcher = lifecycleDispatcher; this.executor = executor; + this.gitHubTokenProviders = gitHubTokenProviders; } /** @@ -92,6 +99,71 @@ void registerHandlers(JsonRpcClient rpc) { (requestId, params) -> handleSystemMessageTransform(rpc, requestId, params)); rpc.registerMethodHandler("providerToken.getToken", (requestId, params) -> handleProviderTokenGetToken(rpc, requestId, params)); + rpc.registerMethodHandler("gitHubToken.getToken", + (requestId, params) -> handleGitHubTokenGetToken(rpc, requestId, params)); + } + + private void handleGitHubTokenGetToken(JsonRpcClient rpc, String requestId, JsonNode params) { + runAsync(() -> { + final long requestIdLong = parseRequestId(requestId, "gitHubToken.getToken"); + if (requestIdLong == -1) { + return; + } + try { + GitHubTokenAcquireRequest request = MAPPER.treeToValue(params, GitHubTokenAcquireRequest.class); + var provider = gitHubTokenProviders.get(request.registrationId()); + if (provider == null) { + rpc.sendErrorResponse(requestIdLong, -32603, "Unknown GitHub token provider registration"); + return; + } + + var resultFuture = provider + .getToken(new GitHubTokenProviderArgs(request.host(), request.sessionId(), request.reason())); + if (resultFuture == null) { + rpc.sendErrorResponse(requestIdLong, -32603, "GitHub token provider returned a null future"); + return; + } + resultFuture.thenAccept(result -> sendGitHubTokenResult(rpc, requestIdLong, result)) + .exceptionally(error -> { + try { + Throwable cause = error instanceof java.util.concurrent.CompletionException + && error.getCause() != null ? error.getCause() : error; + rpc.sendErrorResponse(requestIdLong, -32603, + "GitHub token provider failed: " + cause.getMessage()); + } catch (IOException sendError) { + LOG.log(Level.SEVERE, "Error sending GitHub token provider error", sendError); + } + return null; + }); + } catch (Exception error) { + try { + rpc.sendErrorResponse(requestIdLong, -32603, + "GitHub token provider handler failed: " + error.getMessage()); + } catch (IOException sendError) { + LOG.log(Level.SEVERE, "Error sending GitHub token provider handler error", sendError); + } + } + }); + } + + private void sendGitHubTokenResult(JsonRpcClient rpc, long requestId, GitHubTokenProviderResult result) { + try { + if (result == null) { + rpc.sendErrorResponse(requestId, -32603, "GitHub token provider returned a null result"); + return; + } + if (result.isCancelled()) { + rpc.sendResponse(requestId, new GitHubTokenAcquireResultCancelled()); + return; + } + var wireResult = new GitHubTokenAcquireResultToken(); + wireResult.setAccessToken(result.getAccessToken()); + wireResult.setExpiresIn(result.getExpiresIn()); + wireResult.setTokenType(result.getTokenType()); + rpc.sendResponse(requestId, wireResult); + } catch (IOException error) { + LOG.log(Level.SEVERE, "Error sending GitHub token provider result", error); + } } private void handleSessionEvent(JsonNode params) { diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/CreateSessionRequest.java b/java/sdk/src/main/java/com/github/copilot/rpc/CreateSessionRequest.java index 2eab977db1..403893987d 100644 --- a/java/sdk/src/main/java/com/github/copilot/rpc/CreateSessionRequest.java +++ b/java/sdk/src/main/java/com/github/copilot/rpc/CreateSessionRequest.java @@ -224,6 +224,9 @@ public final class CreateSessionRequest { @JsonProperty("gitHubToken") private String gitHubToken; + @JsonProperty("gitHubTokenProviderRegistrationId") + private String gitHubTokenProviderRegistrationId; + @JsonProperty("remoteSession") private String remoteSession; @@ -1061,6 +1064,21 @@ public void setGitHubToken(String gitHubToken) { this.gitHubToken = gitHubToken; } + /** + * Gets the token-provider registration ID. @return the opaque registration ID + */ + public String getGitHubTokenProviderRegistrationId() { + return gitHubTokenProviderRegistrationId; + } + + /** + * Sets the token-provider registration ID. @param registrationId the opaque + * registration ID + */ + public void setGitHubTokenProviderRegistrationId(String registrationId) { + this.gitHubTokenProviderRegistrationId = registrationId; + } + /** Gets the remote session mode. @return the remote session mode */ public String getRemoteSession() { return remoteSession; diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java new file mode 100644 index 0000000000..20cf11d1cf --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java @@ -0,0 +1,29 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.rpc; + +import java.util.concurrent.CompletableFuture; + +/** + * Acquires rotating GitHub tokens for one session. + *

+ * Implementations return either a token with a positive remaining lifetime or + * an explicit cancellation. Production GitHub tokens typically last eight + * hours. + * + * @since 1.0.0 + */ +@FunctionalInterface +public interface GitHubTokenProvider { + + /** + * Acquires a GitHub token for the supplied host and session context. + * + * @param args + * callback context + * @return a future containing a token or cancellation result + */ + CompletableFuture getToken(GitHubTokenProviderArgs args); +} diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderArgs.java b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderArgs.java new file mode 100644 index 0000000000..283d3c4c7f --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderArgs.java @@ -0,0 +1,22 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.rpc; + +import com.github.copilot.generated.rpc.GitHubTokenAcquireReason; + +/** + * Context supplied when a session needs a GitHub token. + * + * @param host + * effective GitHub host for which a token is required + * @param sessionId + * session receiving the token, or {@code null} before a cloud + * session has been assigned an ID + * @param reason + * whether this is the initial acquisition or a refresh + * @since 1.0.0 + */ +public record GitHubTokenProviderArgs(String host, String sessionId, GitHubTokenAcquireReason reason) { +} diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderResult.java b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderResult.java new file mode 100644 index 0000000000..9425286431 --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProviderResult.java @@ -0,0 +1,113 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.rpc; + +import java.util.Objects; + +/** + * Result of acquiring a session-scoped GitHub token. + *

+ * Token values are redacted from {@link #toString()}. + * + * @since 1.0.0 + */ +public final class GitHubTokenProviderResult { + + private final String accessToken; + private final long expiresIn; + private final String tokenType; + private final boolean cancelled; + + private GitHubTokenProviderResult(String accessToken, long expiresIn, String tokenType, boolean cancelled) { + this.accessToken = accessToken; + this.expiresIn = expiresIn; + this.tokenType = tokenType; + this.cancelled = cancelled; + } + + /** + * Creates a token result. + * + * @param accessToken + * GitHub access token + * @param expiresIn + * positive remaining lifetime in seconds when the callback completes + * @return the token result + */ + public static GitHubTokenProviderResult token(String accessToken, long expiresIn) { + return token(accessToken, expiresIn, null); + } + + /** + * Creates a token result with an explicit OAuth token type. + * + * @param accessToken + * GitHub access token + * @param expiresIn + * positive remaining lifetime in seconds when the callback completes + * @param tokenType + * OAuth token type, or {@code null} to use the runtime's bearer + * default + * @return the token result + */ + public static GitHubTokenProviderResult token(String accessToken, long expiresIn, String tokenType) { + Objects.requireNonNull(accessToken, "accessToken must not be null"); + return new GitHubTokenProviderResult(accessToken, expiresIn, tokenType, false); + } + + /** + * Creates an explicit cancellation result. + * + * @return the cancellation result + */ + public static GitHubTokenProviderResult cancelled() { + return new GitHubTokenProviderResult(null, 0, null, true); + } + + /** + * Gets whether acquisition was cancelled. + * + * @return {@code true} for a cancellation result + */ + public boolean isCancelled() { + return cancelled; + } + + /** + * Gets the access token. + * + * @return the token, or {@code null} for cancellation + */ + public String getAccessToken() { + return accessToken; + } + + /** + * Gets the remaining token lifetime. + * + * @return remaining lifetime in seconds + */ + public long getExpiresIn() { + return expiresIn; + } + + /** + * Gets the OAuth token type. + * + * @return token type, or {@code null} for the runtime default + */ + public String getTokenType() { + return tokenType; + } + + @Override + public String toString() { + if (cancelled) { + return "GitHubTokenProviderResult{cancelled}"; + } + return "GitHubTokenProviderResult{accessToken=, expiresIn=" + expiresIn + ", tokenType=" + tokenType + + "}"; + } +} diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionConfig.java b/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionConfig.java index a188036372..76b6e26a7d 100644 --- a/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionConfig.java +++ b/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionConfig.java @@ -105,6 +105,8 @@ public class ResumeSessionConfig { private boolean enableMcpApps; private GitHubMcpToolConfig githubMcpToolConfig; private String gitHubToken; + @JsonIgnore + private GitHubTokenProvider gitHubTokenProvider; private String remoteSession; private CopilotExpAssignmentResponse expAssignments; private Boolean enableManagedSettings; @@ -1864,6 +1866,33 @@ public ResumeSessionConfig setGitHubToken(String gitHubToken) { return this; } + /** + * Gets the rotating GitHub token provider for the resumed session. + * + * @return the provider, or {@code null} when a static token is used + */ + public GitHubTokenProvider getGitHubTokenProvider() { + return gitHubTokenProvider; + } + + /** + * Sets the rotating GitHub token provider for the resumed session. + *

+ * The provider receives only the effective host, optional assigned session ID, + * and acquisition reason. It must return a positive remaining lifetime in + * seconds when its callback completes. Production GitHub tokens typically last + * eight hours. This option is mutually exclusive with + * {@link #setGitHubToken(String)}. + * + * @param gitHubTokenProvider + * provider used for initial acquisition and refresh + * @return this config instance for method chaining + */ + public ResumeSessionConfig setGitHubTokenProvider(GitHubTokenProvider gitHubTokenProvider) { + this.gitHubTokenProvider = gitHubTokenProvider; + return this; + } + /** * Gets the per-session remote behavior control. *

@@ -2045,6 +2074,7 @@ public ResumeSessionConfig clone() { copy.enableMcpApps = this.enableMcpApps; copy.githubMcpToolConfig = this.githubMcpToolConfig; copy.gitHubToken = this.gitHubToken; + copy.gitHubTokenProvider = this.gitHubTokenProvider; copy.remoteSession = this.remoteSession; copy.expAssignments = this.expAssignments; copy.enableManagedSettings = this.enableManagedSettings; diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionRequest.java b/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionRequest.java index e52892477e..9b8e897fda 100644 --- a/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionRequest.java +++ b/java/sdk/src/main/java/com/github/copilot/rpc/ResumeSessionRequest.java @@ -229,6 +229,9 @@ public final class ResumeSessionRequest { @JsonProperty("gitHubToken") private String gitHubToken; + @JsonProperty("gitHubTokenProviderRegistrationId") + private String gitHubTokenProviderRegistrationId; + @JsonProperty("remoteSession") private String remoteSession; @@ -1086,6 +1089,21 @@ public void setGitHubToken(String gitHubToken) { this.gitHubToken = gitHubToken; } + /** + * Gets the token-provider registration ID. @return the opaque registration ID + */ + public String getGitHubTokenProviderRegistrationId() { + return gitHubTokenProviderRegistrationId; + } + + /** + * Sets the token-provider registration ID. @param registrationId the opaque + * registration ID + */ + public void setGitHubTokenProviderRegistrationId(String registrationId) { + this.gitHubTokenProviderRegistrationId = registrationId; + } + /** Gets the remote session mode. @return the remote session mode */ public String getRemoteSession() { return remoteSession; diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/SessionConfig.java b/java/sdk/src/main/java/com/github/copilot/rpc/SessionConfig.java index 1127e6777b..4e9ede73de 100644 --- a/java/sdk/src/main/java/com/github/copilot/rpc/SessionConfig.java +++ b/java/sdk/src/main/java/com/github/copilot/rpc/SessionConfig.java @@ -105,6 +105,8 @@ public class SessionConfig { private boolean enableMcpApps; private GitHubMcpToolConfig githubMcpToolConfig; private String gitHubToken; + @JsonIgnore + private GitHubTokenProvider gitHubTokenProvider; private String remoteSession; private CloudSessionOptions cloud; private CopilotExpAssignmentResponse expAssignments; @@ -1945,6 +1947,33 @@ public SessionConfig setGitHubToken(String gitHubToken) { return this; } + /** + * Gets the rotating GitHub token provider for this session. + * + * @return the provider, or {@code null} when a static token is used + */ + public GitHubTokenProvider getGitHubTokenProvider() { + return gitHubTokenProvider; + } + + /** + * Sets the rotating GitHub token provider for this session. + *

+ * The provider receives only the effective host, optional assigned session ID, + * and acquisition reason. It must return a positive remaining lifetime in + * seconds when its callback completes. Production GitHub tokens typically last + * eight hours. This option is mutually exclusive with + * {@link #setGitHubToken(String)}. + * + * @param gitHubTokenProvider + * provider used for initial acquisition and refresh + * @return this config instance for method chaining + */ + public SessionConfig setGitHubTokenProvider(GitHubTokenProvider gitHubTokenProvider) { + this.gitHubTokenProvider = gitHubTokenProvider; + return this; + } + /** * Gets the per-session remote behavior control. *

@@ -2063,10 +2092,11 @@ public Optional getEnableManagedSettings() { * (bypass-permissions policy) at session bootstrap. *

* When {@code true}, the runtime self-fetches enterprise managed settings using - * the session's {@link #getGitHubToken() gitHubToken}. Requires - * {@code gitHubToken} to be set; if omitted, the runtime is expected to reject - * session creation (fail-closed). When unset, behaves exactly as before. - * Serialized on the wire as {@code enableManagedSettings}. + * the session's static {@link #getGitHubToken() gitHubToken} or + * {@link #getGitHubTokenProvider() gitHubTokenProvider}. Requires one of those + * credentials; if both are omitted, the runtime is expected to reject session + * creation (fail-closed). When unset, behaves exactly as before. Serialized on + * the wire as {@code enableManagedSettings}. * * @param enableManagedSettings * {@code true} to opt into self-fetching managed settings @@ -2184,6 +2214,7 @@ public SessionConfig clone() { copy.enableMcpApps = this.enableMcpApps; copy.githubMcpToolConfig = this.githubMcpToolConfig; copy.gitHubToken = this.gitHubToken; + copy.gitHubTokenProvider = this.gitHubTokenProvider; copy.remoteSession = this.remoteSession; copy.cloud = this.cloud; copy.expAssignments = this.expAssignments; diff --git a/java/sdk/src/test/java/com/github/copilot/GitHubTokenProviderRegistryTest.java b/java/sdk/src/test/java/com/github/copilot/GitHubTokenProviderRegistryTest.java new file mode 100644 index 0000000000..8f2c7d3314 --- /dev/null +++ b/java/sdk/src/test/java/com/github/copilot/GitHubTokenProviderRegistryTest.java @@ -0,0 +1,74 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; + +import org.junit.jupiter.api.Test; + +import com.github.copilot.rpc.GitHubTokenProviderResult; +import com.github.copilot.rpc.PermissionHandler; +import com.github.copilot.rpc.ResumeSessionConfig; +import com.github.copilot.rpc.SessionConfig; + +class GitHubTokenProviderRegistryTest { + + @Test + void staticTokenAndProviderAreMutuallyExclusive() { + try (var client = new CopilotClient()) { + var create = new SessionConfig().setGitHubToken("static") + .setGitHubTokenProvider( + args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())) + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL); + var createError = assertThrows(CompletionException.class, () -> client.createSession(create).join()); + assertInstanceOf(IllegalArgumentException.class, createError.getCause()); + + var resume = new ResumeSessionConfig().setGitHubToken("static") + .setGitHubTokenProvider( + args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())) + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL); + var resumeError = assertThrows(CompletionException.class, + () -> client.resumeSession("session", resume).join()); + assertInstanceOf(IllegalArgumentException.class, resumeError.getCause()); + } + } + + @Test + void registrationRollbackAndResumeReplacementAreIsolated() { + var registry = new GitHubTokenProviderRegistry(); + var first = registry.register(args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())); + var second = registry + .register(args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())); + + assertNotNull(registry.get(first.id())); + assertNotNull(registry.get(second.id())); + first.claim("session-1"); + second.claim("session-1"); + assertNull(registry.get(first.id())); + assertNotNull(registry.get(second.id())); + + first.close(); + assertNotNull(registry.get(second.id())); + second.close(); + assertNull(registry.get(second.id())); + + var retired = registry + .register(args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())); + retired.claim("session-2"); + registry.retire("session-2"); + assertNull(registry.get(retired.id())); + + var failed = registry + .register(args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())); + failed.close(); + assertNull(registry.get(failed.id())); + } +} diff --git a/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java b/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java index 79aaea10d8..009f15c200 100644 --- a/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java +++ b/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java @@ -115,6 +115,22 @@ void testNotify() throws Exception { } } + @Test + void testCredentialValuesAreRedactedOnlyFromDiagnosticRendering() throws Exception { + String json = """ + {"jsonrpc":"2.0","result":{"accessToken":"secret","nested":{"gitHubToken":"static"}}, + "metadata":{"tokenType":"Bearer"}} + """; + + String rendered = JsonRpcClient.redactCredentialsForLogging(json); + + assertFalse(rendered.contains("secret")); + assertFalse(rendered.contains("static")); + assertEquals("", MAPPER.readTree(rendered).at("/result/accessToken").asText()); + assertEquals("", MAPPER.readTree(rendered).at("/result/nested/gitHubToken").asText()); + assertEquals("Bearer", MAPPER.readTree(rendered).at("/metadata/tokenType").asText()); + } + // ---- isConnected() ---- @Test diff --git a/java/sdk/src/test/java/com/github/copilot/RpcHandlerDispatcherTest.java b/java/sdk/src/test/java/com/github/copilot/RpcHandlerDispatcherTest.java index 76c2d41b29..1cb3ae82ec 100644 --- a/java/sdk/src/test/java/com/github/copilot/RpcHandlerDispatcherTest.java +++ b/java/sdk/src/test/java/com/github/copilot/RpcHandlerDispatcherTest.java @@ -15,6 +15,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; import org.junit.jupiter.api.AfterEach; @@ -32,6 +33,9 @@ import com.github.copilot.rpc.ToolDefinition; import com.github.copilot.rpc.ToolResultObject; import com.github.copilot.rpc.UserInputResponse; +import com.github.copilot.generated.rpc.GitHubTokenAcquireReason; +import com.github.copilot.rpc.GitHubTokenProviderArgs; +import com.github.copilot.rpc.GitHubTokenProviderResult; /** * Unit tests for {@link RpcHandlerDispatcher} focusing on coverage gaps @@ -51,6 +55,7 @@ class RpcHandlerDispatcherTest { private RpcHandlerDispatcher dispatcher; private InputStream responseStream; private Map> handlers; + private GitHubTokenProviderRegistry gitHubTokenProviders; @BeforeEach void setup() throws Exception { @@ -66,7 +71,8 @@ void setup() throws Exception { sessions = new ConcurrentHashMap<>(); lifecycleEvents = new CopyOnWriteArrayList<>(); - dispatcher = new RpcHandlerDispatcher(sessions, lifecycleEvents::add, null); + gitHubTokenProviders = new GitHubTokenProviderRegistry(); + dispatcher = new RpcHandlerDispatcher(sessions, lifecycleEvents::add, null, gitHubTokenProviders); dispatcher.registerHandlers(rpc); // Extract the registered handlers via reflection so we can invoke them directly @@ -119,6 +125,57 @@ private CopilotSession createSession(String sessionId) { return session; } + @Test + void gitHubTokenCallbackMapsRequestAndTokenResult() throws Exception { + AtomicReference received = new AtomicReference<>(); + var registration = gitHubTokenProviders.register(args -> { + received.set(args); + return CompletableFuture.completedFuture(GitHubTokenProviderResult.token("secret", 28_800, "bearer")); + }); + ObjectNode params = MAPPER.createObjectNode(); + params.put("registrationId", registration.id()); + params.put("host", "github.example"); + params.put("sessionId", "session-1"); + params.put("reason", "initial"); + + invokeHandler("gitHubToken.getToken", "80", params); + + JsonNode response = readResponse(); + assertEquals("token", response.at("/result/kind").asText()); + assertEquals("secret", response.at("/result/accessToken").asText()); + assertEquals(28_800, response.at("/result/expiresIn").asLong()); + assertEquals("github.example", received.get().host()); + assertEquals("session-1", received.get().sessionId()); + assertEquals(GitHubTokenAcquireReason.INITIAL, received.get().reason()); + } + + @Test + void gitHubTokenCallbackPreservesCancellationAndErrors() throws Exception { + var cancelled = gitHubTokenProviders + .register(args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())); + ObjectNode cancelledParams = MAPPER.createObjectNode(); + cancelledParams.put("registrationId", cancelled.id()); + cancelledParams.put("host", "github.com"); + cancelledParams.put("reason", "refresh"); + invokeHandler("gitHubToken.getToken", "81", cancelledParams); + assertEquals("cancelled", readResponse().at("/result/kind").asText()); + + var failed = gitHubTokenProviders.register( + args -> CompletableFuture.failedFuture(new IllegalStateException("credential service unavailable"))); + ObjectNode failedParams = cancelledParams.deepCopy(); + failedParams.put("registrationId", failed.id()); + invokeHandler("gitHubToken.getToken", "82", failedParams); + JsonNode error = readResponse(); + assertEquals(-32603, error.at("/error/code").asInt()); + assertTrue(error.at("/error/message").asText().contains("credential service unavailable")); + + failedParams.put("registrationId", "unknown"); + invokeHandler("gitHubToken.getToken", "83", failedParams); + JsonNode unknown = readResponse(); + assertEquals(-32603, unknown.at("/error/code").asInt()); + assertTrue(unknown.at("/error/message").asText().contains("Unknown GitHub token provider registration")); + } + // ===== session.event tests ===== @Test diff --git a/java/sdk/src/test/java/com/github/copilot/SessionRequestBuilderTest.java b/java/sdk/src/test/java/com/github/copilot/SessionRequestBuilderTest.java index 0525786de6..9d76d18ee2 100644 --- a/java/sdk/src/test/java/com/github/copilot/SessionRequestBuilderTest.java +++ b/java/sdk/src/test/java/com/github/copilot/SessionRequestBuilderTest.java @@ -26,6 +26,7 @@ import com.github.copilot.rpc.ExitPlanModeResult; import com.github.copilot.rpc.ExpConfigEntry; import com.github.copilot.rpc.GitHubMcpToolConfig; +import com.github.copilot.rpc.GitHubTokenProviderResult; import com.github.copilot.rpc.LargeToolOutputConfig; import com.github.copilot.rpc.MemoryConfiguration; import com.github.copilot.rpc.ResumeSessionConfig; @@ -56,6 +57,27 @@ void testBuildCreateRequestNullConfig() { assertEquals("direct", request.getEnvValueMode(), "envValueMode should be 'direct' even for null config"); } + @Test + void testGitHubTokenProviderRegistrationWireFieldHasExactCasing() throws Exception { + var create = SessionRequestBuilder.buildCreateRequest(new SessionConfig(), "create-session"); + create.setGitHubTokenProviderRegistrationId("create-registration"); + var resume = SessionRequestBuilder.buildResumeRequest("resume-session", new ResumeSessionConfig()); + resume.setGitHubTokenProviderRegistrationId("resume-registration"); + var mapper = JsonRpcClient.getObjectMapper(); + + assertEquals("create-registration", + mapper.readTree(mapper.writeValueAsBytes(create)).path("gitHubTokenProviderRegistrationId").asText()); + assertEquals("resume-registration", + mapper.readTree(mapper.writeValueAsBytes(resume)).path("gitHubTokenProviderRegistrationId").asText()); + assertFalse(mapper.readTree(mapper.writeValueAsBytes(create)).has("gitHubToken")); + } + + @Test + void testGitHubTokenProviderResultRedactsToken() { + var result = GitHubTokenProviderResult.token("do-not-print", 28_800); + assertFalse(result.toString().contains("do-not-print")); + } + @Test void testBuildCreateRequestHooksNonNullButEmpty() { // Hooks object exists but hasHooks() returns false diff --git a/nodejs/README.md b/nodejs/README.md index eec674ce4e..1e493b7b60 100644 --- a/nodejs/README.md +++ b/nodejs/README.md @@ -137,12 +137,23 @@ Create a new conversation session. - `infiniteSessions?: InfiniteSessionConfig` - Configure automatic context compaction (see below) - `workingDirectory?: string` - Working directory for the session (default: runtime process cwd). - `enableSessionStore?: boolean` - Enables the cross-session store for search and retrieval across sessions. When unset in `"copilot-cli"` mode, the runtime default applies (enabled). In `"empty"` mode, defaults to disabled. +- `gitHubTokenProvider?: GitHubTokenProvider` - Acquires rotating, session-scoped GitHub tokens. Token results require a positive `expiresIn` value in seconds remaining when the callback completes; production tokens typically last eight hours. Cannot be combined with `gitHubToken`. - `provider?: ProviderConfig` - Custom API provider configuration (BYOK - Bring Your Own Key). See [Custom Providers](#custom-providers) section. - `onPermissionRequest?: PermissionHandler` - Optional handler called before each tool execution to approve or deny it. When omitted, permission requests are emitted as events and left pending for manual resolution. `approveAll` approves requests when managed settings are disabled and throws when `enableManagedSettings` is true. Custom handlers can inspect `managedApprovalRequired` for human-facing confirmation logic. See [Permission Handling](#permission-handling) section. - `onUserInputRequest?: UserInputHandler` - Handler for user input requests from the agent. Enables the `ask_user` tool. See [User Input Requests](#user-input-requests) section. - `onElicitationRequest?: ElicitationHandler` - Handler for elicitation requests dispatched by the server. Enables this client to present form-based UI dialogs on behalf of the agent or other session participants. See [Elicitation Requests](#elicitation-requests) section. - `hooks?: SessionHooks` - Hook handlers for session lifecycle events. See [Session Hooks](#session-hooks) section. +```typescript +const session = await client.createSession({ + gitHubTokenProvider: async ({ host }) => ({ + kind: "token", + accessToken: await acquireTokenForHost(host), + expiresIn: 8 * 60 * 60, + }), +}); +``` + ##### `resumeSession(sessionId: string, config?: ResumeSessionConfig): Promise` Resume an existing session. Returns the session with `workspacePath` populated if infinite sessions were enabled. diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts index 1be2cbb941..86f58e6c27 100644 --- a/nodejs/src/client.ts +++ b/nodejs/src/client.ts @@ -35,6 +35,8 @@ import { } from "./generated/rpc.js"; import type { GitHubTelemetryNotification, + GitHubTokenAcquireRequest, + GitHubTokenAcquireResult, OpenCanvasInstance, SessionUpdateOptionsParams, } from "./generated/rpc.js"; @@ -58,6 +60,7 @@ import type { ForegroundSessionInfo, GetAuthStatusResponse, BearerTokenProvider, + GitHubTokenProvider, GetStatusResponse, InternalRuntimeConnection, RuntimeConnection, @@ -524,6 +527,10 @@ export class CopilotClient { private builtinPluginDirectories: string[] = []; private onGitHubTelemetry?: (notification: GitHubTelemetryNotification) => void | Promise; private clientGlobalHandlers: import("./generated/rpc.js").ClientGlobalApiHandlers = {}; + private githubTokenProviders = new Map< + string, + { provider: GitHubTokenProvider; sessionId?: string; committed: boolean } + >(); /** * Typed server-scoped RPC methods. @@ -862,9 +869,65 @@ export class CopilotClient { }, }; } + handlers.gitHubToken = { + getToken: (params) => this.acquireGitHubToken(params), + }; this.clientGlobalHandlers = handlers; } + private async acquireGitHubToken( + params: GitHubTokenAcquireRequest + ): Promise { + const registration = this.githubTokenProviders.get(params.registrationId); + if (!registration) { + throw new Error( + `No GitHub token provider registered for registration ID "${params.registrationId}"` + ); + } + return await registration.provider({ + host: params.host, + sessionId: params.sessionId ?? registration.sessionId, + reason: params.reason, + }); + } + + private registerGitHubTokenProvider( + provider: GitHubTokenProvider | undefined, + sessionId?: string + ): string | undefined { + if (!provider) { + return undefined; + } + const registrationId = randomUUID(); + this.githubTokenProviders.set(registrationId, { provider, sessionId, committed: false }); + return registrationId; + } + + private assignGitHubTokenProvider(registrationId: string | undefined, sessionId: string): void { + if (!registrationId) { + return; + } + const registration = this.githubTokenProviders.get(registrationId); + if (registration) { + registration.sessionId = sessionId; + } + } + + private commitGitHubTokenProvider(sessionId: string, registrationId?: string): void { + for (const [candidateId, registration] of this.githubTokenProviders) { + if (registration.sessionId === sessionId && registration.committed) { + this.githubTokenProviders.delete(candidateId); + } + } + const registration = registrationId + ? this.githubTokenProviders.get(registrationId) + : undefined; + if (registration) { + registration.sessionId = sessionId; + registration.committed = true; + } + } + /** * Starts the CLI server and establishes a connection. * @@ -1015,6 +1078,7 @@ export class CopilotClient { session._markDisconnected(); } this.sessions.clear(); + this.githubTokenProviders.clear(); // Ask SDK-owned runtimes to flush and clean up before we tear down // their transport/process. External runtimes may be shared, so only @@ -1197,6 +1261,7 @@ export class CopilotClient { session._markDisconnected(); } this.sessions.clear(); + this.githubTokenProviders.clear(); // Force close connection. Suppress writer failures first so teardown // write rejections don't surface as unhandled rejections. @@ -1444,6 +1509,9 @@ export class CopilotClient { } async createSession(config: SessionConfig): Promise { + if (config.gitHubToken !== undefined && config.gitHubTokenProvider !== undefined) { + throw new Error("gitHubToken and gitHubTokenProvider are mutually exclusive"); + } if (!this.connection) { await this.start(); } @@ -1463,6 +1531,11 @@ export class CopilotClient { const callerSessionId = config.sessionId; const useServerGeneratedId = config.cloud != null && callerSessionId == null; const localSessionId = useServerGeneratedId ? undefined : (callerSessionId ?? randomUUID()); + const toolFilterOptions = this.resolveToolFilterOptions(config); + const gitHubTokenProviderRegistrationId = this.registerGitHubTokenProvider( + config.gitHubTokenProvider, + localSessionId + ); // Strip non-serializable bearerTokenProvider callbacks from provider configs, // replacing them with a wire flag; keep the callbacks for session-side @@ -1491,6 +1564,13 @@ export class CopilotClient { managedSettingsEnabled: config.enableManagedSettings === true || config.managedSettings !== undefined, + onDisconnected: + gitHubTokenProviderRegistrationId === undefined + ? undefined + : () => + this.githubTokenProviders.delete( + gitHubTokenProviderRegistrationId + ), } ); s.registerTools(config.tools); @@ -1534,12 +1614,17 @@ export class CopilotClient { // processing (e.g. sessionFs.writeFile for workspace metadata) can be // routed to the correct handlers. if (localSessionId !== undefined) { - session = initializeSession(localSessionId); - registeredId = localSessionId; + try { + session = initializeSession(localSessionId); + registeredId = localSessionId; + } catch (error) { + if (gitHubTokenProviderRegistrationId !== undefined) { + this.githubTokenProviders.delete(gitHubTokenProviderRegistrationId); + } + throw error; + } } - const toolFilterOptions = this.resolveToolFilterOptions(config); - try { const response = await this.connection!.sendRequest("session.create", { ...(await getTraceContext(this.onGetTraceContext)), @@ -1628,6 +1713,7 @@ export class CopilotClient { infiniteSessions: config.infiniteSessions, memory: config.memory, gitHubToken: config.gitHubToken, + gitHubTokenProviderRegistrationId, remoteSession: config.remoteSession, cloud: config.cloud, expAssignments: config.expAssignments, @@ -1658,6 +1744,7 @@ export class CopilotClient { session = initializeSession(returnedSessionId); registeredId = returnedSessionId; } + this.assignGitHubTokenProvider(gitHubTokenProviderRegistrationId, returnedSessionId); if (config.onMcpAuthRequest) { await this.connection!.sendRequest("session.eventLog.registerInterest", { sessionId: returnedSessionId, @@ -1668,10 +1755,14 @@ export class CopilotClient { session.setCapabilities(capabilities); await this.updateSessionOptionsForMode(session, config); + this.commitGitHubTokenProvider(returnedSessionId, gitHubTokenProviderRegistrationId); } catch (e) { if (registeredId !== undefined) { this.sessions.delete(registeredId); } + if (gitHubTokenProviderRegistrationId !== undefined) { + this.githubTokenProviders.delete(gitHubTokenProviderRegistrationId); + } throw e; } @@ -1722,6 +1813,9 @@ export class CopilotClient { factories?: FactoryHandle[], extensionOptions?: ExtensionJoinOptions ): Promise { + if (config.gitHubToken !== undefined && config.gitHubTokenProvider !== undefined) { + throw new Error("gitHubToken and gitHubTokenProvider are mutually exclusive"); + } if (!this.connection) { await this.start(); } @@ -1787,6 +1881,15 @@ export class CopilotClient { this.setupSessionFs(session, config); const toolFilterOptions = this.resolveToolFilterOptions(config); + const gitHubTokenProviderRegistrationId = this.registerGitHubTokenProvider( + config.gitHubTokenProvider, + sessionId + ); + if (gitHubTokenProviderRegistrationId !== undefined) { + session._setOnDisconnected(() => + this.githubTokenProviders.delete(gitHubTokenProviderRegistrationId) + ); + } try { const response = await this.connection!.sendRequest("session.resume", { @@ -1880,6 +1983,7 @@ export class CopilotClient { disableResume: config.suppressResumeEvent, continuePendingWork: config.continuePendingWork, gitHubToken: config.gitHubToken, + gitHubTokenProviderRegistrationId, remoteSession: config.remoteSession, openCanvases: config.openCanvases, expAssignments: config.expAssignments, @@ -1927,8 +2031,12 @@ export class CopilotClient { } await this.updateSessionOptionsForMode(session, config); + this.commitGitHubTokenProvider(sessionId, gitHubTokenProviderRegistrationId); } catch (e) { this.sessions.delete(sessionId); + if (gitHubTokenProviderRegistrationId !== undefined) { + this.githubTokenProviders.delete(gitHubTokenProviderRegistrationId); + } throw e; } @@ -2936,6 +3044,7 @@ export class CopilotClient { this.connection.onClose(() => { this.state = "disconnected"; + this.githubTokenProviders.clear(); }); this.connection.onError((_error) => { diff --git a/nodejs/src/index.ts b/nodejs/src/index.ts index ae474eefee..bf7405c87b 100644 --- a/nodejs/src/index.ts +++ b/nodejs/src/index.ts @@ -97,6 +97,11 @@ export type { GitHubTelemetryNotification, GitHubTelemetryEvent, GitHubTelemetryClientInfo, + GitHubTokenAcquireReason, + GitHubTokenAcquireResult, + GitHubTokenProvider, + GitHubTokenProviderArgs, + GitHubTokenProviderResult, InfiniteSessionConfig, LargeToolOutputConfig, MemoryConfiguration, diff --git a/nodejs/src/session.ts b/nodejs/src/session.ts index efe9fa3e91..ca21b19471 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -438,6 +438,7 @@ export class CopilotSession { private _capabilities: SessionCapabilities = {}; private openCanvasInstances: OpenCanvasInstance[] = []; private disconnected = false; + private onDisconnected?: () => void; /** @internal Client session API handlers, populated by CopilotClient during create/resume. */ clientSessionApis: ClientSessionApiHandlers = {}; @@ -617,11 +618,16 @@ export class CopilotSession { private connection: MessageConnection, private _workspacePath?: string, traceContextProvider?: TraceContextProvider, - options?: { mcpAuthHandler?: McpAuthHandler; managedSettingsEnabled?: boolean } + options?: { + mcpAuthHandler?: McpAuthHandler; + managedSettingsEnabled?: boolean; + onDisconnected?: () => void; + } ) { this.traceContextProvider = traceContextProvider; this.mcpAuthHandler = options?.mcpAuthHandler; this.managedSettingsEnabled = options?.managedSettingsEnabled === true; + this.onDisconnected = options?.onDisconnected; } /** @@ -798,7 +804,12 @@ export class CopilotSession { /** @internal */ _markDisconnected(): void { + if (this.disconnected) { + return; + } this.disconnected = true; + this.onDisconnected?.(); + this.onDisconnected = undefined; this.eventHandlers.clear(); this.typedEventHandlers.clear(); this.toolHandlers.clear(); @@ -819,6 +830,11 @@ export class CopilotSession { this.transformCallbacks?.clear(); } + /** @internal */ + _setOnDisconnected(callback: () => void): void { + this.onDisconnected = callback; + } + /** * Subscribes to events from this session. * diff --git a/nodejs/src/types.ts b/nodejs/src/types.ts index 24d23d0826..02c01fac6b 100644 --- a/nodejs/src/types.ts +++ b/nodejs/src/types.ts @@ -21,6 +21,8 @@ import type { import type { CopilotSession } from "./session.js"; import type { FactoryJsonSchema, JsonValue } from "./factory.js"; import type { + GitHubTokenAcquireRequest, + GitHubTokenAcquireResult, GitHubTelemetryNotification, ModelBillingTokenPrices, OpenCanvasInstance, @@ -31,10 +33,37 @@ import type { ToolSet } from "./toolSet.js"; export type { RemoteSessionMode } from "./generated/rpc.js"; export type { CurrentToolMetadata } from "./generated/rpc.js"; export type { + GitHubTokenAcquireReason, + GitHubTokenAcquireResult, GitHubTelemetryNotification, GitHubTelemetryEvent, GitHubTelemetryClientInfo, } from "./generated/rpc.js"; + +/** + * Arguments passed to a session's {@link GitHubTokenProvider}. + * + * The callback registration identifier is intentionally kept inside the SDK. + */ +export type GitHubTokenProviderArgs = Pick< + GitHubTokenAcquireRequest, + "host" | "sessionId" | "reason" +>; + +/** Tagged token or cancellation returned by a {@link GitHubTokenProvider}. */ +export type GitHubTokenProviderResult = GitHubTokenAcquireResult; + +/** + * Acquires a GitHub token for one session. + * + * A token result must include `expiresIn`: the positive number of seconds of + * remaining lifetime when the callback completes. Production GitHub tokens + * typically last eight hours. Return `{ kind: "cancelled" }` to cancel the + * acquisition; thrown errors are returned to the runtime as callback errors. + */ +export type GitHubTokenProvider = ( + args: GitHubTokenProviderArgs +) => GitHubTokenProviderResult | Promise; export type { ModelBillingTokenPrices, ModelBillingTokenPricesLongContext, @@ -2712,6 +2741,16 @@ export interface SessionConfigBase { */ gitHubToken?: string; + /** + * Acquires short-lived GitHub credentials for this session on demand. + * + * Mutually exclusive with {@link SessionConfigBase.gitHubToken}. The + * callback receives the effective GitHub host, the session ID when known, + * and whether this is the initial acquisition or a refresh. Its opaque + * registration ID remains internal to the SDK. + */ + gitHubTokenProvider?: GitHubTokenProvider; + /** * Opt-in: when true, the runtime self-fetches enterprise managed settings * (bypass-permissions policy) at session bootstrap using the session's diff --git a/nodejs/test/github-token-provider.test.ts b/nodejs/test/github-token-provider.test.ts new file mode 100644 index 0000000000..caa8e1bbf4 --- /dev/null +++ b/nodejs/test/github-token-provider.test.ts @@ -0,0 +1,289 @@ +import { describe, expect, it, vi } from "vitest"; +import { CopilotClient, RuntimeConnection, type GitHubTokenProvider } from "../src/index.js"; + +function createMockClient( + request: (method: string, params: Record) => Promise +): CopilotClient { + const client = new CopilotClient({ + connection: RuntimeConnection.forUri("localhost:1234"), + }); + (client as unknown as { connection: unknown }).connection = { + sendRequest: request, + dispose: vi.fn(), + }; + return client; +} + +function getTokenHandler(client: CopilotClient) { + return ( + client as unknown as { + clientGlobalHandlers: { + gitHubToken: { + getToken(params: { + registrationId: string; + host: string; + sessionId?: string; + reason: "initial" | "refresh"; + }): Promise; + }; + }; + } + ).clientGlobalHandlers.gitHubToken.getToken; +} + +describe("session GitHub token providers", () => { + it("rejects a static token and provider together", async () => { + const client = new CopilotClient({ + connection: RuntimeConnection.forUri("localhost:1234"), + }); + + await expect( + client.createSession({ + gitHubToken: "static", + gitHubTokenProvider: async () => ({ + kind: "token", + accessToken: "dynamic", + expiresIn: 28_800, + }), + }) + ).rejects.toThrow("gitHubToken and gitHubTokenProvider are mutually exclusive"); + }); + + it("serializes only the opaque registration and maps token and cancellation results", async () => { + let createPayload: Record | undefined; + const request = vi.fn(async (method: string, params: Record) => { + if (method === "session.create") { + createPayload = params; + return { sessionId: params.sessionId }; + } + throw new Error(`Unexpected method: ${method}`); + }); + const observed: unknown[] = []; + const provider: GitHubTokenProvider = vi + .fn() + .mockImplementationOnce(async (args) => { + observed.push(args); + return { + kind: "token", + accessToken: "secret-token", + tokenType: "Bearer", + expiresIn: 28_800, + }; + }) + .mockImplementationOnce(async (args) => { + observed.push(args); + return { kind: "cancelled" }; + }); + const client = createMockClient(request); + const session = await client.createSession({ + sessionId: "session-one", + gitHubTokenProvider: provider, + }); + + const registrationId = createPayload?.gitHubTokenProviderRegistrationId; + expect(registrationId).toMatch( + /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/ + ); + expect(createPayload).not.toHaveProperty("gitHubTokenProvider"); + expect(createPayload?.gitHubToken).toBeUndefined(); + + const handler = getTokenHandler(client); + await expect( + handler({ + registrationId: registrationId as string, + host: "github.example.com", + reason: "initial", + }) + ).resolves.toEqual({ + kind: "token", + accessToken: "secret-token", + tokenType: "Bearer", + expiresIn: 28_800, + }); + await expect( + handler({ + registrationId: registrationId as string, + host: "github.example.com", + sessionId: session.sessionId, + reason: "refresh", + }) + ).resolves.toEqual({ kind: "cancelled" }); + expect(observed).toEqual([ + { + host: "github.example.com", + sessionId: "session-one", + reason: "initial", + }, + { + host: "github.example.com", + sessionId: "session-one", + reason: "refresh", + }, + ]); + }); + + it("preserves callback errors and rejects unknown registrations", async () => { + const client = createMockClient(async (_method, params) => ({ + sessionId: params.sessionId, + })); + const failure = new Error("credential broker failed"); + await client.createSession({ + sessionId: "error-session", + gitHubTokenProvider: () => { + throw failure; + }, + }); + const registrationId = [ + ...( + client as unknown as { + githubTokenProviders: Map; + } + ).githubTokenProviders.keys(), + ][0]; + const handler = getTokenHandler(client); + + await expect( + handler({ + registrationId, + host: "github.com", + reason: "initial", + }) + ).rejects.toBe(failure); + await expect( + handler({ + registrationId: "unknown", + host: "github.com", + reason: "refresh", + }) + ).rejects.toThrow("No GitHub token provider registered"); + }); + + it("rolls back failed creation and cleans up on session and client close", async () => { + const failingClient = createMockClient(async () => { + throw new Error("create failed"); + }); + await expect( + failingClient.createSession({ + gitHubTokenProvider: async () => ({ kind: "cancelled" }), + }) + ).rejects.toThrow("create failed"); + expect( + ( + failingClient as unknown as { + githubTokenProviders: Map; + } + ).githubTokenProviders + ).toHaveLength(0); + + const client = createMockClient(async (method, params) => { + if (method === "session.create") return { sessionId: params.sessionId }; + if (method === "session.destroy") return {}; + throw new Error(`Unexpected method: ${method}`); + }); + const first = await client.createSession({ + sessionId: "first", + gitHubTokenProvider: async () => ({ kind: "cancelled" }), + }); + await client.createSession({ + sessionId: "second", + gitHubTokenProvider: async () => ({ kind: "cancelled" }), + }); + const registrations = ( + client as unknown as { + githubTokenProviders: Map; + } + ).githubTokenProviders; + expect(registrations).toHaveLength(2); + + await first.disconnect(); + expect(registrations).toHaveLength(1); + await client.forceStop(); + expect(registrations).toHaveLength(0); + }); + + it("rotates a resumed session only after resume succeeds", async () => { + const payloads: Record[] = []; + const client = createMockClient(async (method, params) => { + payloads.push(params); + if (method === "session.create" || method === "session.resume") { + return { sessionId: params.sessionId }; + } + throw new Error(`Unexpected method: ${method}`); + }); + const firstProvider = vi.fn(async () => ({ kind: "cancelled" as const })); + const secondProvider = vi.fn(async () => ({ kind: "cancelled" as const })); + await client.createSession({ + sessionId: "resumed", + gitHubTokenProvider: firstProvider, + }); + const firstRegistration = payloads[0].gitHubTokenProviderRegistrationId as string; + + await client.resumeSession("resumed", { + gitHubTokenProvider: secondProvider, + }); + const secondRegistration = payloads[1].gitHubTokenProviderRegistrationId as string; + const handler = getTokenHandler(client); + + await expect( + handler({ + registrationId: firstRegistration, + host: "github.com", + reason: "refresh", + }) + ).rejects.toThrow("No GitHub token provider registered"); + await expect( + handler({ + registrationId: secondRegistration, + host: "github.com", + reason: "refresh", + }) + ).resolves.toEqual({ kind: "cancelled" }); + expect(secondProvider).toHaveBeenCalledOnce(); + expect(firstProvider).not.toHaveBeenCalled(); + }); + + it("does not retire a concurrent pending registration", async () => { + const resumeResolvers: Array<(value: { sessionId: string }) => void> = []; + const payloads: Record[] = []; + const client = createMockClient(async (method, params) => { + payloads.push(params); + if (method === "session.create") { + return { sessionId: params.sessionId }; + } + if (method === "session.resume") { + return await new Promise<{ sessionId: string }>((resolve) => { + resumeResolvers.push(resolve); + }); + } + throw new Error(`Unexpected method: ${method}`); + }); + await client.createSession({ + sessionId: "concurrent", + gitHubTokenProvider: async () => ({ kind: "cancelled" }), + }); + + const firstResume = client.resumeSession("concurrent", { + gitHubTokenProvider: async () => ({ kind: "cancelled" }), + }); + const secondResume = client.resumeSession("concurrent", { + gitHubTokenProvider: async () => ({ kind: "cancelled" }), + }); + await vi.waitFor(() => expect(resumeResolvers).toHaveLength(2)); + + resumeResolvers[0]({ sessionId: "concurrent" }); + await firstResume; + const registrations = ( + client as unknown as { + githubTokenProviders: Map; + } + ).githubTokenProviders; + expect(registrations).toHaveLength(2); + + resumeResolvers[1]({ sessionId: "concurrent" }); + await secondResume; + expect(registrations).toHaveLength(1); + expect(registrations.has(payloads[2].gitHubTokenProviderRegistrationId as string)).toBe( + true + ); + }); +}); diff --git a/python/README.md b/python/README.md index dc0a6a6794..ada0130b9d 100644 --- a/python/README.md +++ b/python/README.md @@ -281,9 +281,22 @@ These are passed as keyword arguments to `create_session()`: - `infinite_sessions` (InfiniteSessionConfig): Automatic context compaction configuration - `working_directory` (str | None): Working directory for the session (default: runtime process working directory). - `enable_session_store` (bool): Enables the cross-session store for search and retrieval across sessions. When unset in `"copilot-cli"` mode, the runtime default applies (enabled). In `"empty"` mode, defaults to disabled. +- `github_token_provider` (callable): Acquires rotating, session-scoped GitHub tokens. Token results require a positive `expiresIn` value in seconds remaining when the callback completes; production tokens typically last eight hours. Cannot be combined with `github_token`. - `on_permission_request` (callable): Optional handler called before each tool execution to approve or deny it. When omitted, permission requests are emitted as events and left pending for manual resolution. `PermissionHandler.approve_all` approves requests when managed settings are disabled and raises an error when `enable_managed_settings` is true. Custom handlers can inspect `managed_approval_required` for human-facing confirmation logic. See [Permission Handling](#permission-handling) section. - `on_user_input_request` (callable): Handler for user input requests from the agent (enables ask_user tool). See [User Input Requests](#user-input-requests) section. - `hooks` (SessionHooks): Hook handlers for session lifecycle events. See [Session Hooks](#session-hooks) section. + +```python +async def provide_github_token(args): + return { + "kind": "token", + "accessToken": await acquire_token_for_host(args["host"]), + "expiresIn": 8 * 60 * 60, + } + + +session = await client.create_session(github_token_provider=provide_github_token) +``` - `available_tools` / `excluded_tools` / `default_agent.excluded_tools` / custom-agent `tools`: MCP tools registered from `mcp_servers` are exposed to the runtime as `-`. For `available_tools` and `excluded_tools`, prefer `ToolSet().add_mcp("-")` or the raw `mcp:-` form. For custom-agent `tools` and `default_agent.excluded_tools`, use `-` directly. **Session Lifecycle Methods:** diff --git a/python/copilot/__init__.py b/python/copilot/__init__.py index 8f30632e37..5f5bef88bd 100644 --- a/python/copilot/__init__.py +++ b/python/copilot/__init__.py @@ -40,6 +40,11 @@ ExpFlagValue, GetAuthStatusResponse, GetStatusResponse, + GitHubTokenCancelledResult, + GitHubTokenProvider, + GitHubTokenProviderArgs, + GitHubTokenProviderResult, + GitHubTokenResult, InProcessRuntimeConnection, LogLevel, ManagedSettings, @@ -86,6 +91,9 @@ GitHubTelemetryClientInfo, GitHubTelemetryEvent, GitHubTelemetryNotification, + GitHubTokenAcquireReason, + GitHubTokenAcquireResult, + GitHubTokenAcquireResultKind, ModelBillingTokenPrices, ModelBillingTokenPricesLongContext, PermissionDecisionContext, @@ -267,6 +275,14 @@ "GitHubTelemetryClientInfo", "GitHubTelemetryEvent", "GitHubTelemetryNotification", + "GitHubTokenAcquireReason", + "GitHubTokenAcquireResult", + "GitHubTokenAcquireResultKind", + "GitHubTokenProvider", + "GitHubTokenProviderArgs", + "GitHubTokenProviderResult", + "GitHubTokenResult", + "GitHubTokenCancelledResult", "InfiniteSessionConfig", "InProcessRuntimeConnection", "InputOptions", diff --git a/python/copilot/client.py b/python/copilot/client.py index ad4b0fe171..96530c2b9b 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -29,7 +29,7 @@ from dataclasses import dataclass, field from datetime import UTC, datetime from types import TracebackType -from typing import Any, ClassVar, Literal, TypedDict, cast, overload +from typing import Any, ClassVar, Literal, NotRequired, TypedDict, cast, overload from ._diagnostics import log_timing from ._ffi_runtime_host import FfiRuntimeHost @@ -69,6 +69,9 @@ ClientGlobalApiHandlers, ClientSessionApiHandlers, GitHubTelemetryNotification, + GitHubTokenAcquireReason, + GitHubTokenAcquireRequest, + GitHubTokenAcquireResult, ModelBillingTokenPrices, ModelBillingTokenPricesLongContext, # noqa: F401 OpenCanvasInstance, @@ -123,6 +126,49 @@ logger = logging.getLogger(__name__) + +class GitHubTokenProviderArgs(TypedDict): + """Arguments passed to a session-scoped :data:`GitHubTokenProvider`. + + The opaque callback registration identifier is intentionally not exposed. + """ + + host: str + session_id: str | None + reason: GitHubTokenAcquireReason + + +class GitHubTokenResult(TypedDict): + """A GitHub token returned by a session-scoped provider.""" + + kind: Literal["token"] + accessToken: str + expiresIn: int + tokenType: NotRequired[str] + + +class GitHubTokenCancelledResult(TypedDict): + """An explicit cancellation returned by a session-scoped provider.""" + + kind: Literal["cancelled"] + + +GitHubTokenProviderResult = GitHubTokenResult | GitHubTokenCancelledResult +"""Result returned by a session-scoped GitHub token provider.""" + + +GitHubTokenProvider = Callable[ + [GitHubTokenProviderArgs], + GitHubTokenProviderResult | Awaitable[GitHubTokenProviderResult], +] +"""Acquire a GitHub credential for one session. + +Token results require ``expiresIn`` to be the positive number of seconds of +remaining lifetime when the callback completes. Production GitHub tokens +typically last eight hours. Return ``{"kind": "cancelled"}`` to cancel; +exceptions raised by the callback are returned to the runtime unchanged. +""" + # ============================================================================ # Connection Types # ============================================================================ @@ -633,6 +679,43 @@ async def event(self, params: GitHubTelemetryNotification) -> None: logger.warning("Error handling gitHubTelemetry.event notification", exc_info=True) +@dataclass +class _GitHubTokenProviderRegistration: + provider: GitHubTokenProvider + session_id: str | None = None + committed: bool = False + + +class _GitHubTokenProviderAdapter: + """Routes global GitHub token requests to opaque session registrations.""" + + def __init__(self, client: CopilotClient) -> None: + self._client = client + + async def get_token(self, params: GitHubTokenAcquireRequest) -> GitHubTokenAcquireResult: + with self._client._github_token_providers_lock: + registration = self._client._github_token_providers.get(params.registration_id) + if registration is None: + raise JsonRpcError( + -32603, + "No GitHub token provider registered for registration ID " + f"{params.registration_id!r}", + ) + + result = registration.provider( + GitHubTokenProviderArgs( + host=params.host, + session_id=params.session_id or registration.session_id, + reason=params.reason, + ) + ) + if inspect.isawaitable(result): + result = await result + # The generated global-handler wrapper forwards callback results directly, + # so the public tagged dictionary is already in the expected wire shape. + return cast(GitHubTokenAcquireResult, result) + + class _HooksAdapter: """Adapts session-scoped hook dispatch to the generated ``HooksHandler`` protocol. @@ -1622,6 +1705,9 @@ def __init__( self._state: _ConnectionState = "disconnected" self._sessions: dict[str, CopilotSession] = {} self._sessions_lock = threading.Lock() + self._github_token_providers: dict[str, _GitHubTokenProviderRegistration] = {} + self._github_token_providers_lock = threading.Lock() + self._github_token_provider_adapter = _GitHubTokenProviderAdapter(self) self._models_cache: list[ModelInfo] | None = None self._models_cache_lock = asyncio.Lock() self._lifecycle_handlers: list[SessionLifecycleHandler] = [] @@ -1936,6 +2022,8 @@ async def stop(self) -> None: errors.append( StopError(message=f"Failed to disconnect session {session.session_id}: {e}") ) + with self._github_token_providers_lock: + self._github_token_providers.clear() if ( self._rpc is not None @@ -2052,6 +2140,8 @@ async def force_stop(self) -> None: # Clear sessions immediately without trying to destroy them with self._sessions_lock: self._sessions.clear() + with self._github_token_providers_lock: + self._github_token_providers.clear() # Close the transport first to signal the server immediately. # For external servers (TCP), this closes the socket. @@ -2169,6 +2259,7 @@ async def create_session( on_auto_mode_switch_request: AutoModeSwitchHandler | None = None, create_session_fs_handler: CreateSessionFsHandler | None = None, github_token: str | None = None, + github_token_provider: GitHubTokenProvider | None = None, remote_session: RemoteSessionMode | None = None, cloud: CloudSessionOptions | None = None, canvases: list[CanvasDeclaration] | None = None, @@ -2344,6 +2435,11 @@ async def create_session( May be combined with ``enable_managed_settings``. Requires a runtime whose RPC schema includes ``managedSettings``. Sent on the wire as ``managedSettings``. + github_token_provider: Callback that acquires a short-lived GitHub + credential for this session. Mutually exclusive with + ``github_token``. It receives the effective host, session ID + when assigned, and acquisition reason; the registration ID is + kept internal. Returns: A :class:`CopilotSession` instance for the new session. @@ -2365,6 +2461,8 @@ async def create_session( """ if on_permission_request is not None and not callable(on_permission_request): raise ValueError("on_permission_request must be callable when provided.") + if github_token is not None and github_token_provider is not None: + raise ValueError("github_token and github_token_provider are mutually exclusive") if not self._client: await self.start() @@ -2666,6 +2764,11 @@ async def create_session( ) if local_session_id is not None: payload["sessionId"] = local_session_id + github_token_provider_registration_id = self._register_github_token_provider( + github_token_provider, local_session_id + ) + if github_token_provider_registration_id is not None: + payload["gitHubTokenProviderRegistrationId"] = github_token_provider_registration_id # Propagate W3C Trace Context to CLI if OpenTelemetry is active trace_ctx = get_trace_context() @@ -2686,6 +2789,13 @@ def _initialize_session(sid: str) -> CopilotSession: workspace_path=None, managed_settings_enabled=enable_managed_settings is True or managed_settings is not None, + on_disconnect=( + None + if github_token_provider_registration_id is None + else lambda: self._unregister_github_token_provider( + github_token_provider_registration_id + ) + ), ) if self._session_fs_config: if create_session_fs_handler is None: @@ -2747,8 +2857,12 @@ def _initialize_session(sid: str) -> CopilotSession: # processing (e.g. sessionFs.writeFile for workspace metadata) can be # routed to the correct handlers. if local_session_id is not None: - session = _initialize_session(local_session_id) - registered_session_id = local_session_id + try: + session = _initialize_session(local_session_id) + registered_session_id = local_session_id + except BaseException: + self._unregister_github_token_provider(github_token_provider_registration_id) + raise try: rpc_start = time.perf_counter() @@ -2788,6 +2902,9 @@ def _register_inline(raw_response: Any) -> None: f"session.create returned sessionId {response.get('sessionId')} " f"but the caller requested {local_session_id}" ) + self._assign_github_token_provider( + github_token_provider_registration_id, session.session_id + ) if on_mcp_auth_request is not None: await self._client.request( "session.eventLog.registerInterest", @@ -2800,6 +2917,7 @@ def _register_inline(raw_response: Any) -> None: if registered_session_id is not None: with self._sessions_lock: self._sessions.pop(registered_session_id, None) + self._unregister_github_token_provider(github_token_provider_registration_id) if not isinstance(exc, asyncio.CancelledError): log_timing( logger, @@ -2819,6 +2937,9 @@ def _register_inline(raw_response: Any) -> None: coauthor_enabled, manage_schedule_enabled, ) + self._commit_github_token_provider( + session.session_id, github_token_provider_registration_id + ) log_timing( logger, @@ -2897,6 +3018,7 @@ async def resume_session( on_auto_mode_switch_request: AutoModeSwitchHandler | None = None, create_session_fs_handler: CreateSessionFsHandler | None = None, github_token: str | None = None, + github_token_provider: GitHubTokenProvider | None = None, remote_session: RemoteSessionMode | None = None, continue_pending_work: bool | None = None, canvases: list[CanvasDeclaration] | None = None, @@ -3071,6 +3193,10 @@ async def resume_session( injected layer, and omitting it clears that layer so warm and cold resume behave identically. See :meth:`create_session`. Sent on the wire as ``managedSettings``. + github_token_provider: Callback that acquires a short-lived GitHub + credential for this resumed session. Mutually exclusive with + ``github_token``. The new registration replaces the prior + provider only after resume succeeds. Returns: A :class:`CopilotSession` instance for the resumed session. @@ -3094,6 +3220,8 @@ async def resume_session( """ if on_permission_request is not None and not callable(on_permission_request): raise ValueError("on_permission_request must be callable when provided.") + if github_token is not None and github_token_provider is not None: + raise ValueError("github_token and github_token_provider are mutually exclusive") if not self._client: await self.start() @@ -3415,6 +3543,16 @@ async def resume_session( commands_count=len(commands or []), has_hooks=hooks is not None, ) + github_token_provider_registration_id = self._register_github_token_provider( + github_token_provider, session_id + ) + if github_token_provider_registration_id is not None: + payload["gitHubTokenProviderRegistrationId"] = github_token_provider_registration_id + session._set_disconnect_callback( + lambda: self._unregister_github_token_provider( + github_token_provider_registration_id + ) + ) try: rpc_start = time.perf_counter() @@ -3442,6 +3580,7 @@ async def resume_session( except BaseException as exc: with self._sessions_lock: self._sessions.pop(session_id, None) + self._unregister_github_token_provider(github_token_provider_registration_id) if not isinstance(exc, asyncio.CancelledError): log_timing( logger, @@ -3461,6 +3600,7 @@ async def resume_session( coauthor_enabled, manage_schedule_enabled, ) + self._commit_github_token_provider(session_id, github_token_provider_registration_id) log_timing( logger, @@ -4350,7 +4490,7 @@ async def _connect_via_stdio(self) -> None: # Create JSON-RPC client with the process self._client = JsonRpcClient(self._process) - self._client.on_close = lambda: setattr(self, "_state", "disconnected") + self._client.on_close = self._handle_connection_close self._rpc = ServerRpc(self._client) # Set up notification handler for session events @@ -4471,7 +4611,7 @@ def wait(self, timeout=None): self._process = SocketWrapper(sock_file, sock) self._client = JsonRpcClient(self._process) - self._client.on_close = lambda: setattr(self, "_state", "disconnected") + self._client.on_close = self._handle_connection_close self._rpc = ServerRpc(self._client) # Set up notification handler for session events @@ -4592,9 +4732,56 @@ def _register_client_global_handlers(self) -> None: hooks=_HooksAdapter(self._get_session), llm_inference=llm_inference_adapter, git_hub_telemetry=github_telemetry_adapter, + git_hub_token=self._github_token_provider_adapter, ), ) + def _register_github_token_provider( + self, provider: GitHubTokenProvider | None, session_id: str | None + ) -> str | None: + if provider is None: + return None + registration_id = str(uuid.uuid4()) + with self._github_token_providers_lock: + self._github_token_providers[registration_id] = _GitHubTokenProviderRegistration( + provider, session_id + ) + return registration_id + + def _handle_connection_close(self) -> None: + self._state = "disconnected" + with self._github_token_providers_lock: + self._github_token_providers.clear() + + def _assign_github_token_provider(self, registration_id: str | None, session_id: str) -> None: + if registration_id is None: + return + with self._github_token_providers_lock: + registration = self._github_token_providers.get(registration_id) + if registration is not None: + registration.session_id = session_id + + def _unregister_github_token_provider(self, registration_id: str | None) -> None: + if registration_id is None: + return + with self._github_token_providers_lock: + self._github_token_providers.pop(registration_id, None) + + def _commit_github_token_provider(self, session_id: str, registration_id: str | None) -> None: + with self._github_token_providers_lock: + stale = [ + candidate_id + for candidate_id, registration in self._github_token_providers.items() + if registration.session_id == session_id and registration.committed + ] + for candidate_id in stale: + self._github_token_providers.pop(candidate_id, None) + if registration_id is not None: + registration = self._github_token_providers.get(registration_id) + if registration is not None: + registration.session_id = session_id + registration.committed = True + def _get_session(self, session_id: str) -> CopilotSession | None: with self._sessions_lock: return self._sessions.get(session_id) diff --git a/python/copilot/session.py b/python/copilot/session.py index 21c74bcaf5..86d3d7d287 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -1548,6 +1548,7 @@ def __init__( client: Any, workspace_path: os.PathLike[str] | str | None = None, managed_settings_enabled: bool = False, + on_disconnect: Callable[[], None] | None = None, ): """ Initialize a new CopilotSession. @@ -1600,6 +1601,11 @@ def __init__( self._open_canvases_lock = threading.Lock() self._rpc: SessionRpc | None = None self._destroyed = False + self._on_disconnect = on_disconnect + + def _set_disconnect_callback(self, callback: Callable[[], None]) -> None: + """Set the client-owned cleanup callback before the session becomes active.""" + self._on_disconnect = callback @property def rpc(self) -> SessionRpc: @@ -2973,6 +2979,9 @@ async def disconnect(self) -> None: try: await self._client.request("session.destroy", {"sessionId": self.session_id}) finally: + if self._on_disconnect is not None: + self._on_disconnect() + self._on_disconnect = None # Clear handlers even if the request fails. with self._event_handlers_lock: self._event_handlers.clear() diff --git a/python/test_github_token_provider.py b/python/test_github_token_provider.py new file mode 100644 index 0000000000..c4357c5745 --- /dev/null +++ b/python/test_github_token_provider.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import asyncio +from typing import Any, cast + +import pytest + +from copilot import ( + CopilotClient, + GitHubTokenAcquireReason, + RuntimeConnection, +) +from copilot._jsonrpc import JsonRpcClient, JsonRpcError +from copilot.rpc import GitHubTokenAcquireRequest + + +class FakeJsonRpcClient: + def __init__(self, *, fail_method: str | None = None) -> None: + self.fail_method = fail_method + self.requests: list[tuple[str, dict[str, Any]]] = [] + self.request_handlers: dict[str, Any] = {} + self.notification_method_handlers: dict[str, Any] = {} + + async def request(self, method: str, params: dict[str, Any], **kwargs: Any) -> dict[str, Any]: + self.requests.append((method, params)) + if method == self.fail_method: + raise RuntimeError(f"{method} failed") + if method in {"session.create", "session.resume"}: + response = {"sessionId": params["sessionId"]} + callback = kwargs.get("on_response_inline") + if callback is not None: + callback(response) + return response + if method == "session.destroy": + return {} + raise RuntimeError(f"Unexpected method: {method}") + + async def stop(self) -> None: + pass + + def set_request_handler(self, method: str, handler: Any) -> None: + self.request_handlers[method] = handler + + def set_notification_method_handler(self, method: str, handler: Any) -> None: + self.notification_method_handlers[method] = handler + + +class ConcurrentResumeJsonRpcClient(FakeJsonRpcClient): + def __init__(self) -> None: + super().__init__() + self.resume_responses: list[asyncio.Future[dict[str, Any]]] = [] + + async def request(self, method: str, params: dict[str, Any], **kwargs: Any) -> dict[str, Any]: + self.requests.append((method, params)) + if method == "session.create": + response = {"sessionId": params["sessionId"]} + callback = kwargs.get("on_response_inline") + if callback is not None: + callback(response) + return response + if method == "session.resume": + response = asyncio.get_running_loop().create_future() + self.resume_responses.append(response) + return await response + raise RuntimeError(f"Unexpected method: {method}") + + +def make_client(fake: FakeJsonRpcClient) -> CopilotClient: + client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) + client._client = cast(JsonRpcClient, fake) + return client + + +class TestGitHubTokenProvider: + async def test_mutual_exclusion(self) -> None: + client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) + + with pytest.raises( + ValueError, match="github_token and github_token_provider are mutually exclusive" + ): + await client.create_session( + github_token="static", + github_token_provider=lambda _: {"kind": "cancelled"}, + ) + + async def test_wire_mapping_token_and_cancelled(self) -> None: + fake = FakeJsonRpcClient() + client = make_client(fake) + observed: list[dict[str, Any]] = [] + + async def provider(args): + observed.append(dict(args)) + if len(observed) == 1: + return { + "kind": "token", + "accessToken": "secret-token", + "tokenType": "Bearer", + "expiresIn": 28_800, + } + return {"kind": "cancelled"} + + await client.create_session( + session_id="python-session", + github_token_provider=provider, + ) + create_payload = fake.requests[0][1] + registration_id = create_payload["gitHubTokenProviderRegistrationId"] + assert "github_token_provider" not in create_payload + assert "gitHubToken" not in create_payload + client._register_client_global_handlers() + get_token = fake.request_handlers["gitHubToken.getToken"] + + token = await get_token( + { + "registrationId": registration_id, + "host": "github.example.com", + "reason": "initial", + } + ) + cancelled = await get_token( + { + "registrationId": registration_id, + "host": "github.example.com", + "reason": "refresh", + "sessionId": "python-session", + } + ) + + assert token == { + "kind": "token", + "accessToken": "secret-token", + "tokenType": "Bearer", + "expiresIn": 28_800, + } + assert cancelled == {"kind": "cancelled"} + assert observed == [ + { + "host": "github.example.com", + "session_id": "python-session", + "reason": GitHubTokenAcquireReason.INITIAL, + }, + { + "host": "github.example.com", + "session_id": "python-session", + "reason": GitHubTokenAcquireReason.REFRESH, + }, + ] + + async def test_callback_and_unknown_registration_errors(self) -> None: + fake = FakeJsonRpcClient() + client = make_client(fake) + failure = RuntimeError("credential broker failed") + + def provider(_args): + raise failure + + await client.create_session( + session_id="error-session", + github_token_provider=provider, + ) + registration_id = fake.requests[0][1]["gitHubTokenProviderRegistrationId"] + + with pytest.raises(RuntimeError, match="credential broker failed") as exc: + await client._github_token_provider_adapter.get_token( + GitHubTokenAcquireRequest( + registration_id=registration_id, + host="github.com", + reason=GitHubTokenAcquireReason.INITIAL, + ) + ) + assert exc.value is failure + + with pytest.raises(JsonRpcError, match="No GitHub token provider registered"): + await client._github_token_provider_adapter.get_token( + GitHubTokenAcquireRequest( + registration_id="unknown", + host="github.com", + reason=GitHubTokenAcquireReason.REFRESH, + ) + ) + + async def test_failure_session_close_and_client_close_cleanup(self) -> None: + failing = make_client(FakeJsonRpcClient(fail_method="session.create")) + with pytest.raises(RuntimeError, match="session.create failed"): + await failing.create_session(github_token_provider=lambda _: {"kind": "cancelled"}) + assert failing._github_token_providers == {} + + fake = FakeJsonRpcClient() + client = make_client(fake) + first = await client.create_session( + session_id="first", + github_token_provider=lambda _: {"kind": "cancelled"}, + ) + await client.create_session( + session_id="second", + github_token_provider=lambda _: {"kind": "cancelled"}, + ) + assert len(client._github_token_providers) == 2 + + await first.disconnect() + assert len(client._github_token_providers) == 1 + await client.force_stop() + assert client._github_token_providers == {} + + async def test_resume_rotates_provider(self) -> None: + fake = FakeJsonRpcClient() + client = make_client(fake) + calls: list[str] = [] + + await client.create_session( + session_id="resumed", + github_token_provider=lambda _: calls.append("first") or {"kind": "cancelled"}, + ) + first_registration = fake.requests[0][1]["gitHubTokenProviderRegistrationId"] + await client.resume_session( + "resumed", + github_token_provider=lambda _: calls.append("second") or {"kind": "cancelled"}, + ) + second_registration = fake.requests[1][1]["gitHubTokenProviderRegistrationId"] + + with pytest.raises(JsonRpcError): + await client._github_token_provider_adapter.get_token( + GitHubTokenAcquireRequest( + registration_id=first_registration, + host="github.com", + reason=GitHubTokenAcquireReason.REFRESH, + ) + ) + assert await client._github_token_provider_adapter.get_token( + GitHubTokenAcquireRequest( + registration_id=second_registration, + host="github.com", + reason=GitHubTokenAcquireReason.REFRESH, + ) + ) == {"kind": "cancelled"} + assert calls == ["second"] + + async def test_concurrent_resume_keeps_pending_registration(self) -> None: + fake = ConcurrentResumeJsonRpcClient() + client = make_client(fake) + await client.create_session( + session_id="concurrent", + github_token_provider=lambda _: {"kind": "cancelled"}, + ) + + first_resume = asyncio.create_task( + client.resume_session( + "concurrent", + github_token_provider=lambda _: {"kind": "cancelled"}, + ) + ) + second_resume = asyncio.create_task( + client.resume_session( + "concurrent", + github_token_provider=lambda _: {"kind": "cancelled"}, + ) + ) + while len(fake.resume_responses) < 2: + await asyncio.sleep(0) + + fake.resume_responses[0].set_result({"sessionId": "concurrent"}) + await first_resume + assert len(client._github_token_providers) == 2 + + fake.resume_responses[1].set_result({"sessionId": "concurrent"}) + await second_resume + assert len(client._github_token_providers) == 1 + second_registration = fake.requests[2][1]["gitHubTokenProviderRegistrationId"] + assert second_registration in client._github_token_providers diff --git a/rust/README.md b/rust/README.md index 29fe673558..a3083d15d9 100644 --- a/rust/README.md +++ b/rust/README.md @@ -274,6 +274,29 @@ let config = SessionConfig { let session = client.create_session(config).await?; ``` +For rotating per-session GitHub credentials, install a `GitHubTokenProvider` +instead of setting `github_token`: + +```rust,ignore +use github_copilot_sdk::{ + GitHubToken, GitHubTokenProviderArgs, GitHubTokenProviderResult, SessionConfig, +}; + +let provider = Arc::new(|args: GitHubTokenProviderArgs| async move { + let access_token = acquire_for_host(&args.host).await?; + Ok(GitHubTokenProviderResult::Token(GitHubToken::new( + access_token, + 8 * 60 * 60, + ))) +}); +let config = SessionConfig::default().with_github_token_provider(provider); +``` + +The remaining lifetime is required and must be positive when the callback +completes; production GitHub tokens typically last eight hours. Static +`github_token` and a provider are mutually exclusive. The same provider API is +available on `ResumeSessionConfig`. + ### Session Hooks Hooks intercept CLI behavior at lifecycle points — tool use, prompt submission, session start/end, and errors. Install a `SessionHooks` impl with [`SessionConfig::with_hooks`] — the SDK auto-enables `hooks` in `SessionConfig` when one is set. diff --git a/rust/src/errors.rs b/rust/src/errors.rs index 6e05bbfae1..70f4c14ff1 100644 --- a/rust/src/errors.rs +++ b/rust/src/errors.rs @@ -218,6 +218,8 @@ pub enum ErrorKind { }, /// Invalid combination of options or configuration. InvalidConfig, + /// A session-scoped GitHub token provider failed or returned invalid data. + GitHubTokenProvider, } impl fmt::Display for ErrorKind { @@ -238,6 +240,7 @@ impl fmt::Display for ErrorKind { write!(f, "binary not found: {name}") } ErrorKind::InvalidConfig => write!(f, "invalid configuration"), + ErrorKind::GitHubTokenProvider => write!(f, "GitHub token provider error"), } } } diff --git a/rust/src/github_token.rs b/rust/src/github_token.rs new file mode 100644 index 0000000000..7016d0ca01 --- /dev/null +++ b/rust/src/github_token.rs @@ -0,0 +1,375 @@ +//! Session-scoped GitHub token provider callbacks. + +use std::collections::HashMap; +use std::future::Future; +use std::sync::{Arc, OnceLock, Weak}; + +use async_trait::async_trait; +use parking_lot::Mutex; +use serde_json::Value; + +use crate::generated::api_types::{ + GitHubTokenAcquireReason, GitHubTokenAcquireRequest, GitHubTokenAcquireResult, + GitHubTokenAcquireResultCancelled, GitHubTokenAcquireResultToken, +}; +use crate::{Client, ClientInner, JsonRpcError, JsonRpcRequest, JsonRpcResponse, error_codes}; + +/// Why the runtime is requesting a GitHub token. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GitHubTokenRequestReason { + /// The session needs its initial token. + Initial, + /// The session needs a refreshed token. + Refresh, +} + +/// Context supplied when the runtime needs a GitHub token for a session. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GitHubTokenProviderArgs { + /// Effective GitHub host for which a token is required. + pub host: String, + /// Session receiving the token, when the runtime has assigned its ID. + pub session_id: Option, + /// Whether this is the initial token acquisition or a refresh. + pub reason: GitHubTokenRequestReason, +} + +/// A GitHub access token returned by a session token provider. +/// +/// `expires_in_seconds` is the positive remaining lifetime when the callback +/// completes. Production GitHub tokens typically last eight hours. +pub struct GitHubToken { + access_token: String, + expires_in_seconds: i64, + token_type: Option, +} + +impl GitHubToken { + /// Construct a token response with its remaining lifetime in seconds. + pub fn new(access_token: impl Into, expires_in_seconds: i64) -> Self { + Self { + access_token: access_token.into(), + expires_in_seconds, + token_type: None, + } + } + + /// Override the OAuth token type. The runtime defaults to `bearer` when unset. + pub fn with_token_type(mut self, token_type: impl Into) -> Self { + self.token_type = Some(token_type.into()); + self + } + + fn into_wire(self) -> GitHubTokenAcquireResultToken { + GitHubTokenAcquireResultToken { + access_token: self.access_token, + expires_in: self.expires_in_seconds, + kind: Default::default(), + token_type: self.token_type, + } + } +} + +impl std::fmt::Debug for GitHubToken { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GitHubToken") + .field("access_token", &"") + .field("expires_in_seconds", &self.expires_in_seconds) + .field("token_type", &self.token_type) + .finish() + } +} + +/// Result of acquiring a session-scoped GitHub token. +pub enum GitHubTokenProviderResult { + /// A token was acquired. + Token(GitHubToken), + /// The host cancelled acquisition. + Cancelled, +} + +impl std::fmt::Debug for GitHubTokenProviderResult { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Token(token) => f.debug_tuple("Token").field(token).finish(), + Self::Cancelled => f.write_str("Cancelled"), + } + } +} + +/// Async callback used to acquire GitHub tokens for one session. +#[async_trait] +pub trait GitHubTokenProvider: Send + Sync { + /// Acquire a token or explicitly cancel the request. + async fn get_token( + &self, + args: GitHubTokenProviderArgs, + ) -> Result; +} + +#[async_trait] +impl GitHubTokenProvider for F +where + F: Fn(GitHubTokenProviderArgs) -> Fut + Send + Sync, + Fut: Future> + Send, +{ + async fn get_token( + &self, + args: GitHubTokenProviderArgs, + ) -> Result { + (self)(args).await + } +} + +#[derive(Default)] +struct RegistryState { + providers: HashMap>, + session_owners: HashMap, +} + +pub(crate) struct GitHubTokenRegistry { + state: Mutex, + client: OnceLock>, +} + +impl GitHubTokenRegistry { + pub(crate) fn new() -> Self { + Self { + state: Mutex::new(RegistryState::default()), + client: OnceLock::new(), + } + } + + pub(crate) fn set_client(&self, client: Weak) { + let _ = self.client.set(client); + } + + pub(crate) fn register(&self, provider: Arc) -> String { + let registration_id = uuid::Uuid::new_v4().to_string(); + self.state + .lock() + .providers + .insert(registration_id.clone(), provider); + registration_id + } + + pub(crate) fn claim(&self, registration_id: &str, session_id: crate::SessionId) { + let mut state = self.state.lock(); + if let Some(previous) = state + .session_owners + .insert(session_id, registration_id.to_string()) + && previous != registration_id + { + state.providers.remove(&previous); + } + } + + pub(crate) fn unregister(&self, registration_id: &str) { + let mut state = self.state.lock(); + state.providers.remove(registration_id); + state + .session_owners + .retain(|_, owned| owned != registration_id); + } + + pub(crate) fn retire_session(&self, session_id: &crate::SessionId) { + let mut state = self.state.lock(); + if let Some(registration_id) = state.session_owners.remove(session_id) { + state.providers.remove(®istration_id); + } + } + + pub(crate) fn clear(&self) { + let mut state = self.state.lock(); + state.providers.clear(); + state.session_owners.clear(); + } + + pub(crate) async fn dispatch(&self, request: JsonRpcRequest) { + let Some(inner) = self.client.get().and_then(Weak::upgrade) else { + return; + }; + let client = Client::from_inner(inner); + let params = request + .params + .clone() + .unwrap_or(Value::Object(serde_json::Map::new())); + let params: GitHubTokenAcquireRequest = match serde_json::from_value(params) { + Ok(params) => params, + Err(error) => { + send_error( + &client, + request.id, + error_codes::INVALID_PARAMS, + &format!("invalid params: {error}"), + ) + .await; + return; + } + }; + let provider = self + .state + .lock() + .providers + .get(¶ms.registration_id) + .cloned(); + let Some(provider) = provider else { + send_error( + &client, + request.id, + error_codes::INTERNAL_ERROR, + "unknown GitHub token provider registration", + ) + .await; + return; + }; + + let reason = match params.reason { + GitHubTokenAcquireReason::Initial => GitHubTokenRequestReason::Initial, + GitHubTokenAcquireReason::Refresh => GitHubTokenRequestReason::Refresh, + GitHubTokenAcquireReason::Unknown => { + send_error( + &client, + request.id, + error_codes::INVALID_PARAMS, + "unknown GitHub token acquisition reason", + ) + .await; + return; + } + }; + + match provider + .get_token(GitHubTokenProviderArgs { + host: params.host, + session_id: params.session_id, + reason, + }) + .await + { + Ok(GitHubTokenProviderResult::Token(token)) => { + respond( + &client, + request.id, + GitHubTokenAcquireResult::Token(token.into_wire()), + ) + .await; + } + Ok(GitHubTokenProviderResult::Cancelled) => { + respond( + &client, + request.id, + GitHubTokenAcquireResult::Cancelled(GitHubTokenAcquireResultCancelled { + kind: Default::default(), + }), + ) + .await; + } + Err(error) => { + send_error( + &client, + request.id, + error_codes::INTERNAL_ERROR, + &format!("GitHub token provider failed: {error}"), + ) + .await; + } + } + } +} + +pub(crate) struct GitHubTokenRegistration { + registry: Arc, + id: String, +} + +impl GitHubTokenRegistration { + pub(crate) fn new(registry: Arc, id: String) -> Self { + Self { registry, id } + } + + pub(crate) fn id(&self) -> &str { + &self.id + } + + pub(crate) fn claim(&self, session_id: crate::SessionId) { + self.registry.claim(&self.id, session_id); + } +} + +impl Drop for GitHubTokenRegistration { + fn drop(&mut self) { + self.registry.unregister(&self.id); + } +} + +async fn respond(client: &Client, request_id: u64, result: GitHubTokenAcquireResult) { + match serde_json::to_value(result) { + Ok(result) => { + let _ = client + .send_response(&JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: request_id, + result: Some(result), + error: None, + }) + .await; + } + Err(_) => { + send_error( + client, + request_id, + error_codes::INTERNAL_ERROR, + "serialization failure", + ) + .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(JsonRpcError { + code, + message: message.to_string(), + data: None, + }), + }) + .await; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn token_debug_is_redacted() { + let token = GitHubToken::new("do-not-print", 28_800); + assert!(!format!("{token:?}").contains("do-not-print")); + } + + #[test] + fn retiring_session_removes_its_provider() { + let registry = GitHubTokenRegistry::new(); + let provider = Arc::new(|_args: GitHubTokenProviderArgs| async { + Ok(GitHubTokenProviderResult::Cancelled) + }); + let registration_id = registry.register(provider); + let session_id = crate::SessionId::from("session-1"); + registry.claim(®istration_id, session_id.clone()); + + registry.retire_session(&session_id); + + assert!( + !registry + .state + .lock() + .providers + .contains_key(®istration_id) + ); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index bf50adfd48..0517493c2a 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -22,6 +22,8 @@ pub mod copilot_request_handler; /// `#[doc(hidden)]` — re-exports the generated telemetry payload types. #[doc(hidden)] pub mod github_telemetry; +/// Session-scoped GitHub token provider callbacks. +pub mod github_token; /// Event handler traits for session lifecycle. pub mod handler; /// Lifecycle hook callbacks (pre/post tool use, prompt submission, session start/end). @@ -84,6 +86,10 @@ use async_trait::async_trait; pub use indexmap::IndexMap; // JSON-RPC wire types are internal transport details. // External callers interact via Client/Session methods, not raw RPC. +pub use github_token::{ + GitHubToken, GitHubTokenProvider, GitHubTokenProviderArgs, GitHubTokenProviderResult, + GitHubTokenRequestReason, +}; pub(crate) use jsonrpc::{ JsonRpcClient, JsonRpcError, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, error_codes, }; @@ -1020,6 +1026,7 @@ struct ClientInner { request_rx: parking_lot::Mutex>>, notification_tx: broadcast::Sender, router: router::SessionRouter, + github_token_registry: Arc, negotiated_protocol_version: OnceLock, state: parking_lot::Mutex, lifecycle_tx: broadcast::Sender, @@ -1423,6 +1430,7 @@ impl Client { &client.inner.request_rx, Some(dispatcher.clone()), client.inner.on_github_telemetry.clone(), + client.inner.github_token_registry.clone(), ); client.rpc().llm_inference().set_provider().await?; let llm_inference_elapsed = llm_inference_start.elapsed(); @@ -1603,6 +1611,7 @@ impl Client { let pid = child.as_ref().and_then(|c| c.id()); info!(pid = ?pid, "copilot CLI client ready"); + let github_token_registry = Arc::new(github_token::GitHubTokenRegistry::new()); let client = Self { inner: Arc::new(ClientInner { child: parking_lot::Mutex::new(child), @@ -1613,6 +1622,7 @@ impl Client { request_rx: parking_lot::Mutex::new(Some(request_rx)), notification_tx: notification_broadcast_tx, router: router::SessionRouter::new(), + github_token_registry: github_token_registry.clone(), negotiated_protocol_version: OnceLock::new(), state: parking_lot::Mutex::new(ConnectionState::Connected), lifecycle_tx: broadcast::channel(256).0, @@ -1628,6 +1638,7 @@ impl Client { startup_timings: OnceLock::new(), }), }; + github_token_registry.set_client(Arc::downgrade(&client.inner)); client.spawn_lifecycle_dispatcher(); debug!( elapsed_ms = setup_start.elapsed().as_millis(), @@ -2046,6 +2057,7 @@ impl Client { &self.inner.request_rx, self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), + self.inner.github_token_registry.clone(), ); self.inner.router.register(session_id) } @@ -2055,6 +2067,25 @@ impl Client { self.inner.router.unregister(session_id); } + pub(crate) fn register_github_token_provider( + &self, + provider: Arc, + ) -> github_token::GitHubTokenRegistration { + self.inner.router.ensure_started( + &self.inner.notification_tx, + &self.inner.request_rx, + self.inner.llm_inference.get().cloned(), + self.inner.on_github_telemetry.clone(), + self.inner.github_token_registry.clone(), + ); + let id = self.inner.github_token_registry.register(provider); + github_token::GitHubTokenRegistration::new(self.inner.github_token_registry.clone(), id) + } + + pub(crate) fn retire_github_token_provider(&self, session_id: &SessionId) { + self.inner.github_token_registry.retire_session(session_id); + } + /// Returns the protocol version negotiated with the CLI server, if any. /// /// Set during [`start`](Self::start). Returns `None` if the server didn't @@ -2263,6 +2294,7 @@ impl Client { &self.inner.request_rx, self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), + self.inner.github_token_registry.clone(), ); } @@ -2286,6 +2318,7 @@ impl Client { } self.inner.router.unregister(&session_id); } + self.inner.github_token_registry.clear(); match self.list_sessions(None).await { Ok(sessions) => { @@ -2455,6 +2488,7 @@ impl Client { } self.inner.router.unregister(&session_id); } + self.inner.github_token_registry.clear(); let should_shutdown_runtime = self.inner.child.lock().is_some(); #[cfg(feature = "bundled-in-process")] @@ -2581,6 +2615,7 @@ impl Client { // Drop all session channels so any awaiters see a closed channel // instead of waiting for responses that will never arrive. self.inner.router.clear(); + self.inner.github_token_registry.clear(); *self.inner.state.lock() = ConnectionState::Disconnected; *self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new()); } @@ -3388,6 +3423,7 @@ mod tests { request_rx: parking_lot::Mutex::new(None), notification_tx: broadcast::channel(16).0, router: router::SessionRouter::new(), + github_token_registry: Arc::new(github_token::GitHubTokenRegistry::new()), negotiated_protocol_version: OnceLock::new(), state: parking_lot::Mutex::new(ConnectionState::Connected), lifecycle_tx: broadcast::channel(16).0, diff --git a/rust/src/router.rs b/rust/src/router.rs index adc1923824..1dec9d16f3 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -87,6 +87,7 @@ impl SessionRouter { request_rx: &Mutex>>, llm_inference: Option>, github_telemetry: Option, + github_token_registry: Arc, ) { let mut started = self.started.lock(); if *started { @@ -181,6 +182,10 @@ impl SessionRouter { let sessions = self.sessions.clone(); tokio::spawn(async move { while let Some(request) = rx.recv().await { + if request.method == "gitHubToken.getToken" { + github_token_registry.dispatch(request).await; + continue; + } // Client-global `llmInference.*` requests carry no routable // session and are handled by the inference dispatcher. if request.method.starts_with("llmInference.") { diff --git a/rust/src/session.rs b/rust/src/session.rs index 7676c2ad70..d9b4a6bd94 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -190,6 +190,8 @@ pub struct Session { open_canvases: Arc>>, /// Broadcast channel for runtime event subscribers — see [`Session::subscribe`]. event_tx: tokio::sync::broadcast::Sender, + github_token_registration: + ParkingLotMutex>, } impl Session { @@ -584,6 +586,7 @@ impl Session { .await?; self.stop_event_loop().await; self.client.unregister_session(&self.id); + self.github_token_registration.lock().take(); Ok(()) } @@ -661,6 +664,7 @@ impl Drop for Session { // it here because Drop is sync. self.shutdown.cancel(); self.client.unregister_session(&self.id); + self.github_token_registration.lock().take(); } } @@ -933,6 +937,13 @@ impl Client { 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); + let github_token_registration = runtime + .github_token_provider + .take() + .map(|provider| self.register_github_token_provider(provider)); + wire.github_token_provider_registration_id = github_token_registration + .as_ref() + .map(|registration| registration.id().to_string()); let has_mcp_auth_handler = handlers.mcp_auth.is_some(); if self.inner.session_fs_configured && session_fs_provider.is_none() { return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into()); @@ -1090,6 +1101,7 @@ impl Client { capabilities, open_canvases, event_tx, + github_token_registration: ParkingLotMutex::new(github_token_registration), }; apply_mode_post_create_patch( &session, @@ -1100,6 +1112,11 @@ impl Client { opt_manage_schedule_enabled, ) .await?; + if let Some(registration) = session.github_token_registration.lock().as_ref() { + registration.claim(session.id.clone()); + } else { + self.retire_github_token_provider(&session.id); + } Ok(session) } @@ -1206,6 +1223,13 @@ impl Client { 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); + let github_token_registration = runtime + .github_token_provider + .take() + .map(|provider| self.register_github_token_provider(provider)); + wire.github_token_provider_registration_id = github_token_registration + .as_ref() + .map(|registration| registration.id().to_string()); let has_mcp_auth_handler = handlers.mcp_auth.is_some(); if self.inner.session_fs_configured && session_fs_provider.is_none() { return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into()); @@ -1350,6 +1374,7 @@ impl Client { capabilities, open_canvases, event_tx, + github_token_registration: ParkingLotMutex::new(github_token_registration), }; apply_mode_post_create_patch( &session, @@ -1360,6 +1385,11 @@ impl Client { opt_manage_schedule_enabled, ) .await?; + if let Some(registration) = session.github_token_registration.lock().as_ref() { + registration.claim(session.id.clone()); + } else { + self.retire_github_token_provider(&session.id); + } Ok(session) } } diff --git a/rust/src/types.rs b/rust/src/types.rs index 06c0fc4e79..df2d5c5769 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -24,6 +24,7 @@ use crate::generated::api_types::{CurrentToolMetadata, OpenCanvasInstance}; use crate::generated::session_events::ReasoningSummary; /// Context window tier for models that support tiered context windows. pub use crate::generated::session_events::{ContextTier, SessionLimitsConfig}; +use crate::github_token::GitHubTokenProvider; use crate::handler::{ AutoModeSwitchHandler, ElicitationHandler, ExitPlanModeHandler, McpAuthHandler, PermissionHandler, UserInputHandler, @@ -2113,6 +2114,12 @@ pub struct SessionConfig { /// the GitHub identity used for content exclusion, model routing, and /// quota checks for *this session*. pub github_token: Option, + /// Provider used to acquire rotating GitHub tokens for this session. + /// + /// Mutually exclusive with [`github_token`](Self::github_token). The callback + /// receives the effective host, optional assigned session ID, and acquisition + /// reason; its opaque registration ID is never exposed. + pub github_token_provider: Option>, /// Per-session remote behavior control: /// - `Off` — local only, no remote export (default) /// - `Export` — export session events to GitHub without @@ -2140,9 +2147,10 @@ pub struct SessionConfig { pub exp_assignments: Option, /// Opt-in: when `Some(true)`, the runtime self-fetches enterprise managed /// settings (bypass-permissions policy) at session bootstrap using the - /// session's [`github_token`](Self::github_token). Requires `github_token` - /// to be set; if omitted, the runtime is expected to reject session creation - /// (fail-closed). When `None`, behaves exactly as before. Set via + /// session's static [`github_token`](Self::github_token) or + /// [`github_token_provider`](Self::github_token_provider). Requires one of + /// those credentials; if both are omitted, the runtime is expected to reject + /// session creation (fail-closed). When `None`, behaves exactly as before. Set via /// [`with_enable_managed_settings`](Self::with_enable_managed_settings). pub enable_managed_settings: Option, /// Optional managed-settings layer injected at session bootstrap. Currently @@ -2294,6 +2302,10 @@ impl std::fmt::Debug for SessionConfig { "github_token", &self.github_token.as_ref().map(|_| ""), ) + .field( + "github_token_provider", + &self.github_token_provider.as_ref().map(|_| ""), + ) .field("remote_session", &self.remote_session) .field("cloud", &self.cloud) .field( @@ -2411,6 +2423,7 @@ impl Default for SessionConfig { working_directory: None, additional_directories: None, github_token: None, + github_token_provider: None, remote_session: None, cloud: None, include_sub_agent_streaming_events: None, @@ -2456,6 +2469,7 @@ pub(crate) struct SessionConfigRuntime { pub canvas_handler: Option>, pub session_fs_provider: Option>, pub bearer_token_providers: HashMap>, + pub github_token_provider: Option>, pub commands: Option>, } @@ -2475,6 +2489,12 @@ impl SessionConfig { mut self, session_id: Option, ) -> Result<(crate::wire::SessionCreateWire, SessionConfigRuntime), crate::Error> { + if self.github_token.is_some() && self.github_token_provider.is_some() { + return Err(crate::Error::with_message( + crate::ErrorKind::InvalidConfig, + "github_token and github_token_provider are mutually exclusive", + )); + } let permission_active = self.permission_handler.is_some() || self.permission_policy.is_some(); let request_user_input = self.user_input_handler.is_some(); @@ -2576,6 +2596,7 @@ impl SessionConfig { working_directory: self.working_directory, additional_directories: self.additional_directories, github_token: self.github_token, + github_token_provider_registration_id: None, remote_session: self.remote_session, cloud: self.cloud, include_sub_agent_streaming_events: self.include_sub_agent_streaming_events, @@ -2601,6 +2622,7 @@ impl SessionConfig { canvas_handler, session_fs_provider: self.session_fs_provider, bearer_token_providers, + github_token_provider: self.github_token_provider, commands: self.commands, }; @@ -3141,6 +3163,16 @@ impl SessionConfig { self } + /// Install a rotating GitHub token provider for this session. + /// + /// The provider must return a positive remaining lifetime in seconds when + /// its callback completes. Production GitHub tokens typically last eight + /// hours. This option is mutually exclusive with [`with_github_token`](Self::with_github_token). + pub fn with_github_token_provider(mut self, provider: Arc) -> Self { + self.github_token_provider = Some(provider); + self + } + /// Forward sub-agent streaming events to this connection. Defaults /// to true on the CLI when unset. pub fn with_include_sub_agent_streaming_events(mut self, include: bool) -> Self { @@ -3208,9 +3240,10 @@ impl SessionConfig { /// Opt the runtime into self-fetching enterprise managed settings /// (bypass-permissions policy) at session bootstrap using the session's - /// [`github_token`](Self::github_token). Requires `github_token` to be set; - /// if omitted, the runtime is expected to reject session creation - /// (fail-closed). + /// static [`github_token`](Self::github_token) or + /// [`github_token_provider`](Self::github_token_provider). Requires one of + /// those credentials; if both are omitted, the runtime is expected to reject + /// session creation (fail-closed). pub fn with_enable_managed_settings(mut self, enabled: bool) -> Self { self.enable_managed_settings = Some(enabled); self @@ -3399,6 +3432,9 @@ pub struct ResumeSessionConfig { /// Per-session GitHub token on resume. See /// [`SessionConfig::github_token`]. pub github_token: Option, + /// Rotating GitHub token provider on resume. See + /// [`SessionConfig::github_token_provider`]. + pub github_token_provider: Option>, /// Per-session remote behavior control on resume. See /// [`SessionConfig::remote_session`]. pub remote_session: Option, @@ -3560,6 +3596,10 @@ impl std::fmt::Debug for ResumeSessionConfig { "github_token", &self.github_token.as_ref().map(|_| ""), ) + .field( + "github_token_provider", + &self.github_token_provider.as_ref().map(|_| ""), + ) .field("remote_session", &self.remote_session) .field( "include_sub_agent_streaming_events", @@ -3619,6 +3659,12 @@ impl ResumeSessionConfig { pub(crate) fn into_wire( mut self, ) -> Result<(crate::wire::SessionResumeWire, SessionConfigRuntime), crate::Error> { + if self.github_token.is_some() && self.github_token_provider.is_some() { + return Err(crate::Error::with_message( + crate::ErrorKind::InvalidConfig, + "github_token and github_token_provider are mutually exclusive", + )); + } let permission_active = self.permission_handler.is_some() || self.permission_policy.is_some(); let request_user_input = self.user_input_handler.is_some(); @@ -3721,6 +3767,7 @@ impl ResumeSessionConfig { working_directory: self.working_directory, additional_directories: self.additional_directories, github_token: self.github_token, + github_token_provider_registration_id: None, remote_session: self.remote_session, include_sub_agent_streaming_events: self.include_sub_agent_streaming_events, enable_github_telemetry_forwarding: None, @@ -3747,6 +3794,7 @@ impl ResumeSessionConfig { canvas_handler, session_fs_provider: self.session_fs_provider, bearer_token_providers, + github_token_provider: self.github_token_provider, commands: self.commands, }; @@ -3818,6 +3866,7 @@ impl ResumeSessionConfig { working_directory: None, additional_directories: None, github_token: None, + github_token_provider: None, remote_session: None, include_sub_agent_streaming_events: None, commands: None, @@ -4358,6 +4407,16 @@ impl ResumeSessionConfig { self } + /// Install a rotating GitHub token provider for the resumed session. + /// + /// The provider must return a positive remaining lifetime in seconds when + /// its callback completes. Production GitHub tokens typically last eight + /// hours. Mutually exclusive with [`with_github_token`](Self::with_github_token). + pub fn with_github_token_provider(mut self, provider: Arc) -> Self { + self.github_token_provider = Some(provider); + self + } + /// Forward sub-agent streaming events to this connection on resume. pub fn with_include_sub_agent_streaming_events(mut self, include: bool) -> Self { self.include_sub_agent_streaming_events = Some(include); diff --git a/rust/src/wire.rs b/rust/src/wire.rs index 8c08b017b7..f7de33839c 100644 --- a/rust/src/wire.rs +++ b/rust/src/wire.rs @@ -170,6 +170,11 @@ pub(crate) struct SessionCreateWire { pub additional_directories: Option>, #[serde(rename = "gitHubToken", skip_serializing_if = "Option::is_none")] pub github_token: Option, + #[serde( + rename = "gitHubTokenProviderRegistrationId", + skip_serializing_if = "Option::is_none" + )] + pub github_token_provider_registration_id: Option, #[serde(skip_serializing_if = "Option::is_none")] pub remote_session: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -320,6 +325,11 @@ pub(crate) struct SessionResumeWire { pub additional_directories: Option>, #[serde(rename = "gitHubToken", skip_serializing_if = "Option::is_none")] pub github_token: Option, + #[serde( + rename = "gitHubTokenProviderRegistrationId", + skip_serializing_if = "Option::is_none" + )] + pub github_token_provider_registration_id: Option, #[serde(skip_serializing_if = "Option::is_none")] pub remote_session: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index e20d9d0885..83e549f06f 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -7,6 +7,9 @@ use std::time::Duration; use async_trait::async_trait; use github_copilot_sdk::canvas::{CanvasDeclaration, CanvasHandler, CanvasResult}; +use github_copilot_sdk::github_token::{ + GitHubToken, GitHubTokenProviderArgs, GitHubTokenProviderResult, GitHubTokenRequestReason, +}; use github_copilot_sdk::handler::{ ApproveAllHandler, AutoModeSwitchHandler, AutoModeSwitchResponse, ElicitationHandler, ExitPlanModeHandler, ExitPlanModeResult, McpAuthHandler, McpAuthRequest, McpAuthResult, @@ -254,6 +257,211 @@ where (session, server) } +#[tokio::test] +async fn github_token_provider_uses_global_registration_and_maps_results() { + let (client, server_read, server_write) = make_client(); + let mut server = FakeServer { + read: server_read, + write: server_write, + session_id: "github-token-session".to_string(), + }; + let (args_tx, mut args_rx) = tokio::sync::mpsc::unbounded_channel(); + let provider = Arc::new(move |args: GitHubTokenProviderArgs| { + let args_tx = args_tx.clone(); + async move { + args_tx.send(args.clone()).unwrap(); + match args.host.as_str() { + "github.com" => Ok(GitHubTokenProviderResult::Token(GitHubToken::new( + "secret-token", + 8 * 60 * 60, + ))), + "cancel.example" => Ok(GitHubTokenProviderResult::Cancelled), + _ => Err(github_copilot_sdk::Error::with_message( + ErrorKind::GitHubTokenProvider, + "credential service unavailable", + )), + } + } + }); + let create_handle = tokio::spawn({ + let client = client.clone(); + async move { + client + .create_session( + SessionConfig::default() + .with_session_id("github-token-session") + .with_github_token_provider(provider), + ) + .await + .unwrap() + } + }); + + let create_request = server.read_request().await; + assert_eq!(create_request["method"], "session.create"); + assert!(create_request["params"].get("gitHubToken").is_none()); + let registration_id = create_request["params"]["gitHubTokenProviderRegistrationId"] + .as_str() + .unwrap() + .to_string(); + + server + .send_request( + 900, + "gitHubToken.getToken", + serde_json::json!({ + "registrationId": registration_id, + "host": "github.com", + "sessionId": "github-token-session", + "reason": "initial" + }), + ) + .await; + let token_response = server.read_response().await; + assert_eq!(token_response["result"]["kind"], "token"); + assert_eq!(token_response["result"]["accessToken"], "secret-token"); + assert_eq!(token_response["result"]["expiresIn"], 8 * 60 * 60); + let args = args_rx.recv().await.unwrap(); + assert_eq!(args.host, "github.com"); + assert_eq!( + args.session_id.as_ref().map(SessionId::as_str), + Some("github-token-session") + ); + assert_eq!(args.reason, GitHubTokenRequestReason::Initial); + + server + .send_request( + 901, + "gitHubToken.getToken", + serde_json::json!({ + "registrationId": registration_id, + "host": "cancel.example", + "reason": "refresh" + }), + ) + .await; + let cancelled = server.read_response().await; + assert_eq!(cancelled["result"]["kind"], "cancelled"); + + server + .send_request( + 902, + "gitHubToken.getToken", + serde_json::json!({ + "registrationId": registration_id, + "host": "error.example", + "reason": "refresh" + }), + ) + .await; + let provider_error = server.read_response().await; + assert_eq!(provider_error["error"]["code"], -32603); + assert!( + provider_error["error"]["message"] + .as_str() + .unwrap() + .contains("credential service unavailable") + ); + + server + .respond( + &create_request, + serde_json::json!({"sessionId": "github-token-session"}), + ) + .await; + let session = timeout(TIMEOUT, create_handle).await.unwrap().unwrap(); + drop(session); + + server + .send_request( + 903, + "gitHubToken.getToken", + serde_json::json!({ + "registrationId": registration_id, + "host": "github.com", + "reason": "refresh" + }), + ) + .await; + let unknown = server.read_response().await; + assert_eq!(unknown["error"]["code"], -32603); + assert!( + unknown["error"]["message"] + .as_str() + .unwrap() + .contains("unknown GitHub token provider registration") + ); +} + +#[tokio::test] +async fn github_token_provider_is_mutually_exclusive_and_rolls_back_failed_create() { + let provider = Arc::new(|_args: GitHubTokenProviderArgs| async { + Ok(GitHubTokenProviderResult::Cancelled) + }); + let (client, server_read, server_write) = make_client(); + let result = client + .create_session( + SessionConfig::default() + .with_github_token("static") + .with_github_token_provider(provider.clone()), + ) + .await; + let Err(error) = result else { + panic!("mutually exclusive GitHub credentials must be rejected"); + }; + assert!(matches!(error.kind(), ErrorKind::InvalidConfig)); + + let mut server = FakeServer { + read: server_read, + write: server_write, + session_id: "failed-create".to_string(), + }; + let create_handle = tokio::spawn({ + let client = client.clone(); + async move { + client + .create_session( + SessionConfig::default() + .with_session_id("failed-create") + .with_github_token_provider(provider), + ) + .await + } + }); + let create_request = server.read_request().await; + let registration_id = create_request["params"]["gitHubTokenProviderRegistrationId"] + .as_str() + .unwrap() + .to_string(); + let response = serde_json::json!({ + "jsonrpc": "2.0", + "id": create_request["id"], + "error": {"code": -32603, "message": "create failed"} + }); + write_framed(&mut server.write, &serde_json::to_vec(&response).unwrap()).await; + assert!( + timeout(TIMEOUT, create_handle) + .await + .unwrap() + .unwrap() + .is_err() + ); + + server + .send_request( + 904, + "gitHubToken.getToken", + serde_json::json!({ + "registrationId": registration_id, + "host": "github.com", + "reason": "initial" + }), + ) + .await; + let unknown = server.read_response().await; + assert_eq!(unknown["error"]["code"], -32603); +} + fn rand_id() -> u64 { static COUNTER: AtomicUsize = AtomicUsize::new(0); COUNTER.fetch_add(1, Ordering::Relaxed) as u64 From 57ab421bb9896c3aeb2dd3a4ffb4aa0c9c794f2f Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Wed, 26 Aug 2026 18:03:13 +0200 Subject: [PATCH 2/4] Clarify GitHub token provider lifecycle Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: a447a8bc-8687-47ea-a545-4370979e8128 --- CHANGELOG.md | 2 ++ docs/auth/authenticate.md | 4 +++- dotnet/README.md | 2 ++ dotnet/src/Types.cs | 4 +++- go/README.md | 2 ++ go/github_token_provider.go | 4 +++- java/README.md | 5 +++++ .../java/com/github/copilot/rpc/GitHubTokenProvider.java | 4 +++- nodejs/README.md | 2 ++ nodejs/src/types.ts | 5 +++-- python/README.md | 3 +++ python/copilot/client.py | 5 +++-- rust/README.md | 5 +++++ rust/src/github_token.rs | 3 +++ 14 files changed, 42 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bc3b41c71c..d695ed6fb6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,8 @@ All six SDKs can now acquire short-lived GitHub credentials through a session-sc Token responses use the shared tagged token/cancelled shape and require `expiresIn`, expressed as the positive number of seconds remaining when the callback completes. See [github/copilot-agent-runtime#16381](https://github.com/github/copilot-agent-runtime/pull/16381) for the runtime credential-authority implementation. +Initial acquisition occurs during create or resume; cancellation, callback errors, and invalid credentials reject that operation instead of falling back to ambient authentication. Idle sessions refresh only before their next credential-consuming operation. + ### Feature: extensions can request sensitive environment variables Copilot CLI extensions can now ask for named sensitive environment variables when they join a session. `joinSession()` accepts a `requestedEnvironmentVariables` option listing the variable names the extension needs. The CLI shows a permission prompt naming the extension and the exact variables requested. On approval, only those variables reach that extension and their values are written into the extension process's `process.env` before `joinSession()` resolves. On denial, `joinSession()` rejects, the extension does not load, and its tools never reach the model. diff --git a/docs/auth/authenticate.md b/docs/auth/authenticate.md index 1fe3246547..a7582b5510 100644 --- a/docs/auth/authenticate.md +++ b/docs/auth/authenticate.md @@ -383,7 +383,9 @@ let session = client -Return the language's cancelled result when the user or credential broker cancels acquisition. Provider exceptions follow the SDK's normal callback error behavior. The runtime refreshes during async preflight when one hour or less remains; it does not use background timers, rejection-driven replay, 401/403 challenge propagation, or upscope for this callback. +The runtime performs the `initial` acquisition as part of session creation or resume. A cancelled acquisition, provider error, invalid response, or token without a stable account identity rejects the create or resume operation. The runtime does not fall back to ambient authentication. + +After the session is established, the runtime performs async preflight before each credential-consuming operation. It requests a `refresh` when the current token has one hour or less remaining. Idle sessions are not refreshed until their next credential-consuming operation. The runtime does not use background timers, rejection-driven replay, 401/403 challenge propagation, or upscope for this callback. ## Environment variables diff --git a/dotnet/README.md b/dotnet/README.md index ba935c190a..461ff0cf94 100644 --- a/dotnet/README.md +++ b/dotnet/README.md @@ -162,6 +162,8 @@ await using var session = await client.CreateSessionAsync(new SessionConfig }); ``` +Initial acquisition runs during session creation or resume. Cancellation, provider errors, and invalid token responses reject that operation instead of falling back to ambient authentication. Idle sessions refresh only before their next credential-consuming operation; there is no background refresh timer. + ##### `PingAsync(string? message = null): Task` Ping the server to check connectivity. diff --git a/dotnet/src/Types.cs b/dotnet/src/Types.cs index d038721a5e..6f33860e79 100644 --- a/dotnet/src/Types.cs +++ b/dotnet/src/Types.cs @@ -3652,7 +3652,9 @@ protected SessionConfigBase(SessionConfigBase? other) ///

/// Gets or sets a callback that acquires session-scoped GitHub tokens on - /// demand. This cannot be combined with . + /// demand. Initial cancellation, callback errors, and invalid token responses + /// reject session creation or resume instead of falling back to ambient + /// authentication. This cannot be combined with . /// [Experimental(Diagnostics.Experimental)] [JsonIgnore] diff --git a/go/README.md b/go/README.md index 80abbd19c8..ddd74b91aa 100644 --- a/go/README.md +++ b/go/README.md @@ -255,6 +255,8 @@ session, err := client.CreateSession(ctx, &copilot.SessionConfig{ }) ``` +Initial acquisition runs during session creation or resume. Cancellation, provider errors, and invalid token responses reject that operation instead of falling back to ambient authentication. Idle sessions refresh only before their next credential-consuming operation; there is no background refresh timer. + ### Session - `Send(ctx context.Context, options MessageOptions) (string, error)` - Send a message diff --git a/go/github_token_provider.go b/go/github_token_provider.go index fa57e24ca0..8f1233b7cd 100644 --- a/go/github_token_provider.go +++ b/go/github_token_provider.go @@ -76,7 +76,9 @@ func GitHubTokenCancelled() *GitHubTokenProviderResult { return &GitHubTokenProviderResult{Cancelled: true} } -// GitHubTokenProvider acquires session-scoped GitHub tokens on demand. +// GitHubTokenProvider acquires session-scoped GitHub tokens on demand. Initial +// cancellation, errors, and invalid token responses reject session creation or +// resume instead of falling back to ambient authentication. // // Experimental: GitHubTokenProvider may change or be removed. type GitHubTokenProvider func(args GitHubTokenProviderArgs) (*GitHubTokenProviderResult, error) diff --git a/java/README.md b/java/README.md index fd03aa5f6c..4c7fdce462 100644 --- a/java/README.md +++ b/java/README.md @@ -190,6 +190,11 @@ The remaining lifetime is required and must be positive when the callback completes; production GitHub tokens typically last eight hours. A static token and a provider are mutually exclusive. +Initial acquisition runs during session creation or resume. Cancellation, +provider errors, and invalid token responses reject that operation instead of +falling back to ambient authentication. Idle sessions refresh only before their +next credential-consuming operation; there is no background refresh timer. + ## Permission Handling `PermissionHandler.APPROVE_ALL` approves requests when managed settings are disabled. When `enableManagedSettings` is true, it completes exceptionally. Custom handlers can inspect `request.getManagedApprovalRequired()` for human-facing confirmation logic. diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java index 20cf11d1cf..daefc2f8d8 100644 --- a/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java +++ b/java/sdk/src/main/java/com/github/copilot/rpc/GitHubTokenProvider.java @@ -11,7 +11,9 @@ *

* Implementations return either a token with a positive remaining lifetime or * an explicit cancellation. Production GitHub tokens typically last eight - * hours. + * hours. Initial cancellation, callback errors, and invalid token responses + * reject session creation or resume instead of falling back to ambient + * authentication. * * @since 1.0.0 */ diff --git a/nodejs/README.md b/nodejs/README.md index 1e493b7b60..93f9c3fa6b 100644 --- a/nodejs/README.md +++ b/nodejs/README.md @@ -154,6 +154,8 @@ const session = await client.createSession({ }); ``` +Initial acquisition runs during session creation or resume. Cancellation, provider errors, and invalid token responses reject that operation instead of falling back to ambient authentication. Idle sessions refresh only before their next credential-consuming operation; there is no background refresh timer. + ##### `resumeSession(sessionId: string, config?: ResumeSessionConfig): Promise` Resume an existing session. Returns the session with `workspacePath` populated if infinite sessions were enabled. diff --git a/nodejs/src/types.ts b/nodejs/src/types.ts index 02c01fac6b..1fa376d869 100644 --- a/nodejs/src/types.ts +++ b/nodejs/src/types.ts @@ -58,8 +58,9 @@ export type GitHubTokenProviderResult = GitHubTokenAcquireResult; * * A token result must include `expiresIn`: the positive number of seconds of * remaining lifetime when the callback completes. Production GitHub tokens - * typically last eight hours. Return `{ kind: "cancelled" }` to cancel the - * acquisition; thrown errors are returned to the runtime as callback errors. + * typically last eight hours. Initial cancellation, callback errors, and + * invalid token responses reject session creation or resume instead of falling + * back to ambient authentication. */ export type GitHubTokenProvider = ( args: GitHubTokenProviderArgs diff --git a/python/README.md b/python/README.md index ada0130b9d..61608c16a0 100644 --- a/python/README.md +++ b/python/README.md @@ -297,6 +297,9 @@ async def provide_github_token(args): session = await client.create_session(github_token_provider=provide_github_token) ``` + +Initial acquisition runs during session creation or resume. Cancellation, provider errors, and invalid token responses reject that operation instead of falling back to ambient authentication. Idle sessions refresh only before their next credential-consuming operation; there is no background refresh timer. + - `available_tools` / `excluded_tools` / `default_agent.excluded_tools` / custom-agent `tools`: MCP tools registered from `mcp_servers` are exposed to the runtime as `-`. For `available_tools` and `excluded_tools`, prefer `ToolSet().add_mcp("-")` or the raw `mcp:-` form. For custom-agent `tools` and `default_agent.excluded_tools`, use `-` directly. **Session Lifecycle Methods:** diff --git a/python/copilot/client.py b/python/copilot/client.py index 96530c2b9b..c28c343d27 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -165,8 +165,9 @@ class GitHubTokenCancelledResult(TypedDict): Token results require ``expiresIn`` to be the positive number of seconds of remaining lifetime when the callback completes. Production GitHub tokens -typically last eight hours. Return ``{"kind": "cancelled"}`` to cancel; -exceptions raised by the callback are returned to the runtime unchanged. +typically last eight hours. Initial cancellation, callback errors, and invalid +token responses reject session creation or resume instead of falling back to +ambient authentication. """ # ============================================================================ diff --git a/rust/README.md b/rust/README.md index a3083d15d9..323d525d37 100644 --- a/rust/README.md +++ b/rust/README.md @@ -297,6 +297,11 @@ completes; production GitHub tokens typically last eight hours. Static `github_token` and a provider are mutually exclusive. The same provider API is available on `ResumeSessionConfig`. +Initial acquisition runs during session creation or resume. Cancellation, +provider errors, and invalid token responses reject that operation instead of +falling back to ambient authentication. Idle sessions refresh only before their +next credential-consuming operation; there is no background refresh timer. + ### Session Hooks Hooks intercept CLI behavior at lifecycle points — tool use, prompt submission, session start/end, and errors. Install a `SessionHooks` impl with [`SessionConfig::with_hooks`] — the SDK auto-enables `hooks` in `SessionConfig` when one is set. diff --git a/rust/src/github_token.rs b/rust/src/github_token.rs index 7016d0ca01..c5eaa63adf 100644 --- a/rust/src/github_token.rs +++ b/rust/src/github_token.rs @@ -101,6 +101,9 @@ impl std::fmt::Debug for GitHubTokenProviderResult { #[async_trait] pub trait GitHubTokenProvider: Send + Sync { /// Acquire a token or explicitly cancel the request. + /// + /// Initial cancellation, errors, and invalid token responses reject session + /// creation or resume instead of falling back to ambient authentication. async fn get_token( &self, args: GitHubTokenProviderArgs, From 4840175e9cf79310c2f15e7a6be9a596e5a43374 Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Wed, 26 Aug 2026 18:21:39 +0200 Subject: [PATCH 3/4] Fix GitHub token provider cleanup Release session-owned provider registrations after successful session deletion and treat empty static .NET tokens as configured for mutual-exclusion validation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: a447a8bc-8687-47ea-a545-4370979e8128 --- dotnet/src/Client.cs | 12 +++---- .../test/Unit/ClientSessionLifetimeTests.cs | 28 +++++++++++++++-- .../com/github/copilot/CopilotClient.java | 5 ++- .../com/github/copilot/CopilotSession.java | 14 ++++++--- .../com/github/copilot/CopilotClientTest.java | 31 +++++++++++++++++++ nodejs/src/client.ts | 3 +- nodejs/src/session.ts | 9 ++++-- nodejs/test/github-token-provider.test.ts | 3 ++ python/copilot/client.py | 5 +-- python/copilot/session.py | 10 ++++-- python/test_github_token_provider.py | 4 +++ rust/src/lib.rs | 1 + rust/tests/session_test.rs | 19 +++++++++++- 13 files changed, 119 insertions(+), 25 deletions(-) diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index b886a4a42b..e541717595 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -1701,7 +1701,10 @@ public async Task DeleteSessionAsync(string sessionId, CancellationToken cancell throw new InvalidOperationException($"Failed to delete session {sessionId}: {response.Error}"); } - RemoveSession(sessionId); + if (_sessions.TryRemove(sessionId, out var session)) + { + session.ReleaseGitHubTokenProviderRegistration(); + } } ///

@@ -2000,7 +2003,7 @@ await Rpc.SessionFs.SetProviderAsync( private static void ValidateGitHubTokenConfig(SessionConfigBase config) { - if (!string.IsNullOrEmpty(config.GitHubToken) && config.GitHubTokenProvider is not null) + if (config.GitHubToken is not null && config.GitHubTokenProvider is not null) { throw new ArgumentException( $"{nameof(SessionConfigBase.GitHubToken)} and {nameof(SessionConfigBase.GitHubTokenProvider)} cannot be used together.", @@ -2648,11 +2651,6 @@ private void RegisterSession(CopilotSession session) } } - private void RemoveSession(string sessionId) - { - _sessions.TryRemove(sessionId, out _); - } - /// /// Disposes the synchronously. /// diff --git a/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index 356ef605a8..24707a7527 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -20,13 +20,15 @@ public sealed class ClientSessionLifetimeTests { private sealed record RpcRequestRecord(string Method, JsonElement Params); - [Fact] - public async Task GitHubTokenProvider_Is_Mutually_Exclusive_With_Static_Token() + [Theory] + [InlineData("static")] + [InlineData("")] + public async Task GitHubTokenProvider_Is_Mutually_Exclusive_With_Static_Token(string staticToken) { await using var client = new CopilotClient(); var config = new SessionConfig { - GitHubToken = "static", + GitHubToken = staticToken, GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) }; @@ -35,6 +37,26 @@ public async Task GitHubTokenProvider_Is_Mutually_Exclusive_With_Static_Token() Assert.Contains("cannot be used together", error.Message); } + [Fact] + public async Task GitHubTokenProvider_Is_Released_When_Session_Is_Deleted() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + var session = await client.CreateSessionAsync(new SessionConfig + { + GitHubTokenProvider = _ => Task.FromResult(GitHubTokenProviderResult.Cancel()) + }); + var registrationId = Assert.Single(server.Requests, request => request.Method == "session.create") + .Params.GetProperty("gitHubTokenProviderRegistrationId").GetString(); + + await client.DeleteSessionAsync(session.SessionId); + + var error = await Assert.ThrowsAsync(() => + server.SendRequestAsync("gitHubToken.getToken", TokenRequest(registrationId))); + Assert.Contains("Unknown GitHub token provider registration ID", error.Message); + await session.DisposeAsync(); + } + [Fact] public async Task GitHubTokenProvider_Is_Serialized_And_Maps_Callbacks() { diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java index 639562224d..8a4f03fd83 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java @@ -1557,7 +1557,10 @@ public CompletableFuture deleteSession(String sessionId) { if (!response.success()) { throw new RuntimeException("Failed to delete session " + sessionId + ": " + response.error()); } - sessions.remove(sessionId); + CopilotSession session = sessions.remove(sessionId); + if (session != null) { + session.releaseGitHubTokenProviderRegistration(); + } })); } diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java index bf1f17d413..83ce49654e 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java @@ -257,6 +257,14 @@ void setGitHubTokenProviderRegistration(GitHubTokenProviderRegistry.Registration this.gitHubTokenProviderRegistration = registration; } + synchronized void releaseGitHubTokenProviderRegistration() { + GitHubTokenProviderRegistry.Registration registration = gitHubTokenProviderRegistration; + gitHubTokenProviderRegistration = null; + if (registration != null) { + registration.close(); + } + } + /** * Gets the unique identifier for this session. * @@ -2306,11 +2314,7 @@ public void close() { } timeoutScheduler.shutdownNow(); - GitHubTokenProviderRegistry.Registration tokenRegistration = gitHubTokenProviderRegistration; - gitHubTokenProviderRegistration = null; - if (tokenRegistration != null) { - tokenRegistration.close(); - } + releaseGitHubTokenProviderRegistration(); try { rpc.invoke("session.destroy", Map.of("sessionId", sessionId), Void.class).get(5, TimeUnit.SECONDS); diff --git a/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java b/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java index 067571df13..4cfd7f4c48 100644 --- a/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java +++ b/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java @@ -8,6 +8,8 @@ import org.junit.jupiter.api.Test; import com.github.copilot.rpc.CopilotClientOptions; +import com.github.copilot.rpc.DeleteSessionResponse; +import com.github.copilot.rpc.GitHubTokenProviderResult; import com.github.copilot.rpc.PermissionHandler; import com.github.copilot.rpc.PingResponse; import com.github.copilot.rpc.SessionConfig; @@ -111,6 +113,35 @@ void testForceStopAndExternalStopDoNotRequestRuntimeShutdown() throws Exception verify(externalRpc, never()).invoke(eq("runtime.shutdown"), any(), eq(Void.class)); } + @Test + @SuppressWarnings("unchecked") + void testDeleteSessionReleasesGitHubTokenProvider() throws Exception { + var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false)); + var rpc = mock(JsonRpcClient.class); + when(rpc.invoke(eq("session.delete"), any(), eq(DeleteSessionResponse.class))) + .thenReturn(CompletableFuture.completedFuture(new DeleteSessionResponse(true, null))); + when(rpc.invoke(eq("session.destroy"), any(), eq(Void.class))) + .thenReturn(CompletableFuture.completedFuture(null)); + setConnectionFuture(client, rpc, null); + + var registry = new GitHubTokenProviderRegistry(); + var registration = registry + .register(args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled())); + var session = new CopilotSession("delete-session", rpc); + session.setGitHubTokenProviderRegistration(registration); + Field sessionsField = CopilotClient.class.getDeclaredField("sessions"); + sessionsField.setAccessible(true); + ((Map) sessionsField.get(client)).put(session.getSessionId(), session); + + try { + client.deleteSession(session.getSessionId()).join(); + assertNull(registry.get(registration.id())); + } finally { + session.close(); + client.close(); + } + } + @Test void testClientConstruction() { var client = new CopilotClient(); diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts index 86f58e6c27..9bed7e9e0e 100644 --- a/nodejs/src/client.ts +++ b/nodejs/src/client.ts @@ -2280,8 +2280,9 @@ export class CopilotClient { throw new Error(`Failed to delete session ${sessionId}: ${error || "Unknown error"}`); } - // Remove from local sessions map if present + const session = this.sessions.get(sessionId); this.sessions.delete(sessionId); + session?._runOnDisconnected(); } /** diff --git a/nodejs/src/session.ts b/nodejs/src/session.ts index ca21b19471..d8b67133ff 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -808,8 +808,7 @@ export class CopilotSession { return; } this.disconnected = true; - this.onDisconnected?.(); - this.onDisconnected = undefined; + this._runOnDisconnected(); this.eventHandlers.clear(); this.typedEventHandlers.clear(); this.toolHandlers.clear(); @@ -830,6 +829,12 @@ export class CopilotSession { this.transformCallbacks?.clear(); } + /** @internal */ + _runOnDisconnected(): void { + this.onDisconnected?.(); + this.onDisconnected = undefined; + } + /** @internal */ _setOnDisconnected(callback: () => void): void { this.onDisconnected = callback; diff --git a/nodejs/test/github-token-provider.test.ts b/nodejs/test/github-token-provider.test.ts index caa8e1bbf4..2202f466a1 100644 --- a/nodejs/test/github-token-provider.test.ts +++ b/nodejs/test/github-token-provider.test.ts @@ -178,6 +178,7 @@ describe("session GitHub token providers", () => { const client = createMockClient(async (method, params) => { if (method === "session.create") return { sessionId: params.sessionId }; if (method === "session.destroy") return {}; + if (method === "session.delete") return { success: true }; throw new Error(`Unexpected method: ${method}`); }); const first = await client.createSession({ @@ -197,6 +198,8 @@ describe("session GitHub token providers", () => { await first.disconnect(); expect(registrations).toHaveLength(1); + await client.deleteSession("second"); + expect(registrations).toHaveLength(0); await client.forceStop(); expect(registrations).toHaveLength(0); }); diff --git a/python/copilot/client.py b/python/copilot/client.py index c28c343d27..96d57c4d47 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -3821,8 +3821,9 @@ async def delete_session(self, session_id: str) -> None: # Remove from local sessions map if present with self._sessions_lock: - if session_id in self._sessions: - del self._sessions[session_id] + session = self._sessions.pop(session_id, None) + if session is not None: + session._run_disconnect_callback() async def get_last_session_id(self) -> str | None: """ diff --git a/python/copilot/session.py b/python/copilot/session.py index 86d3d7d287..b3d8aa45ce 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -1607,6 +1607,12 @@ def _set_disconnect_callback(self, callback: Callable[[], None]) -> None: """Set the client-owned cleanup callback before the session becomes active.""" self._on_disconnect = callback + def _run_disconnect_callback(self) -> None: + callback = self._on_disconnect + self._on_disconnect = None + if callback is not None: + callback() + @property def rpc(self) -> SessionRpc: """Typed session-scoped RPC methods.""" @@ -2979,9 +2985,7 @@ async def disconnect(self) -> None: try: await self._client.request("session.destroy", {"sessionId": self.session_id}) finally: - if self._on_disconnect is not None: - self._on_disconnect() - self._on_disconnect = None + self._run_disconnect_callback() # Clear handlers even if the request fails. with self._event_handlers_lock: self._event_handlers.clear() diff --git a/python/test_github_token_provider.py b/python/test_github_token_provider.py index c4357c5745..5b203f3e34 100644 --- a/python/test_github_token_provider.py +++ b/python/test_github_token_provider.py @@ -33,6 +33,8 @@ async def request(self, method: str, params: dict[str, Any], **kwargs: Any) -> d return response if method == "session.destroy": return {} + if method == "session.delete": + return {"success": True} raise RuntimeError(f"Unexpected method: {method}") async def stop(self) -> None: @@ -199,6 +201,8 @@ async def test_failure_session_close_and_client_close_cleanup(self) -> None: await first.disconnect() assert len(client._github_token_providers) == 1 + await client.delete_session("second") + assert client._github_token_providers == {} await client.force_stop() assert client._github_token_providers == {} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 0517493c2a..334a395e9b 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2281,6 +2281,7 @@ impl Client { Some(serde_json::json!({ "sessionId": session_id })), ) .await?; + self.retire_github_token_provider(session_id); Ok(()) } diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index 83e549f06f..378d7120bf 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -370,7 +370,23 @@ async fn github_token_provider_uses_global_registration_and_maps_results() { ) .await; let session = timeout(TIMEOUT, create_handle).await.unwrap().unwrap(); - drop(session); + + let delete_handle = tokio::spawn({ + let client = client.clone(); + async move { + client + .delete_session(&SessionId::new("github-token-session")) + .await + } + }); + let delete_request = server.read_request().await; + assert_eq!(delete_request["method"], "session.delete"); + server.respond(&delete_request, serde_json::json!({})).await; + timeout(TIMEOUT, delete_handle) + .await + .unwrap() + .unwrap() + .unwrap(); server .send_request( @@ -391,6 +407,7 @@ async fn github_token_provider_uses_global_registration_and_maps_results() { .unwrap() .contains("unknown GitHub token provider registration") ); + drop(session); } #[tokio::test] From d531885e153f17077630b587dfc612551f7e6594 Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Wed, 26 Aug 2026 20:51:08 +0200 Subject: [PATCH 4/4] Fix Rust import formatting Order the new GitHub token re-export according to the nightly rustfmt configuration used by CI. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: a447a8bc-8687-47ea-a545-4370979e8128 --- rust/src/lib.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 334a395e9b..4a9f73ca4c 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -80,16 +80,16 @@ use std::sync::{Arc, OnceLock}; use std::time::{Duration, Instant}; use async_trait::async_trait; +pub use github_token::{ + GitHubToken, GitHubTokenProvider, GitHubTokenProviderArgs, GitHubTokenProviderResult, + GitHubTokenRequestReason, +}; /// Re-export of [`indexmap::IndexMap`], used for order-preserving maps in the /// public API (e.g. [`Tool::parameters`](types::Tool::parameters) and /// `SessionConfig::mcp_servers`) so serialized key order stays deterministic. pub use indexmap::IndexMap; // JSON-RPC wire types are internal transport details. // External callers interact via Client/Session methods, not raw RPC. -pub use github_token::{ - GitHubToken, GitHubTokenProvider, GitHubTokenProviderArgs, GitHubTokenProviderResult, - GitHubTokenRequestReason, -}; pub(crate) use jsonrpc::{ JsonRpcClient, JsonRpcError, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, error_codes, };