Skip to content

Allow ClientContext.Custom unmarshaling for non-string (JSON) values #620

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
M-Elsaeed wants to merge 5 commits into
base: main
Choose a base branch
from
+105 −0
Open

Allow ClientContext.Custom unmarshaling for non-string (JSON) values #620

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
06a9683
Allow ClientContext.Custom unmarshaling for non-string (JSON) values
Mar 4, 2026
3de2206
Lint
Mar 4, 2026
f9b3690
Adding Tests
Mar 11, 2026
fdcc6b4
More tests
Mar 11, 2026
d1c2110
empty commit to retrigger flakey tests
Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions lambda/invoke_loop_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,46 @@ func TestContextDeserializationErrors(t *testing.T) {
}`, string(record.responses[2]))
}

func TestClientContextWithNestedCustomValues(t *testing.T) {
metadata := defaultInvokeMetadata()
metadata.clientContext = `{
"Client": {
"app_title": "test",
"installation_id": "install1",
"app_version_code": "1.0",
"app_package_name": "com.test"
},
"custom": {
"bedrockAgentCoreTargetId": "target-123",
"bedrockAgentCorePropagatedHeaders": {"x-id": "my-custom-id"}
}
}`

ts, record := runtimeAPIServer(`{}`, 1, metadata)
defer ts.Close()
handler := NewHandler(func(ctx context.Context) (interface{}, error) {
lc, _ := lambdacontext.FromContext(ctx)
return lc.ClientContext, nil
})
endpoint := strings.Split(ts.URL, "://")[1]
_ = startRuntimeAPILoop(endpoint, handler)

expected := `{
"Client": {
"installation_id": "install1",
"app_title": "test",
"app_version_code": "1.0",
"app_package_name": "com.test"
},
"env": null,
"custom": {
"bedrockAgentCoreTargetId": "target-123",
"bedrockAgentCorePropagatedHeaders": "{\"x-id\": \"my-custom-id\"}"
}
}`
assert.JSONEq(t, expected, string(record.responses[0]))
}

type invalidPayload struct{}

func (invalidPayload) MarshalJSON() ([]byte, error) {
Expand Down
30 changes: 30 additions & 0 deletions lambdacontext/context.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ package lambdacontext

import (
"context"
"encoding/json"
"os"
"strconv"
)
Expand Down Expand Up @@ -68,6 +69,35 @@ type ClientContext struct {
Custom map[string]string `json:"custom"`
}

// UnmarshalJSON implements custom JSON unmarshaling for ClientContext.
// This handles the case where values in the "custom" map are not strings
// (e.g. nested JSON objects), by serializing non-string values back to
// their JSON string representation.
func (cc *ClientContext) UnmarshalJSON(data []byte) error {
Copy link
Collaborator

@bmoffatt bmoffatt Mar 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thinking there may need to be a MarshalJSON too, so that a re-serialization produces the same output as what was input. (number, bool values. quote escaping in nested json)

Copy link
Collaborator

@bmoffatt bmoffatt Mar 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see: https://go.dev/play/p/8q6bgWa-Rcw where re marshaling results in a type change

// You can edit this code!
// Click here and start typing.
package main

import (
	"encoding/json"
	"fmt"
	"log"
	"log/slog"
	"os"
)

type X struct {
	Custom map[string]string
}

func (x *X) UnmarshalJSON(b []byte) error {
	var m struct{ Custom map[string]json.RawMessage }
	if err := json.Unmarshal(b, &m); err != nil {
		return err
	}
	x.Custom = make(map[string]string, len(m.Custom))
	for k, v := range m.Custom {
		var s string
		if err := json.Unmarshal(v, &s); err != nil {
			slog.Warn("uhoh", "err", err, "key", k, "val", v)
			s = string(v)
		}
		x.Custom[k] = s
	}
	return nil
}

func main() {
	input := []byte(`{"Custom":{"hello":9001}}`)
	fmt.Println("Input:")
	os.Stdout.Write(input)
	fmt.Println("")

	var x X
	if err := json.Unmarshal(input, &x); err != nil {
		log.Fatal(err)
	}

	output, err := json.Marshal(x)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Output:")
	os.Stdout.Write(output)
}

Input:
{"Custom":{"hello":9001}}
2009/11/10 23:00:00 WARN uhoh err="json: cannot unmarshal number into Go value of type string" key=hello val="9001"
Output:
{"Custom":{"hello":"9001"}}
Program exited.

Copy link
Collaborator

@bmoffatt bmoffatt Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

interesting. Java and .NET already be mangling this with string conversions for numbers and bools. And .NET only will stringify nested values.

Client Context custom — Return Value Test Matrix

All functions return context.clientContext.custom as JSON.

Test 1: {"custom": {"number": 9001, "bool": false}}

Runtime Result
Java ✅ {"number":"9001","bool":"false"}
.NET ✅ {"number":"9001","bool":"False"}
Rust ❌ Runtime crash (logs: invalid type: integer '9001', expected a string)
Go ❌ cannot unmarshal number into Go struct field ClientContext.custom of type string

Test 2: {"custom": {"nested": {"hello": "world"}}}

Runtime Result
Java ❌ Expected a string but was BEGIN_OBJECT
.NET ✅ {"nested":"{\"hello\":\"world\"}"}
Rust ❌ Runtime crash (logs: invalid type: map, expected a string)
Go ❌ cannot unmarshal object into Go struct field ClientContext.custom of type string

var raw struct {
Client ClientApplication `json:"Client"`
Copy link
Collaborator

@bmoffatt bmoffatt Mar 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

typo, should be "client". Well actually, no struct tag. without the json: struct tag, the current ClientContext is lenient to both "Custom" and "custom". The mobile SDKs I belive all strictly send "custom", and .NET, java, Rust, all drop "Custom", so IDK why this was left lenient 🤷‍♂️ - but I guess should also stay lenient for back compat

nevermind

When parsing a JSON object into a Go struct, keys are considered in a case-insensitive fashion.

Env map[string]string `json:"env"`
Custom map[string]json.RawMessage `json:"custom"`
}
if err := json.Unmarshal(data, &raw); err != nil {
return err
}
cc.Client = raw.Client
cc.Env = raw.Env
if raw.Custom != nil {
cc.Custom = make(map[string]string, len(raw.Custom))
for k, v := range raw.Custom {
var s string
if err := json.Unmarshal(v, &s); err == nil {
cc.Custom[k] = s
} else {
Copy link
Collaborator

@bmoffatt bmoffatt Mar 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the nested json in the Client Context like, critical to all use of the Bedrock Agentcore Gateway service? Rather than stringifying, another option is to drop non-string values.

var raw struct { Custom map[string]any }
json.Unmarshal(b, &raw)
for k, v := range raw.Custom {
    if s, ok := v.(string); ok {
        cc.Custom[k] = s
    }
}

Bringing this option up, so as not to be hasty in type coercion decisions that could introduce inconsistencies with other runtimes.

Copy link

@charles-d-burton charles-d-burton Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you elaborate on why you would drop types other than a string? It seems possible that this would bite you down the road if another service was enabled that needed this but had values other than strings.

Copy link
Collaborator

@bmoffatt bmoffatt Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Balancing urgency in (partially?) unblocking a use-case, and introducing warts that'd be hard to undo without semver bumps. Other runtimes have the same blocker, and it'd be prudent to line up a similar strategy for supporting the usecase.

For Go, other options that solve for making the context parse more lenient, without committing to inconsistent type coercion for this novel use of client context:

  • Assign LambdaContext.ClientContext.Custom to nil if parse fails, slog.Warn the error.
  • Move the LambdaContext parse step from pre-invoke to being done lazily only when handlers fetch it using lambdacontext.FromContext(ctx) (and logging errors, returning nil, false)

An option to make the novel use of Client Context readable in a way consistent with node, python, (eg: no type coercion) could be something like stuffing the []byte from the client context header into the context.Context pre-invoke. Combine with swallowing the parse error (set Custom to nil), and a deprecation comment on the ClientContext.Custom field // Deprecated: Custom is 'nil' when client sends something other than a map[string]string. Use `lambdacontext.ClientContextCustomFromContext` instead

and adding new function

ClientContextCustomFromContext[T](ctx context.Context) (*T, error) {
  b, ok := ctx.Value("lambda-client-context-header-value").([]byte)
  if !ok {
      return nil errors.New("ClientContext not found in context)
  }
  var t T
  if err := json.Unmarshal(lc.ClientContext.CustomRaw, &t); err != nil {
      return nil, err  
  }
  return &t, nil
}

^ is my general recommendation to keep back compat while also unblocking the new use-case. But given that .NET and Java are already inconsistent this isn't necessarily the only solution.

cc.Custom[k] = string(v)
}
}
}
return nil
}

// CognitoIdentity is the cognito identity used by the calling application.
type CognitoIdentity struct {
CognitoIdentityID string
Expand Down
35 changes: 35 additions & 0 deletions lambdacontext/context_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
package lambdacontext

import (
"encoding/json"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func TestClientContextUnmarshalJSON(t *testing.T) {
t.Run("non-string custom values are serialized to string", func(t *testing.T) {
input := `{
"Client": {"installation_id": "install1"},
"custom": {
"key1": "stringval",
"key2": {"nested": "object"},
"key3": 42
}
}`
var cc ClientContext
err := json.Unmarshal([]byte(input), &cc)
require.NoError(t, err)
assert.Equal(t, "install1", cc.Client.InstallationID)
assert.Equal(t, "stringval", cc.Custom["key1"])
assert.JSONEq(t, `{"nested":"object"}`, cc.Custom["key2"])
assert.Equal(t, "42", cc.Custom["key3"])
})

t.Run("invalid JSON returns error", func(t *testing.T) {
var cc ClientContext
err := json.Unmarshal([]byte(`not valid json`), &cc)
assert.Error(t, err)
})
}
Loading