Fix OriginalShapeName handling for union types in SDK code generation

SDK code generation functions for unions did not use OriginalShapeName to
restore the correct SDK type name, producing invalid references like
svcsdktypes.MemoryStrategyInput_.

Extend the existing OriginalShapeName handling (used for structures in
varEmptyConstructorSDKType) to the union code paths: setSDKForUnion,
setResourceForUnion, and varEmptyConstructorSDKType when RealType is
union.

Author: jlbutler

pkg/api/passes_test.go

@@ -974,3 +974,91 @@ func TestValidateShapeNameMethod(t *testing.T) {
 		})
 	}
 }
+
+// TestRenamedUnionShapePreservesOriginalName verifies that when a union
+// shape is renamed (e.g. MemoryStrategyInput -> MemoryStrategyInput_) by
+// renameIOSuffixedShapeNames, the OriginalShapeName is preserved. The
+// downstream SDK code generation functions (setSDKForUnion,
+// setResourceForUnion, varEmptyConstructorSDKType) use OriginalShapeName
+// to emit correct SDK type references.
+func TestRenamedUnionShapePreservesOriginalName(t *testing.T) {
+	a := &API{
+		name: "testapi",
+		Metadata: Metadata{
+			APIVersion:          "0000-00-00",
+			EndpointPrefix:      "testapi",
+			JSONVersion:         "1.1",
+			Protocol:            "json",
+			ServiceAbbreviation: "TestAPI",
+			ServiceFullName:     "Test API",
+			SignatureVersion:    "v4",
+		},
+		Operations: map[string]*Operation{
+			"CreateMemory": {
+				Name:      "CreateMemory",
+				InputRef:  ShapeRef{ShapeName: "CreateMemoryInput"},
+				OutputRef: ShapeRef{ShapeName: "CreateMemoryOutput"},
+			},
+		},
+		Shapes: map[string]*Shape{
+			"CreateMemoryInput": {
+				ShapeName: "CreateMemoryInput",
+				Type:      "structure",
+				MemberRefs: map[string]*ShapeRef{
+					"Strategies": {ShapeName: "MemoryStrategyInputList"},
+				},
+			},
+			"CreateMemoryOutput": {
+				ShapeName: "CreateMemoryOutput",
+				Type:      "structure",
+			},
+			"MemoryStrategyInputList": {
+				ShapeName: "MemoryStrategyInputList",
+				Type:      "list",
+				MemberRef: ShapeRef{ShapeName: "MemoryStrategyInput"},
+			},
+			"MemoryStrategyInput": {
+				ShapeName: "MemoryStrategyInput",
+				Type:      "structure",
+				RealType:  "union",
+			},
+		},
+	}
+
+	// Wire up shape refs
+	for _, op := range a.Operations {
+		op.InputRef.API = a
+		op.InputRef.Shape = a.Shapes[op.InputRef.ShapeName]
+		op.OutputRef.API = a
+		op.OutputRef.Shape = a.Shapes[op.OutputRef.ShapeName]
+	}
+	for _, s := range a.Shapes {
+		s.API = a
+		for k := range s.MemberRefs {
+			ref := s.MemberRefs[k]
+			ref.API = a
+			ref.Shape = a.Shapes[ref.ShapeName]
+			s.MemberRefs[k] = ref
+		}
+		if s.MemberRef.ShapeName != "" {
+			s.MemberRef.API = a
+			s.MemberRef.Shape = a.Shapes[s.MemberRef.ShapeName]
+		}
+	}
+
+	unionShape := a.Shapes["MemoryStrategyInput"]
+	if unionShape.OriginalShapeName != "" {
+		t.Fatalf("expected empty OriginalShapeName before rename, got %q", unionShape.OriginalShapeName)
+	}
+
+	a.renameIOSuffixedShapeNames()
+
+	// ShapeName gets underscore suffix, OriginalShapeName preserves the
+	// pre-rename value for downstream SDK code generation.
+	if unionShape.ShapeName != "MemoryStrategyInput_" {
+		t.Fatalf("expected ShapeName %q after rename, got %q", "MemoryStrategyInput_", unionShape.ShapeName)
+	}
+	if unionShape.OriginalShapeName != "MemoryStrategyInput" {
+		t.Fatalf("expected OriginalShapeName %q after rename, got %q", "MemoryStrategyInput", unionShape.OriginalShapeName)
+	}
+}

