handle eventHeader and eventPayload traits

kuhe · kuhe · commit f6bfd6691764 · 2025-08-28T11:24:28.000-04:00
diff --git a/packages/core/src/submodules/eventStreams/EventStreamSerde.spec.ts b/packages/core/src/submodules/eventStreams/EventStreamSerde.spec.ts
@@ -1,5 +1,5 @@
 import { cbor, CborCodec, dateToTag } from "@smithy/core/cbor";
-import { NormalizedSchema, SCHEMA, struct } from "@smithy/core/schema";
+import { NormalizedSchema, SCHEMA, sim,struct } from "@smithy/core/schema";
 import { EventStreamMarshaller } from "@smithy/eventstream-serde-node";
 import { HttpResponse } from "@smithy/protocol-http";
 import { Message as EventMessage } from "@smithy/types";
@@ -42,9 +42,27 @@ describe(EventStreamSerde.name, () => {
       "ns",
       "EventStreamStructure",
       { streaming: 1 },
-      ["A", "B", "C"],
+      ["A", "B", "C", "Payload", "CustomHeaders"],
       // D is omitted to represent an unknown event.
-      [struct("ns", "A", 0, ["name"], [0]), struct("ns", "B", 0, ["name"], [0]), struct("ns", "C", 0, ["name"], [0])]
+      [
+        struct("ns", "A", 0, ["name"], [0]),
+        struct("ns", "B", 0, ["name"], [0]),
+        struct("ns", "C", 0, ["name"], [0]),
+        struct(
+          "ns",
+          "Payload",
+          0,
+          ["payload"],
+          [sim("ns", "StreamingBlobPayload", SCHEMA.STREAMING_BLOB, { eventPayload: 1 })]
+        ),
+        struct(
+          "ns",
+          "CustomHeaders",
+          0,
+          ["header1", "header2"],
+          [sim("ns", "EventHeader", 0, { eventHeader: 1 }), sim("ns", "EventHeader", 0, { eventHeader: 1 })]
+        ),
+      ]
     );
 
     const eventStreamContainerSchema = struct(
@@ -68,6 +86,8 @@ describe(EventStreamSerde.name, () => {
           yield { B: { name: "b" } };
           yield { C: { name: "c" } };
           yield { $unknown: ["D", { name: "d" }] };
+          yield { Payload: { payload: new Uint8Array([0, 1, 2, 3, 4, 5, 6]) } };
+          yield { CustomHeaders: { header1: "h1", header2: "h2" } };
         },
       };
 
@@ -104,6 +124,24 @@ describe(EventStreamSerde.name, () => {
           },
           body: { name: "d" },
         },
+        {
+          headers: {
+            ":event-type": { type: "string", value: "Payload" },
+            ":message-type": { type: "string", value: "event" },
+            ":content-type": { type: "string", value: "application/octet-stream" },
+          },
+          body: new Uint8Array([0, 1, 2, 3, 4, 5, 6]),
+        },
+        {
+          headers: {
+            ":event-type": { type: "string", value: "CustomHeaders" },
+            ":message-type": { type: "string", value: "event" },
+            ":content-type": { type: "string", value: "application/cbor" },
+            header1: { type: "string", value: "h1" },
+            header2: { type: "string", value: "h2" },
+          },
+          body: {},
+        },
       ];
 
       /**
@@ -188,6 +226,8 @@ describe(EventStreamSerde.name, () => {
           yield { B: { name: "b" } };
           yield { C: { name: "c" } };
           yield { D: { name: "d" } };
+          yield { Payload: { payload: new Uint8Array([0, 1, 2, 3, 4, 5, 6]) } };
+          yield { CustomHeaders: { header1: "h1", header2: "h2" } };
         },
       };
 
@@ -230,6 +270,8 @@ describe(EventStreamSerde.name, () => {
           { C: { name: `c` } },
           // todo(schema) getMessageUnmarshaller.ts must be patched to return unknown events.
           // $unknownEvent,
+          { Payload: { payload: new Uint8Array([0, 1, 2, 3, 4, 5, 6]) } },
+          { CustomHeaders: { header1: "h1", header2: "h2" } },
         ]);
       });
 
@@ -265,6 +307,8 @@ describe(EventStreamSerde.name, () => {
           { C: { name: `c` } },
           // todo(schema) getMessageUnmarshaller.ts must be patched to return unknown events.
           // $unknownEvent,
+          { Payload: { payload: new Uint8Array([0, 1, 2, 3, 4, 5, 6]) } },
+          { CustomHeaders: { header1: "h1", header2: "h2" } },
         ]);
 
         expect(initialResponseContainer).toEqual({
diff --git a/packages/core/src/submodules/eventStreams/EventStreamSerde.ts b/packages/core/src/submodules/eventStreams/EventStreamSerde.ts
@@ -109,65 +109,19 @@ export class EventStreamSerde {
         Object.keys(event).find((key) => {
           return key !== "__type";
         }) ?? "";
-      const eventStreamSchema = memberSchemas[unionMember] ?? NormalizedSchema.of(SCHEMA.DOCUMENT);
-
-      let messageSerialization: string | Uint8Array;
-      let eventType = unionMember;
-
-      if (eventStreamSchema.isStructSchema()) {
-        serializer.write(eventStreamSchema, event[unionMember]);
-        messageSerialization = serializer.flush();
-      } else {
-        // $unknown member
-        const [type, value] = event[unionMember];
-        eventType = type;
-        serializer.write(NormalizedSchema.of(SCHEMA.DOCUMENT), value);
-        messageSerialization = serializer.flush();
-      }
-
-      const body =
-        typeof messageSerialization === "string"
-          ? (this.serdeContext?.utf8Decoder ?? fromUtf8)(messageSerialization)
-          : messageSerialization;
+      const { additionalHeaders, body, eventType, explicitPayloadContentType } = this.writeEventBody(
+        unionMember,
+        unionSchema,
+        event
+      );
 
       const headers: MessageHeaders = {
         ":event-type": { type: "string", value: eventType },
         ":message-type": { type: "string", value: "event" },
-        ":content-type": { type: "string", value: defaultContentType },
+        ":content-type": { type: "string", value: explicitPayloadContentType ?? defaultContentType },
+        ...additionalHeaders,
       };
 
-      // additional trait-annotated event headers.
-      if (eventStreamSchema.isStructSchema()) {
-        for (const [memberName, memberSchema] of eventStreamSchema.structIterator()) {
-          const isHeader = !!memberSchema.getMergedTraits().eventHeader;
-          if (!isHeader) {
-            continue;
-          }
-          const value = event[memberName];
-          let type = "binary" as MessageHeaderValue["type"];
-          if (memberSchema.isNumericSchema()) {
-            if ((-2) ** 31 <= value && value <= 2 ** 31 - 1) {
-              type = "integer";
-            } else {
-              type = "long";
-            }
-          } else if (memberSchema.isTimestampSchema()) {
-            type = "timestamp";
-          } else if (memberSchema.isStringSchema()) {
-            type = "string";
-          } else if (memberSchema.isBooleanSchema()) {
-            type = "boolean";
-          }
-
-          if (isHeader && value != undefined) {
-            headers[memberName] = {
-              type,
-              value,
-            };
-          }
-        }
-      }
-
       return {
         headers,
         body,
@@ -263,4 +217,93 @@ export class EventStreamSerde {
       },
     };
   }
+
+  /**
+   * @param unionMember - member name within the structure that contains an event stream union.
+   * @param unionSchema - schema of the union.
+   * @param event
+   *
+   * @returns the event body (bytes) and event type (string).
+   */
+  private writeEventBody(unionMember: string, unionSchema: NormalizedSchema, event: any) {
+    const serializer = this.serializer;
+    let eventType = unionMember;
+    let explicitPayloadMember = null as null | string;
+    let explicitPayloadContentType: undefined | string;
+
+    const isKnownSchema = unionSchema.hasMemberSchema(unionMember);
+    const additionalHeaders: MessageHeaders = {};
+
+    if (!isKnownSchema) {
+      // $unknown member
+      const [type, value] = event[unionMember];
+      eventType = type;
+      serializer.write(SCHEMA.DOCUMENT, value);
+    } else {
+      const eventSchema = unionSchema.getMemberSchema(unionMember);
+
+      if (eventSchema.isStructSchema()) {
+        for (const [memberName, memberSchema] of eventSchema.structIterator()) {
+          const { eventHeader, eventPayload } = memberSchema.getMergedTraits();
+
+          if (eventPayload) {
+            explicitPayloadMember = memberName;
+            break;
+          } else if (eventHeader) {
+            const value = event[unionMember][memberName];
+            let type = "binary" as MessageHeaderValue["type"];
+            if (memberSchema.isNumericSchema()) {
+              if ((-2) ** 31 <= value && value <= 2 ** 31 - 1) {
+                type = "integer";
+              } else {
+                type = "long";
+              }
+            } else if (memberSchema.isTimestampSchema()) {
+              type = "timestamp";
+            } else if (memberSchema.isStringSchema()) {
+              type = "string";
+            } else if (memberSchema.isBooleanSchema()) {
+              type = "boolean";
+            }
+
+            if (value != null) {
+              additionalHeaders[memberName] = {
+                type,
+                value,
+              };
+              delete event[unionMember][memberName];
+            }
+          }
+        }
+
+        if (explicitPayloadMember !== null) {
+          const payloadSchema = eventSchema.getMemberSchema(explicitPayloadMember);
+          if (payloadSchema.isBlobSchema()) {
+            explicitPayloadContentType = "application/octet-stream";
+          } else if (payloadSchema.isStringSchema()) {
+            explicitPayloadContentType = "text/plain";
+          }
+          serializer.write(payloadSchema, event[unionMember][explicitPayloadMember]);
+        } else {
+          serializer.write(eventSchema, event[unionMember]);
+        }
+      } else {
+        throw new Error("@smithy/core/eventStreams - non-struct member not supported in event stream union.");
+      }
+    }
+
+    const messageSerialization: string | Uint8Array = serializer.flush();
+
+    const body =
+      typeof messageSerialization === "string"
+        ? (this.serdeContext?.utf8Decoder ?? fromUtf8)(messageSerialization)
+        : messageSerialization;
+
+    return {
+      body,
+      eventType,
+      explicitPayloadContentType,
+      additionalHeaders,
+    };
+  }
 }
