feat: notify servers when credentials are updated

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<T>(request: UpdateCredentialsParams): Promise<T> {
         this.connection.console.info('Runtime: Decoding encrypted credentials token')
         if (!this.key) {

runtimes/runtimes/base-runtime.ts

@@ -325,6 +325,7 @@ export const baseRuntime = (connections: { reader: MessageReader; writer: Messag
         const sdkInitializator: SDKInitializator = <T, P>(Ctor: SDKClientConstructorV3<T, P>, current_config: P): T =>
             new Ctor({ ...current_config })
         credentialsProvider.onCredentialsDeleted = lspServer.setCredentialsDeleteHandler
+        credentialsProvider.onCredentialsUpdated = lspServer.setCredentialsUpdateHandler
 
         const agent = newAgent()
 

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)

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) {

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'))
+        })
     })
 })

runtimes/runtimes/lsp/router/lspServer.ts

@@ -42,6 +42,7 @@ export class LspServer {
     private initializedHandler?: NotificationHandler<InitializedParams>
     private updateConfigurationHandler?: RequestHandler<UpdateConfigurationParams, void, void>
     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<InitializedParams>): 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}`)
+        }
+    }
 }

runtimes/runtimes/standalone.ts

@@ -488,6 +488,7 @@ export const standalone = (props: RuntimeProps) => {
             }
 
             credentialsProvider.onCredentialsDeleted = lspServer.setCredentialsDeleteHandler
+            credentialsProvider.onCredentialsUpdated = lspServer.setCredentialsUpdateHandler
 
             return s({
                 chat,

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
 }