pkg/generate/code/set_resource.go

@@ -2454,6 +2454,11 @@ func setResourceForUnion(
 
 	sdkGoType := sourceShape.GoTypeWithPkgName()
 	sdkGoType = model.ReplacePkgName(sdkGoType, r.SDKAPIPackageName(), "svcsdktypes", true)
+	// Use the original shape name for SDK type references when the shape
+	// was renamed (e.g. Input/Output suffix collision avoidance).
+	if sourceShape.OriginalShapeName != "" {
+		sdkGoType = "*svcsdktypes." + sourceShape.OriginalShapeName
+	}
 
 	out += fmt.Sprintf("%sswitch %s.(type) {\n", indent, sourceVarName)
 	for _, targetMemberName := range targetShape.MemberNames() {

pkg/generate/code/set_resource_test.go

@@ -5579,3 +5579,31 @@ func TestSetResource_BedrockAgentCoreControl_GatewayTarget_NestedUnionTypeSwitch
 	assert.Contains(got, "case *svcsdktypes.ApiSchemaConfigurationMemberInlinePayload:")
 	assert.Contains(got, "case *svcsdktypes.ApiSchemaConfigurationMemberS3:")
 }
+
+// TestSetResource_BedrockAgentCoreControl_Memory_InputSuffixUnion tests that
+// union types whose names end in "Input" (e.g. MemoryStrategyInput) use the
+// original SDK shape name in generated set-output code, not the renamed
+// version with an underscore suffix (MemoryStrategyInput_).
+func TestSetResource_BedrockAgentCoreControl_Memory_InputSuffixUnion(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	g := testutil.NewModelForServiceWithOptions(t, "bedrock-agentcore-control", &testutil.TestingModelOptions{
+		GeneratorConfigFile: "generator-with-input-suffix-union.yaml",
+	})
+
+	crd := testutil.GetCRDByName(t, g, "Memory")
+	require.NotNil(crd)
+	assert.NotNil(crd.Ops.ReadOne)
+
+	got, err := code.SetResource(crd.Config(), crd, model.OpTypeGet, "resp", "ko", 1)
+	require.NoError(err)
+
+	// On the read path, union type switches must use the original SDK name
+	assert.NotContains(got, "MemoryStrategyInput_",
+		"Should use original SDK shape name MemoryStrategyInput, not renamed MemoryStrategyInput_")
+	assert.NotContains(got, "CustomConfigurationInput_",
+		"Should use original SDK shape name CustomConfigurationInput, not renamed version")
+	assert.NotContains(got, "TriggerConditionInput_",
+		"Should use original SDK shape name TriggerConditionInput, not renamed version")
+}

pkg/generate/code/set_sdk.go

@@ -1533,7 +1533,7 @@ func varEmptyConstructorSDKType(
 	// For SDK types, we need to use the original shape name (before stutter
 	// removal) since the AWS SDK uses the original names. The stutter removal
 	// renames are only for CRD types.
-	if shape.Type == "structure" && shape.OriginalShapeName != "" {
+	if (shape.Type == "structure" || shape.RealType == "union") && shape.OriginalShapeName != "" {
 		// Replace the renamed shape name with the original SDK shape name
 		goType = "svcsdktypes." + shape.OriginalShapeName
 	}
@@ -1863,6 +1863,11 @@ func setSDKForUnion(
 
 	sdkGoType := targetShape.GoTypeWithPkgName()
 	sdkGoType = model.ReplacePkgName(sdkGoType, r.SDKAPIPackageName(), "svcsdktypes", false)
+	// Use the original shape name for SDK type references when the shape
+	// was renamed (e.g. Input/Output suffix collision avoidance).
+	if targetShape.OriginalShapeName != "" {
+		sdkGoType = "svcsdktypes." + targetShape.OriginalShapeName
+	}
 
 	out += fmt.Sprintf("%sisInterfaceSet := false\n", indent)
 

pkg/generate/code/set_sdk_test.go

@@ -6769,3 +6769,40 @@ func TestSetSDK_BedrockAgentCoreControl_GatewayTarget_NestedUnionNoDereference(t
 	// dereferenced (e.g., ApiGatewayTargetConfiguration is a regular struct)
 	assert.Contains(got, "f7f0f0Parent.Value = *f7f0f0\n")
 }
+
+// TestSetSDK_BedrockAgentCoreControl_Memory_InputSuffixUnion tests that
+// union types whose names end in "Input" (e.g. MemoryStrategyInput) use
+// the original SDK shape name in generated code, not the renamed version
+// with an underscore suffix (MemoryStrategyInput_). The rename pass adds
+// the underscore to avoid collisions with operation input types, but SDK
+// type references must use the original name.
+func TestSetSDK_BedrockAgentCoreControl_Memory_InputSuffixUnion(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	g := testutil.NewModelForServiceWithOptions(t, "bedrock-agentcore-control", &testutil.TestingModelOptions{
+		GeneratorConfigFile: "generator-with-input-suffix-union.yaml",
+	})
+
+	crd := testutil.GetCRDByName(t, g, "Memory")
+	require.NotNil(crd)
+	assert.NotNil(crd.Ops.Create)
+
+	got, err := code.SetSDK(crd.Config(), crd, model.OpTypeCreate, "r.ko", "res", 1)
+	require.NoError(err)
+
+	// MemoryStrategyInput is a union whose name ends in "Input". The code-gen
+	// renames it to MemoryStrategyInput_ to avoid collision with operation
+	// input types. SDK references must use the original name.
+	assert.Contains(got, "svcsdktypes.MemoryStrategyInput")
+	assert.NotContains(got, "MemoryStrategyInput_",
+		"Should use original SDK shape name MemoryStrategyInput, not renamed MemoryStrategyInput_")
+
+	// CustomConfigurationInput is also a union ending in "Input"
+	assert.NotContains(got, "CustomConfigurationInput_",
+		"Should use original SDK shape name CustomConfigurationInput, not renamed version")
+
+	// TriggerConditionInput is also a union ending in "Input"
+	assert.NotContains(got, "TriggerConditionInput_",
+		"Should use original SDK shape name TriggerConditionInput, not renamed version")
+}

pkg/testdata/models/apis/bedrock-agentcore-control/0000-00-00/generator-with-input-suffix-union.yaml

@@ -0,0 +1,33 @@
+ignore:
+  resource_names:
+    - AgentRuntime
+    - AgentRuntimeEndpoint
+    - ApiKeyCredentialProvider
+    - Browser
+    - BrowserProfile
+    - CodeInterpreter
+    - Evaluator
+    - Gateway
+    - GatewayTarget
+    # Memory
+    - Oauth2CredentialProvider
+    - OnlineEvaluationConfig
+    - WorkloadIdentity
+    - Policy
+    - PolicyEngine
+  field_paths:
+    - MemoryStrategy.Configuration
+    - CreateMemoryInput.ClientToken
+    - UpdateMemoryInput.ClientToken
+    - DeleteMemoryInput.ClientToken
+sdk_names:
+  model_name: bedrock-agentcore-control
+
+resources:
+  Memory:
+    ignore_idempotency_token: true
+    fields:
+      MemoryId:
+        is_primary_key: true
+    tags:
+      ignore: true