fix: account ID based endpoints

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/KotlinJmespathExpressionVisitor.kt

@@ -19,16 +19,26 @@ import software.amazon.smithy.kotlin.codegen.model.isEnum
 import software.amazon.smithy.kotlin.codegen.model.targetOrSelf
 import software.amazon.smithy.kotlin.codegen.model.traits.OperationInput
 import software.amazon.smithy.kotlin.codegen.model.traits.OperationOutput
+import software.amazon.smithy.kotlin.codegen.utils.doubleQuote
 import software.amazon.smithy.kotlin.codegen.utils.dq
 import software.amazon.smithy.kotlin.codegen.utils.getOrNull
 import software.amazon.smithy.kotlin.codegen.utils.toCamelCase
-import software.amazon.smithy.model.shapes.ListShape
-import software.amazon.smithy.model.shapes.MapShape
-import software.amazon.smithy.model.shapes.MemberShape
-import software.amazon.smithy.model.shapes.Shape
+import software.amazon.smithy.model.shapes.*
 
 private val suffixSequence = sequenceOf("") + generateSequence(2) { it + 1 }.map(Int::toString) // "", "2", "3", etc.
 
+/*
+JMESPath has the concept of objects.
+This visitor assumes JMESPath objects will be Kotlin objects/classes.
+The smithy spec contains an instance where it's assumed a JMESPath object will be a Kotlin map.
+
+Specifically it's the keys function
+Smithy spec: https://smithy.io/2.0/additional-specs/rules-engine/parameters.html#smithy-rules-operationcontextparams-trait
+JMESPath spec: https://jmespath.org/specification.html#keys
+
+TODO: Test relevant uses of JMESPath to determine if we should support JMESPath objects as Kotlin maps throughout the entire visitor
+ */
+
 /**
  * An [ExpressionVisitor] used for traversing a JMESPath expression to generate code for traversing an equivalent
  * modeled object. This visitor is passed to [JmespathExpression.accept], at which point specific expression methods
@@ -526,11 +536,22 @@ class KotlinJmespathExpressionVisitor(
             ".$expr"
         }
 
-    private fun VisitedExpression.getKeys(): String {
-        val keys = this.shape?.targetOrSelf(ctx.model)?.allMembers
-            ?.keys?.joinToString(", ", "listOf(", ")") { "\"$it\"" }
-        return keys ?: "listOf<String>()"
-    }
+    /*
+    Smithy spec expects a map, JMESPath spec expects an object
+    Smithy spec: https://smithy.io/2.0/additional-specs/rules-engine/parameters.html#smithy-rules-operationcontextparams-trait
+    JMESPath spec: https://jmespath.org/specification.html#keys
+     */
+    private fun VisitedExpression.getKeys(): String =
+        if (shape?.targetOrSelf(ctx.model)?.type == ShapeType.MAP) {
+            "$identifier?.keys?.map { it.toString() }?.toList()"
+        } else {
+            shape
+                ?.targetOrSelf(ctx.model)
+                ?.allMembers
+                ?.keys
+                ?.joinToString(", ", "listOf(", ")") { it.doubleQuote() }
+                ?: "listOf<String>()"
+        }
 
     private fun VisitedExpression.getValues(): String {
         val values = this.shape?.targetOrSelf(ctx.model)?.allMembers?.keys

codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/endpoints/OperationContextParamsTest.kt

@@ -88,7 +88,7 @@ class OperationContextParamsTest {
     }
 
     @Test
-    fun testKeysFunctionPath() {
+    fun testKeysFunctionPathWithStructure() {
         val input = """
             structure TestOperationRequest {
                 Object: Object
@@ -114,4 +114,31 @@ class OperationContextParamsTest {
 
         codegen(pathResultType, path, input).shouldContainOnlyOnceWithDiff(expected)
     }
+
+    @Test
+    fun testKeysFunctionPathWithMap() {
+        val input = """
+            structure TestOperationRequest {
+                Object: StringMap
+            }
+
+            map StringMap {
+                key: String
+                value: String
+            }
+        """.trimIndent()
+
+        val path = "keys(Object)"
+        val pathResultType = "stringArray"
+
+        val expected = """
+            @Suppress("UNCHECKED_CAST")
+            val input = request.context[HttpOperationContext.OperationInput] as TestOperationRequest
+            val object = input.object
+            val keys = object?.keys?.map { it.toString() }?.toList()
+            builder.foo = keys
+        """.formatForTest("    ")
+
+        codegen(pathResultType, path, input).shouldContainOnlyOnceWithDiff(expected)
+    }
 }

tests/codegen/waiter-tests/model/function-keys.smithy

@@ -32,6 +32,20 @@ use smithy.waiters#waitable
             }
         ]
     },
+    KeysFunctionMapStringEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "keys(maps.strings)",
+                        expected: "key",
+                        comparator: "anyStringEquals"
+                    }
+                }
+            }
+        ]
+    },
 )
 @readonly
 @http(method: "GET", uri: "/keys/{name}", code: 200)

tests/codegen/waiter-tests/src/test/kotlin/com/test/FunctionKeysTest.kt

@@ -5,10 +5,12 @@
 
 package com.test
 
+import com.test.model.EntityMaps
 import com.test.model.EntityPrimitives
 import com.test.model.GetFunctionKeysEqualsRequest
 import com.test.model.GetFunctionKeysEqualsResponse
 import com.test.utils.successTest
+import com.test.waiters.waitUntilKeysFunctionMapStringEquals
 import com.test.waiters.waitUntilKeysFunctionPrimitivesIntegerEquals
 import com.test.waiters.waitUntilKeysFunctionPrimitivesStringEquals
 import kotlin.test.Test
@@ -27,4 +29,11 @@ class FunctionKeysTest {
         WaitersTestClient::waitUntilKeysFunctionPrimitivesIntegerEquals,
         GetFunctionKeysEqualsResponse { primitives = EntityPrimitives { } },
     )
+
+    @Test
+    fun testKeysFunctionMapStringEquals() = successTest(
+        GetFunctionKeysEqualsRequest { name = "test" },
+        WaitersTestClient::waitUntilKeysFunctionMapStringEquals,
+        GetFunctionKeysEqualsResponse { maps = EntityMaps { strings = mapOf("key" to "value") } },
+    )
 }