diff --git a/coderd/x/chatd/chatprompt/chatprompt.go b/coderd/x/chatd/chatprompt/chatprompt.go index 628494c7431..8188edc1fca 100644 --- a/coderd/x/chatd/chatprompt/chatprompt.go +++ b/coderd/x/chatd/chatprompt/chatprompt.go @@ -1175,18 +1175,21 @@ func sanitizeToolCallID(id string) string { } // MarshalParts encodes SDK chat message parts for persistence. -// NUL characters in string fields are encoded as PUA sentinel -// pairs (U+E000 U+E001) before marshaling so the resulting JSON -// never contains \u0000 (rejected by PostgreSQL jsonb). The -// encoding operates on Go string values, not JSON bytes, so it -// survives jsonb text normalization. +// NUL characters in free-form and opaque fields are encoded as PUA +// sentinel pairs (U+E000 U+E001) so PostgreSQL jsonb can store them. +// Structured fields such as identifiers, paths, and media types reject +// NUL because their downstream representations cannot consume it. func MarshalParts(parts []codersdk.ChatMessagePart) (pqtype.NullRawMessage, error) { if len(parts) == 0 { return pqtype.NullRawMessage{}, nil } - data, err := json.Marshal(encodeNulInParts(parts)) + encoded, err := encodeNulInParts(parts) if err != nil { - return pqtype.NullRawMessage{}, xerrors.Errorf("encode chat message parts: %w", err) + return pqtype.NullRawMessage{}, xerrors.Errorf("validate chat message parts: %w", err) + } + data, err := json.Marshal(encoded) + if err != nil { + return pqtype.NullRawMessage{}, xerrors.Errorf("marshal chat message parts: %w", err) } return pqtype.NullRawMessage{RawMessage: data, Valid: true}, nil } @@ -1797,16 +1800,12 @@ func decodeNulInValue(v any) any { } // encodeNulInJSON walks all string values (and keys) inside a -// json.RawMessage and applies encodeNulInString. Returns the -// original unchanged when the raw message does not contain NUL -// escapes or U+E000 bytes, or when parsing fails. +// json.RawMessage and applies encodeNulInString. It also normalizes +// unpaired UTF-16 surrogate escapes, which PostgreSQL jsonb rejects. +// Returns the original unchanged when no transformation is needed or +// when parsing fails. func encodeNulInJSON(raw json.RawMessage) json.RawMessage { - if len(raw) == 0 { - return raw - } - // Quick exit: no \u0000 escape and no U+E000 UTF-8 bytes. - if !bytes.Contains(raw, []byte(`\u0000`)) && - !bytes.Contains(raw, []byte{0xEE, 0x80, 0x80}) { + if len(raw) == 0 || !needsNulEncodingInJSON(raw) || !json.Valid(raw) { return raw } dec := json.NewDecoder(bytes.NewReader(raw)) @@ -1822,6 +1821,80 @@ func encodeNulInJSON(raw json.RawMessage) json.RawMessage { return result } +// needsNulEncodingInJSON reports whether valid JSON contains a string +// value or object key that encodeNulInJSON must transform. It only +// detects candidates; encoding/json owns parsing and normalization. +func needsNulEncodingInJSON(raw json.RawMessage) bool { + // encoding/json emits U+E000 literally, while json.RawMessage may + // contain either the literal UTF-8 bytes or a Unicode escape. + if bytes.Contains(raw, []byte{0xEE, 0x80, 0x80}) { + return true + } + + inString := false + for i := 0; i < len(raw); { + switch raw[i] { + case '"': + inString = !inString + i++ + case '\\': + if !inString || i+1 >= len(raw) { + i++ + continue + } + if raw[i+1] != 'u' { + i += 2 + continue + } + value, ok := parseJSONHexEscape(raw[i+2:]) + if !ok { + return true + } + switch { + case value == 0 || value == 0xE000: + return true + case 0xD800 <= value && value <= 0xDBFF: + if i+12 <= len(raw) && raw[i+6] == '\\' && raw[i+7] == 'u' { + low, lowOK := parseJSONHexEscape(raw[i+8:]) + if lowOK && 0xDC00 <= low && low <= 0xDFFF { + i += 12 + continue + } + } + return true + case 0xDC00 <= value && value <= 0xDFFF: + return true + default: + i += 6 + } + default: + i++ + } + } + return false +} + +func parseJSONHexEscape(raw []byte) (uint16, bool) { + if len(raw) < 4 { + return 0, false + } + var value uint16 + for _, c := range raw[:4] { + value <<= 4 + switch { + case '0' <= c && c <= '9': + value |= uint16(c - '0') + case 'a' <= c && c <= 'f': + value |= uint16(c-'a') + 10 + case 'A' <= c && c <= 'F': + value |= uint16(c-'A') + 10 + default: + return 0, false + } + } + return value, true +} + // decodeNulInJSON walks all string values (and keys) inside a // json.RawMessage and applies decodeNulInString. func decodeNulInJSON(raw json.RawMessage) json.RawMessage { @@ -1845,33 +1918,154 @@ func decodeNulInJSON(raw json.RawMessage) json.RawMessage { return result } -// encodeNulInParts returns a shallow copy of parts with all -// string and json.RawMessage fields NUL-encoded. The caller's -// slice is not modified. -func encodeNulInParts(parts []codersdk.ChatMessagePart) []codersdk.ChatMessagePart { +// nulPolicy classifies how MarshalParts treats NUL in one persisted +// ChatMessagePart field. +type nulPolicy int + +const ( + // nulEncode reversibly encodes NUL as PUA sentinel pairs. Used for + // free-form text and opaque provider data that must survive + // persistence byte-for-byte, e.g. Gemini binary thought signatures + // in ProviderMetadata. + nulEncode nulPolicy = iota + // nulReject fails MarshalParts with an actionable error. Used for + // structured values (enums, identifiers, paths, filenames, URLs, + // media types, OS labels, and parsed commands) whose downstream + // consumers cannot handle NUL, so persisting an encoded form would + // only hide invalid input. + nulReject +) + +// partNulField binds one string-bearing ChatMessagePart field to its +// NUL policy. Pointer accessors let encoding, decoding, and validation +// share the table. Exactly one accessor is set per entry; grid is only +// supported with nulReject and raw only with nulEncode. Unsupported +// combinations and missing entries for new string-bearing fields are +// caught by TestChatMessagePartNULCoverage. +type partNulField struct { + name string + policy nulPolicy + str func(*codersdk.ChatMessagePart) *string + raw func(*codersdk.ChatMessagePart) *json.RawMessage + grid func(*codersdk.ChatMessagePart) [][]string +} + +// partNulFields classifies every persisted string-bearing field of +// codersdk.ChatMessagePart. Adding a string-bearing field to the struct +// requires an entry here. +var partNulFields = []partNulField{ + {name: "Text", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.Text }}, + {name: "Args", policy: nulEncode, raw: func(p *codersdk.ChatMessagePart) *json.RawMessage { return &p.Args }}, + {name: "ArgsDelta", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.ArgsDelta }}, + {name: "Result", policy: nulEncode, raw: func(p *codersdk.ChatMessagePart) *json.RawMessage { return &p.Result }}, + {name: "ResultDelta", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.ResultDelta }}, + {name: "Title", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.Title }}, + {name: "Content", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.Content }}, + {name: "ProviderMetadata", policy: nulEncode, raw: func(p *codersdk.ChatMessagePart) *json.RawMessage { return &p.ProviderMetadata }}, + {name: "ContextFileContent", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileContent }}, + {name: "SkillDescription", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.SkillDescription }}, + {name: "Type", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return (*string)(&p.Type) }}, + {name: "ToolCallID", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ToolCallID }}, + {name: "ToolName", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ToolName }}, + {name: "ParsedCommands", policy: nulReject, grid: func(p *codersdk.ChatMessagePart) [][]string { return p.ParsedCommands }}, + {name: "SourceID", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.SourceID }}, + {name: "URL", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.URL }}, + {name: "MediaType", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.MediaType }}, + {name: "Name", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.Name }}, + {name: "FileName", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.FileName }}, + {name: "ContextFilePath", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFilePath }}, + {name: "ContextFileOS", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileOS }}, + {name: "ContextFileDirectory", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileDirectory }}, + {name: "ContextFileSkillMetaFile", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileSkillMetaFile }}, + {name: "SkillName", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.SkillName }}, + {name: "SkillDir", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.SkillDir }}, +} + +// encode applies the reversible NUL encoding to a nulEncode field. +func (f partNulField) encode(p *codersdk.ChatMessagePart) { + switch { + case f.str != nil: + *f.str(p) = encodeNulInString(*f.str(p)) + case f.raw != nil: + *f.raw(p) = encodeNulInJSON(*f.raw(p)) + } +} + +// decode reverses encode. +func (f partNulField) decode(p *codersdk.ChatMessagePart) { + switch { + case f.str != nil: + *f.str(p) = decodeNulInString(*f.str(p)) + case f.raw != nil: + *f.raw(p) = decodeNulInJSON(*f.raw(p)) + } +} + +// validateNulFree returns an actionable error when a nulReject field +// contains NUL. +func (f partNulField) validateNulFree(partIndex int, p *codersdk.ChatMessagePart) error { + switch { + case f.str != nil: + if nul := strings.IndexByte(*f.str(p), 0); nul >= 0 { + return xerrors.Errorf( + "chat message part %d field %s contains NUL at byte %d", + partIndex, + f.name, + nul, + ) + } + case f.grid != nil: + for i, values := range f.grid(p) { + for j, value := range values { + if nul := strings.IndexByte(value, 0); nul >= 0 { + return xerrors.Errorf( + "chat message part %d field %s[%d][%d] contains NUL at byte %d", + partIndex, + f.name, + i, + j, + nul, + ) + } + } + } + } + return nil +} + +// encodeNulInParts returns a copy of parts with every nulEncode field +// encoded and every nulReject field validated per partNulFields. The +// caller's values are not modified. +// +// Content version 1 cannot distinguish sentinel pairs written by this +// codec from matching natural text stored before a field was covered, so +// matching historical values in encoded fields decode as NUL. +func encodeNulInParts(parts []codersdk.ChatMessagePart) ([]codersdk.ChatMessagePart, error) { encoded := make([]codersdk.ChatMessagePart, len(parts)) copy(encoded, parts) for i := range encoded { p := &encoded[i] - p.Text = encodeNulInString(p.Text) - p.Content = encodeNulInString(p.Content) - p.Args = encodeNulInJSON(p.Args) - p.ArgsDelta = encodeNulInString(p.ArgsDelta) - p.Result = encodeNulInJSON(p.Result) - p.ResultDelta = encodeNulInString(p.ResultDelta) + for _, f := range partNulFields { + switch f.policy { + case nulEncode: + f.encode(p) + case nulReject: + if err := f.validateNulFree(i, p); err != nil { + return nil, err + } + } + } } - return encoded + return encoded, nil } // decodeNulInParts reverses encodeNulInParts in place. func decodeNulInParts(parts []codersdk.ChatMessagePart) { for i := range parts { - p := &parts[i] - p.Text = decodeNulInString(p.Text) - p.Content = decodeNulInString(p.Content) - p.Args = decodeNulInJSON(p.Args) - p.ArgsDelta = decodeNulInString(p.ArgsDelta) - p.Result = decodeNulInJSON(p.Result) - p.ResultDelta = decodeNulInString(p.ResultDelta) + for _, f := range partNulFields { + if f.policy == nulEncode { + f.decode(&parts[i]) + } + } } } diff --git a/coderd/x/chatd/chatprompt/chatprompt_test.go b/coderd/x/chatd/chatprompt/chatprompt_test.go index 8529356babe..1899a0c5d3a 100644 --- a/coderd/x/chatd/chatprompt/chatprompt_test.go +++ b/coderd/x/chatd/chatprompt/chatprompt_test.go @@ -10,6 +10,7 @@ import ( "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" + fantasygoogle "charm.land/fantasy/providers/google" "github.com/google/uuid" "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/assert" @@ -2007,7 +2008,7 @@ func extractToolResultIDs(t *testing.T, msgs ...fantasy.Message) []string { func TestNulEscapeRoundTrip(t *testing.T) { t.Parallel() - db, _ := dbtestutil.NewDB(t) + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) // Seed minimal dependencies for the DB round-trip path: // user, provider, model config, chat. @@ -2157,6 +2158,97 @@ func TestNulEscapeRoundTrip(t *testing.T) { assert.JSONEq(t, string(resultJSON), string(decoded[0].Result)) }) + // Gemini thought signatures are binary data converted directly to a + // string by the Google provider. NUL bytes in the signature must survive + // the full persistence and provider-metadata round-trip. + t.Run("GeminiReasoningMetadataWithNul", func(t *testing.T) { + t.Parallel() + + const signature = "binary\x00thought\x00signature" + part := chatprompt.PartFromContent(fantasy.ReasoningContent{ + Text: "reasoning", + ProviderMetadata: fantasy.ProviderMetadata{ + fantasygoogle.Name: &fantasygoogle.ReasoningMetadata{ + Signature: signature, + ToolID: "tool-1", + }, + }, + }) + encoded, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part}) + require.NoError(t, err) + require.NotContains(t, string(encoded.RawMessage), `\u0000`) + + ctx := testutil.Context(t, testutil.WaitShort) + dbMsg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: encoded, + ContentVersion: chatprompt.CurrentContentVersion, + }) + readBack, err := db.GetChatMessageByID(ctx, dbMsg.ID) + require.NoError(t, err) + + // Physical evidence of the encoding: the stored jsonb row contains + // one sentinel pair per NUL byte and no \u0000 escape. + var stored string + require.NoError(t, sqlDB.QueryRowContext(ctx, + "SELECT content::text FROM chat_messages WHERE id = $1", dbMsg.ID, + ).Scan(&stored)) + require.NotContains(t, stored, `\u0000`) + require.Equal(t, 2, strings.Count(stored, "\uE000\uE001")) + + prompt, err := chatprompt.ConvertMessagesWithFiles( + ctx, + []database.ChatMessage{readBack}, + nil, + slogtest.Make(t, nil), + nil, + ) + require.NoError(t, err) + require.Len(t, prompt, 1) + require.Len(t, prompt[0].Content, 1) + reasoning, ok := fantasy.AsMessagePart[fantasy.ReasoningPart](prompt[0].Content[0]) + require.True(t, ok) + metadata := fantasygoogle.GetReasoningMetadata(reasoning.ProviderOptions) + require.NotNil(t, metadata) + require.Equal(t, signature, metadata.Signature) + }) + + // PostgreSQL rejects lone UTF-16 surrogate escapes. The RawMessage + // encoder normalizes them through encoding/json before persistence. + t.Run("ToolArgsWithLoneSurrogate", func(t *testing.T) { + t.Parallel() + + parts := []codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolCall, + Args: json.RawMessage(`{"value":"before\uD800after","pair":"\uD83D\uDE00","literal":"\\uD800"}`), + }} + encoded, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + require.NotContains(t, string(encoded.RawMessage), `\uD800after`) + + ctx := testutil.Context(t, testutil.WaitShort) + dbMsg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: encoded, + ContentVersion: chatprompt.CurrentContentVersion, + }) + readBack, err := db.GetChatMessageByID(ctx, dbMsg.ID) + require.NoError(t, err) + decoded, err := chatprompt.ParseContent(readBack) + require.NoError(t, err) + require.Len(t, decoded, 1) + assert.JSONEq(t, + `{"value":"before�after","pair":"😀","literal":"\\uD800"}`, + string(decoded[0].Args), + ) + }) + // Multiple parts in one message: one with NUL, one without. t.Run("MultiPartMixed", func(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatprompt/export_test.go b/coderd/x/chatd/chatprompt/export_test.go index 16b3090c2b5..d546449b31d 100644 --- a/coderd/x/chatd/chatprompt/export_test.go +++ b/coderd/x/chatd/chatprompt/export_test.go @@ -15,6 +15,12 @@ const SyntheticPasteTitleBudgetForTest = syntheticPasteTitleBudget // for external tests. var ToolResultPartToMessagePartForTest = toolResultPartToMessagePart +// EncodeNulInJSONForTest exposes encodeNulInJSON for external tests. +var EncodeNulInJSONForTest = encodeNulInJSON + +// NeedsNulEncodingInJSONForTest exposes needsNulEncodingInJSON for external tests. +var NeedsNulEncodingInJSONForTest = needsNulEncodingInJSON + // ToolResultContentToPartForTest exposes toolResultContentToPart // for external tests. var ToolResultContentToPartForTest = func(logger slog.Logger, content fantasy.ToolResultContent) codersdk.ChatMessagePart { diff --git a/coderd/x/chatd/chatprompt/nul_coverage_test.go b/coderd/x/chatd/chatprompt/nul_coverage_test.go new file mode 100644 index 00000000000..87d94c999d0 --- /dev/null +++ b/coderd/x/chatd/chatprompt/nul_coverage_test.go @@ -0,0 +1,302 @@ +package chatprompt_test + +import ( + "encoding/json" + "reflect" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/codersdk" +) + +// Probe values injected into string-bearing fields. The NUL probe's +// prefix "before" puts the NUL at byte 6, which rejection tests assert. +const ( + nulProbe = "before\x00middle\uE000\uE001after" + naturalSentinelProbe = "before\uE000\uE001after" +) + +func TestNeedsNulEncodingInJSON(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + raw json.RawMessage + want bool + }{ + {name: "clean", raw: json.RawMessage(`{"value":"clean"}`)}, + {name: "nul", raw: json.RawMessage(`{"value":"\u0000"}`), want: true}, + {name: "escaped nul text", raw: json.RawMessage(`{"value":"\\u0000"}`)}, + {name: "escaped backslash then nul", raw: json.RawMessage(`{"value":"\\\u0000"}`), want: true}, + {name: "ordinary unicode escape", raw: json.RawMessage(`{"value":"\u1234"}`)}, + {name: "truncated unicode escape", raw: json.RawMessage(`{"value":"\u12"}`), want: true}, + {name: "unicode escape without digits", raw: json.RawMessage(`{"value":"\u`), want: true}, + {name: "invalid unicode escape", raw: json.RawMessage(`{"value":"\uZZZZ"}`), want: true}, + {name: "backslash outside string", raw: json.RawMessage(`\`)}, + {name: "trailing backslash in string", raw: json.RawMessage(`{"value":"\`)}, + + {name: "two escaped backslashes then nul text", raw: json.RawMessage(`{"value":"\\\\u0000"}`)}, + {name: "escaped sentinel", raw: json.RawMessage(`{"value":"\uE000"}`), want: true}, + {name: "literal sentinel", raw: json.RawMessage(`{"value":""}`), want: true}, + {name: "valid surrogate pair", raw: json.RawMessage(`{"value":"\uD83D\uDE00"}`)}, + {name: "mixed case surrogate pair", raw: json.RawMessage(`{"value":"\ud83D\uDe00"}`)}, + {name: "lone high surrogate", raw: json.RawMessage(`{"value":"\uD800"}`), want: true}, + {name: "lone low surrogate", raw: json.RawMessage(`{"value":"\udc00"}`), want: true}, + {name: "reversed surrogate pair", raw: json.RawMessage(`{"value":"\uDC00\uD800"}`), want: true}, + {name: "surrogate literal text", raw: json.RawMessage(`{"value":"\\uD800"}`)}, + {name: "object key", raw: json.RawMessage(`{"key\u0000":"value"}`), want: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := chatprompt.NeedsNulEncodingInJSONForTest(tt.raw) + require.Equalf(t, tt.want, got, "needsNulEncodingInJSON(%s)", tt.raw) + }) + } +} + +func TestEncodeNulInJSONPreservesCleanInput(t *testing.T) { + t.Parallel() + + raw := json.RawMessage(`{"number":1.2300,"pair":"\uD83D\uDE00","literal":"\\u0000"}`) + encoded := chatprompt.EncodeNulInJSONForTest(raw) + require.Equal(t, raw, encoded, "clean JSON changed") + require.Same(t, &raw[0], &encoded[0], "clean JSON did not preserve its backing array") +} + +func TestEncodeNulInJSONNormalizesSurrogates(t *testing.T) { + t.Parallel() + + raw := json.RawMessage(`{"high":"\uD800","low":"\uDC00","pair":"\uD83D\uDE00","literal":"\\uD800"}`) + encoded := chatprompt.EncodeNulInJSONForTest(raw) + require.Truef(t, json.Valid(encoded), "encoded invalid JSON: %s", encoded) + var got map[string]string + require.NoError(t, json.Unmarshal(encoded, &got)) + want := map[string]string{ + "high": "�", + "low": "�", + "pair": "😀", + "literal": `\uD800`, + } + require.Equal(t, want, got, "normalized JSON mismatch") +} + +func TestEncodeNulInJSONPreservesInvalidInput(t *testing.T) { + t.Parallel() + + raw := json.RawMessage(`{"value":"\uD800"} trailing`) + encoded := chatprompt.EncodeNulInJSONForTest(raw) + require.Equal(t, raw, encoded, "invalid JSON changed") +} + +func TestChatMessagePartNaturalSentinelCoverage(t *testing.T) { + t.Parallel() + + partType := reflect.TypeFor[codersdk.ChatMessagePart]() + for i := range partType.NumField() { + field := partType.Field(i) + if _, transforms := chatMessagePartNULProbe(t, field); !transforms { + continue + } + + t.Run(field.Name, func(t *testing.T) { + t.Parallel() + + part := reflect.New(partType).Elem() + part.Field(i).Set(chatMessagePartNaturalSentinelProbeValue(t, field.Type)) + parts := []codersdk.ChatMessagePart{part.Interface().(codersdk.ChatMessagePart)} + + encoded, err := chatprompt.MarshalParts(parts) + require.NoErrorf(t, err, "MarshalParts rejected natural sentinels in ChatMessagePart.%s (%s)", field.Name, field.Type) + decoded, err := chatprompt.ParseContent(database.ChatMessage{ + Role: database.ChatMessageRoleAssistant, + Content: encoded, + ContentVersion: chatprompt.CurrentContentVersion, + }) + require.NoErrorf(t, err, "ParseContent with natural sentinels in ChatMessagePart.%s (%s)", field.Name, field.Type) + require.Equalf(t, parts, decoded, "natural sentinels in ChatMessagePart.%s did not round-trip", field.Name) + }) + } +} + +func TestChatMessagePartNULCoverage(t *testing.T) { + t.Parallel() + + partType := reflect.TypeFor[codersdk.ChatMessagePart]() + for i := range partType.NumField() { + field := partType.Field(i) + probe, transforms := chatMessagePartNULProbe(t, field) + + t.Run(field.Name, func(t *testing.T) { + t.Parallel() + + part := reflect.New(partType).Elem() + part.Field(i).Set(probe) + parts := []codersdk.ChatMessagePart{part.Interface().(codersdk.ChatMessagePart)} + + want := reflect.New(partType).Elem() + want.Field(i).Set(chatMessagePartNULProbeValue(t, field.Type)) + wantParts := []codersdk.ChatMessagePart{want.Interface().(codersdk.ChatMessagePart)} + + encoded, err := chatprompt.MarshalParts(parts) + require.Equalf(t, wantParts, parts, "MarshalParts mutated caller-owned ChatMessagePart.%s data", field.Name) + if err != nil { + require.Truef(t, transforms, "MarshalParts rejected NUL-free ChatMessagePart.%s (%s): %v", field.Name, field.Type, err) + require.ErrorContainsf(t, err, "chat message part 0", "MarshalParts rejection for ChatMessagePart.%s does not name the part", field.Name) + require.ErrorContainsf(t, err, "field "+field.Name, "MarshalParts rejection for ChatMessagePart.%s does not name the field", field.Name) + require.ErrorContainsf(t, err, "contains NUL at byte 6", "MarshalParts rejection for ChatMessagePart.%s does not locate the NUL", field.Name) + return + } + if transforms { + require.NotContainsf(t, string(encoded.RawMessage), `\u0000`, "ChatMessagePart.%s is neither rejected nor handled by encodeNulInParts", field.Name) + } + + decoded, err := chatprompt.ParseContent(database.ChatMessage{ + Role: database.ChatMessageRoleAssistant, + Content: encoded, + ContentVersion: chatprompt.CurrentContentVersion, + }) + require.NoErrorf(t, err, "ParseContent with ChatMessagePart.%s (%s)", field.Name, field.Type) + require.Equalf(t, wantParts, decoded, "ChatMessagePart.%s is not handled by decodeNulInParts", field.Name) + require.Equalf(t, wantParts, parts, "ParseContent mutated caller-owned ChatMessagePart.%s data through an encoded alias", field.Name) + }) + } +} + +func chatMessagePartNaturalSentinelProbeValue(t *testing.T, typ reflect.Type) reflect.Value { + t.Helper() + + if typ == reflect.TypeFor[json.RawMessage]() { + return reflect.ValueOf(json.RawMessage(`{"key":"value"}`)) + } + if typ.Kind() == reflect.String { + value := reflect.New(typ).Elem() + value.SetString(naturalSentinelProbe) + return value + } + if chatMessagePartNULTransforms(typ) { + return chatMessagePartStringContainerValue(t, typ, naturalSentinelProbe) + } + + t.Fatalf("no natural sentinel probe for ChatMessagePart field type %s", typ) + return reflect.Value{} +} + +func chatMessagePartNULProbe(t *testing.T, field reflect.StructField) (reflect.Value, bool) { + t.Helper() + + transforms := chatMessagePartNULTransforms(field.Type) + if transforms || chatMessagePartNULSafeType(field.Type) { + return chatMessagePartNULProbeValue(t, field.Type), transforms + } + + t.Fatalf( + "ChatMessagePart.%s has unknown NUL coverage shape %s; classify the type and add a partNulFields entry if it can contain strings", + field.Name, + field.Type, + ) + return reflect.Value{}, false +} + +func chatMessagePartNULTransforms(typ reflect.Type) bool { + if typ == reflect.TypeFor[json.RawMessage]() { + return true + } + if typ.Kind() == reflect.String { + return true + } + if typ.Kind() != reflect.Slice && typ.Kind() != reflect.Array { + return false + } + return chatMessagePartNULStringContainer(typ.Elem()) +} + +func chatMessagePartNULStringContainer(typ reflect.Type) bool { + if typ.Kind() == reflect.String { + return true + } + if typ.Kind() != reflect.Slice && typ.Kind() != reflect.Array { + return false + } + return chatMessagePartNULStringContainer(typ.Elem()) +} + +func chatMessagePartNULSafeType(typ reflect.Type) bool { + switch typ { + case reflect.TypeFor[bool](), + reflect.TypeFor[int](), + reflect.TypeFor[[]byte](), + reflect.TypeFor[uuid.NullUUID](), + reflect.TypeFor[*time.Time](): + return true + default: + return false + } +} + +func chatMessagePartNULProbeValue(t *testing.T, typ reflect.Type) reflect.Value { + t.Helper() + + if typ == reflect.TypeFor[json.RawMessage]() { + return reflect.ValueOf(json.RawMessage(`{"key\u0000":"value\u0000"}`)) + } + if typ == reflect.TypeFor[[][]string]() { + return reflect.ValueOf([][]string{{nulProbe}, nil}) + } + if typ.Kind() == reflect.String { + value := reflect.New(typ).Elem() + value.SetString(nulProbe) + return value + } + if chatMessagePartNULTransforms(typ) { + return chatMessagePartStringContainerValue(t, typ, nulProbe) + } + + switch typ { + case reflect.TypeFor[bool](): + return reflect.ValueOf(true) + case reflect.TypeFor[int](): + return reflect.ValueOf(42) + case reflect.TypeFor[[]byte](): + return reflect.ValueOf([]byte{0, 1, 2}) + case reflect.TypeFor[uuid.NullUUID](): + return reflect.ValueOf(uuid.NullUUID{UUID: uuid.MustParse("00000000-0000-0000-0000-000000000001"), Valid: true}) + case reflect.TypeFor[*time.Time](): + return reflect.ValueOf(new(time.Date(2026, time.August, 25, 1, 2, 3, 0, time.UTC))) + default: + t.Fatalf("no probe value for recognized ChatMessagePart field type %s", typ) + return reflect.Value{} + } +} + +func chatMessagePartStringContainerValue(t *testing.T, typ reflect.Type, probe string) reflect.Value { + t.Helper() + + if typ.Kind() == reflect.String { + value := reflect.New(typ).Elem() + value.SetString(probe) + return value + } + + switch typ.Kind() { + case reflect.Slice: + value := reflect.MakeSlice(typ, 1, 1) + value.Index(0).Set(chatMessagePartStringContainerValue(t, typ.Elem(), probe)) + return value + case reflect.Array: + if typ.Len() == 0 { + t.Fatalf("cannot inject a string probe into zero-length string container %s", typ) + } + value := reflect.New(typ).Elem() + value.Index(0).Set(chatMessagePartStringContainerValue(t, typ.Elem(), probe)) + return value + default: + t.Fatalf("unsupported recursive string container %s", typ) + return reflect.Value{} + } +}