feat(client): allow overriding unions

Author: stainless-app[bot]

packages/param/encoder.go

@@ -68,7 +68,7 @@ func MarshalWithExtras[T ParamStruct, R any](f T, underlying any, extras map[str
 // MarshalUnion uses a shimmed 'encoding/json' from Go 1.24, to support the 'omitzero' tag
 //
 // Stability for the API of MarshalUnion is not guaranteed.
-func MarshalUnion[T any](variants ...any) ([]byte, error) {
+func MarshalUnion[T ParamStruct](metadata T, variants ...any) ([]byte, error) {
 	nPresent := 0
 	presentIdx := -1
 	for i, variant := range variants {
@@ -78,6 +78,9 @@ func MarshalUnion[T any](variants ...any) ([]byte, error) {
 		}
 	}
 	if nPresent == 0 || presentIdx == -1 {
+		if ovr, ok := metadata.Overrides(); ok {
+			return shimjson.Marshal(ovr)
+		}
 		return []byte(`null`), nil
 	} else if nPresent > 1 {
 		return nil, &json.MarshalerError{

packages/param/encoder_test.go

@@ -218,10 +218,11 @@ func TestExtraFieldsForceOmitted(t *testing.T) {
 type UnionWithDates struct {
 	OfDate param.Opt[time.Time]
 	OfTime param.Opt[time.Time]
+	param.APIUnion
 }
 
 func (r UnionWithDates) MarshalJSON() (data []byte, err error) {
-	return param.MarshalUnion[UnionWithDates](param.EncodedAsDate(r.OfDate), r.OfTime)
+	return param.MarshalUnion(r, param.EncodedAsDate(r.OfDate), r.OfTime)
 }
 
 func TestUnionDateMarshal(t *testing.T) {
@@ -324,3 +325,41 @@ func TestOptionalInterfaceAssignability(t *testing.T) {
 
 	notOpt.implOpt() // silence the warning
 }
+
+type PrimitiveUnion struct {
+	OfString param.Opt[string]
+	OfInt    param.Opt[int]
+	param.APIUnion
+}
+
+func (p PrimitiveUnion) MarshalJSON() (data []byte, err error) {
+	return param.MarshalUnion(p, p.OfString, p.OfInt)
+}
+
+func TestOverriddenUnion(t *testing.T) {
+	tests := map[string]struct {
+		value    PrimitiveUnion
+		expected string
+	}{
+		"string": {
+			param.Override[PrimitiveUnion](json.RawMessage(`"hello"`)),
+			`"hello"`,
+		},
+		"int": {
+			param.Override[PrimitiveUnion](json.RawMessage(`42`)),
+			`42`,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			b, err := json.Marshal(test.value)
+			if err != nil {
+				t.Fatalf("didn't expect error %v, expected %s", err, test.expected)
+			}
+			if string(b) != test.expected {
+				t.Fatalf("expected %s, received %s", test.expected, string(b))
+			}
+		})
+	}
+}