From 4a79af1a0de17a726dd617df1a8bbe701a103c4b Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 13 Mar 2026 18:47:36 +0200 Subject: [PATCH] refactor: add chat_message_role enum and content_version column (#23042) Migration 000434 converts chat_messages.role from text to a Postgres enum, rebuilds the partial index, and adds content_version smallint. The column is backfilled with DEFAULT 0, then the default is dropped so future inserts must set it explicitly. Version 0 uses the role-aware heuristic from #22958. Version 1 (all new inserts) stores []ChatMessagePart JSON for all roles, including system messages. ParseContent takes database.ChatMessage directly and dispatches on version internally. Unknown versions error. All string(codersdk.ChatMessageRole*) casts at DB write sites are replaced with database.ChatMessageRole* constants from sqlc. Refs #22958 --- coderd/chatd/chatd.go | 97 +++++------ coderd/chatd/chatd_test.go | 45 ++++-- coderd/chatd/chatprompt/chatprompt.go | 58 +++++-- coderd/chatd/chatprompt/chatprompt_test.go | 152 ++++++++++++------ coderd/chatd/quickgen.go | 8 +- coderd/chatd/subagent.go | 4 +- coderd/chats_test.go | 12 +- coderd/database/db2sdk/db2sdk.go | 18 ++- coderd/database/db2sdk/db2sdk_test.go | 41 ++--- coderd/database/dbauthz/dbauthz_test.go | 2 +- coderd/database/dump.sql | 14 +- ..._message_role_and_content_version.down.sql | 15 ++ ...at_message_role_and_content_version.up.sql | 32 ++++ coderd/database/models.go | 67 +++++++- coderd/database/querier_test.go | 54 ++++--- coderd/database/queries.sql.go | 38 +++-- coderd/database/queries/chats.sql | 6 +- 17 files changed, 455 insertions(+), 208 deletions(-) create mode 100644 coderd/database/migrations/000434_chat_message_role_and_content_version.down.sql create mode 100644 coderd/database/migrations/000434_chat_message_role_and_content_version.up.sql diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 227eea9c4f..fbdb59b267 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -268,7 +268,9 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C systemPrompt := strings.TrimSpace(opts.SystemPrompt) if systemPrompt != "" { - systemContent, err := json.Marshal(systemPrompt) + systemContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(systemPrompt), + }) if err != nil { return xerrors.Errorf("marshal system prompt: %w", err) } @@ -279,11 +281,9 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C UUID: opts.ModelConfigID, Valid: true, }, - Role: string(codersdk.ChatMessageRoleSystem), - Content: pqtype.NullRawMessage{ - RawMessage: systemContent, - Valid: len(systemContent) > 0, - }, + Role: database.ChatMessageRoleSystem, + ContentVersion: chatprompt.CurrentContentVersion, + Content: systemContent, Visibility: database.ChatMessageVisibilityModel, InputTokens: sql.NullInt64{}, OutputTokens: sql.NullInt64{}, @@ -309,7 +309,8 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C UUID: opts.ModelConfigID, Valid: true, }, - Role: string(codersdk.ChatMessageRoleUser), + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, Content: userContent, CreatedBy: uuid.NullUUID{UUID: opts.OwnerID, Valid: opts.OwnerID != uuid.Nil}, Visibility: database.ChatMessageVisibilityBoth, @@ -523,7 +524,7 @@ func (p *Server) EditMessage( if existing.ChatID != opts.ChatID { return ErrEditedMessageNotFound } - if existing.Role != "user" { + if existing.Role != database.ChatMessageRoleUser { return ErrEditedMessageNotUser } @@ -897,7 +898,8 @@ func insertUserMessageAndSetPending( message, err := insertChatMessageWithStore(ctx, store, database.InsertChatMessageParams{ ChatID: lockedChat.ID, ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: string(codersdk.ChatMessageRoleUser), + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, Content: content, CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: createdBy != uuid.Nil}, Visibility: database.ChatMessageVisibilityBoth, @@ -1950,9 +1952,10 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { nextQueued, popErr := tx.PopNextQueuedMessage(cleanupCtx, chat.ID) if popErr == nil { msg, insertErr := tx.InsertChatMessage(cleanupCtx, database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: latestChat.LastModelConfigID, Valid: true}, - Role: string(codersdk.ChatMessageRoleUser), + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: latestChat.LastModelConfigID, Valid: true}, + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, Content: pqtype.NullRawMessage{ RawMessage: nextQueued.Content, Valid: len(nextQueued.Content) > 0, @@ -2345,15 +2348,16 @@ func (p *Server) runChat( hasUsage := step.Usage != (fantasy.Usage{}) assistantMessage, insertErr := tx.InsertChatMessage(persistCtx, database.InsertChatMessageParams{ - ChatID: chat.ID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: string(codersdk.ChatMessageRoleAssistant), - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: usageNullInt64(step.Usage.InputTokens, hasUsage), - OutputTokens: usageNullInt64(step.Usage.OutputTokens, hasUsage), - TotalTokens: usageNullInt64(step.Usage.TotalTokens, hasUsage), + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{}, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + InputTokens: usageNullInt64(step.Usage.InputTokens, hasUsage), + OutputTokens: usageNullInt64(step.Usage.OutputTokens, hasUsage), + TotalTokens: usageNullInt64(step.Usage.TotalTokens, hasUsage), ReasoningTokens: usageNullInt64( step.Usage.ReasoningTokens, hasUsage, @@ -2383,7 +2387,8 @@ func (p *Server) runChat( ChatID: chat.ID, CreatedBy: uuid.NullUUID{}, ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: string(codersdk.ChatMessageRoleTool), + Role: database.ChatMessageRoleTool, + ContentVersion: chatprompt.CurrentContentVersion, Content: resultContent, Visibility: database.ChatMessageVisibilityBoth, InputTokens: sql.NullInt64{}, @@ -2676,7 +2681,9 @@ func (p *Server) persistChatContextSummary( return nil } - systemContent, err := json.Marshal(result.SystemSummary) + systemContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(result.SystemSummary), + }) if err != nil { return xerrors.Errorf("encode system summary: %w", err) } @@ -2718,14 +2725,12 @@ func (p *Server) persistChatContextSummary( txErr := p.db.InTx(func(tx database.Store) error { _, txErr := tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: string(codersdk.ChatMessageRoleUser), - Content: pqtype.NullRawMessage{ - RawMessage: systemContent, - Valid: len(systemContent) > 0, - }, + ChatID: chatID, + CreatedBy: uuid.NullUUID{}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, + Content: systemContent, Visibility: database.ChatMessageVisibilityModel, Compressed: sql.NullBool{Bool: true, Valid: true}, InputTokens: sql.NullInt64{}, @@ -2741,12 +2746,13 @@ func (p *Server) persistChatContextSummary( } assistantMessage, txErr := tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: string(codersdk.ChatMessageRoleAssistant), - Content: assistantContent, - Visibility: database.ChatMessageVisibilityUser, + ChatID: chatID, + CreatedBy: uuid.NullUUID{}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityUser, Compressed: sql.NullBool{ Bool: true, Valid: true, @@ -2765,12 +2771,13 @@ func (p *Server) persistChatContextSummary( insertedMessages = append(insertedMessages, assistantMessage) toolMessage, txErr := tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: string(codersdk.ChatMessageRoleTool), - Content: toolResult, - Visibility: database.ChatMessageVisibilityBoth, + ChatID: chatID, + CreatedBy: uuid.NullUUID{}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleTool, + ContentVersion: chatprompt.CurrentContentVersion, + Content: toolResult, + Visibility: database.ChatMessageVisibilityBoth, Compressed: sql.NullBool{ Bool: true, Valid: true, @@ -3134,10 +3141,10 @@ func (p *Server) maybeSendPushNotification( msg, err := p.db.GetLastChatMessageByRole(pushCtx, database.GetLastChatMessageByRoleParams{ ChatID: chat.ID, - Role: string(codersdk.ChatMessageRoleAssistant), + Role: database.ChatMessageRoleAssistant, }) if err == nil { - content, parseErr := chatprompt.ParseContent(codersdk.ChatMessageRoleAssistant, msg.Content) + content, parseErr := chatprompt.ParseContent(msg) if parseErr == nil { assistantText := strings.TrimSpace(contentBlocksToText(content)) if assistantText != "" { diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index ac9d5bd750..5bfe159b55 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -13,12 +13,12 @@ import ( "time" "github.com/google/uuid" - "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/agent/agenttest" "github.com/coder/coder/v2/coderd/chatd" + "github.com/coder/coder/v2/coderd/chatd/chatprompt" "github.com/coder/coder/v2/coderd/chatd/chattest" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" @@ -463,9 +463,13 @@ func TestEditMessageUpdatesAndTruncatesAndClearsQueue(t *testing.T) { }) require.NoError(t, err) + queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("queued"), + }) + require.NoError(t, err) _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ ChatID: chat.ID, - Content: json.RawMessage(`"queued"`), + Content: queuedContent, }) require.NoError(t, err) @@ -556,14 +560,17 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { }) require.NoError(t, err) + assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("assistant"), + }) + require.NoError(t, err) + assistantMessage, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: "assistant", - Content: pqtype.NullRawMessage{ - RawMessage: json.RawMessage(`"assistant"`), - Valid: true, - }, + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, Visibility: database.ChatMessageVisibilityBoth, InputTokens: sql.NullInt64{}, OutputTokens: sql.NullInt64{}, @@ -901,11 +908,17 @@ func TestSubscribeAfterMessageID(t *testing.T) { // Insert two more messages so we have three total visible // messages (the initial user message plus these two). + secondContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("second"), + }) + require.NoError(t, err) + msg2, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{ ChatID: chat.ID, ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: "assistant", - Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`"second"`), Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: secondContent, Visibility: database.ChatMessageVisibilityBoth, InputTokens: sql.NullInt64{}, OutputTokens: sql.NullInt64{}, @@ -918,11 +931,17 @@ func TestSubscribeAfterMessageID(t *testing.T) { }) require.NoError(t, err) + thirdContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("third"), + }) + require.NoError(t, err) + _, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{ ChatID: chat.ID, ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: "user", - Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`"third"`), Valid: true}, + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, + Content: thirdContent, Visibility: database.ChatMessageVisibilityBoth, InputTokens: sql.NullInt64{}, OutputTokens: sql.NullInt64{}, diff --git a/coderd/chatd/chatprompt/chatprompt.go b/coderd/chatd/chatprompt/chatprompt.go index 0085254ad6..5295026e7b 100644 --- a/coderd/chatd/chatprompt/chatprompt.go +++ b/coderd/chatd/chatprompt/chatprompt.go @@ -91,15 +91,14 @@ func ConvertMessagesWithFiles( continue } - role := codersdk.ChatMessageRole(msg.Role) - parts, err := ParseContent(role, msg.Content) + parts, err := ParseContent(msg) if err != nil { return nil, err } - parsed[i] = parsedMessage{role: role, parts: parts} + parsed[i] = parsedMessage{role: codersdk.ChatMessageRole(msg.Role), parts: parts} // Collect file IDs from user messages for resolution. - if resolver != nil && msg.Role == string(codersdk.ChatMessageRoleUser) { + if resolver != nil && msg.Role == database.ChatMessageRoleUser { for _, part := range parts { if part.Type == codersdk.ChatMessagePartTypeFile && part.FileID.Valid { if _, seen := seenFileIDs[part.FileID.UUID]; !seen { @@ -262,16 +261,45 @@ func AppendUser(prompt []fantasy.Message, instruction string) []fantasy.Message return out } +const ( + // ContentVersionV0 is the legacy content format. Parsing uses + // role-aware heuristics to distinguish fantasy envelope format + // from SDK parts. + ContentVersionV0 int16 = 0 + // ContentVersionV1 stores content as []codersdk.ChatMessagePart + // JSON for all roles. + ContentVersionV1 int16 = 1 + + // CurrentContentVersion is the version used for new inserts. + CurrentContentVersion = ContentVersionV1 +) + // ParseContent decodes persisted chat message content blocks into -// SDK parts. Role-aware: system messages are JSON strings, -// assistant and user messages use a structural heuristic -// (isFantasyEnvelopeFormat) to distinguish legacy fantasy envelope -// from SDK parts, and tool messages use try/fallback. -func ParseContent(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) ([]codersdk.ChatMessagePart, error) { - if !raw.Valid || len(raw.RawMessage) == 0 { +// SDK parts. Dispatches on content version: version 0 (legacy) uses +// a role-aware heuristic to distinguish fantasy envelope format +// from SDK parts, version 1 (current) unmarshals SDK-format +// []ChatMessagePart directly. +func ParseContent(msg database.ChatMessage) ([]codersdk.ChatMessagePart, error) { + if !msg.Content.Valid || len(msg.Content.RawMessage) == 0 { return nil, nil } + role := codersdk.ChatMessageRole(msg.Role) + + switch msg.ContentVersion { + case ContentVersionV0: + return parseLegacyContent(role, msg.Content) + case ContentVersionV1: + return parseContentV1(role, msg.Content) + default: + return nil, xerrors.Errorf("unsupported content version %d", msg.ContentVersion) + } +} + +// parseLegacyContent handles content version 0, where the format +// varies by role and era. Uses structural heuristics to distinguish +// fantasy envelope format from SDK parts. +func parseLegacyContent(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) ([]codersdk.ChatMessagePart, error) { switch role { case codersdk.ChatMessageRoleSystem: return parseSystemRole(raw) @@ -286,6 +314,16 @@ func ParseContent(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) ([]c } } +// parseContentV1 handles content version 1. Content is a JSON +// array of ChatMessagePart structs. +func parseContentV1(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) ([]codersdk.ChatMessagePart, error) { + var parts []codersdk.ChatMessagePart + if err := json.Unmarshal(raw.RawMessage, &parts); err != nil { + return nil, xerrors.Errorf("parse %s content: %w", role, err) + } + return parts, nil +} + // parseSystemRole decodes a system message (JSON string) into a // single text part. func parseSystemRole(raw pqtype.NullRawMessage) ([]codersdk.ChatMessagePart, error) { diff --git a/coderd/chatd/chatprompt/chatprompt_test.go b/coderd/chatd/chatprompt/chatprompt_test.go index 5a4e41beff..da2acbbbcb 100644 --- a/coderd/chatd/chatprompt/chatprompt_test.go +++ b/coderd/chatd/chatprompt/chatprompt_test.go @@ -20,6 +20,25 @@ import ( "github.com/coder/coder/v2/codersdk" ) +// testMsg builds a database.ChatMessage for ParseContent tests. +// ContentVersion defaults to 0 (legacy), which exercises the +// heuristic detection path. +func testMsg(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) database.ChatMessage { + return database.ChatMessage{ + Role: database.ChatMessageRole(role), + Content: raw, + } +} + +// testMsgV1 builds a database.ChatMessage with ContentVersion 1. +func testMsgV1(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) database.ChatMessage { + return database.ChatMessage{ + Role: database.ChatMessageRole(role), + Content: raw, + ContentVersion: chatprompt.CurrentContentVersion, + } +} + func TestConvertMessages_NormalizesAssistantToolCallInput(t *testing.T) { t.Parallel() @@ -75,12 +94,12 @@ func TestConvertMessages_NormalizesAssistantToolCallInput(t *testing.T) { prompt, err := chatprompt.ConvertMessages([]database.ChatMessage{ { - Role: string(fantasy.MessageRoleAssistant), + Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent, }, { - Role: string(fantasy.MessageRoleTool), + Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: toolContent, }, @@ -135,7 +154,7 @@ func TestConvertMessagesWithFiles_ResolvesFileData(t *testing.T) { context.Background(), []database.ChatMessage{ { - Role: string(fantasy.MessageRoleUser), + Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: pqtype.NullRawMessage{RawMessage: rawContent, Valid: true}, }, @@ -192,7 +211,7 @@ func TestConvertMessagesWithFiles_BackwardCompat(t *testing.T) { context.Background(), []database.ChatMessage{ { - Role: string(fantasy.MessageRoleUser), + Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: pqtype.NullRawMessage{RawMessage: rawContent, Valid: true}, }, @@ -279,12 +298,12 @@ func TestInjectMissingToolResults_SkipsProviderExecuted(t *testing.T) { prompt, err := chatprompt.ConvertMessages([]database.ChatMessage{ { - Role: "assistant", + Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent, }, { - Role: "tool", + Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: localResult, }, @@ -381,16 +400,16 @@ func TestInjectMissingToolUses_DropsProviderExecutedOrphans(t *testing.T) { prompt, err := chatprompt.ConvertMessages([]database.ChatMessage{ // Step 1 - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: step1Assistant}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: resultA}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: resultB}, + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: step1Assistant}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: resultA}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: resultB}, // Step 2 - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: step2Assistant}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: resultC}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: resultD}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: resultE}, + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: step2Assistant}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: resultC}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: resultD}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: resultE}, // User follow-up - {Role: "user", Visibility: database.ChatMessageVisibilityBoth, Content: mustMarshalContent(t, []fantasy.Content{ + {Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: mustMarshalContent(t, []fantasy.Content{ fantasy.TextContent{Text: "?"}, })}, }) @@ -468,10 +487,10 @@ func TestInjectMissingToolUses_DropsOnlyProviderExecutedMessage(t *testing.T) { ) prompt, err := chatprompt.ConvertMessages([]database.ChatMessage{ - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: localResult}, - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: assistant2Content}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: peResult}, + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: localResult}, + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: assistant2Content}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: peResult}, }) require.NoError(t, err) @@ -513,8 +532,8 @@ func TestProviderExecutedResultInAssistantContent(t *testing.T) { }) prompt, err := chatprompt.ConvertMessages([]database.ChatMessage{ - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent}, - {Role: "user", Visibility: database.ChatMessageVisibilityBoth, Content: mustMarshalContent(t, []fantasy.Content{ + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent}, + {Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: mustMarshalContent(t, []fantasy.Content{ fantasy.TextContent{Text: "Thanks!"}, })}, }) @@ -586,10 +605,10 @@ func TestProviderExecutedResult_LegacyToolRow(t *testing.T) { ) prompt, err := chatprompt.ConvertMessages([]database.ChatMessage{ - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: peResult}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: execResult}, - {Role: "user", Visibility: database.ChatMessageVisibilityBoth, Content: mustMarshalContent(t, []fantasy.Content{ + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: assistantContent}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: peResult}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: execResult}, + {Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: mustMarshalContent(t, []fantasy.Content{ fantasy.TextContent{Text: "next"}, })}, }) @@ -921,13 +940,50 @@ func TestParseContent_BackwardCompat(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() - parts, err := chatprompt.ParseContent(tc.role, tc.raw) + parts, err := chatprompt.ParseContent(testMsg(tc.role, tc.raw)) require.NoError(t, err) tc.check(t, parts) }) } } +func TestParseContent_V1(t *testing.T) { + t.Parallel() + + t.Run("system", func(t *testing.T) { + t.Parallel() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("You are helpful."), + }) + require.NoError(t, err) + + parts, err := chatprompt.ParseContent(testMsgV1(codersdk.ChatMessageRoleSystem, raw)) + require.NoError(t, err) + require.Len(t, parts, 1) + assert.Equal(t, codersdk.ChatMessagePartTypeText, parts[0].Type) + assert.Equal(t, "You are helpful.", parts[0].Text) + }) + + t.Run("system_bare_string_errors", func(t *testing.T) { + t.Parallel() + // A bare JSON string is not valid V1 content. + _, err := chatprompt.ParseContent(testMsgV1( + codersdk.ChatMessageRoleSystem, + nullRaw(json.RawMessage(`"You are helpful."`)), + )) + require.Error(t, err) + }) + + t.Run("unknown_version_errors", func(t *testing.T) { + t.Parallel() + msg := testMsgV1(codersdk.ChatMessageRoleUser, nullRaw(json.RawMessage(`[{"type":"text","text":"hi"}]`))) + msg.ContentVersion = 99 + _, err := chatprompt.ParseContent(msg) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported content version") + }) +} + // TestProviderMetadataRoundTrip verifies that Anthropic cache // control hints survive the full path: legacy fantasy DB row → // ParseContent → SDK part (ProviderMetadata) → partsToMessageParts @@ -948,7 +1004,7 @@ func TestProviderMetadataRoundTrip(t *testing.T) { require.NoError(t, err) // Step 1: ParseContent preserves metadata on the SDK part. - parts, err := chatprompt.ParseContent(codersdk.ChatMessageRoleAssistant, legacyContent) + parts, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleAssistant, legacyContent)) require.NoError(t, err) require.Len(t, parts, 1) require.NotNil(t, parts[0].ProviderMetadata, @@ -959,7 +1015,7 @@ func TestProviderMetadataRoundTrip(t *testing.T) { prompt, err := chatprompt.ConvertMessagesWithFiles( context.Background(), []database.ChatMessage{{ - Role: "assistant", + Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: legacyContent, }}, @@ -994,7 +1050,7 @@ func TestFileReferencePreservation(t *testing.T) { require.NoError(t, err) // Storage round-trip: all fields intact. - parts, err := chatprompt.ParseContent(codersdk.ChatMessageRoleUser, raw) + parts, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleUser, raw)) require.NoError(t, err) require.Len(t, parts, 1) assert.Equal(t, codersdk.ChatMessagePartTypeFileReference, parts[0].Type) @@ -1007,7 +1063,7 @@ func TestFileReferencePreservation(t *testing.T) { prompt, err := chatprompt.ConvertMessagesWithFiles( context.Background(), []database.ChatMessage{{ - Role: "user", + Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: raw, }}, @@ -1052,7 +1108,7 @@ func TestAssistantWriteRoundTrip(t *testing.T) { // Read back via ParseContent (takes the new SDK path, not // the legacy fallback, because the stored format is flat). - parts, err := chatprompt.ParseContent(codersdk.ChatMessageRoleAssistant, raw) + parts, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleAssistant, raw)) require.NoError(t, err) require.Len(t, parts, 1) assert.Equal(t, "response with cache hints", parts[0].Text) @@ -1062,7 +1118,7 @@ func TestAssistantWriteRoundTrip(t *testing.T) { prompt, err := chatprompt.ConvertMessagesWithFiles( context.Background(), []database.ChatMessage{{ - Role: "assistant", + Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: raw, }}, @@ -1153,12 +1209,12 @@ func TestMixedFormatConversation(t *testing.T) { require.NoError(t, err) messages := []database.ChatMessage{ - {Role: "system", Visibility: database.ChatMessageVisibilityModel, Content: pqtype.NullRawMessage{RawMessage: systemRaw, Valid: true}}, - {Role: "user", Visibility: database.ChatMessageVisibilityBoth, Content: pqtype.NullRawMessage{RawMessage: oldUserRaw, Valid: true}}, - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: oldAssistantRaw}, - {Role: "tool", Visibility: database.ChatMessageVisibilityBoth, Content: oldToolRaw}, - {Role: "user", Visibility: database.ChatMessageVisibilityBoth, Content: newUserRaw}, - {Role: "assistant", Visibility: database.ChatMessageVisibilityBoth, Content: newAssistantRaw}, + {Role: database.ChatMessageRoleSystem, Visibility: database.ChatMessageVisibilityModel, Content: pqtype.NullRawMessage{RawMessage: systemRaw, Valid: true}}, + {Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: pqtype.NullRawMessage{RawMessage: oldUserRaw, Valid: true}}, + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: oldAssistantRaw}, + {Role: database.ChatMessageRoleTool, Visibility: database.ChatMessageVisibilityBoth, Content: oldToolRaw}, + {Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: newUserRaw}, + {Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth, Content: newAssistantRaw}, } prompt, err := chatprompt.ConvertMessagesWithFiles( @@ -1258,10 +1314,10 @@ func TestQueuedMessageRoundTrip(t *testing.T) { // Step 2: PromoteQueued copies the raw bytes into // chat_messages. ParseContent must handle them identically. - promoted, err := chatprompt.ParseContent(codersdk.ChatMessageRoleUser, pqtype.NullRawMessage{ + promoted, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleUser, pqtype.NullRawMessage{ RawMessage: raw.RawMessage, Valid: true, - }) + })) require.NoError(t, err) require.Len(t, promoted, 2) assert.Equal(t, codersdk.ChatMessagePartTypeText, promoted[0].Type) @@ -1277,7 +1333,7 @@ func TestQueuedMessageRoundTrip(t *testing.T) { prompt, err := chatprompt.ConvertMessagesWithFiles( context.Background(), []database.ChatMessage{{ - Role: "user", + Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, Content: pqtype.NullRawMessage{RawMessage: raw.RawMessage, Valid: true}, }}, @@ -1303,50 +1359,50 @@ func TestParseContent_ErrorPaths(t *testing.T) { t.Run("null_content_returns_nil", func(t *testing.T) { t.Parallel() - parts, err := chatprompt.ParseContent(codersdk.ChatMessageRoleUser, pqtype.NullRawMessage{}) + parts, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleUser, pqtype.NullRawMessage{})) require.NoError(t, err) assert.Nil(t, parts) }) t.Run("empty_content_returns_nil", func(t *testing.T) { t.Parallel() - parts, err := chatprompt.ParseContent(codersdk.ChatMessageRoleAssistant, pqtype.NullRawMessage{ + parts, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleAssistant, pqtype.NullRawMessage{ RawMessage: []byte{}, Valid: true, - }) + })) require.NoError(t, err) assert.Nil(t, parts) }) t.Run("unknown_role", func(t *testing.T) { t.Parallel() - _, err := chatprompt.ParseContent(codersdk.ChatMessageRole("banana"), nullRaw(json.RawMessage(`"hello"`))) + _, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRole("banana"), nullRaw(json.RawMessage(`"hello"`)))) require.Error(t, err) assert.Contains(t, err.Error(), "unsupported chat message role") }) t.Run("system/malformed_json", func(t *testing.T) { t.Parallel() - _, err := chatprompt.ParseContent(codersdk.ChatMessageRoleSystem, nullRaw(json.RawMessage(`not json`))) + _, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleSystem, nullRaw(json.RawMessage(`not json`)))) require.Error(t, err) assert.Contains(t, err.Error(), "parse system content") }) t.Run("user/malformed_json", func(t *testing.T) { t.Parallel() - _, err := chatprompt.ParseContent(codersdk.ChatMessageRoleUser, nullRaw(json.RawMessage(`{not json`))) + _, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleUser, nullRaw(json.RawMessage(`{not json`)))) require.Error(t, err) }) t.Run("assistant/malformed_json", func(t *testing.T) { t.Parallel() - _, err := chatprompt.ParseContent(codersdk.ChatMessageRoleAssistant, nullRaw(json.RawMessage(`{not json`))) + _, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleAssistant, nullRaw(json.RawMessage(`{not json`)))) require.Error(t, err) }) t.Run("tool/malformed_json", func(t *testing.T) { t.Parallel() - _, err := chatprompt.ParseContent(codersdk.ChatMessageRoleTool, nullRaw(json.RawMessage(`{not json`))) + _, err := chatprompt.ParseContent(testMsg(codersdk.ChatMessageRoleTool, nullRaw(json.RawMessage(`{not json`)))) require.Error(t, err) }) } diff --git a/coderd/chatd/quickgen.go b/coderd/chatd/quickgen.go index f5747404f3..4b9d3dc9b8 100644 --- a/coderd/chatd/quickgen.go +++ b/coderd/chatd/quickgen.go @@ -159,14 +159,12 @@ func titleInput( } switch message.Role { - case string(codersdk.ChatMessageRoleAssistant), string(codersdk.ChatMessageRoleTool): + case database.ChatMessageRoleAssistant, database.ChatMessageRoleTool: return "", false - case string(codersdk.ChatMessageRoleUser): + case database.ChatMessageRoleUser: userCount++ if firstUserText == "" { - parsed, err := chatprompt.ParseContent( - codersdk.ChatMessageRoleUser, message.Content, - ) + parsed, err := chatprompt.ParseContent(message) if err != nil { return "", false } diff --git a/coderd/chatd/subagent.go b/coderd/chatd/subagent.go index 8288e8ef3c..2a07960a00 100644 --- a/coderd/chatd/subagent.go +++ b/coderd/chatd/subagent.go @@ -482,12 +482,12 @@ func latestSubagentAssistantMessage( for i := len(messages) - 1; i >= 0; i-- { message := messages[i] - if message.Role != string(codersdk.ChatMessageRoleAssistant) || + if message.Role != database.ChatMessageRoleAssistant || message.Visibility == database.ChatMessageVisibilityModel { continue } - content, parseErr := chatprompt.ParseContent(codersdk.ChatMessageRole(message.Role), message.Content) + content, parseErr := chatprompt.ParseContent(message) if parseErr != nil { continue } diff --git a/coderd/chats_test.go b/coderd/chats_test.go index d765120e2f..17c31202e5 100644 --- a/coderd/chats_test.go +++ b/coderd/chats_test.go @@ -3056,11 +3056,15 @@ func TestDeleteChatQueuedMessage(t *testing.T) { }) require.NoError(t, err) + deleteContent, err := json.Marshal([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("queued message for delete route"), + }) + require.NoError(t, err) queuedMessage, err := db.InsertChatQueuedMessage( dbauthz.AsSystemRestricted(ctx), database.InsertChatQueuedMessageParams{ ChatID: chat.ID, - Content: []byte(`"queued message for delete route"`), + Content: deleteContent, }, ) require.NoError(t, err) @@ -3138,11 +3142,15 @@ func TestPromoteChatQueuedMessage(t *testing.T) { require.NoError(t, err) const queuedText = "queued message for promote route" + queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(queuedText), + }) + require.NoError(t, err) queuedMessage, err := db.InsertChatQueuedMessage( dbauthz.AsSystemRestricted(ctx), database.InsertChatQueuedMessageParams{ ChatID: chat.ID, - Content: []byte(fmt.Sprintf("%q", queuedText)), + Content: queuedContent, }, ) require.NoError(t, err) diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index ee4cca1129..91c839bac6 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1071,7 +1071,7 @@ func ChatMessage(m database.ChatMessage) codersdk.ChatMessage { Role: codersdk.ChatMessageRole(m.Role), } if m.Content.Valid { - parts, err := chatMessageParts(codersdk.ChatMessageRole(m.Role), m.Content) + parts, err := chatMessageParts(m) if err == nil { msg.Content = parts } @@ -1113,9 +1113,15 @@ func chatMessageUsage(m database.ChatMessage) *codersdk.ChatMessageUsage { // ChatQueuedMessage converts a queued message to its SDK representation. func ChatQueuedMessage(message database.ChatQueuedMessage) codersdk.ChatQueuedMessage { - parts, err := chatMessageParts(codersdk.ChatMessageRoleUser, pqtype.NullRawMessage{ - RawMessage: message.Content, - Valid: len(message.Content) > 0, + // Queued messages are always written by current code via + // MarshalParts, so they are always current content version. + parts, err := chatMessageParts(database.ChatMessage{ + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{ + RawMessage: message.Content, + Valid: len(message.Content) > 0, + }, + ContentVersion: chatprompt.CurrentContentVersion, }) if err != nil { parts = nil @@ -1139,8 +1145,8 @@ func ChatQueuedMessages(messages []database.ChatQueuedMessage) []codersdk.ChatQu return out } -func chatMessageParts(role codersdk.ChatMessageRole, raw pqtype.NullRawMessage) ([]codersdk.ChatMessagePart, error) { - parts, err := chatprompt.ParseContent(role, raw) +func chatMessageParts(m database.ChatMessage) ([]codersdk.ChatMessagePart, error) { + parts, err := chatprompt.ParseContent(m) if err != nil { return nil, err } diff --git a/coderd/database/db2sdk/db2sdk_test.go b/coderd/database/db2sdk/db2sdk_test.go index 95cc8cf2da..3b98e185ff 100644 --- a/coderd/database/db2sdk/db2sdk_test.go +++ b/coderd/database/db2sdk/db2sdk_test.go @@ -467,7 +467,7 @@ func TestChatMessage_PreservesProviderExecutedOnToolResults(t *testing.T) { dbMsg := database.ChatMessage{ ID: 1, ChatID: uuid.New(), - Role: string(codersdk.ChatMessageRoleAssistant), + Role: database.ChatMessageRoleAssistant, Content: pqtype.NullRawMessage{ RawMessage: rawContent, Valid: true, @@ -495,8 +495,9 @@ func TestChatMessage_PreservesProviderExecutedOnToolResults(t *testing.T) { func TestChatQueuedMessage_ParsesUserContentParts(t *testing.T) { t.Parallel() - rawContent, err := json.Marshal([]fantasy.Content{ - fantasy.TextContent{Text: "queued text"}, + // Queued messages are always written via MarshalParts (SDK format). + rawContent, err := json.Marshal([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("queued text"), }) require.NoError(t, err) @@ -512,35 +513,15 @@ func TestChatQueuedMessage_ParsesUserContentParts(t *testing.T) { require.Equal(t, "queued text", queued.Content[0].Text) } -func TestChatQueuedMessage_FallsBackToTextForLegacyContent(t *testing.T) { +func TestChatQueuedMessage_MalformedContent(t *testing.T) { t.Parallel() - t.Run("legacy_string", func(t *testing.T) { - t.Parallel() - - queued := db2sdk.ChatQueuedMessage(database.ChatQueuedMessage{ - ID: 1, - ChatID: uuid.New(), - Content: json.RawMessage(`"legacy queued text"`), - CreatedAt: time.Now(), - }) - - require.Len(t, queued.Content, 1) - require.Equal(t, codersdk.ChatMessagePartTypeText, queued.Content[0].Type) - require.Equal(t, "legacy queued text", queued.Content[0].Text) + queued := db2sdk.ChatQueuedMessage(database.ChatQueuedMessage{ + ID: 1, + ChatID: uuid.New(), + Content: json.RawMessage(`{"unexpected":"shape"}`), + CreatedAt: time.Now(), }) - t.Run("malformed_payload", func(t *testing.T) { - t.Parallel() - - raw := json.RawMessage(`{"unexpected":"shape"}`) - queued := db2sdk.ChatQueuedMessage(database.ChatQueuedMessage{ - ID: 1, - ChatID: uuid.New(), - Content: raw, - CreatedAt: time.Now(), - }) - - require.Empty(t, queued.Content) - }) + require.Empty(t, queued.Content) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index a9cb62ecd5..4035f78cd5 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -486,7 +486,7 @@ func (s *MethodTestSuite) TestChats() { s.Run("GetLastChatMessageByRole", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) msg := testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID}) - arg := database.GetLastChatMessageByRoleParams{ChatID: chat.ID, Role: "assistant"} + arg := database.GetLastChatMessageByRoleParams{ChatID: chat.ID, Role: database.ChatMessageRoleAssistant} dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().GetLastChatMessageByRole(gomock.Any(), arg).Return(msg, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msg) diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index a8c2b23549..784b6628cc 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -265,6 +265,13 @@ CREATE TYPE build_reason AS ENUM ( 'task_resume' ); +CREATE TYPE chat_message_role AS ENUM ( + 'system', + 'user', + 'assistant', + 'tool' +); + CREATE TYPE chat_message_visibility AS ENUM ( 'user', 'model', @@ -1207,7 +1214,7 @@ CREATE TABLE chat_messages ( chat_id uuid NOT NULL, model_config_id uuid, created_at timestamp with time zone DEFAULT now() NOT NULL, - role text NOT NULL, + role chat_message_role NOT NULL, content jsonb, visibility chat_message_visibility DEFAULT 'both'::chat_message_visibility NOT NULL, input_tokens bigint, @@ -1218,7 +1225,8 @@ CREATE TABLE chat_messages ( cache_read_tokens bigint, context_limit bigint, compressed boolean DEFAULT false NOT NULL, - created_by uuid + created_by uuid, + content_version smallint NOT NULL ); CREATE SEQUENCE chat_messages_id_seq @@ -3524,7 +3532,7 @@ CREATE INDEX idx_chat_messages_chat ON chat_messages USING btree (chat_id); CREATE INDEX idx_chat_messages_chat_created ON chat_messages USING btree (chat_id, created_at); -CREATE INDEX idx_chat_messages_compressed_summary_boundary ON chat_messages USING btree (chat_id, created_at DESC, id DESC) WHERE ((compressed = true) AND (role = 'system'::text) AND (visibility = ANY (ARRAY['model'::chat_message_visibility, 'both'::chat_message_visibility]))); +CREATE INDEX idx_chat_messages_compressed_summary_boundary ON chat_messages USING btree (chat_id, created_at DESC, id DESC) WHERE ((compressed = true) AND (role = 'system'::chat_message_role) AND (visibility = ANY (ARRAY['model'::chat_message_visibility, 'both'::chat_message_visibility]))); CREATE INDEX idx_chat_model_configs_enabled ON chat_model_configs USING btree (enabled); diff --git a/coderd/database/migrations/000434_chat_message_role_and_content_version.down.sql b/coderd/database/migrations/000434_chat_message_role_and_content_version.down.sql new file mode 100644 index 0000000000..223ca278fb --- /dev/null +++ b/coderd/database/migrations/000434_chat_message_role_and_content_version.down.sql @@ -0,0 +1,15 @@ +ALTER TABLE chat_messages DROP COLUMN content_version; + +DROP INDEX idx_chat_messages_compressed_summary_boundary; + +ALTER TABLE chat_messages + ALTER COLUMN role TYPE text + USING (role::text); + +CREATE INDEX idx_chat_messages_compressed_summary_boundary + ON chat_messages(chat_id, created_at DESC, id DESC) + WHERE compressed = TRUE + AND role = 'system' + AND visibility IN ('model', 'both'); + +DROP TYPE chat_message_role; diff --git a/coderd/database/migrations/000434_chat_message_role_and_content_version.up.sql b/coderd/database/migrations/000434_chat_message_role_and_content_version.up.sql new file mode 100644 index 0000000000..8612aba41b --- /dev/null +++ b/coderd/database/migrations/000434_chat_message_role_and_content_version.up.sql @@ -0,0 +1,32 @@ +-- Add chat_message_role enum. +CREATE TYPE chat_message_role AS ENUM ( + 'system', + 'user', + 'assistant', + 'tool' +); + +-- Drop the partial index that references role as text before +-- converting the column type. +DROP INDEX idx_chat_messages_compressed_summary_boundary; + +-- Convert role column from text to enum. +ALTER TABLE chat_messages + ALTER COLUMN role TYPE chat_message_role + USING (role::chat_message_role); + +-- Recreate the partial index with enum-typed comparison. +CREATE INDEX idx_chat_messages_compressed_summary_boundary + ON chat_messages(chat_id, created_at DESC, id DESC) + WHERE compressed = TRUE + AND role = 'system' + AND visibility IN ('model', 'both'); + +-- Add content_version column. Default 0 backfills existing rows. +-- The default is then dropped so future inserts must specify the +-- version explicitly. +ALTER TABLE chat_messages + ADD COLUMN content_version smallint NOT NULL DEFAULT 0; + +ALTER TABLE chat_messages + ALTER COLUMN content_version DROP DEFAULT; diff --git a/coderd/database/models.go b/coderd/database/models.go index d59a5eda1f..948ab02739 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -1049,6 +1049,70 @@ func AllBuildReasonValues() []BuildReason { } } +type ChatMessageRole string + +const ( + ChatMessageRoleSystem ChatMessageRole = "system" + ChatMessageRoleUser ChatMessageRole = "user" + ChatMessageRoleAssistant ChatMessageRole = "assistant" + ChatMessageRoleTool ChatMessageRole = "tool" +) + +func (e *ChatMessageRole) Scan(src interface{}) error { + switch s := src.(type) { + case []byte: + *e = ChatMessageRole(s) + case string: + *e = ChatMessageRole(s) + default: + return fmt.Errorf("unsupported scan type for ChatMessageRole: %T", src) + } + return nil +} + +type NullChatMessageRole struct { + ChatMessageRole ChatMessageRole `json:"chat_message_role"` + Valid bool `json:"valid"` // Valid is true if ChatMessageRole is not NULL +} + +// Scan implements the Scanner interface. +func (ns *NullChatMessageRole) Scan(value interface{}) error { + if value == nil { + ns.ChatMessageRole, ns.Valid = "", false + return nil + } + ns.Valid = true + return ns.ChatMessageRole.Scan(value) +} + +// Value implements the driver Valuer interface. +func (ns NullChatMessageRole) Value() (driver.Value, error) { + if !ns.Valid { + return nil, nil + } + return string(ns.ChatMessageRole), nil +} + +func (e ChatMessageRole) Valid() bool { + switch e { + case ChatMessageRoleSystem, + ChatMessageRoleUser, + ChatMessageRoleAssistant, + ChatMessageRoleTool: + return true + } + return false +} + +func AllChatMessageRoleValues() []ChatMessageRole { + return []ChatMessageRole{ + ChatMessageRoleSystem, + ChatMessageRoleUser, + ChatMessageRoleAssistant, + ChatMessageRoleTool, + } +} + type ChatMessageVisibility string const ( @@ -3943,7 +4007,7 @@ type ChatMessage struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` CreatedAt time.Time `db:"created_at" json:"created_at"` - Role string `db:"role" json:"role"` + Role ChatMessageRole `db:"role" json:"role"` Content pqtype.NullRawMessage `db:"content" json:"content"` Visibility ChatMessageVisibility `db:"visibility" json:"visibility"` InputTokens sql.NullInt64 `db:"input_tokens" json:"input_tokens"` @@ -3955,6 +4019,7 @@ type ChatMessage struct { ContextLimit sql.NullInt64 `db:"context_limit" json:"context_limit"` Compressed bool `db:"compressed" json:"compressed"` CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"` + ContentVersion int16 `db:"content_version" json:"content_version"` } type ChatModelConfig struct { diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 1e5accda37..c186e2e11e 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -21,6 +21,7 @@ import ( "github.com/stretchr/testify/require" "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/chatd/chatprompt" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" @@ -9044,17 +9045,18 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { insertMsg := func( t *testing.T, chatID uuid.UUID, - role string, + role database.ChatMessageRole, vis database.ChatMessageVisibility, compressed bool, content string, ) database.ChatMessage { t.Helper() msg, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - Role: role, - Visibility: vis, - Compressed: sql.NullBool{Bool: compressed, Valid: true}, + ChatID: chatID, + Role: role, + ContentVersion: chatprompt.CurrentContentVersion, + Visibility: vis, + Compressed: sql.NullBool{Bool: compressed, Valid: true}, Content: pqtype.NullRawMessage{ RawMessage: json.RawMessage(`"` + content + `"`), Valid: true, @@ -9076,9 +9078,9 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { t.Parallel() chat := newChat(t) - sys := insertMsg(t, chat.ID, "system", database.ChatMessageVisibilityModel, false, "system prompt") - usr := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityBoth, false, "hello") - ast := insertMsg(t, chat.ID, "assistant", database.ChatMessageVisibilityBoth, false, "hi there") + sys := insertMsg(t, chat.ID, database.ChatMessageRoleSystem, database.ChatMessageVisibilityModel, false, "system prompt") + usr := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "hello") + ast := insertMsg(t, chat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, false, "hi there") got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) @@ -9091,9 +9093,9 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { // Messages with visibility=user should NOT appear in the // prompt (they are only for the UI). - insertMsg(t, chat.ID, "system", database.ChatMessageVisibilityModel, false, "system prompt") - insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityUser, false, "user-only msg") - usr := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityBoth, false, "hello") + insertMsg(t, chat.ID, database.ChatMessageRoleSystem, database.ChatMessageVisibilityModel, false, "system prompt") + insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityUser, false, "user-only msg") + usr := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "hello") got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) @@ -9109,21 +9111,21 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { chat := newChat(t) // Pre-compaction conversation. - sys := insertMsg(t, chat.ID, "system", database.ChatMessageVisibilityModel, false, "system prompt") - preUser := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityBoth, false, "old question") - preAsst := insertMsg(t, chat.ID, "assistant", database.ChatMessageVisibilityBoth, false, "old answer") + sys := insertMsg(t, chat.ID, database.ChatMessageRoleSystem, database.ChatMessageVisibilityModel, false, "system prompt") + preUser := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "old question") + preAsst := insertMsg(t, chat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, false, "old answer") // Compaction messages: // 1. Summary (role=user, visibility=model, compressed=true). - summary := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityModel, true, "compaction summary") + summary := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, true, "compaction summary") // 2. Compressed assistant tool-call (visibility=user). - insertMsg(t, chat.ID, "assistant", database.ChatMessageVisibilityUser, true, "tool call") + insertMsg(t, chat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, true, "tool call") // 3. Compressed tool result (visibility=both). - insertMsg(t, chat.ID, "tool", database.ChatMessageVisibilityBoth, true, "tool result") + insertMsg(t, chat.ID, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, true, "tool result") // Post-compaction messages. - postUser := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityBoth, false, "new question") - postAsst := insertMsg(t, chat.ID, "assistant", database.ChatMessageVisibilityBoth, false, "new answer") + postUser := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "new question") + postAsst := insertMsg(t, chat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, false, "new answer") got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) @@ -9151,9 +9153,9 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { // After compaction the summary must appear as role=user so // that LLM APIs (e.g. Anthropic) see at least one // non-system message in the prompt. - insertMsg(t, chat.ID, "system", database.ChatMessageVisibilityModel, false, "system prompt") - summary := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityModel, true, "summary text") - newUsr := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityBoth, false, "new question") + insertMsg(t, chat.ID, database.ChatMessageRoleSystem, database.ChatMessageVisibilityModel, false, "system prompt") + summary := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, true, "summary text") + newUsr := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "new question") got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) @@ -9179,10 +9181,10 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { // used IN ('model','both'), the compressed tool result // (visibility=both) would be picked as the "summary" // instead of the actual summary. - insertMsg(t, chat.ID, "system", database.ChatMessageVisibilityModel, false, "system prompt") - summary := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityModel, true, "real summary") - compressedTool := insertMsg(t, chat.ID, "tool", database.ChatMessageVisibilityBoth, true, "tool result") - postUser := insertMsg(t, chat.ID, "user", database.ChatMessageVisibilityBoth, false, "follow-up") + insertMsg(t, chat.ID, database.ChatMessageRoleSystem, database.ChatMessageVisibilityModel, false, "system prompt") + summary := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, true, "real summary") + compressedTool := insertMsg(t, chat.ID, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, true, "tool result") + postUser := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "follow-up") got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 5b84ad764d..24de7b15b5 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -3311,7 +3311,7 @@ func (q *sqlQuerier) GetChatDiffStatusesByChatIDs(ctx context.Context, chatIds [ const getChatMessageByID = `-- name: GetChatMessageByID :one SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version FROM chat_messages WHERE @@ -3338,13 +3338,14 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess &i.ContextLimit, &i.Compressed, &i.CreatedBy, + &i.ContentVersion, ) return i, err } const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version FROM chat_messages WHERE @@ -3386,6 +3387,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes &i.ContextLimit, &i.Compressed, &i.CreatedBy, + &i.ContentVersion, ); err != nil { return nil, err } @@ -3417,7 +3419,7 @@ WITH latest_compressed_summary AS ( 1 ) SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version FROM chat_messages WHERE @@ -3483,6 +3485,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI &i.ContextLimit, &i.Compressed, &i.CreatedBy, + &i.ContentVersion, ); err != nil { return nil, err } @@ -3626,12 +3629,12 @@ func (q *sqlQuerier) GetChatsByOwnerID(ctx context.Context, arg GetChatsByOwnerI const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version FROM chat_messages WHERE chat_id = $1::uuid - AND role = $2::text + AND role = $2::chat_message_role ORDER BY created_at DESC, id DESC LIMIT @@ -3639,8 +3642,8 @@ LIMIT ` type GetLastChatMessageByRoleParams struct { - ChatID uuid.UUID `db:"chat_id" json:"chat_id"` - Role string `db:"role" json:"role"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Role ChatMessageRole `db:"role" json:"role"` } func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastChatMessageByRoleParams) (ChatMessage, error) { @@ -3663,6 +3666,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh &i.ContextLimit, &i.Compressed, &i.CreatedBy, + &i.ContentVersion, ) return i, err } @@ -3793,6 +3797,7 @@ INSERT INTO chat_messages ( model_config_id, role, content, + content_version, visibility, input_tokens, output_tokens, @@ -3806,28 +3811,30 @@ INSERT INTO chat_messages ( $1::uuid, $2::uuid, $3::uuid, - $4::text, + $4::chat_message_role, $5::jsonb, - $6::chat_message_visibility, - $7::bigint, + $6::smallint, + $7::chat_message_visibility, $8::bigint, $9::bigint, $10::bigint, $11::bigint, $12::bigint, $13::bigint, - COALESCE($14::boolean, FALSE) + $14::bigint, + COALESCE($15::boolean, FALSE) ) RETURNING - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version ` type InsertChatMessageParams struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"` ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` - Role string `db:"role" json:"role"` + Role ChatMessageRole `db:"role" json:"role"` Content pqtype.NullRawMessage `db:"content" json:"content"` + ContentVersion int16 `db:"content_version" json:"content_version"` Visibility ChatMessageVisibility `db:"visibility" json:"visibility"` InputTokens sql.NullInt64 `db:"input_tokens" json:"input_tokens"` OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` @@ -3846,6 +3853,7 @@ func (q *sqlQuerier) InsertChatMessage(ctx context.Context, arg InsertChatMessag arg.ModelConfigID, arg.Role, arg.Content, + arg.ContentVersion, arg.Visibility, arg.InputTokens, arg.OutputTokens, @@ -3874,6 +3882,7 @@ func (q *sqlQuerier) InsertChatMessage(ctx context.Context, arg InsertChatMessag &i.ContextLimit, &i.Compressed, &i.CreatedBy, + &i.ContentVersion, ) return i, err } @@ -4008,7 +4017,7 @@ SET WHERE id = $3::bigint RETURNING - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version ` type UpdateChatMessageByIDParams struct { @@ -4037,6 +4046,7 @@ func (q *sqlQuerier) UpdateChatMessageByID(ctx context.Context, arg UpdateChatMe &i.ContextLimit, &i.Compressed, &i.CreatedBy, + &i.ContentVersion, ) return i, err } diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 80435f24dd..d341b9d931 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -170,6 +170,7 @@ INSERT INTO chat_messages ( model_config_id, role, content, + content_version, visibility, input_tokens, output_tokens, @@ -183,8 +184,9 @@ INSERT INTO chat_messages ( @chat_id::uuid, sqlc.narg('created_by')::uuid, sqlc.narg('model_config_id')::uuid, - @role::text, + @role::chat_message_role, sqlc.narg('content')::jsonb, + @content_version::smallint, @visibility::chat_message_visibility, sqlc.narg('input_tokens')::bigint, sqlc.narg('output_tokens')::bigint, @@ -422,7 +424,7 @@ FROM chat_messages WHERE chat_id = @chat_id::uuid - AND role = @role::text + AND role = @role::chat_message_role ORDER BY created_at DESC, id DESC LIMIT