Merge master into feature/v2-to-v3-migration

Author: aws-toolkit-automation

packages/core/src/auth/deprecated/loginManager.ts

@@ -30,7 +30,11 @@ import { isAutomation } from '../../shared/vscode/env'
 import { Credentials } from '@aws-sdk/types'
 import { ToolkitError } from '../../shared/errors'
 import * as localizedText from '../../shared/localizedText'
-import { DefaultStsClient, type GetCallerIdentityResponse } from '../../shared/clients/stsClient'
+import {
+    DefaultStsClient,
+    type GetCallerIdentityResponse,
+    type GetCallerIdentityResponseWithHeaders,
+} from '../../shared/clients/stsClient'
 import { findAsync } from '../../shared/utilities/collectionUtils'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { withTelemetryContext } from '../../shared/telemetry/util'
@@ -135,9 +139,11 @@ export class LoginManager {
         return accountId
     }
 
-    private async detectExternalConnection(callerIdentity: GetCallerIdentityResponse): Promise<void> {
-        // @ts-ignore
-        const headers = callerIdentity.$response?.httpResponse?.headers
+    private async detectExternalConnection(
+        callerIdentity: GetCallerIdentityResponse | GetCallerIdentityResponseWithHeaders
+    ): Promise<void> {
+        // SDK v3: Headers are captured via middleware and attached as $httpHeaders
+        const headers = (callerIdentity as GetCallerIdentityResponseWithHeaders).$httpHeaders
         if (headers !== undefined && localStackConnectionHeader in headers) {
             await globals.globalState.update('aws.toolkit.externalConnection', localStackConnectionString)
             telemetry.auth_localstackEndpoint.emit({ source: 'validateCredentials', result: 'Succeeded' })

packages/core/src/auth/sso/ssoAccessTokenProvider.ts

@@ -59,6 +59,15 @@ export abstract class SsoAccessTokenProvider {
     private static logIfChanged = onceChanged((s: string) => getLogger().info(s))
     private readonly className = 'SsoAccessTokenProvider'
 
+    /**
+     * Prevents concurrent token refresh operations.
+     * Maps tokenCacheKey to an in-flight refresh promise.
+     */
+    private static refreshPromises = new Map<
+        string,
+        Promise<{ token: SsoToken; registration: ClientRegistration; region: string; startUrl: string }>
+    >()
+
     public static set authSource(val: string) {
         SsoAccessTokenProvider._authSource = val
     }
@@ -108,15 +117,43 @@ export abstract class SsoAccessTokenProvider {
                 true
             )
         )
+
         if (!data || !isExpired(data.token)) {
+            getLogger().debug('Auth: token is valid, returning cached token (key=%s)', this.tokenCacheKey)
             return data?.token
         }
 
+        getLogger().info(
+            `Auth: bearer token expired (expires at ${data.token.expiresAt}), attempting refresh (key=${this.tokenCacheKey})`
+        )
+
         if (data.registration && !isExpired(data.registration) && hasProps(data.token, 'refreshToken')) {
-            const refreshed = await this.refreshToken(data.token, data.registration)
+            getLogger().debug(`Auth: refresh token available, calling refreshToken() (key=${this.tokenCacheKey})`)
+            // Check if a refresh is already in progress for this token
+            const existingRefresh = SsoAccessTokenProvider.refreshPromises.get(this.tokenCacheKey)
+            if (existingRefresh) {
+                getLogger().debug(
+                    'SsoAccessTokenProvider: Token refresh already in progress, waiting for existing refresh'
+                )
+                const refreshed = await existingRefresh
+                return refreshed.token
+            }
+
+            // Start a new refresh and store the promise
+            const refreshPromise = this.refreshToken(data.token, data.registration)
+            SsoAccessTokenProvider.refreshPromises.set(this.tokenCacheKey, refreshPromise)
 
-            return refreshed.token
+            try {
+                const refreshed = await refreshPromise
+                return refreshed.token
+            } finally {
+                // Clean up the promise from the map once complete (success or failure)
+                SsoAccessTokenProvider.refreshPromises.delete(this.tokenCacheKey)
+            }
         } else {
+            getLogger().warn(
+                `getToken: cannot refresh - registration expired or no refresh token available (key=${this.tokenCacheKey})`
+            )
             await this.invalidate('allCacheExpired')
         }
     }
@@ -172,10 +209,18 @@ export abstract class SsoAccessTokenProvider {
 
         try {
             const clientInfo = selectFrom(registration, 'clientId', 'clientSecret')
+            getLogger().debug(`Auth refreshToken: calling OIDC createToken API (key=${this.tokenCacheKey})`)
             const response = await this.oidc.createToken({ ...clientInfo, ...token, grantType: refreshGrantType })
+
+            getLogger().debug(`Auth refreshToken: got response, now saving to cache...`)
+
             const refreshed = this.formatToken(response, registration)
+            getLogger().debug(`refreshToken: saving refreshed token to cache (key=${this.tokenCacheKey})`)
             await this.cache.token.save(this.tokenCacheKey, refreshed)
 
+            getLogger().info(
+                `Auth refreshToken: token refresh successful (key=${this.tokenCacheKey}, new expiry=${response.expiresAt})`
+            )
             telemetry.aws_refreshCredentials.emit({
                 result: 'Succeeded',
                 requestId: response.requestId,
@@ -184,6 +229,10 @@ export abstract class SsoAccessTokenProvider {
 
             return refreshed
         } catch (err) {
+            getLogger().error(
+                `Auth refreshToken: token refresh failed (key=${this.tokenCacheKey}): ${getErrorMsg(err as unknown as Error)}`
+            )
+
             if (err instanceof DiskCacheError) {
                 /**
                  * Background:
@@ -197,6 +246,9 @@ export abstract class SsoAccessTokenProvider {
                  *   to the logs where the error was logged. Hopefully they can use this information to fix the issue,
                  *   or at least hint for them to provide the logs in a bug report.
                  */
+                getLogger().warn(
+                    `Auth refreshToken: DiskCacheError during refresh, not invalidating session (key=${this.tokenCacheKey})`
+                )
                 void DiskCacheErrorMessage.instance.showMessageThrottled(err)
             } else if (!isNetworkError(err)) {
                 const reason = getTelemetryReason(err)

packages/core/src/lambda/remoteDebugging/ldkClient.ts

@@ -17,7 +17,7 @@ import { DefaultLambdaClient } from '../../shared/clients/lambdaClient'
 import { LocalProxy } from './localProxy'
 import globals from '../../shared/extensionGlobals'
 import { getLogger } from '../../shared/logger/logger'
-import { getIoTSTClientWithAgent, getLambdaClientWithAgent, getLambdaDebugUserAgent } from './utils'
+import { getIoTSTClientWithAgent, getLambdaClientWithAgent, getLambdaDebugUserAgentPairs } from './utils'
 import { ToolkitError } from '../../shared/errors'
 import * as nls from 'vscode-nls'
 
@@ -99,7 +99,7 @@ export class LdkClient {
      */
     private getLambdaClient(region: string): DefaultLambdaClient {
         if (!this.lambdaClientCache.has(region)) {
-            this.lambdaClientCache.set(region, getLambdaClientWithAgent(region, getLambdaDebugUserAgent()))
+            this.lambdaClientCache.set(region, getLambdaClientWithAgent(region, getLambdaDebugUserAgentPairs()))
         }
         return this.lambdaClientCache.get(region)!
     }

packages/core/src/lambda/remoteDebugging/utils.ts

@@ -5,36 +5,50 @@
 
 import { IoTSecureTunnelingClient } from '@aws-sdk/client-iotsecuretunneling'
 import { DefaultLambdaClient } from '../../shared/clients/lambdaClient'
-import { getUserAgent } from '../../shared/telemetry/util'
+import { getUserAgentPairs, userAgentPairsToString } from '../../shared/telemetry/util'
 import globals from '../../shared/extensionGlobals'
+import type { UserAgent } from '@aws-sdk/types'
 
-const customUserAgentBase = 'LAMBDA-DEBUG/1.0.0'
+const customUserAgentName = 'LAMBDA-DEBUG'
+const customUserAgentVersion = '1.0.0'
 
-export function getLambdaClientWithAgent(region: string, customUserAgent?: string): DefaultLambdaClient {
+export function getLambdaClientWithAgent(region: string, customUserAgent?: UserAgent): DefaultLambdaClient {
     if (!customUserAgent) {
-        customUserAgent = getLambdaUserAgent()
+        customUserAgent = getLambdaUserAgentPairs()
     }
     return new DefaultLambdaClient(region, customUserAgent)
 }
 
-// Example user agent:
-// LAMBDA-DEBUG/1.0.0 AWS-Toolkit-For-VSCode/testPluginVersion Visual-Studio-Code/1.102.2 ClientId/11111111-1111-1111-1111-111111111111
-export function getLambdaDebugUserAgent(): string {
-    return `${customUserAgentBase} ${getLambdaUserAgent()}`
+/**
+ * Returns properly formatted UserAgent pairs for AWS SDK v3
+ */
+export function getLambdaDebugUserAgentPairs(): UserAgent {
+    return [
+        [customUserAgentName, customUserAgentVersion],
+        ...getUserAgentPairs({ includePlatform: true, includeClientId: true }),
+    ]
 }
 
-// Example user agent:
-// AWS-Toolkit-For-VSCode/testPluginVersion Visual-Studio-Code/1.102.2 ClientId/11111111-1111-1111-1111-111111111111
-export function getLambdaUserAgent(): string {
-    return `${getUserAgent({ includePlatform: true, includeClientId: true })}`
+/**
+ * Returns properly formatted UserAgent pairs for AWS SDK v3
+ */
+export function getLambdaUserAgentPairs(): UserAgent {
+    return getUserAgentPairs({ includePlatform: true, includeClientId: true })
+}
+
+/**
+ * Returns user agent string for Lambda debugging in traditional format.
+ * Example: "LAMBDA-DEBUG/1.0.0 AWS-Toolkit-For-VSCode/testPluginVersion Visual-Studio-Code/1.105.1 ClientId/11111111-1111-1111-1111-111111111111"
+ */
+export function getLambdaDebugUserAgent(): string {
+    return userAgentPairsToString(getLambdaDebugUserAgentPairs())
 }
 
 export function getIoTSTClientWithAgent(region: string): IoTSecureTunnelingClient {
-    const customUserAgent = `${customUserAgentBase} ${getUserAgent({ includePlatform: true, includeClientId: true })}`
     return globals.sdkClientBuilderV3.createAwsService({
         serviceClient: IoTSecureTunnelingClient,
         clientOptions: {
-            userAgent: [[customUserAgent]],
+            customUserAgent: getLambdaDebugUserAgentPairs(),
             region,
         },
         userAgent: false,

packages/core/src/lambda/vue/remoteInvoke/invokeLambda.ts

@@ -36,7 +36,7 @@ import { getLambdaHandlerFile } from '../../../awsService/appBuilder/utils'
 import { runUploadDirectory } from '../../commands/uploadLambda'
 import fs from '../../../shared/fs/fs'
 import { showConfirmationMessage, showMessage } from '../../../shared/utilities/messages'
-import { getLambdaClientWithAgent, getLambdaDebugUserAgent } from '../../remoteDebugging/utils'
+import { getLambdaClientWithAgent, getLambdaDebugUserAgentPairs } from '../../remoteDebugging/utils'
 import { isLocalStackConnection } from '../../../auth/utils'
 import { getRemoteDebugLayer } from '../../remoteDebugging/remoteLambdaDebugger'
 
@@ -906,7 +906,7 @@ export async function invokeRemoteLambda(
     const resource: LambdaFunctionNode = params.functionNode
     const source: string = params.source || 'AwsExplorerRemoteInvoke'
     const client = getLambdaClientWithAgent(resource.regionCode)
-    const clientDebug = getLambdaClientWithAgent(resource.regionCode, getLambdaDebugUserAgent())
+    const clientDebug = getLambdaClientWithAgent(resource.regionCode, getLambdaDebugUserAgentPairs())
 
     const Panel = VueWebview.compilePanel(RemoteInvokeWebview)
 

packages/core/src/shared/awsClientBuilderV3.ts

@@ -12,7 +12,7 @@ import {
     TokenIdentity,
     TokenIdentityProvider,
 } from '@smithy/types'
-import { getUserAgent } from './telemetry/util'
+import { getUserAgentPairs } from './telemetry/util'
 import { DevSettings } from './settings'
 import {
     BuildHandler,
@@ -37,7 +37,6 @@ import { HttpResponse, HttpRequest } from '@aws-sdk/protocol-http'
 import { ConfiguredRetryStrategy } from '@smithy/util-retry'
 import { telemetry } from './telemetry/telemetry'
 import { getRequestId, getTelemetryReason, getTelemetryReasonDesc, getTelemetryResult } from './errors'
-import { extensionVersion } from './vscode/env'
 import { getLogger } from './logger/logger'
 import { partialClone } from './utilities/collectionUtils'
 import { selectFrom } from './utilities/tsUtils'
@@ -72,7 +71,7 @@ export interface AwsCommand<InputType extends object, OutputType extends object>
 export interface AwsClientOptions {
     credentials: AwsCredentialIdentityProvider
     region: string | Provider<string>
-    userAgent: UserAgent
+    customUserAgent: UserAgent
     requestHandler: {
         metadata?: RequestHandlerMetadata
         handle: (req: any, options?: any) => Promise<RequestHandlerOutput<any>>
@@ -155,8 +154,8 @@ export class AWSClientBuilderV3 {
             opt.region = serviceOptions.region
         }
 
-        if (!opt.userAgent && userAgent) {
-            opt.userAgent = [[getUserAgent({ includePlatform: true, includeClientId: true }), extensionVersion]]
+        if (!opt.customUserAgent && userAgent) {
+            opt.customUserAgent = getUserAgentPairs({ includePlatform: true, includeClientId: true })
         }
 
         if (!opt.retryStrategy) {
@@ -196,6 +195,7 @@ export class AWSClientBuilderV3 {
         }
         const service = new serviceOptions.serviceClient(opt)
         service.middlewareStack.add(telemetryMiddleware, { step: 'deserialize' })
+        service.middlewareStack.add(captureHeadersMiddleware, { step: 'deserialize' })
         service.middlewareStack.add(loggingMiddleware, { step: 'finalizeRequest' })
         service.middlewareStack.add(getEndpointMiddleware(serviceOptions.settings), { step: 'build' })
 
@@ -254,6 +254,26 @@ function getEndpointMiddleware(settings: DevSettings = DevSettings.instance): Bu
 const keepAliveMiddleware: BuildMiddleware<any, any> = (next: BuildHandler<any, any>) => async (args: any) =>
     addKeepAliveHeader(next, args)
 
+/**
+ * Middleware that captures HTTP response headers and attaches them to the output object.
+ * This makes headers accessible via `response.$httpHeaders` for all AWS SDK v3 operations.
+ * Useful for detecting custom headers from services like LocalStack.
+ */
+const captureHeadersMiddleware: DeserializeMiddleware<any, any> =
+    (next: DeserializeHandler<any, any>) => async (args: any) => {
+        const result = await next(args)
+
+        // Extract headers from HTTP response and attach to output for easy access
+        if (HttpResponse.isInstance(result.response)) {
+            const headers = result.response.headers
+            if (headers && result.output) {
+                result.output.$httpHeaders = headers
+            }
+        }
+
+        return result
+    }
+
 export async function emitOnRequest(next: DeserializeHandler<any, any>, context: HandlerExecutionContext, args: any) {
     if (!HttpResponse.isInstance(args.request)) {
         return next(args)

packages/core/src/shared/clients/lambdaClient.ts

@@ -41,6 +41,7 @@ import { CancellationError } from '../utilities/timeoutUtils'
 import { fromSSO } from '@aws-sdk/credential-provider-sso'
 import { getIAMConnection } from '../../auth/utils'
 import { NodeHttpHandler } from '@smithy/node-http-handler'
+import type { UserAgent } from '@aws-sdk/types'
 
 export type LambdaClient = ClassToInterfaceType<DefaultLambdaClient>
 
@@ -49,7 +50,7 @@ export class DefaultLambdaClient {
 
     public constructor(
         public readonly regionCode: string,
-        public readonly userAgent: string | undefined = undefined
+        public readonly userAgent: UserAgent | undefined = undefined
     ) {
         this.defaultTimeoutInMs = 5 * 60 * 1000 // 5 minutes (SDK default is 2 minutes)
     }
@@ -322,7 +323,7 @@ export class DefaultLambdaClient {
             serviceClient: LambdaSdkClient,
             userAgent: !this.userAgent,
             clientOptions: {
-                userAgent: this.userAgent ? [[this.userAgent]] : undefined,
+                customUserAgent: this.userAgent,
                 region: this.regionCode,
                 requestHandler: new NodeHttpHandler({
                     requestTimeout: this.defaultTimeoutInMs,

packages/core/src/shared/clients/stsClient.ts

@@ -10,6 +10,11 @@ import { Credentials } from '@aws-sdk/types'
 import globals from '../extensionGlobals'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
 
+// Extended response type that includes captured HTTP headers (added by global middleware)
+export interface GetCallerIdentityResponseWithHeaders extends GetCallerIdentityResponse {
+    $httpHeaders?: Record<string, string>
+}
+
 export type { GetCallerIdentityResponse }
 export type StsClient = ClassToInterfaceType<DefaultStsClient>
 
@@ -35,8 +40,9 @@ export class DefaultStsClient {
         return response
     }
 
-    public async getCallerIdentity(): Promise<GetCallerIdentityResponse> {
+    public async getCallerIdentity(): Promise<GetCallerIdentityResponseWithHeaders> {
         const sdkClient = this.createSdkClient()
+        // Note: $httpHeaders are added by global middleware in awsClientBuilderV3
         const response = await sdkClient.send(new GetCallerIdentityCommand({}))
         return response
     }