feat(client): optimize json encoder for internal types

Author: stainless-app[bot]

internal/encoding/json/encode.go

@@ -173,15 +173,21 @@ import (
 // JSON cannot represent cyclic data structures and Marshal does not
 // handle them. Passing cyclic structures to Marshal will result in
 // an error.
-func Marshal(v any) ([]byte, error) {
+// EDIT(begin): add optimization options
+func Marshal(v any, opts ...Option) ([]byte, error) {
+	// EDIT(end): add optimization options
 	e := newEncodeState()
 	defer encodeStatePool.Put(e)
 
-	// SHIM(begin): don't escape HTML by default
-	err := e.marshal(v, encOpts{escapeHTML: shims.EscapeHTMLByDefault})
+	// EDIT(begin): don't escape HTML by default, and apply options
+	encOpts := encOpts{escapeHTML: shims.EscapeHTMLByDefault}
+	if opts != nil {
+		encOpts = encOpts.apply(opts...)
+	}
+	err := e.marshal(v, encOpts)
 	// ORIGINAL:
 	//  err := e.marshal(v, encOpts{escapeHTML: true})
-	// SHIM(end)
+	// EDIT(end)
 	if err != nil {
 		return nil, err
 	}
@@ -352,6 +358,9 @@ type encOpts struct {
 	// EDIT(begin): save the timefmt
 	timefmt string
 	// EDIT(end)
+	// EDIT(begin): add optimization to skip compaction
+	skipCompaction bool
+	// EDIT(end)
 }
 
 type encoderFunc func(e *encodeState, v reflect.Value, opts encOpts)
@@ -483,7 +492,7 @@ func marshalerEncoder(e *encodeState, v reflect.Value, opts encOpts) {
 	if err == nil {
 		e.Grow(len(b))
 		out := e.AvailableBuffer()
-		out, err = appendCompact(out, b, opts.escapeHTML)
+		out, err = appendCompact(out, b, opts)
 		e.Buffer.Write(out)
 	}
 	if err != nil {
@@ -509,7 +518,7 @@ func addrMarshalerEncoder(e *encodeState, v reflect.Value, opts encOpts) {
 	if err == nil {
 		e.Grow(len(b))
 		out := e.AvailableBuffer()
-		out, err = appendCompact(out, b, opts.escapeHTML)
+		out, err = appendCompact(out, b, opts)
 		e.Buffer.Write(out)
 	}
 	if err != nil {

internal/encoding/json/indent.go

@@ -4,7 +4,9 @@
 
 package json
 
-import "bytes"
+import (
+	"bytes"
+)
 
 // HTMLEscape appends to dst the JSON-encoded src with <, >, &, U+2028 and U+2029
 // characters inside string literals changed to \u003c, \u003e, \u0026, \u2028, \u2029
@@ -41,12 +43,21 @@ func appendHTMLEscape(dst, src []byte) []byte {
 func Compact(dst *bytes.Buffer, src []byte) error {
 	dst.Grow(len(src))
 	b := dst.AvailableBuffer()
-	b, err := appendCompact(b, src, false)
+	b, err := appendCompact(b, src, encOpts{})
 	dst.Write(b)
 	return err
 }
 
-func appendCompact(dst, src []byte, escape bool) ([]byte, error) {
+func appendCompact(dst, src []byte, opts encOpts) ([]byte, error) {
+	// EDIT(begin): optimize for skipCompaction
+	if opts.skipCompaction {
+		dst = append(dst, src...)
+		return dst, nil
+	}
+
+	escape := opts.escapeHTML
+	// EDIT(end)
+
 	origLen := len(dst)
 	scan := newScanner()
 	defer freeScanner(scan)

internal/encoding/json/opt.go

@@ -0,0 +1,24 @@
+// EDIT(begin): add custom options for JSON encoding
+package json
+
+type Option func(*encOpts)
+
+// Every time a sub-type of [json.Marshaler] is encountered,
+// skip a redundant and costly compaction step, trust it to self-compact.
+//
+// This is a divergence from the standard library behavior, and is only guaranteed
+// safe with SDK types.
+func WithSkipCompaction(b bool) Option {
+	return func(eos *encOpts) {
+		eos.skipCompaction = true
+	}
+}
+
+func (eos encOpts) apply(opts ...Option) encOpts {
+	for _, opt := range opts {
+		opt(&eos)
+	}
+	return eos
+}
+
+// EDIT(end)

internal/encoding/json/stream.go

@@ -6,7 +6,6 @@ package json
 
 import (
 	"bytes"
-	"errors"
 	"io"
 )
 
@@ -253,30 +252,34 @@ func (enc *Encoder) SetEscapeHTML(on bool) {
 	enc.escapeHTML = on
 }
 
-// RawMessage is a raw encoded JSON value.
-// It implements [Marshaler] and [Unmarshaler] and can
-// be used to delay JSON decoding or precompute a JSON encoding.
-type RawMessage []byte
-
-// MarshalJSON returns m as the JSON encoding of m.
-func (m RawMessage) MarshalJSON() ([]byte, error) {
-	if m == nil {
-		return []byte("null"), nil
-	}
-	return m, nil
-}
-
-// UnmarshalJSON sets *m to a copy of data.
-func (m *RawMessage) UnmarshalJSON(data []byte) error {
-	if m == nil {
-		return errors.New("json.RawMessage: UnmarshalJSON on nil pointer")
-	}
-	*m = append((*m)[0:0], data...)
-	return nil
-}
-
-var _ Marshaler = (*RawMessage)(nil)
-var _ Unmarshaler = (*RawMessage)(nil)
+// EDIT(begin): remove RawMessage
+//
+// // RawMessage is a raw encoded JSON value.
+// // It implements [Marshaler] and [Unmarshaler] and can
+// // be used to delay JSON decoding or precompute a JSON encoding.
+// type RawMessage []byte
+//
+// // MarshalJSON returns m as the JSON encoding of m.
+// func (m RawMessage) MarshalJSON() ([]byte, error) {
+// 	if m == nil {
+// 		return []byte("null"), nil
+// 	}
+// 	return m, nil
+// }
+//
+// // UnmarshalJSON sets *m to a copy of data.
+// func (m *RawMessage) UnmarshalJSON(data []byte) error {
+// 	if m == nil {
+// 		return errors.New("json.RawMessage: UnmarshalJSON on nil pointer")
+// 	}
+// 	*m = append((*m)[0:0], data...)
+// 	return nil
+// }
+//
+// var _ Marshaler = (*RawMessage)(nil)
+// var _ Unmarshaler = (*RawMessage)(nil)
+//
+// EDIT(end)
 
 // A Token holds a value of one of these types:
 //

internal/encoding/json/time.go

@@ -50,7 +50,7 @@ func timeMarshalEncoder(e *encodeState, v reflect.Value, opts encOpts) bool {
 	if b != nil {
 		e.Grow(len(b))
 		out := e.AvailableBuffer()
-		out, _ = appendCompact(out, b, opts.escapeHTML)
+		out, _ = appendCompact(out, b, opts)
 		e.Buffer.Write(out)
 		return true
 	}

packages/param/encoder.go

@@ -66,7 +66,7 @@ func MarshalWithExtras[T ParamStruct, R any](f T, underlying any, extras map[str
 	} else if ovr, ok := f.Overrides(); ok {
 		return shimjson.Marshal(ovr)
 	} else {
-		return shimjson.Marshal(underlying)
+		return shimjson.Marshal(underlying, shimjson.WithSkipCompaction(true))
 	}
 }
 
@@ -96,7 +96,7 @@ func MarshalUnion[T ParamStruct](metadata T, variants ...any) ([]byte, error) {
 			Err:  fmt.Errorf("expected union to have only one present variant, got %d", nPresent),
 		}
 	}
-	return shimjson.Marshal(variants[presentIdx])
+	return shimjson.Marshal(variants[presentIdx], shimjson.WithSkipCompaction(true))
 }
 
 // typeFor is shimmed from Go 1.23 "reflect" package

packages/param/encoder_test.go

@@ -1,10 +1,13 @@
 package param_test
 
 import (
+	"bytes"
 	"encoding/json"
+	"reflect"
 	"testing"
 	"time"
 
+	shimjson "github.com/browserbase/stagehand-go/v3/internal/encoding/json"
 	"github.com/browserbase/stagehand-go/v3/packages/param"
 )
 
@@ -375,3 +378,176 @@ func TestNullStructUnion(t *testing.T) {
 		t.Fatalf("expected null, received %s", string(b))
 	}
 }
+
+//
+// Compaction optimization
+//
+
+type NonCompactedDoubleParent struct {
+	Prop   string             `json:"prop"`
+	Parent NonCompactedParent `json:"parent"`
+
+	param.APIObject
+}
+
+type NonCompactedParent struct {
+	BadChild NonCompacted `json:"bad_child"`
+
+	param.APIObject
+}
+
+type NonCompacted struct {
+	Raw string
+
+	param.APIObject
+}
+
+func (a NonCompactedDoubleParent) MarshalJSON() ([]byte, error) {
+	type shadow NonCompactedDoubleParent
+	return param.MarshalObject(a, (*shadow)(&a))
+}
+
+func (a NonCompactedParent) MarshalJSON() ([]byte, error) {
+	type shadow NonCompactedParent
+	return param.MarshalObject(a, (*shadow)(&a))
+}
+
+func (a NonCompacted) MarshalJSON() ([]byte, error) {
+	if a.Raw == "" {
+		a.Raw = nonCompactedRaw
+	}
+	return []byte(a.Raw), nil
+}
+
+var nonCompactedRaw string = ` { "foo": "bar" } `
+
+func TestAppendCompactBroken(t *testing.T) {
+	tests := map[string]struct {
+		value json.Marshaler
+	}{
+		"red/illegal-json": {
+			NonCompacted{Raw: `{ "broken": "json" `},
+		},
+		"red/nested-with-illegal-json": {
+			NonCompactedParent{BadChild: NonCompacted{
+				Raw: `{ "broken": "json" `,
+			}},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			v, err := json.Marshal(test.value)
+			if err == nil {
+				t.Fatal("expected error got", v)
+			}
+		})
+	}
+}
+
+// TestAppendCompact validates an optimization for internal SDK types to
+// avoid O(keys^2) iteration over each JSON object.
+//
+// It's possible to intentionally trigger this behavior as both a user and
+// SDK developer. However, the edge case is quite pathological and requires
+// calling [json.Marshaler.MarshalJSON] rather than [json.Marshal].
+func TestAppendCompact(t *testing.T) {
+
+	tests := map[string]struct {
+		value    json.Marshaler
+		expected string
+	}{
+		//
+		// Non-compacted cases
+		//
+		// Note this is how to exploit the compacter to fail, you must call [json.Marshaler.MarshalJSON] rather than [json.Marshal].
+		// The type must also embed [param.APIObject] and return non-compacted JSON.
+		//
+
+		"no-compact/fails-compaction": {
+			NonCompacted{Raw: nonCompactedRaw},
+			nonCompactedRaw,
+		},
+		"no-compact/nested-with-bad-child": {
+			NonCompactedParent{BadChild: NonCompacted{
+				Raw: nonCompactedRaw,
+			}},
+			`{"bad_child":` + nonCompactedRaw + `}`,
+		},
+		"no-compact/double-nested-with-bad-child": {
+			NonCompactedDoubleParent{Prop: "1", Parent: NonCompactedParent{BadChild: NonCompacted{
+				Raw: nonCompactedRaw,
+			}}},
+			`{"prop":"1","parent":{"bad_child":` + nonCompactedRaw + `}}`,
+		},
+
+		//
+		// Compacted cases
+		//
+
+		"override/spaces-within": {
+			param.Override[NonCompactedDoubleParent](json.RawMessage(`{"com": "pact"}`)),
+			`{"com":"pact"}`,
+		},
+		"override/spaces-after": {
+			param.Override[NonCompactedDoubleParent](json.RawMessage(`{"com":"pact"}  `)),
+			`{"com":"pact"}`,
+		},
+		"override/spaces-before": {
+			param.Override[NonCompactedDoubleParent](json.RawMessage(`  {"com":"pact"}`)),
+			`{"com":"pact"}`,
+		},
+		"override/spaces-around": {
+			param.Override[NonCompactedDoubleParent](json.RawMessage(` { "com": "pact"   }`)),
+			`{"com":"pact"}`,
+		},
+		"override/override-with-nested": {
+			param.Override[NonCompactedDoubleParent](NonCompactedParent{}),
+			`{"bad_child":{"foo":"bar"}}`,
+		},
+		"override/override-with-non-compacted": {
+			param.Override[NonCompactedDoubleParent](NonCompacted{}),
+			`{"foo":"bar"}`,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name+"/marshal-json", func(t *testing.T) {
+			b, err := test.value.MarshalJSON()
+			if err != nil {
+				t.Fatalf("didn't expect error %v, expected %s", err, test.expected)
+			}
+			if string(b) != test.expected {
+				t.Fatalf("expected %s (%s), received %s", test.expected, reflect.TypeOf(test.value), string(b))
+			}
+		})
+
+		t.Run(name+"/json-marshal", 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)
+			}
+
+			// expected output of JSON Marshal should always be compacted
+			var compactedExpected bytes.Buffer
+			err = json.Compact(&compactedExpected, []byte(test.expected))
+			if err != nil {
+				t.Fatalf("didn't expect error %v, expected %s", err, test.expected)
+			}
+
+			if string(b) != compactedExpected.String() {
+				t.Fatalf("expected %s (%s), received %s", test.expected, reflect.TypeOf(test.value), string(b))
+			}
+		})
+
+		t.Run(name+"/shimjson-marshal", func(t *testing.T) {
+			b, err := shimjson.Marshal(test.value)
+			if err != nil {
+				t.Fatalf("didn't expect error %v, expected %s", err, test.expected)
+			}
+			if string(b) != test.expected {
+				t.Logf("expected %s (%s), received %s", test.expected, reflect.TypeOf(test.value), string(b))
+			}
+		})
+	}
+}