From 00d1aec38416c42873670ff3c46509c04a22f78b Mon Sep 17 00:00:00 2001
From: George Fu <kuhe@users.noreply.github.com>
Date: Wed, 2 Oct 2024 19:00:21 +0000
Subject: [PATCH] feat: allow feature identification on the smithy context

---
 .changeset/sharp-horses-fry.md       |  6 +++++
 packages/core/src/setFeature.spec.ts | 10 ++++++++
 packages/core/src/setFeature.ts      | 26 +++++++++++++++++++
 packages/types/src/feature-ids.ts    | 18 +++++++++++++
 packages/types/src/index.ts          |  1 +
 packages/types/src/middleware.ts     | 38 +++++++++++++++++++++++-----
 6 files changed, 92 insertions(+), 7 deletions(-)
 create mode 100644 .changeset/sharp-horses-fry.md
 create mode 100644 packages/core/src/setFeature.spec.ts
 create mode 100644 packages/core/src/setFeature.ts
 create mode 100644 packages/types/src/feature-ids.ts

diff --git a/.changeset/sharp-horses-fry.md b/.changeset/sharp-horses-fry.md
new file mode 100644
index 00000000000..d08c81e5e61
--- /dev/null
+++ b/.changeset/sharp-horses-fry.md
@@ -0,0 +1,6 @@
+---
+"@smithy/types": minor
+"@smithy/core": minor
+---
+
+add feature identification map to smithy context
diff --git a/packages/core/src/setFeature.spec.ts b/packages/core/src/setFeature.spec.ts
new file mode 100644
index 00000000000..bfe44eb65f4
--- /dev/null
+++ b/packages/core/src/setFeature.spec.ts
@@ -0,0 +1,10 @@
+import { HandlerExecutionContext } from "@smithy/types";
+
+import { setFeature } from "./setFeature";
+
+describe(setFeature.name, () => {
+  it("creates the context object path if needed", () => {
+    const context: HandlerExecutionContext = {};
+    setFeature(context, "RETRY_MODE_STANDARD", "E");
+  });
+});
diff --git a/packages/core/src/setFeature.ts b/packages/core/src/setFeature.ts
new file mode 100644
index 00000000000..8bfaf30c28d
--- /dev/null
+++ b/packages/core/src/setFeature.ts
@@ -0,0 +1,26 @@
+import type { HandlerExecutionContext, SmithyFeatures } from "@smithy/types";
+
+/**
+ * @internal
+ * Indicates to the request context that a given feature is active.
+ *
+ * @param context - handler execution context.
+ * @param feature - readable name of feature.
+ * @param value - encoding value of feature. This is required because the
+ * specification asks the library not to include a runtime lookup of all
+ * the feature identifiers.
+ */
+export function setFeature<F extends keyof SmithyFeatures>(
+  context: HandlerExecutionContext,
+  feature: F,
+  value: SmithyFeatures[F]
+) {
+  if (!context.__smithy_context) {
+    context.__smithy_context = {
+      features: {},
+    };
+  } else if (!context.__smithy_context.features) {
+    context.__smithy_context.features = {};
+  }
+  context.__smithy_context.features![feature] = value;
+}
diff --git a/packages/types/src/feature-ids.ts b/packages/types/src/feature-ids.ts
new file mode 100644
index 00000000000..219193559aa
--- /dev/null
+++ b/packages/types/src/feature-ids.ts
@@ -0,0 +1,18 @@
+/**
+ * @internal
+ */
+export type SmithyFeatures = Partial<{
+  RESOURCE_MODEL: "A";
+  WAITER: "B";
+  PAGINATOR: "C";
+  RETRY_MODE_LEGACY: "D";
+  RETRY_MODE_STANDARD: "E";
+  RETRY_MODE_ADAPTIVE: "F";
+  GZIP_REQUEST_COMPRESSION: "L";
+  PROTOCOL_RPC_V2_CBOR: "M";
+  ENDPOINT_OVERRIDE: "N";
+  SIGV4A_SIGNING: "S";
+  CREDENTIALS_CODE: "e";
+  CREDENTIALS_HTTP: "z";
+  CREDENTIALS_IMDS: "0";
+}>;
diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts
index 17c2c945d56..c370335c2ab 100644
--- a/packages/types/src/index.ts
+++ b/packages/types/src/index.ts
@@ -11,6 +11,7 @@ export * from "./endpoint";
 export * from "./endpoints";
 export * from "./eventStream";
 export * from "./extensions";
+export * from "./feature-ids";
 export * from "./http";
 export * from "./http/httpHandlerInitialization";
 export * from "./identity";
diff --git a/packages/types/src/middleware.ts b/packages/types/src/middleware.ts
index c31a123f198..4ff74e1d4a1 100644
--- a/packages/types/src/middleware.ts
+++ b/packages/types/src/middleware.ts
@@ -1,7 +1,10 @@
-import { AuthScheme, HttpAuthDefinition } from "./auth/auth";
-import { EndpointV2 } from "./endpoint";
-import { Logger } from "./logger";
-import { UserAgent } from "./util";
+import type { AuthScheme, HttpAuthDefinition } from "./auth/auth";
+import type { SelectedHttpAuthScheme } from "./auth/HttpAuthScheme";
+import type { Command } from "./command";
+import type { EndpointV2 } from "./endpoint";
+import type { SmithyFeatures } from "./feature-ids";
+import type { Logger } from "./logger";
+import type { UserAgent } from "./util";
 
 /**
  * @public
@@ -554,6 +557,7 @@ export interface HandlerExecutionContext {
   currentAuthConfig?: HttpAuthDefinition;
 
   /**
+   * @deprecated do not extend this field, it is a carryover from AWS SDKs.
    * Used by DynamoDbDocumentClient.
    */
   dynamoDbDocumentClientOptions?: Partial<{
@@ -563,10 +567,30 @@ export interface HandlerExecutionContext {
 
   /**
    * @internal
-   * Context for Smithy properties
+   * Context for Smithy properties.
+   */
+  [SMITHY_CONTEXT_KEY]?: {
+    service?: string;
+    operation?: string;
+    commandInstance?: Command<any, any, any, any, any>;
+    selectedHttpAuthScheme?: SelectedHttpAuthScheme;
+    features?: SmithyFeatures;
+    /**
+     * @deprecated
+     * Do not assign arbitrary members to the Smithy Context,
+     * fields should be explicitly declared here to avoid collisions.
+     */
+    [key: string]: unknown;
+  };
+
+  /**
+   * @deprecated
+   * Do not assign arbitrary members to the context, since
+   * they can interfere with existing functionality.
+   *
+   * Additional members should instead be declared on the SMITHY_CONTEXT_KEY
+   * or other reserved keys.
    */
-  [SMITHY_CONTEXT_KEY]?: Record<string, unknown>;
-
   [key: string]: any;
 }