fix: bugfix for setting JSON keys with special characters

Author: stainless-app[bot]

internal/apijson/encoder.go

@@ -16,6 +16,10 @@ import (
 
 var encoders sync.Map // map[encoderEntry]encoderFunc
 
+// If we want to set a literal key value into JSON using sjson, we need to make sure it doesn't have
+// special characters that sjson interprets as a path.
+var EscapeSJSONKey = strings.NewReplacer("\\", "\\\\", "|", "\\|", "#", "\\#", "@", "\\@", "*", "\\*", ".", "\\.", ":", "\\:", "?", "\\?").Replace
+
 func Marshal(value any) ([]byte, error) {
 	e := &encoder{dateFormat: time.RFC3339}
 	return e.marshal(value)
@@ -270,7 +274,7 @@ func (e *encoder) newStructTypeEncoder(t reflect.Type) encoderFunc {
 			if encoded == nil {
 				continue
 			}
-			json, err = sjson.SetRawBytes(json, ef.tag.name, encoded)
+			json, err = sjson.SetRawBytes(json, EscapeSJSONKey(ef.tag.name), encoded)
 			if err != nil {
 				return nil, err
 			}
@@ -348,7 +352,7 @@ func (e *encoder) encodeMapEntries(json []byte, v reflect.Value) ([]byte, error)
 			}
 			encodedKeyString = string(encodedKeyBytes)
 		}
-		encodedKey := []byte(sjsonReplacer.Replace(encodedKeyString))
+		encodedKey := []byte(encodedKeyString)
 		pairs = append(pairs, mapPair{key: encodedKey, value: iter.Value()})
 	}
 
@@ -366,7 +370,7 @@ func (e *encoder) encodeMapEntries(json []byte, v reflect.Value) ([]byte, error)
 		if len(encodedValue) == 0 {
 			continue
 		}
-		json, err = sjson.SetRawBytes(json, string(p.key), encodedValue)
+		json, err = sjson.SetRawBytes(json, EscapeSJSONKey(string(p.key)), encodedValue)
 		if err != nil {
 			return nil, err
 		}
@@ -386,7 +390,3 @@ func (e *encoder) newMapEncoder(_ reflect.Type) encoderFunc {
 		return json, nil
 	}
 }
-
-// If we want to set a literal key value into JSON using sjson, we need to make sure it doesn't have
-// special characters that sjson interprets as a path.
-var sjsonReplacer *strings.Replacer = strings.NewReplacer(".", "\\.", ":", "\\:", "*", "\\*")

internal/apijson/union.go

@@ -78,7 +78,7 @@ func (d *decoderBuilder) newStructUnionDecoder(t reflect.Type) decoderFunc {
 
 	return func(n gjson.Result, v reflect.Value, state *decoderState) error {
 		if discriminated && n.Type == gjson.JSON && len(unionEntry.discriminatorKey) != 0 {
-			discriminator := n.Get(unionEntry.discriminatorKey).Value()
+			discriminator := n.Get(EscapeSJSONKey(unionEntry.discriminatorKey)).Value()
 			for _, decoder := range discriminatedDecoders {
 				if discriminator == decoder.discriminator {
 					inner := v.FieldByIndex(decoder.field.Index)
@@ -162,7 +162,7 @@ func (d *decoderBuilder) newUnionDecoder(t reflect.Type) decoderFunc {
 			}
 
 			if len(unionEntry.discriminatorKey) != 0 {
-				discriminatorValue := n.Get(unionEntry.discriminatorKey).Value()
+				discriminatorValue := n.Get(EscapeSJSONKey(unionEntry.discriminatorKey)).Value()
 				if discriminatorValue == variant.DiscriminatorValue {
 					inner := reflect.New(variant.Type).Elem()
 					err := decoder(n, inner, state)

packages/param/encoder.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"reflect"
+	"strings"
 	"time"
 
 	shimjson "github.com/llamastack/llama-stack-client-go/internal/encoding/json"
@@ -14,6 +15,10 @@ import (
 // EncodedAsDate is not be stable and shouldn't be relied upon
 type EncodedAsDate Opt[time.Time]
 
+// If we want to set a literal key value into JSON using sjson, we need to make sure it doesn't have
+// special characters that sjson interprets as a path.
+var EscapeSJSONKey = strings.NewReplacer("\\", "\\\\", "|", "\\|", "#", "\\#", "@", "\\@", "*", "\\*", ".", "\\.", ":", "\\:", "?", "\\?").Replace
+
 type forceOmit int
 
 func (m EncodedAsDate) MarshalJSON() ([]byte, error) {
@@ -52,7 +57,7 @@ func MarshalWithExtras[T ParamStruct, R any](f T, underlying any, extras map[str
 				}
 				continue
 			}
-			bytes, err = sjson.SetBytes(bytes, k, v)
+			bytes, err = sjson.SetBytes(bytes, EscapeSJSONKey(k), v)
 			if err != nil {
 				return nil, err
 			}