diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 3d732726a0..37105e2955 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -3,6 +3,7 @@ package chatd_test import ( "context" "database/sql" + "encoding/base64" "encoding/json" "errors" "fmt" @@ -143,8 +144,13 @@ func TestSubagentChatExcludesWorkspaceProvisioningTools(t *testing.T) { ) } // Subsequent calls (including the subagent): just reply. + // Include literal \u0000 in the response text, which is + // what a real LLM writes when explaining binary output. + // json.Marshal encodes the backslash as \\, producing + // \\u0000 in the JSON bytes. The sanitizer must not + // corrupt this into invalid JSON. return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("Done.")..., + chattest.OpenAITextChunks("The file contains \\u0000 null bytes.")..., ) }) @@ -1479,6 +1485,179 @@ func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) { require.Equal(t, codersdk.ChatStatusPending, snapshot[0].Status.Status) } +func TestPersistToolResultWithBinaryData(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + const binaryOutputBase64 = "SEVBREVSAAAAc29tZSBkYXRhAABtb3JlIGRhdGEARU5E" + binaryOutput, err := io.ReadAll(base64.NewDecoder( + base64.StdEncoding, + strings.NewReader(binaryOutputBase64), + )) + require.NoError(t, err) + + var streamedCallCount atomic.Int32 + var streamedCallsMu sync.Mutex + streamedCalls := make([][]chattest.OpenAIMessage, 0, 2) + + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("Binary tool result test") + } + + streamedCallsMu.Lock() + streamedCalls = append(streamedCalls, append([]chattest.OpenAIMessage(nil), req.Messages...)) + streamedCallsMu.Unlock() + + if streamedCallCount.Add(1) == 1 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk( + "execute", + `{"command":"cat /home/coder/binary_file.bin"}`, + ), + ) + } + // Include literal \u0000 in the response text, which is + // what a real LLM writes when explaining binary output. + // json.Marshal encodes the backslash as \\, producing + // \\u0000 in the JSON bytes. The sanitizer must not + // corrupt this into invalid JSON. + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("The file contains \\u0000 null bytes.")..., + ) + }) + + // Use "openai-compat" provider so the chatd framework uses the + // /chat/completions endpoint, where the mock server supports + // streaming tool calls. The default "openai" provider routes to + // /responses which only handles text deltas in the mock. + user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + mockConn.EXPECT(). + SetExtraHeaders(gomock.Any()). + AnyTimes() + mockConn.EXPECT(). + LS(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.LSResponse{}, nil). + AnyTimes() + mockConn.EXPECT(). + ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(io.NopCloser(strings.NewReader("")), "", nil). + AnyTimes() + mockConn.EXPECT(). + StartProcess(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, req workspacesdk.StartProcessRequest) (workspacesdk.StartProcessResponse, error) { + require.Equal(t, "cat /home/coder/binary_file.bin", req.Command) + return workspacesdk.StartProcessResponse{ID: "proc-binary", Started: true}, nil + }). + Times(1) + mockConn.EXPECT(). + ProcessOutput(gomock.Any(), "proc-binary"). + Return(workspacesdk.ProcessOutputResponse{ + Output: string(binaryOutput), + Running: false, + ExitCode: ptrRef(0), + }, nil). + AnyTimes() + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + Title: "binary-tool-result", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("Read /home/coder/binary_file.bin."), + }, + }) + require.NoError(t, err) + + var chatResult database.Chat + require.Eventually(t, func() bool { + got, getErr := db.GetChatByID(ctx, chat.ID) + if getErr != nil { + return false + } + chatResult = got + return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError + }, testutil.WaitLong, testutil.IntervalFast) + + if chatResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat run failed", "last_error=%q", chatResult.LastError.String) + } + + var toolMessage *database.ChatMessage + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + messages, dbErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + if dbErr != nil { + return false + } + for i := range messages { + if messages[i].Role == database.ChatMessageRoleTool { + toolMessage = &messages[i] + return true + } + } + return false + }, testutil.IntervalFast) + require.NotNil(t, toolMessage) + + parts, err := chatprompt.ParseContent(*toolMessage) + require.NoError(t, err) + require.Len(t, parts, 1) + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, parts[0].Type) + require.Equal(t, "execute", parts[0].ToolName) + + var result chattool.ExecuteResult + require.NoError(t, json.Unmarshal(parts[0].Result, &result)) + require.True(t, result.Success) + require.Equal(t, string(binaryOutput), result.Output) + require.Equal(t, 0, result.ExitCode) + + require.GreaterOrEqual(t, streamedCallCount.Load(), int32(2)) + streamedCallsMu.Lock() + recordedStreamCalls := append([][]chattest.OpenAIMessage(nil), streamedCalls...) + streamedCallsMu.Unlock() + require.GreaterOrEqual(t, len(recordedStreamCalls), 2) + + var foundToolResultInSecondCall bool + for _, message := range recordedStreamCalls[1] { + if message.Role != "tool" { + continue + } + if !json.Valid([]byte(message.Content)) { + continue + } + var result chattool.ExecuteResult + if err := json.Unmarshal([]byte(message.Content), &result); err != nil { + continue + } + if result.Output == string(binaryOutput) { + foundToolResultInSecondCall = true + break + } + } + require.True(t, foundToolResultInSecondCall, "expected second streamed model call to include execute tool output") +} + +func ptrRef[T any](v T) *T { + return &v +} + func TestSubscribeNoPubsubNoDuplicateMessageParts(t *testing.T) { t.Parallel() @@ -2172,26 +2351,69 @@ func newTestServer( return server } +// newActiveTestServer creates a chatd server that actively polls for +// and processes pending chats. Use this instead of newTestServer when +// the test needs the chat loop to actually run. Optional config +// overrides are applied after the defaults. +func newActiveTestServer( + t *testing.T, + db database.Store, + ps dbpubsub.Pubsub, + overrides ...func(*chatd.Config), +) *chatd.Server { + t.Helper() + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + cfg := chatd.Config{ + Logger: logger, + Database: db, + ReplicaID: uuid.New(), + Pubsub: ps, + PendingChatAcquireInterval: 10 * time.Millisecond, + InFlightChatStaleAfter: testutil.WaitSuperLong, + } + for _, o := range overrides { + o(&cfg) + } + server := chatd.New(cfg) + t.Cleanup(func() { + require.NoError(t, server.Close()) + }) + return server +} + func seedChatDependencies( ctx context.Context, t *testing.T, db database.Store, ) (database.User, database.ChatModelConfig) { t.Helper() + return seedChatDependenciesWithProvider(ctx, t, db, "openai", "") +} + +// seedChatDependenciesWithProvider creates a user, chat provider, and +// model config for the given provider type and base URL. +func seedChatDependenciesWithProvider( + ctx context.Context, + t *testing.T, + db database.Store, + provider string, + baseURL string, +) (database.User, database.ChatModelConfig) { + t.Helper() user := dbgen.User(t, db, database.User{}) _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", + Provider: provider, + DisplayName: provider, APIKey: "test-key", - BaseUrl: "", - ApiKeyKeyID: sql.NullString{}, + BaseUrl: baseURL, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, Enabled: true, }) require.NoError(t, err) model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", + Provider: provider, Model: "gpt-4o-mini", DisplayName: "Test Model", CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, @@ -2206,6 +2428,50 @@ func seedChatDependencies( return user, model } +// seedWorkspaceWithAgent creates a full workspace chain with a connected +// agent. This is the common setup needed by tests that exercise tool +// execution against a workspace. +func seedWorkspaceWithAgent( + t *testing.T, + db database.Store, + userID uuid.UUID, +) (database.WorkspaceTable, database.WorkspaceAgent) { + t.Helper() + + org := dbgen.Organization(t, db, database.Organization{}) + tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: userID, + }) + tpl := dbgen.Template(t, db, database.Template{ + CreatedBy: userID, + OrganizationID: org.ID, + ActiveVersionID: tv.ID, + }) + ws := dbgen.Workspace(t, db, database.WorkspaceTable{ + TemplateID: tpl.ID, + OwnerID: userID, + OrganizationID: org.ID, + }) + pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + InitiatorID: userID, + OrganizationID: org.ID, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + TemplateVersionID: tv.ID, + WorkspaceID: ws.ID, + JobID: pj.ID, + }) + res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + Transition: database.WorkspaceTransitionStart, + JobID: pj.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: res.ID, + }) + return ws, agent +} + func setOpenAIProviderBaseURL( ctx context.Context, t *testing.T, @@ -2782,38 +3048,21 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { ), ) } + // Include literal \u0000 in the response text, which is + // what a real LLM writes when explaining binary output. + // json.Marshal encodes the backslash as \\, producing + // \\u0000 in the JSON bytes. The sanitizer must not + // corrupt this into invalid JSON. return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("Done.")..., + chattest.OpenAITextChunks("The file contains \\u0000 null bytes.")..., ) }) // Seed the DB: user, openai-compat provider, model config. - user := dbgen.User(t, db, database.User{}) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai-compat", - DisplayName: "OpenAI Compat", - APIKey: "test-key", - BaseUrl: openAIURL, - CreatedBy: uuid.NullUUID{}, - Enabled: true, - }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai-compat", - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{}, - UpdatedBy: uuid.NullUUID{}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), - }) - require.NoError(t, err) + user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) // Add an Anthropic provider pointing to our mock server. - _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ + _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: "anthropic", DisplayName: "Anthropic", APIKey: "test-anthropic-key", @@ -2828,37 +3077,7 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { // Build workspace + agent records so getWorkspaceConn can // resolve the agent for the computer use child. - org := dbgen.Organization(t, db, database.Organization{}) - tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ - OrganizationID: org.ID, - CreatedBy: user.ID, - }) - tpl := dbgen.Template(t, db, database.Template{ - CreatedBy: user.ID, - OrganizationID: org.ID, - ActiveVersionID: tv.ID, - }) - ws := dbgen.Workspace(t, db, database.WorkspaceTable{ - TemplateID: tpl.ID, - OwnerID: user.ID, - OrganizationID: org.ID, - }) - pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ - InitiatorID: user.ID, - OrganizationID: org.ID, - }) - _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ - TemplateVersionID: tv.ID, - WorkspaceID: ws.ID, - JobID: pj.ID, - }) - res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ - Transition: database.WorkspaceTransitionStart, - JobID: pj.ID, - }) - dbAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ - ResourceID: res.ID, - }) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) // Mock agent connection that returns valid display dimensions // for the initial screenshot check in the computer use path. @@ -2880,25 +3099,11 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { Return(workspacesdk.LSResponse{}, xerrors.New("not found")). AnyTimes() - agentConnFn := func( - _ context.Context, agentID uuid.UUID, - ) (workspacesdk.AgentConn, func(), error) { - require.Equal(t, dbAgent.ID, agentID) - return mockConn, func() {}, nil - } - - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: 10 * time.Millisecond, - InFlightChatStaleAfter: testutil.WaitSuperLong, - AgentConn: agentConnFn, - }) - t.Cleanup(func() { - require.NoError(t, server.Close()) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } }) // Create a root chat with a workspace so the child inherits it. diff --git a/coderd/chatd/chatprompt/chatprompt.go b/coderd/chatd/chatprompt/chatprompt.go index 5295026e7b..8530d9f1ce 100644 --- a/coderd/chatd/chatprompt/chatprompt.go +++ b/coderd/chatd/chatprompt/chatprompt.go @@ -321,6 +321,7 @@ func parseContentV1(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) ([ if err := json.Unmarshal(raw.RawMessage, &parts); err != nil { return nil, xerrors.Errorf("parse %s content: %w", role, err) } + decodeNulInParts(parts) return parts, nil } @@ -1018,11 +1019,16 @@ 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. func MarshalParts(parts []codersdk.ChatMessagePart) (pqtype.NullRawMessage, error) { if len(parts) == 0 { return pqtype.NullRawMessage{}, nil } - data, err := json.Marshal(parts) + data, err := json.Marshal(encodeNulInParts(parts)) if err != nil { return pqtype.NullRawMessage{}, xerrors.Errorf("encode chat message parts: %w", err) } @@ -1216,3 +1222,186 @@ func partsToMessageParts( } return result } + +// encodeNulInString replaces NUL (U+0000) characters in s with +// the sentinel pair U+E000 U+E001, and doubles any pre-existing +// U+E000 to U+E000 U+E000 so the encoding is reversible. +// Operates on Unicode code points, not JSON escape sequences, +// making it safe through jsonb round-trips (jsonb stores parsed +// characters, not original escape text). +func encodeNulInString(s string) string { + if !strings.ContainsRune(s, 0) && !strings.ContainsRune(s, '\uE000') { + return s + } + var b strings.Builder + b.Grow(len(s)) + for _, r := range s { + switch r { + case '\uE000': + _, _ = b.WriteRune('\uE000') + _, _ = b.WriteRune('\uE000') + case 0: + _, _ = b.WriteRune('\uE000') + _, _ = b.WriteRune('\uE001') + default: + _, _ = b.WriteRune(r) + } + } + return b.String() +} + +// decodeNulInString reverses encodeNulInString: U+E000 U+E000 +// becomes U+E000, and U+E000 U+E001 becomes NUL. +func decodeNulInString(s string) string { + if !strings.ContainsRune(s, '\uE000') { + return s + } + var b strings.Builder + b.Grow(len(s)) + runes := []rune(s) + for i := 0; i < len(runes); i++ { + if runes[i] == '\uE000' && i+1 < len(runes) { + switch runes[i+1] { + case '\uE000': + _, _ = b.WriteRune('\uE000') + i++ + case '\uE001': + _, _ = b.WriteRune(0) + i++ + default: + // Unpaired sentinel — preserve as-is. + _, _ = b.WriteRune(runes[i]) + } + } else { + _, _ = b.WriteRune(runes[i]) + } + } + return b.String() +} + +// encodeNulInValue recursively walks a JSON value (as produced +// by json.Unmarshal with UseNumber) and applies +// encodeNulInString to every string, including map keys. +func encodeNulInValue(v any) any { + switch val := v.(type) { + case string: + return encodeNulInString(val) + case map[string]any: + out := make(map[string]any, len(val)) + for k, elem := range val { + out[encodeNulInString(k)] = encodeNulInValue(elem) + } + return out + case []any: + out := make([]any, len(val)) + for i, elem := range val { + out[i] = encodeNulInValue(elem) + } + return out + default: + return v // numbers, bools, nil + } +} + +// decodeNulInValue recursively walks a JSON value and applies +// decodeNulInString to every string, including map keys. +func decodeNulInValue(v any) any { + switch val := v.(type) { + case string: + return decodeNulInString(val) + case map[string]any: + out := make(map[string]any, len(val)) + for k, elem := range val { + out[decodeNulInString(k)] = decodeNulInValue(elem) + } + return out + case []any: + out := make([]any, len(val)) + for i, elem := range val { + out[i] = decodeNulInValue(elem) + } + return out + default: + return v + } +} + +// 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. +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}) { + return raw + } + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.UseNumber() + var v any + if err := dec.Decode(&v); err != nil { + return raw + } + result, err := json.Marshal(encodeNulInValue(v)) + if err != nil { + return raw + } + return result +} + +// decodeNulInJSON walks all string values (and keys) inside a +// json.RawMessage and applies decodeNulInString. +func decodeNulInJSON(raw json.RawMessage) json.RawMessage { + if len(raw) == 0 { + return raw + } + // U+E000 encoded as UTF-8 is 0xEE 0x80 0x80. + if !bytes.Contains(raw, []byte{0xEE, 0x80, 0x80}) { + return raw + } + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.UseNumber() + var v any + if err := dec.Decode(&v); err != nil { + return raw + } + result, err := json.Marshal(decodeNulInValue(v)) + if err != nil { + return raw + } + 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 { + 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) + } + return encoded +} + +// 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) + } +} diff --git a/coderd/chatd/chatprompt/chatprompt_test.go b/coderd/chatd/chatprompt/chatprompt_test.go index da2acbbbcb..515684c948 100644 --- a/coderd/chatd/chatprompt/chatprompt_test.go +++ b/coderd/chatd/chatprompt/chatprompt_test.go @@ -17,7 +17,10 @@ import ( "github.com/coder/coder/v2/coderd/chatd/chatprompt" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" ) // testMsg builds a database.ChatMessage for ParseContent tests. @@ -1441,3 +1444,205 @@ func extractToolResultIDs(t *testing.T, msgs ...fantasy.Message) []string { } return ids } + +func TestNulEscapeRoundTrip(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + // Seed minimal dependencies for the DB round-trip path: + // user, provider, model config, chat. + user := dbgen.User(t, db, database.User{}) + + _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ + Provider: "openai", + DisplayName: "openai", + APIKey: "test-key", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + Enabled: true, + }) + require.NoError(t, err) + + model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ + Provider: "openai", + Model: "gpt-4o-mini", + DisplayName: "Test Model", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + Enabled: true, + IsDefault: true, + ContextLimit: 128000, + CompressionThreshold: 70, + Options: json.RawMessage(`{}`), + }) + require.NoError(t, err) + + chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "nul-roundtrip-test", + }) + require.NoError(t, err) + + textTests := []struct { + name string + input string + hasNul bool // Whether the input contains actual NUL bytes. + }{ + // --- basic --- + {"NoNul", "hello world", false}, + {"SingleNul", "a\x00b", true}, + {"MultipleNuls", "a\x00b\x00c", true}, + {"ConsecutiveNuls", "\x00\x00\x00", true}, + + // --- boundaries --- + {"EmptyString", "", false}, + {"NulOnly", "\x00", true}, + {"NulAtStart", "\x00hello", true}, + {"NulAtEnd", "hello\x00", true}, + + // --- sentinel / marker in original data --- + // U+E000 is the sentinel character. The encoder must + // double it so it round-trips without being mistaken + // for an encoded NUL. + {"SentinelInOriginal", "a\uE000b", false}, + {"ConsecutiveSentinels", "\uE000\uE000\uE000", false}, + // U+E001 is the marker character used in the NUL pair. + {"MarkerCharInOriginal", "a\uE001b", false}, + // U+E000 followed by U+E001 looks exactly like an + // encoded NUL in the encoded form, so the encoder must + // double the U+E000 to avoid confusion. + {"SentinelThenMarkerChar", "\uE000\uE001", false}, + {"NulAndSentinel", "a\x00b\uE000c", true}, + // Both orders: sentinel adjacent to NUL. + {"SentinelThenNul", "\uE000\x00", true}, + {"NulThenSentinel", "\x00\uE000", true}, + {"AlternatingSentinelNul", "\x00\uE000\x00\uE000", true}, + + // --- strings containing backslashes --- + // Backslashes are normal characters at the Go string + // level; no special handling needed (unlike the old + // JSON-byte approach). + {"BackslashU0000Text", "\\u0000", false}, + {"BackslashThenNul", "\\\x00", true}, + + // --- literal text that looks like escape patterns --- + {"LiteralTextU0000", "the value is u0000 here", false}, + {"LiteralTextUE000", "sentinel uE000 text", false}, + + // --- other control characters mixed with NUL --- + {"ControlCharsMixedWithNul", "\x01\x00\x02\x00\x1f", true}, + + // --- long / stress --- + {"LongNulRun", "\x00\x00\x00\x00\x00\x00\x00\x00", true}, + // Simulated find -print0 output. + {"FindPrint0", "/usr/bin/ls\x00/usr/bin/cat\x00/usr/bin/grep\x00", true}, + } + + for _, tc := range textTests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + parts := []codersdk.ChatMessagePart{ + codersdk.ChatMessageText(tc.input), + } + + encoded, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + + // When the input has real NUL bytes, the stored JSON + // must not contain the \u0000 escape sequence. + if tc.hasNul { + require.NotContains(t, string(encoded.RawMessage), `\u0000`, + "encoded JSON must not contain \\u0000") + } + + // In-memory round-trip through ParseContent. + msg := testMsgV1(codersdk.ChatMessageRoleAssistant, encoded) + decoded, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + + require.Len(t, decoded, 1) + require.Equal(t, tc.input, decoded[0].Text) + + // Full DB round-trip: write to PostgreSQL jsonb, read + // back, and verify the value survives storage. + ctx := testutil.Context(t, testutil.WaitShort) + dbMsgs, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{user.ID}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + Content: []string{string(encoded.RawMessage)}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, + }) + require.NoError(t, err) + require.Len(t, dbMsgs, 1) + + readBack, err := db.GetChatMessageByID(ctx, dbMsgs[0].ID) + require.NoError(t, err) + + dbDecoded, err := chatprompt.ParseContent(readBack) + require.NoError(t, err) + require.Len(t, dbDecoded, 1) + require.Equal(t, tc.input, dbDecoded[0].Text) + }) + } + + // Tool result with NUL in the result JSON value. + t.Run("ToolResultWithNul", func(t *testing.T) { + t.Parallel() + + resultJSON := json.RawMessage(`"output:\u0000done"`) + parts := []codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("call-1", "my_tool", resultJSON, false), + } + + encoded, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + require.NotContains(t, string(encoded.RawMessage), `\u0000`, + "encoded JSON must not contain \\u0000") + + msg := testMsgV1(codersdk.ChatMessageRoleTool, encoded) + decoded, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + require.Len(t, decoded, 1) + // JSON re-serialization may reformat, so compare + // semantically. + assert.JSONEq(t, string(resultJSON), string(decoded[0].Result)) + }) + + // Multiple parts in one message: one with NUL, one without. + t.Run("MultiPartMixed", func(t *testing.T) { + t.Parallel() + + parts := []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("clean text"), + codersdk.ChatMessageText("has\x00nul"), + } + + encoded, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + require.NotContains(t, string(encoded.RawMessage), `\u0000`, + "encoded JSON must not contain \\u0000") + + msg := testMsgV1(codersdk.ChatMessageRoleAssistant, encoded) + decoded, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + require.Len(t, decoded, 2) + require.Equal(t, "clean text", decoded[0].Text) + require.Equal(t, "has\x00nul", decoded[1].Text) + }) +}