Skip to content

Commit 8b8ce9c

Browse files
SteveSandersonMSCopilot
SteveSandersonMS
and
Copilot
committed
Add Java BYOK get_bearer_token glue + e2e
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent e92b4b5 commit 8b8ce9c

8 files changed

Lines changed: 576 additions & 0 deletions

File tree

‎java/src/main/java/com/github/copilot/CopilotSession.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@
7575
import com.github.copilot.rpc.ExitPlanModeRequest;
7676
import com.github.copilot.rpc.ExitPlanModeResult;
7777
import com.github.copilot.rpc.ElicitationSchema;
78+
import com.github.copilot.rpc.GetBearerToken;
7879
import com.github.copilot.rpc.GetMessagesResponse;
7980
import com.github.copilot.rpc.HookInvocation;
8081
import com.github.copilot.rpc.InputOptions;
@@ -169,6 +170,7 @@ public final class CopilotSession implements AutoCloseable {
169170
private final Set<Consumer<SessionEvent>> eventHandlers = ConcurrentHashMap.newKeySet();
170171
private final Map<String, ToolDefinition> toolHandlers = new ConcurrentHashMap<>();
171172
private final Map<String, CommandHandler> commandHandlers = new ConcurrentHashMap<>();
173+
private final Map<String, GetBearerToken> bearerTokenProviders = new ConcurrentHashMap<>();
172174
private final AtomicReference<PermissionHandler> permissionHandler = new AtomicReference<>();
173175
private final AtomicReference<UserInputHandler> userInputHandler = new AtomicReference<>();
174176
private final AtomicReference<ElicitationHandler> elicitationHandler = new AtomicReference<>();
@@ -1348,6 +1350,33 @@ void registerElicitationHandler(ElicitationHandler handler) {
13481350
elicitationHandler.set(handler);
13491351
}
13501352

1353+
/**
1354+
* Registers bearer-token provider callbacks for this session.
1355+
* <p>
1356+
* Called internally when creating or resuming a session with BYOK providers
1357+
* that use managed-identity token callbacks.
1358+
*
1359+
* @param providers
1360+
* the callbacks keyed by provider name
1361+
*/
1362+
void registerBearerTokenProviders(Map<String, GetBearerToken> providers) {
1363+
bearerTokenProviders.clear();
1364+
if (providers != null) {
1365+
bearerTokenProviders.putAll(providers);
1366+
}
1367+
}
1368+
1369+
/**
1370+
* Gets the bearer-token provider callback for the given provider name.
1371+
*
1372+
* @param providerName
1373+
* the provider name
1374+
* @return the registered callback, or {@code null} if none is registered
1375+
*/
1376+
GetBearerToken getBearerTokenProvider(String providerName) {
1377+
return bearerTokenProviders.get(providerName);
1378+
}
1379+
13511380
/**
13521381
* Registers an exit-plan-mode handler for this session.
13531382
* <p>

‎java/src/main/java/com/github/copilot/RpcHandlerDispatcher.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import com.github.copilot.generated.SessionEvent;
2020
import com.github.copilot.rpc.AutoModeSwitchRequest;
2121
import com.github.copilot.rpc.ExitPlanModeRequest;
22+
import com.github.copilot.rpc.GetBearerToken;
23+
import com.github.copilot.rpc.ProviderTokenArgs;
2224
import com.github.copilot.rpc.PermissionRequestResult;
2325
import com.github.copilot.rpc.PermissionRequestResultKind;
2426
import com.github.copilot.rpc.SessionLifecycleEvent;
@@ -88,6 +90,8 @@ void registerHandlers(JsonRpcClient rpc) {
8890
rpc.registerMethodHandler("hooks.invoke", (requestId, params) -> handleHooksInvoke(rpc, requestId, params));
8991
rpc.registerMethodHandler("systemMessage.transform",
9092
(requestId, params) -> handleSystemMessageTransform(rpc, requestId, params));
93+
rpc.registerMethodHandler("providerToken.getToken",
94+
(requestId, params) -> handleProviderTokenGetToken(rpc, requestId, params));
9195
}
9296

9397
private void handleSessionEvent(JsonNode params) {
@@ -300,6 +304,68 @@ private void handleUserInputRequest(JsonRpcClient rpc, String requestId, JsonNod
300304
});
301305
}
302306

307+
private void handleProviderTokenGetToken(JsonRpcClient rpc, String requestId, JsonNode params) {
308+
LOG.fine("Received providerToken.getToken: " + params);
309+
runAsync(() -> {
310+
final long requestIdLong = parseRequestId(requestId, "providerToken.getToken");
311+
if (requestIdLong == -1) {
312+
return;
313+
}
314+
try {
315+
String sessionId = params.get("sessionId").asText();
316+
String providerName = params.get("providerName").asText();
317+
318+
CopilotSession session = sessions.get(sessionId);
319+
if (session == null) {
320+
rpc.sendErrorResponse(requestIdLong, -32602, "Unknown session " + sessionId);
321+
return;
322+
}
323+
324+
GetBearerToken provider = session.getBearerTokenProvider(providerName);
325+
if (provider == null) {
326+
rpc.sendErrorResponse(requestIdLong, -32603,
327+
"No bearer-token provider registered for provider " + providerName);
328+
return;
329+
}
330+
331+
CompletableFuture<String> tokenFuture = provider.getToken(new ProviderTokenArgs(providerName));
332+
if (tokenFuture == null) {
333+
rpc.sendErrorResponse(requestIdLong, -32603,
334+
"Bearer-token provider returned null future for provider " + providerName);
335+
return;
336+
}
337+
338+
tokenFuture.thenAccept(token -> {
339+
try {
340+
if (token == null) {
341+
rpc.sendErrorResponse(requestIdLong, -32603,
342+
"Bearer-token provider returned null token for provider " + providerName);
343+
return;
344+
}
345+
rpc.sendResponse(requestIdLong, Map.of("token", token));
346+
} catch (IOException e) {
347+
LOG.log(Level.SEVERE, "Error sending provider token response", e);
348+
}
349+
}).exceptionally(ex -> {
350+
LOG.log(Level.WARNING, "Bearer-token provider exception", ex);
351+
try {
352+
rpc.sendErrorResponse(requestIdLong, -32603, "Bearer-token provider error: " + ex.getMessage());
353+
} catch (IOException e) {
354+
LOG.log(Level.SEVERE, "Error sending provider token error", e);
355+
}
356+
return null;
357+
});
358+
} catch (Exception e) {
359+
LOG.log(Level.SEVERE, "Error handling providerToken.getToken", e);
360+
try {
361+
rpc.sendErrorResponse(requestIdLong, -32603, "Provider token handler error: " + e.getMessage());
362+
} catch (IOException ioException) {
363+
LOG.log(Level.SEVERE, "Error sending provider token handler error", ioException);
364+
}
365+
}
366+
});
367+
}
368+
303369
private void handleExitPlanModeRequest(JsonRpcClient rpc, String requestId, JsonNode params) {
304370
runAsync(() -> {
305371
final long requestIdLong = parseRequestId(requestId, "exitPlanMode.request");

‎java/src/main/java/com/github/copilot/SessionRequestBuilder.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,15 @@
55
package com.github.copilot;
66

77
import java.util.HashMap;
8+
import java.util.List;
89
import java.util.Map;
910
import java.util.concurrent.CompletableFuture;
1011
import java.util.function.Function;
1112

1213
import com.github.copilot.rpc.CreateSessionRequest;
14+
import com.github.copilot.rpc.ProviderConfig;
15+
import com.github.copilot.rpc.NamedProviderConfig;
16+
import com.github.copilot.rpc.GetBearerToken;
1317
import com.github.copilot.rpc.CommandWireDefinition;
1418
import com.github.copilot.rpc.ResumeSessionConfig;
1519
import com.github.copilot.rpc.ResumeSessionRequest;
@@ -329,6 +333,11 @@ static void configureSession(CopilotSession session, SessionConfig config) {
329333
if (config.getOnElicitationRequest() != null) {
330334
session.registerElicitationHandler(config.getOnElicitationRequest());
331335
}
336+
Map<String, GetBearerToken> bearerTokenProviders = collectBearerTokenProviders(config.getProvider(),
337+
config.getProviders());
338+
if (!bearerTokenProviders.isEmpty()) {
339+
session.registerBearerTokenProviders(bearerTokenProviders);
340+
}
332341
if (config.getOnExitPlanMode() != null) {
333342
session.registerExitPlanModeHandler(config.getOnExitPlanMode());
334343
}
@@ -371,6 +380,11 @@ static void configureSession(CopilotSession session, ResumeSessionConfig config)
371380
if (config.getOnElicitationRequest() != null) {
372381
session.registerElicitationHandler(config.getOnElicitationRequest());
373382
}
383+
Map<String, GetBearerToken> bearerTokenProviders = collectBearerTokenProviders(config.getProvider(),
384+
config.getProviders());
385+
if (!bearerTokenProviders.isEmpty()) {
386+
session.registerBearerTokenProviders(bearerTokenProviders);
387+
}
374388
if (config.getOnExitPlanMode() != null) {
375389
session.registerExitPlanModeHandler(config.getOnExitPlanMode());
376390
}
@@ -381,4 +395,21 @@ static void configureSession(CopilotSession session, ResumeSessionConfig config)
381395
session.on(config.getOnEvent());
382396
}
383397
}
398+
399+
private static Map<String, GetBearerToken> collectBearerTokenProviders(ProviderConfig provider,
400+
List<NamedProviderConfig> providers) {
401+
Map<String, GetBearerToken> bearerTokenProviders = new HashMap<>();
402+
if (provider != null && provider.getGetBearerToken() != null) {
403+
bearerTokenProviders.put("default", provider.getGetBearerToken());
404+
}
405+
if (providers != null) {
406+
for (NamedProviderConfig namedProvider : providers) {
407+
if (namedProvider != null && namedProvider.getName() != null
408+
&& namedProvider.getGetBearerToken() != null) {
409+
bearerTokenProviders.put(namedProvider.getName(), namedProvider.getGetBearerToken());
410+
}
411+
}
412+
}
413+
return bearerTokenProviders;
414+
}
384415
}

‎java/src/main/java/com/github/copilot/rpc/GetBearerToken.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
*--------------------------------------------------------------------------------------------*/
4+
5+
package com.github.copilot.rpc;
6+
7+
import java.util.concurrent.CompletableFuture;
8+
9+
import com.github.copilot.CopilotExperimental;
10+
11+
/**
12+
* Functional interface for supplying per-provider bearer tokens for BYOK
13+
* provider requests.
14+
* <p>
15+
* The callback returns the raw token without a {@code Bearer } prefix. The SDK
16+
* keeps this callback client-side and the runtime requests a token via the
17+
* session-scoped {@code providerToken.getToken} RPC before each outbound model
18+
* request.
19+
* <p>
20+
* <strong>Experimental.</strong> This managed-identity surface may change or be
21+
* removed in future SDK or CLI releases.
22+
*
23+
* @see ProviderConfig#setGetBearerToken(GetBearerToken)
24+
* @see NamedProviderConfig#setGetBearerToken(GetBearerToken)
25+
* @since 1.0.0
26+
*/
27+
@CopilotExperimental
28+
@FunctionalInterface
29+
public interface GetBearerToken {
30+
31+
/**
32+
* Gets a bearer token for the provider identified by {@code args}.
33+
*
34+
* @param args
35+
* the provider token request arguments
36+
* @return a future that completes with the raw token, without a {@code Bearer }
37+
* prefix
38+
*/
39+
CompletableFuture<String> getToken(ProviderTokenArgs args);
40+
}

‎java/src/main/java/com/github/copilot/rpc/NamedProviderConfig.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.util.Collections;
88
import java.util.Map;
99

10+
import com.fasterxml.jackson.annotation.JsonIgnore;
1011
import com.fasterxml.jackson.annotation.JsonInclude;
1112
import com.fasterxml.jackson.annotation.JsonProperty;
1213

@@ -59,6 +60,9 @@ public class NamedProviderConfig {
5960
@JsonProperty("bearerToken")
6061
private String bearerToken;
6162

63+
@JsonIgnore
64+
private GetBearerToken getBearerToken;
65+
6266
@JsonProperty("azure")
6367
private AzureOptions azure;
6468

@@ -212,6 +216,39 @@ public NamedProviderConfig setBearerToken(String bearerToken) {
212216
return this;
213217
}
214218

219+
/**
220+
* Gets the bearer-token provider callback.
221+
*
222+
* @return the bearer-token provider callback, or {@code null} if not set
223+
*/
224+
public GetBearerToken getGetBearerToken() {
225+
return getBearerToken;
226+
}
227+
228+
/**
229+
* Sets a callback that supplies bearer tokens for outbound provider requests.
230+
* <p>
231+
* <strong>Experimental.</strong> The callback stays SDK-side and is not
232+
* serialized. Instead, the runtime receives a {@code hasBearerTokenProvider}
233+
* flag and calls back over the session-scoped {@code providerToken.getToken}
234+
* RPC before each model request. Return the raw token without a {@code Bearer }
235+
* prefix.
236+
*
237+
* @param getBearerToken
238+
* the bearer-token provider callback
239+
* @return this config for method chaining
240+
*/
241+
public NamedProviderConfig setGetBearerToken(GetBearerToken getBearerToken) {
242+
this.getBearerToken = getBearerToken;
243+
return this;
244+
}
245+
246+
@JsonProperty("hasBearerTokenProvider")
247+
@JsonInclude(JsonInclude.Include.NON_NULL)
248+
Boolean hasBearerTokenProviderWireFlag() {
249+
return getBearerToken != null ? Boolean.TRUE : null;
250+
}
251+
215252
/**
216253
* Gets the Azure-specific options.
217254
*

‎java/src/main/java/com/github/copilot/rpc/ProviderConfig.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ public class ProviderConfig {
5656
@JsonProperty("bearerToken")
5757
private String bearerToken;
5858

59+
@JsonIgnore
60+
private GetBearerToken getBearerToken;
61+
5962
@JsonProperty("azure")
6063
private AzureOptions azure;
6164

@@ -222,6 +225,39 @@ public ProviderConfig setBearerToken(String bearerToken) {
222225
return this;
223226
}
224227

228+
/**
229+
* Gets the bearer-token provider callback.
230+
*
231+
* @return the bearer-token provider callback, or {@code null} if not set
232+
*/
233+
public GetBearerToken getGetBearerToken() {
234+
return getBearerToken;
235+
}
236+
237+
/**
238+
* Sets a callback that supplies bearer tokens for outbound provider requests.
239+
* <p>
240+
* <strong>Experimental.</strong> The callback stays SDK-side and is not
241+
* serialized. Instead, the runtime receives a {@code hasBearerTokenProvider}
242+
* flag and calls back over the session-scoped {@code providerToken.getToken}
243+
* RPC before each model request. Return the raw token without a {@code Bearer }
244+
* prefix.
245+
*
246+
* @param getBearerToken
247+
* the bearer-token provider callback
248+
* @return this config for method chaining
249+
*/
250+
public ProviderConfig setGetBearerToken(GetBearerToken getBearerToken) {
251+
this.getBearerToken = getBearerToken;
252+
return this;
253+
}
254+
255+
@JsonProperty("hasBearerTokenProvider")
256+
@JsonInclude(JsonInclude.Include.NON_NULL)
257+
Boolean hasBearerTokenProviderWireFlag() {
258+
return getBearerToken != null ? Boolean.TRUE : null;
259+
}
260+
225261
/**
226262
* Gets the Azure-specific options.
227263
*

0 commit comments

Comments
 (0)