From 06a9683896abfa1cd78b78c151dd5d6b724f24c3 Mon Sep 17 00:00:00 2001 From: Mohammed Ehab Date: Wed, 4 Mar 2026 12:42:41 +0000 Subject: [PATCH 1/5] Allow ClientContext.Custom unmarshaling for non-string (JSON) values --- lambda/invoke_loop_test.go | 40 ++++++++++++++++++++++++++++++++++++++ lambdacontext/context.go | 30 ++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/lambda/invoke_loop_test.go b/lambda/invoke_loop_test.go index af0e0c0c..c91df347 100644 --- a/lambda/invoke_loop_test.go +++ b/lambda/invoke_loop_test.go @@ -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) { diff --git a/lambdacontext/context.go b/lambdacontext/context.go index d75d8282..58dc5e99 100644 --- a/lambdacontext/context.go +++ b/lambdacontext/context.go @@ -11,6 +11,7 @@ package lambdacontext import ( "context" + "encoding/json" "os" "strconv" ) @@ -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 { + var raw struct { + Client ClientApplication `json:"Client"` + 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 { + cc.Custom[k] = string(v) + } + } + } + return nil +} + // CognitoIdentity is the cognito identity used by the calling application. type CognitoIdentity struct { CognitoIdentityID string From 3de2206e228122fc53eea9c866d7bd3fc5be89fa Mon Sep 17 00:00:00 2001 From: Mohammed Ehab Date: Wed, 4 Mar 2026 12:51:23 +0000 Subject: [PATCH 2/5] Lint --- lambda/invoke_loop_test.go | 6 +++--- lambdacontext/context.go | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/lambda/invoke_loop_test.go b/lambda/invoke_loop_test.go index c91df347..4c305564 100644 --- a/lambda/invoke_loop_test.go +++ b/lambda/invoke_loop_test.go @@ -389,7 +389,7 @@ func TestContextDeserializationErrors(t *testing.T) { }`, string(record.responses[2])) } - func TestClientContextWithNestedCustomValues(t *testing.T) { +func TestClientContextWithNestedCustomValues(t *testing.T) { metadata := defaultInvokeMetadata() metadata.clientContext = `{ "Client": { @@ -403,7 +403,7 @@ func TestContextDeserializationErrors(t *testing.T) { "bedrockAgentCorePropagatedHeaders": {"x-id": "my-custom-id"} } }` - + ts, record := runtimeAPIServer(`{}`, 1, metadata) defer ts.Close() handler := NewHandler(func(ctx context.Context) (interface{}, error) { @@ -412,7 +412,7 @@ func TestContextDeserializationErrors(t *testing.T) { }) endpoint := strings.Split(ts.URL, "://")[1] _ = startRuntimeAPILoop(endpoint, handler) - + expected := `{ "Client": { "installation_id": "install1", diff --git a/lambdacontext/context.go b/lambdacontext/context.go index 58dc5e99..f3e70399 100644 --- a/lambdacontext/context.go +++ b/lambdacontext/context.go @@ -75,7 +75,7 @@ type ClientContext struct { // their JSON string representation. func (cc *ClientContext) UnmarshalJSON(data []byte) error { var raw struct { - Client ClientApplication `json:"Client"` + Client ClientApplication `json:"Client"` Env map[string]string `json:"env"` Custom map[string]json.RawMessage `json:"custom"` } From f9b3690c20ccc107d54130ff167bfef5f1c6b104 Mon Sep 17 00:00:00 2001 From: Mohammed Ehab Date: Wed, 11 Mar 2026 11:06:41 +0000 Subject: [PATCH 3/5] Adding Tests --- lambdacontext/context_test.go | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 lambdacontext/context_test.go diff --git a/lambdacontext/context_test.go b/lambdacontext/context_test.go new file mode 100644 index 00000000..82d204bd --- /dev/null +++ b/lambdacontext/context_test.go @@ -0,0 +1,15 @@ +package lambdacontext + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestClientContextUnmarshalJSON_InvalidJSON(t *testing.T) { + var cc ClientContext + err := json.Unmarshal([]byte(`not valid json`), &cc) + assert.Error(t, err) + assert.Empty(t, cc.Custom) +} From fdcc6b4d522389dea904a96cae1e3f53b372ed2a Mon Sep 17 00:00:00 2001 From: Mohammed Ehab Date: Wed, 11 Mar 2026 11:22:23 +0000 Subject: [PATCH 4/5] More tests --- lambdacontext/context_test.go | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/lambdacontext/context_test.go b/lambdacontext/context_test.go index 82d204bd..28266373 100644 --- a/lambdacontext/context_test.go +++ b/lambdacontext/context_test.go @@ -5,11 +5,31 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) -func TestClientContextUnmarshalJSON_InvalidJSON(t *testing.T) { - var cc ClientContext - err := json.Unmarshal([]byte(`not valid json`), &cc) - assert.Error(t, err) - assert.Empty(t, cc.Custom) +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) + }) } From d1c21107bcbfc44be13dfb9f7cd9b4dceb3aa744 Mon Sep 17 00:00:00 2001 From: Mohammed Ehab Date: Wed, 11 Mar 2026 11:28:43 +0000 Subject: [PATCH 5/5] empty commit to retrigger flakey tests