@@ -26,6 +26,9 @@ import software.amazon.smithy.model.traits.TimestampFormatTrait
26
26
import software.amazon.smithy.model.traits.UnitTypeTrait
27
27
import software.amazon.smithy.protocol.traits.Rpcv2CborTrait
28
28
29
+ private const val ACCEPT_HEADER = " application/cbor"
30
+ private const val ACCEPT_HEADER_EVENT_STREAM = " application/vnd.amazon.eventstream"
31
+
29
32
class RpcV2Cbor : AwsHttpBindingProtocolGenerator () {
30
33
override val protocol: ShapeId = Rpcv2CborTrait .ID
31
34
@@ -59,13 +62,16 @@ class RpcV2Cbor : AwsHttpBindingProtocolGenerator() {
59
62
}
60
63
}
61
64
62
- // Requests with event stream responses MUST include an `Accept` header set to the value `application/vnd.amazon.eventstream`
63
- val eventStreamsAcceptHeaderMiddleware = object : ProtocolMiddleware {
64
- private val mutateHeadersMiddleware = MutateHeadersMiddleware (extraHeaders = mapOf (" Accept" to " application/vnd.amazon.eventstream" ))
65
-
66
- override fun isEnabledFor (ctx : ProtocolGenerator .GenerationContext , op : OperationShape ): Boolean = op.isOutputEventStream(ctx.model)
67
- override val name: String = " RpcV2CborEventStreamsAcceptHeaderMiddleware"
68
- override fun render (ctx : ProtocolGenerator .GenerationContext , op : OperationShape , writer : KotlinWriter ) = mutateHeadersMiddleware.render(ctx, op, writer)
65
+ // Add `Accept` header with value `application/cbor` for standard responses
66
+ // and `application/vnd.amazon.eventstream` for event stream responses
67
+ val acceptHeaderMiddleware = object : ProtocolMiddleware {
68
+ override val name: String = " RpcV2CborAcceptHeaderMiddleware"
69
+ override fun render (ctx : ProtocolGenerator .GenerationContext , op : OperationShape , writer : KotlinWriter ) {
70
+ val acceptHeaderValue = if (op.isOutputEventStream(ctx.model)) ACCEPT_HEADER_EVENT_STREAM else ACCEPT_HEADER
71
+ MutateHeadersMiddleware (
72
+ extraHeaders = mapOf (" Accept" to acceptHeaderValue)
73
+ ).render(ctx, op, writer)
74
+ }
69
75
}
70
76
71
77
// Emit a metric to track usage of RpcV2Cbor
@@ -79,7 +85,7 @@ class RpcV2Cbor : AwsHttpBindingProtocolGenerator() {
79
85
return super .getDefaultHttpMiddleware(ctx) + listOf (
80
86
smithyProtocolHeaderMiddleware,
81
87
validateSmithyProtocolHeaderMiddleware,
82
- eventStreamsAcceptHeaderMiddleware ,
88
+ acceptHeaderMiddleware ,
83
89
businessMetricsMiddleware,
84
90
)
85
91
}
0 commit comments