diff --git a/src/replit_river/codegen/client.py b/src/replit_river/codegen/client.py
index a760acee..977c85a3 100644
--- a/src/replit_river/codegen/client.py
+++ b/src/replit_river/codegen/client.py
@@ -342,7 +342,7 @@ def flatten_union(tpe: RiverType) -> list[RiverType]:
                         )
                     typeddict_encoder.append(
                         f"""
-                            if x[{repr(discriminator_name)}]
+                            if x.get({repr(discriminator_name)})
                             == {repr(discriminator_value)}
                             else
                         """,
diff --git a/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py b/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py
index 0f0b0d8b..81e98a6d 100644
--- a/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py
+++ b/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py
@@ -75,7 +75,7 @@ def encode_NeedsenumobjectInput(
         encode_NeedsenumobjectInputOneOf_in_first(
             cast("NeedsenumobjectInputOneOf_in_first", x)
         )
-        if x["kind"] == "in_first"
+        if x.get("kind") == "in_first"
         else encode_NeedsenumobjectInputOneOf_in_second(
             cast("NeedsenumobjectInputOneOf_in_second", x)
         )