deps(sdkv3): start migration to v3 with new client builder

packages/core/package.json

@@ -496,6 +496,7 @@
     "dependencies": {
         "@amzn/amazon-q-developer-streaming-client": "file:../../src.gen/@amzn/amazon-q-developer-streaming-client",
         "@amzn/codewhisperer-streaming": "file:../../src.gen/@amzn/codewhisperer-streaming",
+        "@aws-sdk/protocol-http": "^3.370.0",
         "@aws-sdk/client-cloudformation": "^3.667.0",
         "@aws-sdk/client-cloudwatch-logs": "^3.666.0",
         "@aws-sdk/client-cognito-identity": "^3.637.0",

packages/core/src/shared/awsClientBuilderV3.ts

@@ -0,0 +1,151 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { CredentialsShim } from '../auth/deprecated/loginManager'
+import { AwsContext } from './awsContext'
+import { AwsCredentialIdentityProvider, RetryStrategyV2 } from '@smithy/types'
+import { getUserAgent } from './telemetry/util'
+import { DevSettings } from './settings'
+import {
+    DeserializeHandler,
+    DeserializeHandlerOptions,
+    DeserializeMiddleware,
+    HandlerExecutionContext,
+    MetadataBearer,
+    MiddlewareStack,
+    Provider,
+    RequestHandlerMetadata,
+    RequestHandlerOutput,
+    RetryStrategy,
+    UserAgent,
+} from '@aws-sdk/types'
+import { HttpResponse } from '@aws-sdk/protocol-http'
+import { ConfiguredRetryStrategy } from '@smithy/util-retry'
+import { telemetry } from './telemetry'
+import { getRequestId, getTelemetryReason, getTelemetryReasonDesc, getTelemetryResult } from './errors'
+import { extensionVersion } from '.'
+import { getLogger } from './logger'
+import { partialClone } from './utilities/collectionUtils'
+
+export type AwsClientConstructor<C> = new (o: AwsClientOptions) => C
+
+// AWS-SDKv3 does not export generic types for clients so we need to build them as needed
+// https://github.com/aws/aws-sdk-js-v3/issues/5856#issuecomment-2096950979
+interface AwsClient {
+    middlewareStack: {
+        add: MiddlewareStack<any, MetadataBearer>['add']
+    }
+}
+
+interface AwsClientOptions {
+    credentials: AwsCredentialIdentityProvider
+    region: string | Provider<string>
+    userAgent: UserAgent
+    requestHandler: {
+        metadata?: RequestHandlerMetadata
+        handle: (req: any, options?: any) => Promise<RequestHandlerOutput<any>>
+        destroy?: () => void
+    }
+    apiVersion: string
+    endpoint: string
+    retryStrategy: RetryStrategy | RetryStrategyV2
+}
+
+export class AWSClientBuilderV3 {
+    public constructor(private readonly context: AwsContext) {}
+
+    private getShim(): CredentialsShim {
+        const shim = this.context.credentialsShim
+        if (!shim) {
+            throw new Error('Toolkit is not logged-in.')
+        }
+        return shim
+    }
+
+    public async createAwsService<C extends AwsClient>(
+        type: AwsClientConstructor<C>,
+        options?: Partial<AwsClientOptions>,
+        region?: string,
+        userAgent: boolean = true,
+        settings?: DevSettings
+    ): Promise<C> {
+        const shim = this.getShim()
+        const opt = (options ?? {}) as AwsClientOptions
+
+        if (!opt.region && region) {
+            opt.region = region
+        }
+
+        if (!opt.userAgent && userAgent) {
+            opt.userAgent = [[getUserAgent({ includePlatform: true, includeClientId: true }), extensionVersion]]
+        }
+
+        if (!opt.retryStrategy) {
+            // Simple exponential backoff strategy as default.
+            opt.retryStrategy = new ConfiguredRetryStrategy(5, (attempt: number) => 1000 * 2 ** attempt)
+        }
+        // TODO: add tests for refresh logic.
+        opt.credentials = async () => {
+            const creds = await shim.get()
+            if (creds.expiration && creds.expiration.getTime() < Date.now()) {
+                return shim.refresh()
+            }
+            return creds
+        }
+
+        const service = new type(opt)
+        // TODO: add middleware for logging, telemetry, endpoints.
+        service.middlewareStack.add(telemetryMiddleware, { step: 'deserialize' } as DeserializeHandlerOptions)
+        return service
+    }
+}
+
+export function getServiceId(context: { clientName?: string; commandName?: string }): string {
+    return context.clientName?.toLowerCase().replace(/client$/, '') ?? 'unknown-service'
+}
+
+/**
+ * Record request IDs to the current context, potentially overriding the field if
+ * multiple API calls are made in the same context. We only do failures as successes are generally uninteresting and noisy.
+ */
+export function recordErrorTelemetry(err: Error, serviceName?: string) {
+    telemetry.record({
+        requestId: getRequestId(err),
+        requestServiceType: serviceName,
+        reasonDesc: getTelemetryReasonDesc(err),
+        reason: getTelemetryReason(err),
+        result: getTelemetryResult(err),
+    })
+}
+
+function logAndThrow(e: any, serviceId: string, errorMessageAppend: string): never {
+    if (e instanceof Error) {
+        recordErrorTelemetry(e, serviceId)
+        const err = { ...e }
+        delete err['stack']
+        getLogger().error('API Response %s: %O', errorMessageAppend, err)
+    }
+    throw e
+}
+/**
+ * Telemetry logic to be added to all created clients. Adds logging and emitting metric on errors.
+ */
+const telemetryMiddleware: DeserializeMiddleware<any, any> =
+    (next: DeserializeHandler<any, any>, context: HandlerExecutionContext) => async (args: any) => {
+        if (!HttpResponse.isInstance(args.request)) {
+            return next(args)
+        }
+        const serviceId = getServiceId(context as object)
+        const { hostname, path } = args.request
+        const logTail = `(${hostname} ${path})`
+        const result = await next(args).catch((e: any) => logAndThrow(e, serviceId, logTail))
+        if (HttpResponse.isInstance(result.response)) {
+            // TODO: omit credentials / sensitive info from the logs / telemetry.
+            const output = partialClone(result.output, 3)
+            getLogger().debug('API Response %s: %O', logTail, output)
+        }
+
+        return result
+    }

packages/core/src/shared/utilities/tsUtils.ts

@@ -151,3 +151,13 @@ export type FactoryFunction<T extends abstract new (...args: any[]) => any> = (
 
 /** Can be used to isolate all number fields of a record `T` */
 export type NumericKeys<T> = { [P in keyof T]-?: T[P] extends number | undefined ? P : never }[keyof T]
+
+export function omitIfPresent<T extends Record<string, unknown>>(obj: T, keys: string[]): T {
+    const objCopy = { ...obj }
+    for (const key of keys) {
+        if (key in objCopy) {
+            ;(objCopy as any)[key] = '[omitted]'
+        }
+    }
+    return objCopy
+}

packages/core/src/test/shared/awsClientBuilderV3.test.ts

@@ -0,0 +1,79 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import { version } from 'vscode'
+import { getClientId } from '../../shared/telemetry/util'
+import { FakeMemento } from '../fakeExtensionContext'
+import { FakeAwsContext } from '../utilities/fakeAwsContext'
+import { GlobalState } from '../../shared/globalState'
+import { AWSClientBuilderV3, getServiceId, recordErrorTelemetry } from '../../shared/awsClientBuilderV3'
+import { Client } from '@aws-sdk/smithy-client'
+import { extensionVersion } from '../../shared'
+import { assertTelemetry } from '../testUtil'
+import { telemetry } from '../../shared/telemetry'
+
+describe('AwsClientBuilderV3', function () {
+    let builder: AWSClientBuilderV3
+
+    beforeEach(async function () {
+        builder = new AWSClientBuilderV3(new FakeAwsContext())
+    })
+
+    describe('createAndConfigureSdkClient', function () {
+        it('includes Toolkit user-agent if no options are specified', async function () {
+            const service = await builder.createAwsService(Client)
+            const clientId = getClientId(new GlobalState(new FakeMemento()))
+
+            assert.ok(service.config.userAgent)
+            assert.strictEqual(
+                service.config.userAgent![0][0].replace('---Insiders', ''),
+                `AWS-Toolkit-For-VSCode/testPluginVersion Visual-Studio-Code/${version} ClientId/${clientId}`
+            )
+            assert.strictEqual(service.config.userAgent![0][1], extensionVersion)
+        })
+
+        it('adds region to client', async function () {
+            const service = await builder.createAwsService(Client, { region: 'us-west-2' })
+
+            assert.ok(service.config.region)
+            assert.strictEqual(service.config.region, 'us-west-2')
+        })
+
+        it('adds Client-Id to user agent', async function () {
+            const service = await builder.createAwsService(Client)
+            const clientId = getClientId(new GlobalState(new FakeMemento()))
+            const regex = new RegExp(`ClientId/${clientId}`)
+            assert.ok(service.config.userAgent![0][0].match(regex))
+        })
+
+        it('does not override custom user-agent if specified in options', async function () {
+            const service = await builder.createAwsService(Client, {
+                userAgent: [['CUSTOM USER AGENT']],
+            })
+
+            assert.strictEqual(service.config.userAgent[0][0], 'CUSTOM USER AGENT')
+        })
+    })
+})
+
+describe('getServiceId', function () {
+    it('returns the service ID', function () {
+        assert.strictEqual(getServiceId({ clientName: 'ec2' }), 'ec2')
+        assert.strictEqual(getServiceId({ clientName: 'ec2client' }), 'ec2')
+        assert.strictEqual(getServiceId({ clientName: 's3client' }), 's3')
+    })
+})
+
+describe('recordErrorTelemetry', function () {
+    it('includes requestServiceType in span', function () {
+        const e = new Error('test error')
+        // Using vscode_executeCommand as general span to test functionality. This metric is unrelated to what is done here.
+        telemetry.vscode_executeCommand.run((span) => {
+            recordErrorTelemetry(e, 'aws-service')
+        })
+        assertTelemetry('vscode_executeCommand', { requestServiceType: 'aws-service' })
+    })
+})