diff --git a/packages/core/src/submodules/schema/schemas/NormalizedSchema.spec.ts b/packages/core/src/submodules/schema/schemas/NormalizedSchema.spec.ts
@@ -173,6 +173,15 @@ describe(NormalizedSchema.name, () => {
     });
   });
 
+  it("can identify whether a member exists", () => {
+    expect(ns.hasMemberSchema("list")).toBe(true);
+    expect(ns.hasMemberSchema("map")).toBe(true);
+    expect(ns.hasMemberSchema("struct")).toBe(true);
+    expect(ns.hasMemberSchema("a")).toBe(false);
+    expect(ns.hasMemberSchema("b")).toBe(false);
+    expect(ns.hasMemberSchema("c")).toBe(false);
+  });
+
   describe("iteration", () => {
     it("iterates over member schemas", () => {
       const iteration = Array.from(ns.structIterator());
diff --git a/packages/core/src/submodules/schema/schemas/NormalizedSchema.ts b/packages/core/src/submodules/schema/schemas/NormalizedSchema.ts
@@ -36,10 +36,7 @@ export class NormalizedSchema implements INormalizedSchema {
    * @param ref - a polymorphic SchemaRef to be dereferenced/normalized.
    * @param memberName - optional memberName if this NormalizedSchema should be considered a member schema.
    */
-  public constructor(
-    private readonly ref: SchemaRef,
-    private memberName?: string
-  ) {
+  public constructor(private readonly ref: SchemaRef, private memberName?: string) {
     const traitStack = [] as SchemaTraits[];
     let _ref = ref;
     let schema = ref;
@@ -373,6 +370,18 @@ export class NormalizedSchema implements INormalizedSchema {
     throw new Error(`@smithy/core/schema - the schema ${this.getName(true)} does not have a value member.`);
   }
 
+  /**
+   * @param member - to query.
+   * @returns whether there is a memberSchema with the given member name. False if not a structure (or union).
+   */
+  public hasMemberSchema(member: string): boolean {
+    if (this.isStructSchema()) {
+      const struct = this.getSchema() as StructureSchema;
+      return member in struct.members;
+    }
+    return false;
+  }
+
   /**
    * @returns the NormalizedSchema for the given member name. The returned instance will return true for `isMemberSchema()`
    * and will have the member name given.
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/schema/SchemaTraitFilterIndex.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/schema/SchemaTraitFilterIndex.java
@@ -74,8 +74,8 @@ final class SchemaTraitFilterIndex implements KnowledgeIndex {
             ErrorTrait.ID, // set by the ServiceException runtime classes.
             RequiresLengthTrait.ID, // unhandled
 
-            EventHeaderTrait.ID,
-            EventPayloadTrait.ID,
+            EventHeaderTrait.ID, // @smithy/core/EventStreams::EventStreamSerde
+            EventPayloadTrait.ID, // @smithy/core/EventStreams::EventStreamSerde
 
             // afaict, HttpErrorTrait is ignored by the client. The discriminator selects the error structure
             // but the actual HTTP response status code is used with no particular comparison
