From 0062e965b972b248a47d3129812d0cba5762f61c Mon Sep 17 00:00:00 2001 From: Ashish Reddy Podduturi Date: Fri, 14 Aug 2026 15:30:17 -0700 Subject: [PATCH] feat: notify servers when credentials are updated Servers are told when credentials are deleted but not when they arrive, so anything that depends on credentials can only start when something else happens to ask for it. In the Q server that means the first service call waits for the chat webview to finish booting -- measured at ~0.66s after the credentials landed, on top of the call's own latency. Adds onCredentialsUpdated, mirroring the existing deletion path exactly: the handler is registered per server, and Auth routes the event through lspRouter so every server gets it rather than only the last one to subscribe. Optional on the interface rather than required. CredentialsProvider is implemented by test doubles across several consumers, and requiring a new member would fail their builds on upgrade. Optional also means a server built against this version still runs on an older runtime; it simply never receives the event. Handlers are invoked inside the client's credentials request, so a throwing handler is caught and logged instead of failing that request -- otherwise the client would be left believing the credentials never landed. Tests cover the fan-out to all servers, a registered handler, no handler registered, and a handler that throws. 78 passing. --- runtimes/runtimes/auth/auth.ts | 18 +++++++++++++ runtimes/runtimes/base-runtime.ts | 1 + .../runtimes/lsp/router/lspRouter.test.ts | 19 ++++++++++++++ runtimes/runtimes/lsp/router/lspRouter.ts | 4 +++ .../runtimes/lsp/router/lspServer.test.ts | 25 +++++++++++++++++++ runtimes/runtimes/lsp/router/lspServer.ts | 19 ++++++++++++++ runtimes/runtimes/standalone.ts | 1 + runtimes/server-interface/auth.ts | 14 +++++++++++ 8 files changed, 101 insertions(+) diff --git a/runtimes/runtimes/auth/auth.ts b/runtimes/runtimes/auth/auth.ts index 117e8ac4..d4498de6 100644 --- a/runtimes/runtimes/auth/auth.ts +++ b/runtimes/runtimes/auth/auth.ts @@ -42,6 +42,7 @@ export class Auth { private credentialsEncoding: CredentialsEncoding | undefined private lspRouter: LspRouter private credentialsDeleteHandler?: (type: CredentialsType) => void + private credentialsUpdateHandler?: (type: CredentialsType) => void constructor( private readonly connection: Connection, @@ -116,6 +117,9 @@ export class Auth { onCredentialsDeleted: (handler: (type: CredentialsType) => void) => { this.credentialsDeleteHandler = handler }, + onCredentialsUpdated: (handler: (type: CredentialsType) => void) => { + this.credentialsUpdateHandler = handler + }, } this.registerLspCredentialsUpdateHandlers() @@ -152,6 +156,9 @@ export class Auth { onCredentialsDeleted: (handler: (type: CredentialsType) => void) => { this.credentialsDeleteHandler = handler }, + onCredentialsUpdated: (handler: (type: CredentialsType) => void) => { + this.credentialsUpdateHandler = handler + }, } } @@ -239,13 +246,24 @@ export class Auth { this.iamCredentials = creds as IamCredentials // Prevent modifying credentials by implementors Object.freeze(this.iamCredentials) + this.notifyCredentialsUpdate('iam') } else { this.bearerCredentials = creds as BearerCredentials Object.freeze(this.bearerCredentials) + this.notifyCredentialsUpdate('bearer') } } } + /** + * Announces stored credentials so servers can start dependent work immediately instead of waiting + * for whatever happens to ask for a service first. Fired after the credentials are readable, and + * routed to every server the same way deletion is, so more than one server can react. + */ + private notifyCredentialsUpdate(type: CredentialsType) { + this.lspRouter.onCredentialsUpdate(type) + } + private async decodeCredentialsRequestToken(request: UpdateCredentialsParams): Promise { this.connection.console.info('Runtime: Decoding encrypted credentials token') if (!this.key) { diff --git a/runtimes/runtimes/base-runtime.ts b/runtimes/runtimes/base-runtime.ts index 1810cc6c..112ca653 100644 --- a/runtimes/runtimes/base-runtime.ts +++ b/runtimes/runtimes/base-runtime.ts @@ -325,6 +325,7 @@ export const baseRuntime = (connections: { reader: MessageReader; writer: Messag const sdkInitializator: SDKInitializator = (Ctor: SDKClientConstructorV3, current_config: P): T => new Ctor({ ...current_config }) credentialsProvider.onCredentialsDeleted = lspServer.setCredentialsDeleteHandler + credentialsProvider.onCredentialsUpdated = lspServer.setCredentialsUpdateHandler const agent = newAgent() diff --git a/runtimes/runtimes/lsp/router/lspRouter.test.ts b/runtimes/runtimes/lsp/router/lspRouter.test.ts index 8d382c60..b9268e38 100644 --- a/runtimes/runtimes/lsp/router/lspRouter.test.ts +++ b/runtimes/runtimes/lsp/router/lspRouter.test.ts @@ -415,6 +415,22 @@ describe('LspRouter', () => { }) }) + describe('onCredentialsUpdate', () => { + it('should send notifyCredentialsUpdate to all servers', () => { + const params: CredentialsType = 'bearer' + + const spy1 = sandbox.spy() + const spy2 = sandbox.spy() + const server1 = newServer({ credentialsUpdateHandler: spy1 }) + const server2 = newServer({ credentialsUpdateHandler: spy2 }) + + lspRouter.servers = [server1, server2] + lspRouter.onCredentialsUpdate(params) + assert(spy1.calledWith(params)) + assert(spy2.calledWith(params)) + }) + }) + describe('handleGetConfigurationFromServer', () => { it(`should return the result from initializeResult cache for key: ${SERVER_CAPABILITES_CONFIGURATION_SECTION}`, async () => { const initHandler1 = () => { @@ -1000,6 +1016,7 @@ describe('LspRouter', () => { initializedHandler, updateConfigurationHandler, credentialsDeleteHandler, + credentialsUpdateHandler, didChangeWorkspaceFoldersHandler, didCreateFilesHandler, didDeleteFilesHandler, @@ -1014,6 +1031,7 @@ describe('LspRouter', () => { initializedHandler?: any updateConfigurationHandler?: any credentialsDeleteHandler?: any + credentialsUpdateHandler?: any didChangeWorkspaceFoldersHandler?: any didCreateFilesHandler?: any didDeleteFilesHandler?: any @@ -1028,6 +1046,7 @@ describe('LspRouter', () => { server.setInitializedHandler(initializedHandler) server.setUpdateConfigurationHandler(updateConfigurationHandler) server.setCredentialsDeleteHandler(credentialsDeleteHandler) + server.setCredentialsUpdateHandler(credentialsUpdateHandler) server.setDidChangeWorkspaceFoldersHandler(didChangeWorkspaceFoldersHandler) server.setDidCreateFilesHandler(didCreateFilesHandler) server.setDidDeleteFilesHandler(didDeleteFilesHandler) diff --git a/runtimes/runtimes/lsp/router/lspRouter.ts b/runtimes/runtimes/lsp/router/lspRouter.ts index a1641a4e..2be135d4 100644 --- a/runtimes/runtimes/lsp/router/lspRouter.ts +++ b/runtimes/runtimes/lsp/router/lspRouter.ts @@ -221,6 +221,10 @@ ${JSON.stringify({ ...result.capabilities, ...result.awsServerCapabilities })}` this.routeNotificationToAllServers((server, type) => server.notifyCredentialsDeletion(type), type) } + onCredentialsUpdate = (type: CredentialsType): void => { + this.routeNotificationToAllServers((server, type) => server.notifyCredentialsUpdate(type), type) + } + onInitialized = (params: InitializedParams): void => { const workspaceCapabilities = this.clientInitializeParams?.capabilities.workspace if (workspaceCapabilities?.didChangeConfiguration?.dynamicRegistration) { diff --git a/runtimes/runtimes/lsp/router/lspServer.test.ts b/runtimes/runtimes/lsp/router/lspServer.test.ts index f128d557..0f24e3eb 100644 --- a/runtimes/runtimes/lsp/router/lspServer.test.ts +++ b/runtimes/runtimes/lsp/router/lspServer.test.ts @@ -343,5 +343,30 @@ describe('LspServer', () => { assert(credentialsDeleteHandler.calledOnce) assert(credentialsDeleteHandler.calledWith(params)) }) + + it('should handle credential update notifications', async () => { + const credentialsUpdateHandler = sandbox.stub() + const params = 'bearer' + + lspServer.setCredentialsUpdateHandler(credentialsUpdateHandler) + lspServer.notifyCredentialsUpdate(params) + + assert(credentialsUpdateHandler.calledOnce) + assert(credentialsUpdateHandler.calledWith(params)) + }) + + it('should not throw when no credential update handler is registered', async () => { + assert.doesNotThrow(() => lspServer.notifyCredentialsUpdate('bearer')) + }) + + it('should swallow errors thrown by a credential update handler', async () => { + // This runs inside the client's credentials request. A server that throws must not fail + // that request, or the client is left believing the credentials never landed. + lspServer.setCredentialsUpdateHandler(() => { + throw new Error('handler exploded') + }) + + assert.doesNotThrow(() => lspServer.notifyCredentialsUpdate('bearer')) + }) }) }) diff --git a/runtimes/runtimes/lsp/router/lspServer.ts b/runtimes/runtimes/lsp/router/lspServer.ts index 11a10ad5..6ae6ccc9 100644 --- a/runtimes/runtimes/lsp/router/lspServer.ts +++ b/runtimes/runtimes/lsp/router/lspServer.ts @@ -42,6 +42,7 @@ export class LspServer { private initializedHandler?: NotificationHandler private updateConfigurationHandler?: RequestHandler private credentialsDeleteHandler?: (type: CredentialsType) => void + private credentialsUpdateHandler?: (type: CredentialsType) => void private clientSupportsNotifications?: boolean private initializeResult?: PartialInitializeResult @@ -81,6 +82,10 @@ export class LspServer { this.credentialsDeleteHandler = handler } + setCredentialsUpdateHandler = (handler: (type: CredentialsType) => void): void => { + this.credentialsUpdateHandler = handler + } + public setInitializedHandler = (handler: NotificationHandler): void => { this.initializedHandler = handler } @@ -246,4 +251,18 @@ export class LspServer { this.credentialsDeleteHandler(type) } } + + notifyCredentialsUpdate = (type: CredentialsType): void => { + if (!this.credentialsUpdateHandler) { + return + } + + try { + this.credentialsUpdateHandler(type) + } catch (error) { + // This runs inside the client's credentials request. A server that throws here must not + // fail that request, or the client is left believing the credentials never landed. + this.logger.log(`Credentials update handler threw, ignoring: ${(error as Error)?.message}`) + } + } } diff --git a/runtimes/runtimes/standalone.ts b/runtimes/runtimes/standalone.ts index 5286c8f2..37be8f2c 100644 --- a/runtimes/runtimes/standalone.ts +++ b/runtimes/runtimes/standalone.ts @@ -488,6 +488,7 @@ export const standalone = (props: RuntimeProps) => { } credentialsProvider.onCredentialsDeleted = lspServer.setCredentialsDeleteHandler + credentialsProvider.onCredentialsUpdated = lspServer.setCredentialsUpdateHandler return s({ chat, diff --git a/runtimes/server-interface/auth.ts b/runtimes/server-interface/auth.ts index 11b2a7a4..5c46493e 100644 --- a/runtimes/server-interface/auth.ts +++ b/runtimes/server-interface/auth.ts @@ -13,4 +13,18 @@ export interface CredentialsProvider { getConnectionMetadata: () => ConnectionMetadata | undefined getConnectionType: () => SsoConnectionType onCredentialsDeleted: (handler: (type: CredentialsType) => void) => void + /** + * Notifies the server that credentials of the given type have been stored, so work that depends on + * them can start immediately rather than waiting for the first consumer to ask. + * + * Fired after the credentials are available from {@link getCredentials}, so a handler can read them + * synchronously. Handlers must not throw; a throwing handler is logged and ignored so it cannot + * fail the client's credentials request. + * + * Optional so that adding it is not a breaking change: this interface is implemented by test + * doubles across several consumers, and requiring a new member would fail their builds on upgrade. + * Servers should call it with optional chaining, which also means a server built against this + * version still runs on an older runtime -- it simply never gets the event. + */ + onCredentialsUpdated?: (handler: (type: CredentialsType) => void) => void }