feat: update getUsageLimits API and add service client generator script

server/aws-lsp-codewhisperer/package.json

@@ -18,6 +18,7 @@
         "copyServiceClient": "copyfiles -u 1 --error ./src/client/sigv4/*.json out && copyfiles -u 1 --error ./src/client/token/*.json out",
         "fix": "npm run fix:prettier",
         "fix:prettier": "prettier . --write",
+        "generateClients": "ts-node ./script/generateServiceClient.ts",
         "lint": "npm run lint:src",
         "lint:bundle:webworker": "webpack --config webpack.lint.config.js && eslint bundle/aws-lsp-codewhisperer-webworker.js # Verify compatibility with web runtime target",
         "lint:src": "eslint src/ --ext .ts,.tsx",

server/aws-lsp-codewhisperer/script/generateServiceClient.ts

@@ -0,0 +1,223 @@
+import * as child_process from 'child_process'
+import * as fs from 'fs'
+import * as path from 'path'
+
+const repoRoot = path.join(process.cwd(), '../..')
+/**
+ * This script uses the AWS JS SDK to generate service clients where the client definition is contained within
+ * this repo. Client definitions are added at the bottom of this script.
+ */
+
+interface ServiceClientDefinition {
+    serviceName: string
+    serviceJsonPath: string
+}
+
+async function generateServiceClients(serviceClientDefinitions: ServiceClientDefinition[]): Promise<void> {
+    const tempJsSdkPath = path.join(repoRoot, 'node_modules', '.zzz-awssdk2')
+    fs.rmSync(tempJsSdkPath, { recursive: true, force: true })
+    fs.mkdirSync(tempJsSdkPath)
+    console.log(`Temp JS SDK Repo location: ${tempJsSdkPath}`)
+    console.log('Service Clients to Generate: ', serviceClientDefinitions.map(x => x.serviceName).join(', '))
+
+    await cloneJsSdk(tempJsSdkPath)
+    await insertServiceClientsIntoJsSdk(tempJsSdkPath, serviceClientDefinitions)
+    await runTypingsGenerator(tempJsSdkPath)
+    await integrateServiceClients(tempJsSdkPath, serviceClientDefinitions)
+
+    console.log('Done generating service client(s)')
+}
+
+/** When cloning aws-sdk-js, we want to pull the version actually used in package-lock.json. */
+function getJsSdkVersion(): string {
+    const packageLockPath = path.resolve(repoRoot, 'package-lock.json')
+    console.log(`packageLockPath: ${packageLockPath}`)
+    const json = fs.readFileSync(packageLockPath).toString()
+    const packageLock = JSON.parse(json)
+
+    return packageLock['packages']['node_modules/aws-sdk']['version']
+}
+
+async function cloneJsSdk(dir: string): Promise<void> {
+    // Output stderr while it clones so it doesn't look frozen
+    return new Promise<void>((resolve, reject) => {
+        const tag = `v${getJsSdkVersion()}`
+
+        const msg = `'Cloning' AWS JS SDK...
+    tag: ${tag}`
+        console.log(msg)
+
+        const gitArgs = [
+            '-c',
+            'advice.detachedHead=false',
+            'clone',
+            '--quiet',
+            '-b',
+            tag,
+            '--depth',
+            '1',
+            'https://github.com/aws/aws-sdk-js.git',
+            dir,
+        ]
+
+        const gitCmd = child_process.execFile('git', gitArgs, { encoding: 'utf8' })
+
+        gitCmd.stderr?.on('data', (data: any) => {
+            console.log(data)
+        })
+        gitCmd.once('close', (code, signal) => {
+            gitCmd.stdout?.removeAllListeners()
+
+            // Only needed for the "update" case, but harmless for "clone".
+            const gitCheckout = child_process.spawnSync('git', [
+                '-c',
+                'advice.detachedHead=false',
+                '-C',
+                dir,
+                'checkout',
+                '--force',
+                tag,
+            ])
+            if (gitCheckout.status !== undefined && gitCheckout.status !== 0) {
+                console.log(`error: git: status=${gitCheckout.status} output=${gitCheckout.output.toString()}`)
+            }
+
+            resolve()
+        })
+    })
+}
+
+async function insertServiceClientsIntoJsSdk(
+    jsSdkPath: string,
+    serviceClientDefinitions: ServiceClientDefinition[]
+): Promise<void> {
+    serviceClientDefinitions.forEach(serviceClientDefinition => {
+        const apiVersion = getApiVersion(serviceClientDefinition.serviceJsonPath)
+
+        // Copy the Service Json into the JS SDK for generation
+        const jsSdkServiceJsonPath = path.join(
+            jsSdkPath,
+            'apis',
+            `${serviceClientDefinition.serviceName.toLowerCase()}-${apiVersion}.normal.json`
+        )
+        fs.copyFileSync(serviceClientDefinition.serviceJsonPath, jsSdkServiceJsonPath)
+    })
+
+    const apiMetadataPath = path.join(jsSdkPath, 'apis', 'metadata.json')
+    await patchServicesIntoApiMetadata(
+        apiMetadataPath,
+        serviceClientDefinitions.map(x => x.serviceName)
+    )
+}
+
+interface ServiceJsonSchema {
+    metadata: {
+        apiVersion: string
+    }
+}
+
+function getApiVersion(serviceJsonPath: string): string {
+    const json = fs.readFileSync(serviceJsonPath).toString()
+    const serviceJson = JSON.parse(json) as ServiceJsonSchema
+
+    return serviceJson.metadata.apiVersion
+}
+
+interface ApiMetadata {
+    [key: string]: { name: string }
+}
+
+/**
+ * Updates the JS SDK's api metadata to contain the provided services
+ */
+async function patchServicesIntoApiMetadata(apiMetadataPath: string, serviceNames: string[]): Promise<void> {
+    console.log(`Patching services (${serviceNames.join(', ')}) into API Metadata...`)
+
+    const apiMetadataJson = fs.readFileSync(apiMetadataPath).toString()
+    const apiMetadata = JSON.parse(apiMetadataJson) as ApiMetadata
+
+    serviceNames.forEach(serviceName => {
+        apiMetadata[serviceName.toLowerCase()] = { name: serviceName }
+    })
+
+    fs.writeFileSync(apiMetadataPath, JSON.stringify(apiMetadata, undefined, 4))
+}
+
+/**
+ * Generates service clients
+ */
+async function runTypingsGenerator(repoPath: string): Promise<void> {
+    console.log('Generating service client typings...')
+
+    const stdout = child_process.execFileSync('node', ['scripts/typings-generator.js'], {
+        encoding: 'utf8',
+        cwd: repoPath,
+    })
+    console.log(stdout)
+}
+
+/**
+ * Copies the generated service clients into the repo
+ */
+async function integrateServiceClients(
+    repoPath: string,
+    serviceClientDefinitions: ServiceClientDefinition[]
+): Promise<void> {
+    for (const serviceClientDefinition of serviceClientDefinitions) {
+        await integrateServiceClient(
+            repoPath,
+            serviceClientDefinition.serviceJsonPath,
+            serviceClientDefinition.serviceName
+        )
+    }
+}
+
+/**
+ * Copies the generated service client into the repo
+ */
+async function integrateServiceClient(repoPath: string, serviceJsonPath: string, serviceName: string): Promise<void> {
+    const typingsFilename = `${serviceName.toLowerCase()}.d.ts`
+    const sourceClientPath = path.join(repoPath, 'clients', typingsFilename)
+    const destinationClientPath = path.join(path.dirname(serviceJsonPath), typingsFilename)
+
+    console.log(`Integrating ${typingsFilename} ...`)
+
+    fs.copyFileSync(sourceClientPath, destinationClientPath)
+
+    await sanitizeServiceClient(destinationClientPath)
+}
+
+/**
+ * Patches the type file imports to be relative to the SDK module
+ */
+async function sanitizeServiceClient(generatedClientPath: string): Promise<void> {
+    console.log('Altering Service Client to fit the codebase...')
+
+    let fileContents = fs.readFileSync(generatedClientPath).toString()
+
+    // Add a header stating the file is autogenerated
+    fileContents = `
+/**
+ * THIS FILE IS AUTOGENERATED BY 'generateServiceClient.ts'.
+ * DO NOT EDIT BY HAND.
+ */
+
+${fileContents}
+    `
+
+    fileContents = fileContents.replace(/(import .* from.*)\.\.(.*)/g, '$1aws-sdk$2')
+
+    fs.writeFileSync(generatedClientPath, fileContents)
+}
+
+// ---------------------------------------------------------------------------------------------------------------------
+
+;(async () => {
+    const serviceClientDefinitions: ServiceClientDefinition[] = [
+        {
+            serviceJsonPath: 'src/client/token/bearer-token-service.json',
+            serviceName: 'CodeWhispererBearerTokenClient',
+        },
+    ]
+    await generateServiceClients(serviceClientDefinitions)
+})()

