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 }