server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json

@@ -1318,6 +1318,9 @@
             "min": 1,
             "pattern": "(?:[A-Za-z0-9\\+/]{4})*(?:[A-Za-z0-9\\+/]{2}\\=\\=|[A-Za-z0-9\\+/]{3}\\=)?"
         },
+        "BigDecimal": {
+            "type": "bigdecimal"
+        },
         "Boolean": {
             "type": "boolean",
             "box": true
@@ -3241,6 +3244,9 @@
                 "profileArn": {
                     "shape": "ProfileArn",
                     "documentation": "<p>The ARN of the Q Developer profile. Required for enterprise customers, optional for Builder ID users.</p>"
+                },
+                "resourceType": {
+                    "shape": "ResourceType"
                 }
             }
         },
@@ -3254,6 +3260,10 @@
                 "daysUntilReset": {
                     "shape": "Integer",
                     "documentation": "<p>Number of days remaining until the usage metrics reset</p>"
+                },
+                "usageBreakdown": {
+                    "shape": "UsageBreakdown",
+                    "documentation": "<p>Usage breakdown by SKU type</p>"
                 }
             }
         },
@@ -4448,6 +4458,10 @@
             "type": "string",
             "enum": ["ALLOW", "DENY"]
         },
+        "ResourceType": {
+            "type": "string",
+            "enum": ["AGENTIC_REQUEST"]
+        },
         "ResumeTransformationRequest": {
             "type": "structure",
             "required": ["transformationJobId"],
@@ -5597,7 +5611,7 @@
         "ThrottlingExceptionReason": {
             "type": "string",
             "documentation": "<p>Reason for ThrottlingException</p>",
-            "enum": ["MONTHLY_REQUEST_COUNT"]
+            "enum": ["MONTHLY_REQUEST_COUNT", "INSUFFICIENT_MODEL_CAPACITY"]
         },
         "Timestamp": {
             "type": "timestamp"
@@ -6075,6 +6089,9 @@
                 "accountlessUserId": {
                     "shape": "String"
                 },
+                "directoryId": {
+                    "shape": "String"
+                },
                 "featureType": {
                     "shape": "UsageLimitType"
                 },
@@ -6147,6 +6164,28 @@
             "max": 1024,
             "min": 1
         },
+        "UsageBreakdown": {
+            "type": "structure",
+            "required": ["currentUsage", "currentOverages", "usageLimit", "overageCharges"],
+            "members": {
+                "currentUsage": {
+                    "shape": "Integer",
+                    "documentation": "<p>Current usage count for the billing period</p>"
+                },
+                "currentOverages": {
+                    "shape": "Integer",
+                    "documentation": "<p>Current overages count for the billing period</p>"
+                },
+                "usageLimit": {
+                    "shape": "Integer",
+                    "documentation": "<p>Usage limit based on subscription tier</p>"
+                },
+                "overageCharges": {
+                    "shape": "BigDecimal",
+                    "documentation": "<p>Total overage charges</p>"
+                }
+            }
+        },
         "UsageLimitList": {
             "type": "structure",
             "required": ["type", "currentUsage", "totalUsageLimit"],

server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts

@@ -3,6 +3,7 @@
  * THIS FILE IS AUTOGENERATED BY 'generateServiceClient.ts'.
  * DO NOT EDIT BY HAND.
  */
+
 import {Request} from 'aws-sdk/lib/request';
 import {Response} from 'aws-sdk/lib/response';
 import {AWSError} from 'aws-sdk/lib/error';
@@ -378,6 +379,7 @@ declare namespace CodeWhispererBearerTokenClient {
   export type AttributesMap = {[key: string]: StringList};
   export type AttributesMapKeyString = string;
   export type Base64EncodedPaginationToken = string;
+  export type BigDecimal = number;
   export type Boolean = boolean;
   export interface ByUserAnalytics {
     s3Uri?: S3Uri;
@@ -1000,13 +1002,18 @@ declare namespace CodeWhispererBearerTokenClient {
      * The ARN of the Q Developer profile. Required for enterprise customers, optional for Builder ID users.
      */
     profileArn?: ProfileArn;
+    resourceType?: ResourceType;
   }
   export interface GetUsageLimitsResponse {
     limits: UsageLimits;
     /**
      * Number of days remaining until the usage metrics reset
      */
     daysUntilReset: Integer;
+    /**
+     * Usage breakdown by SKU type
+     */
+    usageBreakdown?: UsageBreakdown;
   }
   export interface GitState {
     /**
@@ -1406,6 +1413,7 @@ declare namespace CodeWhispererBearerTokenClient {
     effect: ResourcePolicyEffect;
   }
   export type ResourcePolicyEffect = "ALLOW"|"DENY"|string;
+  export type ResourceType = "AGENTIC_REQUEST"|string;
   export interface ResumeTransformationRequest {
     transformationJobId: TransformationJobId;
     userActionStatus?: TransformationUserActionStatus;
@@ -1944,6 +1952,7 @@ declare namespace CodeWhispererBearerTokenClient {
   export interface UpdateUsageLimitsRequest {
     accountId: String;
     accountlessUserId?: String;
+    directoryId?: String;
     featureType: UsageLimitType;
     requestedLimit: Long;
     justification?: String;
@@ -1963,6 +1972,24 @@ declare namespace CodeWhispererBearerTokenClient {
   export type UploadId = string;
   export type UploadIntent = "TRANSFORMATION"|"TASK_ASSIST_PLANNING"|"AUTOMATIC_FILE_SECURITY_SCAN"|"FULL_PROJECT_SECURITY_SCAN"|"UNIT_TESTS_GENERATION"|"CODE_FIX_GENERATION"|"WORKSPACE_CONTEXT"|"AGENTIC_CODE_REVIEW"|string;
   export type Url = string;
+  export interface UsageBreakdown {
+    /**
+     * Current usage count for the billing period
+     */
+    currentUsage: Integer;
+    /**
+     * Current overages count for the billing period
+     */
+    currentOverages: Integer;
+    /**
+     * Usage limit based on subscription tier
+     */
+    usageLimit: Integer;
+    /**
+     * Total overage charges
+     */
+    overageCharges: BigDecimal;
+  }
   export interface UsageLimitList {
     type: UsageLimitType;
     currentUsage: Long;