From 21f726957af88bbad9c39cdeb4e9779a33c3fb13 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Fri, 10 Jul 2026 17:35:29 +0000 Subject: [PATCH] refactor(coderd/x/chatd): stop storing gateway key IDs --- coderd/database/dbauthz/dbauthz.go | 153 +++- coderd/database/dbauthz/dbauthz_test.go | 76 +- coderd/database/dbgen/dbgen.go | 10 - coderd/database/dbmetrics/querymetrics.go | 142 +++- coderd/database/dbmock/dbmock.go | 148 +++- coderd/database/modelmethods.go | 83 ++ coderd/database/modelqueries.go | 84 ++ coderd/database/querier.go | 30 +- coderd/database/querier_test.go | 6 - coderd/database/queries.sql.go | 722 ++++++++++++++---- coderd/database/queries/chats.sql | 234 +++++- coderd/exp_chats.go | 6 +- coderd/exp_chats_test.go | 61 +- coderd/x/chatd/ARCHITECTURE.md | 2 +- coderd/x/chatd/chatd.go | 104 +-- coderd/x/chatd/chatd_internal_test.go | 13 +- coderd/x/chatd/chatd_test.go | 230 +----- coderd/x/chatd/chatstate/machine_test.go | 10 - coderd/x/chatd/chatstate/messages.go | 5 - .../chatstate/synthetic_cancellation_test.go | 1 - coderd/x/chatd/chatstate/transitions.go | 33 +- .../chatstate/transitions_helpers_test.go | 11 +- .../chatstate/transitions_matrix_test.go | 1 - coderd/x/chatd/chatstate_bridge.go | 5 +- coderd/x/chatd/generation.go | 1 - .../generation_preparer_internal_test.go | 72 +- coderd/x/chatd/helpers_test.go | 1 - coderd/x/chatd/message_conversion.go | 2 - coderd/x/chatd/model_routing_internal_test.go | 4 - coderd/x/chatd/stream_loop.go | 6 +- coderd/x/chatd/subagent.go | 7 +- coderd/x/chatd/subagent_internal_test.go | 77 -- coderd/x/chatd/synthetickey_internal_test.go | 16 +- coderd/x/chatd/tasks_test.go | 1 - coderd/x/chatd/turn_summary_internal_test.go | 2 - 35 files changed, 1541 insertions(+), 818 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index f2ce964249f..acaae654c62 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -3356,17 +3356,14 @@ func (q *querier) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, return q.db.GetChatIncludeDefaultSystemPrompt(ctx) } -func (q *querier) GetChatMessageByID(ctx context.Context, id int64) (database.ChatMessage, error) { - // ChatMessages are authorized through their parent Chat. - // We need to fetch the message first to get its chat_id. - msg, err := q.db.GetChatMessageByID(ctx, id) +func (q *querier) GetChatMessageByIDRaw(ctx context.Context, id int64) (database.GetChatMessageByIDRawRow, error) { + msg, err := q.db.GetChatMessageByIDRaw(ctx, id) if err != nil { - return database.ChatMessage{}, err + return database.GetChatMessageByIDRawRow{}, err } - // Authorize read on the parent chat. _, err = q.GetChatByID(ctx, msg.ChatID) if err != nil { - return database.ChatMessage{}, err + return database.GetChatMessageByIDRawRow{}, err } return msg, nil } @@ -3379,46 +3376,44 @@ func (q *querier) GetChatMessageSummariesPerChat(ctx context.Context, createdAft return q.db.GetChatMessageSummariesPerChat(ctx, createdAfter) } -func (q *querier) GetChatMessagesByChatID(ctx context.Context, arg database.GetChatMessagesByChatIDParams) ([]database.ChatMessage, error) { - // Authorize read on the parent chat. +func (q *querier) GetChatMessagesByChatIDAscPaginatedRaw(ctx context.Context, arg database.GetChatMessagesByChatIDAscPaginatedRawParams) ([]database.GetChatMessagesByChatIDAscPaginatedRawRow, error) { _, err := q.GetChatByID(ctx, arg.ChatID) if err != nil { return nil, err } - return q.db.GetChatMessagesByChatID(ctx, arg) + return q.db.GetChatMessagesByChatIDAscPaginatedRaw(ctx, arg) } -func (q *querier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDAscPaginatedParams) ([]database.ChatMessage, error) { +func (q *querier) GetChatMessagesByChatIDDescPaginatedRaw(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedRawParams) ([]database.GetChatMessagesByChatIDDescPaginatedRawRow, error) { _, err := q.GetChatByID(ctx, arg.ChatID) if err != nil { return nil, err } - return q.db.GetChatMessagesByChatIDAscPaginated(ctx, arg) + return q.db.GetChatMessagesByChatIDDescPaginatedRaw(ctx, arg) } -func (q *querier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedParams) ([]database.ChatMessage, error) { +func (q *querier) GetChatMessagesByChatIDRaw(ctx context.Context, arg database.GetChatMessagesByChatIDRawParams) ([]database.GetChatMessagesByChatIDRawRow, error) { _, err := q.GetChatByID(ctx, arg.ChatID) if err != nil { return nil, err } - return q.db.GetChatMessagesByChatIDDescPaginated(ctx, arg) + return q.db.GetChatMessagesByChatIDRaw(ctx, arg) } -func (q *querier) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { +func (q *querier) GetChatMessagesByRevisionForStreamRaw(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamRawParams) ([]database.GetChatMessagesByRevisionForStreamRawRow, error) { _, err := q.GetChatByID(ctx, arg.ChatID) if err != nil { return nil, err } - return q.db.GetChatMessagesByRevisionForStream(ctx, arg) + return q.db.GetChatMessagesByRevisionForStreamRaw(ctx, arg) } -func (q *querier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { - // Authorize read on the parent chat. +func (q *querier) GetChatMessagesForPromptByChatIDRaw(ctx context.Context, chatID uuid.UUID) ([]database.GetChatMessagesForPromptByChatIDRawRow, error) { _, err := q.GetChatByID(ctx, chatID) if err != nil { return nil, err } - return q.db.GetChatMessagesForPromptByChatID(ctx, chatID) + return q.db.GetChatMessagesForPromptByChatIDRaw(ctx, chatID) } func (q *querier) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) { @@ -3460,23 +3455,23 @@ func (q *querier) GetChatPlanModeInstructions(ctx context.Context) (string, erro return q.db.GetChatPlanModeInstructions(ctx) } -func (q *querier) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { +func (q *querier) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.GetChatQueuedMessageByIDRow, error) { _, err := q.GetChatByID(ctx, arg.ChatID) if err != nil { - return database.ChatQueuedMessage{}, err + return database.GetChatQueuedMessageByIDRow{}, err } return q.db.GetChatQueuedMessageByID(ctx, arg) } -func (q *querier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { +func (q *querier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.GetChatQueuedMessageHeadRow, error) { _, err := q.GetChatByID(ctx, chatID) if err != nil { - return database.ChatQueuedMessage{}, err + return database.GetChatQueuedMessageHeadRow{}, err } return q.db.GetChatQueuedMessageHead(ctx, chatID) } -func (q *querier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { +func (q *querier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.GetChatQueuedMessagesRow, error) { _, err := q.GetChatByID(ctx, chatID) if err != nil { return nil, err @@ -3484,7 +3479,7 @@ func (q *querier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ( return q.db.GetChatQueuedMessages(ctx, chatID) } -func (q *querier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { +func (q *querier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.GetChatQueuedMessagesByPositionRow, error) { _, err := q.GetChatByID(ctx, chatID) if err != nil { return nil, err @@ -3933,13 +3928,12 @@ func (q *querier) GetInboxNotificationsByUserID(ctx context.Context, userID data return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetInboxNotificationsByUserID)(ctx, userID) } -func (q *querier) GetLastChatMessageByRole(ctx context.Context, arg database.GetLastChatMessageByRoleParams) (database.ChatMessage, error) { - // Authorize read on the parent chat. +func (q *querier) GetLastChatMessageByRoleRaw(ctx context.Context, arg database.GetLastChatMessageByRoleRawParams) (database.GetLastChatMessageByRoleRawRow, error) { _, err := q.GetChatByID(ctx, arg.ChatID) if err != nil { - return database.ChatMessage{}, err + return database.GetLastChatMessageByRoleRawRow{}, err } - return q.db.GetLastChatMessageByRole(ctx, arg) + return q.db.GetLastChatMessageByRoleRaw(ctx, arg) } func (q *querier) GetLastUpdateCheck(ctx context.Context) (string, error) { @@ -5984,8 +5978,7 @@ func (q *querier) InsertChatFile(ctx context.Context, arg database.InsertChatFil return insert(q.log, q.auth, rbac.ResourceChat.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), q.db.InsertChatFile)(ctx, arg) } -func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { - // Authorize create on the parent chat (using update permission). +func (q *querier) InsertChatMessagesRaw(ctx context.Context, arg database.InsertChatMessagesRawParams) ([]database.InsertChatMessagesRawRow, error) { chat, err := q.db.GetChatByID(ctx, arg.ChatID) if err != nil { return nil, err @@ -5993,7 +5986,7 @@ func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertCha if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { return nil, err } - return q.db.InsertChatMessages(ctx, arg) + return q.db.InsertChatMessagesRaw(ctx, arg) } func (q *querier) InsertChatModelConfig(ctx context.Context, arg database.InsertChatModelConfigParams) (database.ChatModelConfig, error) { @@ -6003,24 +5996,24 @@ func (q *querier) InsertChatModelConfig(ctx context.Context, arg database.Insert return q.db.InsertChatModelConfig(ctx, arg) } -func (q *querier) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.ChatQueuedMessage, error) { +func (q *querier) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.InsertChatQueuedMessageRow, error) { chat, err := q.db.GetChatByID(ctx, arg.ChatID) if err != nil { - return database.ChatQueuedMessage{}, err + return database.InsertChatQueuedMessageRow{}, err } if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { - return database.ChatQueuedMessage{}, err + return database.InsertChatQueuedMessageRow{}, err } return q.db.InsertChatQueuedMessage(ctx, arg) } -func (q *querier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { +func (q *querier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.InsertChatQueuedMessageWithCreatorRow, error) { chat, err := q.db.GetChatByID(ctx, arg.ChatID) if err != nil { - return database.ChatQueuedMessage{}, err + return database.InsertChatQueuedMessageWithCreatorRow{}, err } if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { - return database.ChatQueuedMessage{}, err + return database.InsertChatQueuedMessageWithCreatorRow{}, err } _ = chat return q.db.InsertChatQueuedMessageWithCreator(ctx, arg) @@ -6952,13 +6945,13 @@ func (q *querier) PinChatByID(ctx context.Context, id uuid.UUID) error { return q.db.PinChatByID(ctx, id) } -func (q *querier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { +func (q *querier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.PopNextQueuedMessageRow, error) { chat, err := q.db.GetChatByID(ctx, chatID) if err != nil { - return database.ChatQueuedMessage{}, err + return database.PopNextQueuedMessageRow{}, err } if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { - return database.ChatQueuedMessage{}, err + return database.PopNextQueuedMessageRow{}, err } return q.db.PopNextQueuedMessage(ctx, chatID) } @@ -9303,3 +9296,81 @@ func (q *querier) GetAuthorizedChats(ctx context.Context, arg database.GetChatsP func (q *querier) GetAuthorizedChatsByChatFileID(ctx context.Context, fileID uuid.UUID, prepared rbac.PreparedAuthorized) ([]database.Chat, error) { return q.db.GetAuthorizedChatsByChatFileID(ctx, fileID, prepared) } + +func (q *querier) GetChatMessageByID(ctx context.Context, id int64) (database.ChatMessage, error) { + // ChatMessages are authorized through their parent Chat. + // We need to fetch the message first to get its chat_id. + msg, err := q.db.GetChatMessageByID(ctx, id) + if err != nil { + return database.ChatMessage{}, err + } + // Authorize read on the parent chat. + _, err = q.GetChatByID(ctx, msg.ChatID) + if err != nil { + return database.ChatMessage{}, err + } + return msg, nil +} + +func (q *querier) GetChatMessagesByChatID(ctx context.Context, arg database.GetChatMessagesByChatIDParams) ([]database.ChatMessage, error) { + // Authorize read on the parent chat. + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return nil, err + } + return q.db.GetChatMessagesByChatID(ctx, arg) +} + +func (q *querier) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return nil, err + } + return q.db.GetChatMessagesByRevisionForStream(ctx, arg) +} + +func (q *querier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDAscPaginatedParams) ([]database.ChatMessage, error) { + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return nil, err + } + return q.db.GetChatMessagesByChatIDAscPaginated(ctx, arg) +} + +func (q *querier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedParams) ([]database.ChatMessage, error) { + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return nil, err + } + return q.db.GetChatMessagesByChatIDDescPaginated(ctx, arg) +} + +func (q *querier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { + // Authorize read on the parent chat. + _, err := q.GetChatByID(ctx, chatID) + if err != nil { + return nil, err + } + return q.db.GetChatMessagesForPromptByChatID(ctx, chatID) +} + +func (q *querier) GetLastChatMessageByRole(ctx context.Context, arg database.GetLastChatMessageByRoleParams) (database.ChatMessage, error) { + // Authorize read on the parent chat. + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return database.ChatMessage{}, err + } + return q.db.GetLastChatMessageByRole(ctx, arg) +} + +func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { + // Authorize create on the parent chat (using update permission). + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if err != nil { + return nil, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return nil, err + } + return q.db.InsertChatMessages(ctx, arg) +} diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 0d5c8e2c38a..d61097d8d87 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -1129,6 +1129,60 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().GetChatMessagesForPromptByChatID(gomock.Any(), chat.ID).Return(msgs, nil).AnyTimes() check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(msgs) })) + s.Run("GetChatMessageByIDRaw", 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.GetChatMessageByIDRawRow{ChatID: chat.ID}) + dbm.EXPECT().GetChatMessageByIDRaw(gomock.Any(), msg.ID).Return(msg, nil).AnyTimes() + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + check.Args(msg.ID).Asserts(chat, policy.ActionRead).Returns(msg) + })) + s.Run("GetChatMessagesByChatIDRaw", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + msgs := []database.GetChatMessagesByChatIDRawRow{testutil.Fake(s.T(), faker, database.GetChatMessagesByChatIDRawRow{ChatID: chat.ID})} + arg := database.GetChatMessagesByChatIDRawParams{ChatID: chat.ID, AfterID: 0} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatMessagesByChatIDRaw(gomock.Any(), arg).Return(msgs, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs) + })) + s.Run("GetChatMessagesByChatIDAscPaginatedRaw", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + msgs := []database.GetChatMessagesByChatIDAscPaginatedRawRow{testutil.Fake(s.T(), faker, database.GetChatMessagesByChatIDAscPaginatedRawRow{ChatID: chat.ID})} + arg := database.GetChatMessagesByChatIDAscPaginatedRawParams{ChatID: chat.ID, AfterID: 0, LimitVal: 50} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatMessagesByChatIDAscPaginatedRaw(gomock.Any(), arg).Return(msgs, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs) + })) + s.Run("GetChatMessagesByChatIDDescPaginatedRaw", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + msgs := []database.GetChatMessagesByChatIDDescPaginatedRawRow{testutil.Fake(s.T(), faker, database.GetChatMessagesByChatIDDescPaginatedRawRow{ChatID: chat.ID})} + arg := database.GetChatMessagesByChatIDDescPaginatedRawParams{ChatID: chat.ID, BeforeID: 0, LimitVal: 50} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatMessagesByChatIDDescPaginatedRaw(gomock.Any(), arg).Return(msgs, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs) + })) + s.Run("GetChatMessagesByRevisionForStreamRaw", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + msgs := []database.GetChatMessagesByRevisionForStreamRawRow{testutil.Fake(s.T(), faker, database.GetChatMessagesByRevisionForStreamRawRow{ChatID: chat.ID})} + arg := database.GetChatMessagesByRevisionForStreamRawParams{ChatID: chat.ID, AfterRevision: 1} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatMessagesByRevisionForStreamRaw(gomock.Any(), arg).Return(msgs, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs) + })) + s.Run("GetLastChatMessageByRoleRaw", 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.GetLastChatMessageByRoleRawRow{ChatID: chat.ID}) + arg := database.GetLastChatMessageByRoleRawParams{ChatID: chat.ID, Role: database.ChatMessageRoleAssistant} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetLastChatMessageByRoleRaw(gomock.Any(), arg).Return(msg, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msg) + })) + s.Run("GetChatMessagesForPromptByChatIDRaw", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + msgs := []database.GetChatMessagesForPromptByChatIDRawRow{testutil.Fake(s.T(), faker, database.GetChatMessagesForPromptByChatIDRawRow{ChatID: chat.ID})} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatMessagesForPromptByChatIDRaw(gomock.Any(), chat.ID).Return(msgs, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(msgs) + })) s.Run("GetChatModelConfigByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { config := testutil.Fake(s.T(), faker, database.ChatModelConfig{}) dbm.EXPECT().GetChatModelConfigByID(gomock.Any(), config.ID).Return(config, nil).AnyTimes() @@ -1195,7 +1249,7 @@ func (s *MethodTestSuite) TestChats() { })) s.Run("GetChatQueuedMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) - qms := []database.ChatQueuedMessage{testutil.Fake(s.T(), faker, database.ChatQueuedMessage{})} + qms := []database.GetChatQueuedMessagesRow{testutil.Fake(s.T(), faker, database.GetChatQueuedMessagesRow{})} dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().GetChatQueuedMessages(gomock.Any(), chat.ID).Return(qms, nil).AnyTimes() check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(qms) @@ -1288,10 +1342,18 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().InsertChatMessages(gomock.Any(), arg).Return(msgs, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(msgs) })) + s.Run("InsertChatMessagesRaw", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := testutil.Fake(s.T(), faker, database.InsertChatMessagesRawParams{ChatID: chat.ID}) + msgs := []database.InsertChatMessagesRawRow{testutil.Fake(s.T(), faker, database.InsertChatMessagesRawRow{ChatID: chat.ID})} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().InsertChatMessagesRaw(gomock.Any(), arg).Return(msgs, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(msgs) + })) s.Run("InsertChatQueuedMessage", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) arg := testutil.Fake(s.T(), faker, database.InsertChatQueuedMessageParams{ChatID: chat.ID}) - qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{}) + qm := testutil.Fake(s.T(), faker, database.InsertChatQueuedMessageRow{}) dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().InsertChatQueuedMessage(gomock.Any(), arg).Return(qm, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(qm) @@ -1309,7 +1371,7 @@ func (s *MethodTestSuite) TestChats() { s.Run("PopNextQueuedMessage", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) - qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{}) + qm := testutil.Fake(s.T(), faker, database.PopNextQueuedMessageRow{}) dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().PopNextQueuedMessage(gomock.Any(), chat.ID).Return(qm, nil).AnyTimes() check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(qm) @@ -1368,14 +1430,14 @@ func (s *MethodTestSuite) TestChats() { s.Run("InsertChatQueuedMessageWithCreator", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) arg := testutil.Fake(s.T(), faker, database.InsertChatQueuedMessageWithCreatorParams{ChatID: chat.ID}) - qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{}) + qm := testutil.Fake(s.T(), faker, database.InsertChatQueuedMessageWithCreatorRow{}) dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().InsertChatQueuedMessageWithCreator(gomock.Any(), arg).Return(qm, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(qm) })) s.Run("GetChatQueuedMessagesByPosition", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) - qms := []database.ChatQueuedMessage{} + qms := []database.GetChatQueuedMessagesByPositionRow{} dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().GetChatQueuedMessagesByPosition(gomock.Any(), chat.ID).Return(qms, nil).AnyTimes() check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(qms) @@ -1388,14 +1450,14 @@ func (s *MethodTestSuite) TestChats() { })) s.Run("GetChatQueuedMessageHead", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) - qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{ChatID: chat.ID}) + qm := testutil.Fake(s.T(), faker, database.GetChatQueuedMessageHeadRow{ChatID: chat.ID}) dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().GetChatQueuedMessageHead(gomock.Any(), chat.ID).Return(qm, nil).AnyTimes() check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(qm) })) s.Run("GetChatQueuedMessageByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) - qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{ChatID: chat.ID}) + qm := testutil.Fake(s.T(), faker, database.GetChatQueuedMessageByIDRow{ChatID: chat.ID}) arg := database.GetChatQueuedMessageByIDParams{ID: qm.ID, ChatID: chat.ID} dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().GetChatQueuedMessageByID(gomock.Any(), arg).Return(qm, nil).AnyTimes() diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index fc071c5f11a..7efa2fafe11 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -116,20 +116,10 @@ func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) dat content = string(seed.Content.RawMessage) } role := takeFirst(seed.Role, database.ChatMessageRoleUser) - apiKeyID := seed.APIKeyID.String - // Mint a real API key for user turns so the api_key_id foreign key is - // satisfied. Without a creator we leave it empty, which the insert query - // stores as NULL. - if role == database.ChatMessageRoleUser && apiKeyID == "" && - seed.CreatedBy.Valid && seed.CreatedBy.UUID != uuid.Nil { - key, _ := APIKey(t, db, database.APIKey{UserID: seed.CreatedBy.UUID}) - apiKeyID = key.ID - } msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{ ChatID: seed.ChatID, CreatedBy: []uuid.UUID{seed.CreatedBy.UUID}, - APIKeyID: []string{apiKeyID}, ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID}, ReasoningEffort: []string{string(seed.ReasoningEffort.ChatReasoningEffort)}, Role: []database.ChatMessageRole{role}, diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index b2789a53f25..86fa3a70fbb 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1609,11 +1609,11 @@ func (m queryMetricsStore) GetChatIncludeDefaultSystemPrompt(ctx context.Context return r0, r1 } -func (m queryMetricsStore) GetChatMessageByID(ctx context.Context, id int64) (database.ChatMessage, error) { +func (m queryMetricsStore) GetChatMessageByIDRaw(ctx context.Context, id int64) (database.GetChatMessageByIDRawRow, error) { start := time.Now() - r0, r1 := m.s.GetChatMessageByID(ctx, id) - m.queryLatencies.WithLabelValues("GetChatMessageByID").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessageByID").Inc() + r0, r1 := m.s.GetChatMessageByIDRaw(ctx, id) + m.queryLatencies.WithLabelValues("GetChatMessageByIDRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessageByIDRaw").Inc() return r0, r1 } @@ -1625,43 +1625,43 @@ func (m queryMetricsStore) GetChatMessageSummariesPerChat(ctx context.Context, c return r0, r1 } -func (m queryMetricsStore) GetChatMessagesByChatID(ctx context.Context, chatID database.GetChatMessagesByChatIDParams) ([]database.ChatMessage, error) { +func (m queryMetricsStore) GetChatMessagesByChatIDAscPaginatedRaw(ctx context.Context, arg database.GetChatMessagesByChatIDAscPaginatedRawParams) ([]database.GetChatMessagesByChatIDAscPaginatedRawRow, error) { start := time.Now() - r0, r1 := m.s.GetChatMessagesByChatID(ctx, chatID) - m.queryLatencies.WithLabelValues("GetChatMessagesByChatID").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatID").Inc() + r0, r1 := m.s.GetChatMessagesByChatIDAscPaginatedRaw(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByChatIDAscPaginatedRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatIDAscPaginatedRaw").Inc() return r0, r1 } -func (m queryMetricsStore) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDAscPaginatedParams) ([]database.ChatMessage, error) { +func (m queryMetricsStore) GetChatMessagesByChatIDDescPaginatedRaw(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedRawParams) ([]database.GetChatMessagesByChatIDDescPaginatedRawRow, error) { start := time.Now() - r0, r1 := m.s.GetChatMessagesByChatIDAscPaginated(ctx, arg) - m.queryLatencies.WithLabelValues("GetChatMessagesByChatIDAscPaginated").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatIDAscPaginated").Inc() + r0, r1 := m.s.GetChatMessagesByChatIDDescPaginatedRaw(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByChatIDDescPaginatedRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatIDDescPaginatedRaw").Inc() return r0, r1 } -func (m queryMetricsStore) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedParams) ([]database.ChatMessage, error) { +func (m queryMetricsStore) GetChatMessagesByChatIDRaw(ctx context.Context, arg database.GetChatMessagesByChatIDRawParams) ([]database.GetChatMessagesByChatIDRawRow, error) { start := time.Now() - r0, r1 := m.s.GetChatMessagesByChatIDDescPaginated(ctx, arg) - m.queryLatencies.WithLabelValues("GetChatMessagesByChatIDDescPaginated").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatIDDescPaginated").Inc() + r0, r1 := m.s.GetChatMessagesByChatIDRaw(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByChatIDRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatIDRaw").Inc() return r0, r1 } -func (m queryMetricsStore) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { +func (m queryMetricsStore) GetChatMessagesByRevisionForStreamRaw(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamRawParams) ([]database.GetChatMessagesByRevisionForStreamRawRow, error) { start := time.Now() - r0, r1 := m.s.GetChatMessagesByRevisionForStream(ctx, arg) - m.queryLatencies.WithLabelValues("GetChatMessagesByRevisionForStream").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByRevisionForStream").Inc() + r0, r1 := m.s.GetChatMessagesByRevisionForStreamRaw(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByRevisionForStreamRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByRevisionForStreamRaw").Inc() return r0, r1 } -func (m queryMetricsStore) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { +func (m queryMetricsStore) GetChatMessagesForPromptByChatIDRaw(ctx context.Context, chatID uuid.UUID) ([]database.GetChatMessagesForPromptByChatIDRawRow, error) { start := time.Now() - r0, r1 := m.s.GetChatMessagesForPromptByChatID(ctx, chatID) - m.queryLatencies.WithLabelValues("GetChatMessagesForPromptByChatID").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesForPromptByChatID").Inc() + r0, r1 := m.s.GetChatMessagesForPromptByChatIDRaw(ctx, chatID) + m.queryLatencies.WithLabelValues("GetChatMessagesForPromptByChatIDRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesForPromptByChatIDRaw").Inc() return r0, r1 } @@ -1705,7 +1705,7 @@ func (m queryMetricsStore) GetChatPlanModeInstructions(ctx context.Context) (str return r0, r1 } -func (m queryMetricsStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { +func (m queryMetricsStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.GetChatQueuedMessageByIDRow, error) { start := time.Now() r0, r1 := m.s.GetChatQueuedMessageByID(ctx, arg) m.queryLatencies.WithLabelValues("GetChatQueuedMessageByID").Observe(time.Since(start).Seconds()) @@ -1713,7 +1713,7 @@ func (m queryMetricsStore) GetChatQueuedMessageByID(ctx context.Context, arg dat return r0, r1 } -func (m queryMetricsStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { +func (m queryMetricsStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.GetChatQueuedMessageHeadRow, error) { start := time.Now() r0, r1 := m.s.GetChatQueuedMessageHead(ctx, chatID) m.queryLatencies.WithLabelValues("GetChatQueuedMessageHead").Observe(time.Since(start).Seconds()) @@ -1721,7 +1721,7 @@ func (m queryMetricsStore) GetChatQueuedMessageHead(ctx context.Context, chatID return r0, r1 } -func (m queryMetricsStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { +func (m queryMetricsStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.GetChatQueuedMessagesRow, error) { start := time.Now() r0, r1 := m.s.GetChatQueuedMessages(ctx, chatID) m.queryLatencies.WithLabelValues("GetChatQueuedMessages").Observe(time.Since(start).Seconds()) @@ -1729,7 +1729,7 @@ func (m queryMetricsStore) GetChatQueuedMessages(ctx context.Context, chatID uui return r0, r1 } -func (m queryMetricsStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { +func (m queryMetricsStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.GetChatQueuedMessagesByPositionRow, error) { start := time.Now() r0, r1 := m.s.GetChatQueuedMessagesByPosition(ctx, chatID) m.queryLatencies.WithLabelValues("GetChatQueuedMessagesByPosition").Observe(time.Since(start).Seconds()) @@ -2217,11 +2217,11 @@ func (m queryMetricsStore) GetInboxNotificationsByUserID(ctx context.Context, ar return r0, r1 } -func (m queryMetricsStore) GetLastChatMessageByRole(ctx context.Context, arg database.GetLastChatMessageByRoleParams) (database.ChatMessage, error) { +func (m queryMetricsStore) GetLastChatMessageByRoleRaw(ctx context.Context, arg database.GetLastChatMessageByRoleRawParams) (database.GetLastChatMessageByRoleRawRow, error) { start := time.Now() - r0, r1 := m.s.GetLastChatMessageByRole(ctx, arg) - m.queryLatencies.WithLabelValues("GetLastChatMessageByRole").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetLastChatMessageByRole").Inc() + r0, r1 := m.s.GetLastChatMessageByRoleRaw(ctx, arg) + m.queryLatencies.WithLabelValues("GetLastChatMessageByRoleRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetLastChatMessageByRoleRaw").Inc() return r0, r1 } @@ -4081,11 +4081,11 @@ func (m queryMetricsStore) InsertChatFile(ctx context.Context, arg database.Inse return r0, r1 } -func (m queryMetricsStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { +func (m queryMetricsStore) InsertChatMessagesRaw(ctx context.Context, arg database.InsertChatMessagesRawParams) ([]database.InsertChatMessagesRawRow, error) { start := time.Now() - r0, r1 := m.s.InsertChatMessages(ctx, arg) - m.queryLatencies.WithLabelValues("InsertChatMessages").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatMessages").Inc() + r0, r1 := m.s.InsertChatMessagesRaw(ctx, arg) + m.queryLatencies.WithLabelValues("InsertChatMessagesRaw").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatMessagesRaw").Inc() return r0, r1 } @@ -4097,7 +4097,7 @@ func (m queryMetricsStore) InsertChatModelConfig(ctx context.Context, arg databa return r0, r1 } -func (m queryMetricsStore) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.ChatQueuedMessage, error) { +func (m queryMetricsStore) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.InsertChatQueuedMessageRow, error) { start := time.Now() r0, r1 := m.s.InsertChatQueuedMessage(ctx, arg) m.queryLatencies.WithLabelValues("InsertChatQueuedMessage").Observe(time.Since(start).Seconds()) @@ -4105,7 +4105,7 @@ func (m queryMetricsStore) InsertChatQueuedMessage(ctx context.Context, arg data return r0, r1 } -func (m queryMetricsStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { +func (m queryMetricsStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.InsertChatQueuedMessageWithCreatorRow, error) { start := time.Now() r0, r1 := m.s.InsertChatQueuedMessageWithCreator(ctx, arg) m.queryLatencies.WithLabelValues("InsertChatQueuedMessageWithCreator").Observe(time.Since(start).Seconds()) @@ -4897,7 +4897,7 @@ func (m queryMetricsStore) PinChatByID(ctx context.Context, id uuid.UUID) error return r0 } -func (m queryMetricsStore) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { +func (m queryMetricsStore) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.PopNextQueuedMessageRow, error) { start := time.Now() r0, r1 := m.s.PopNextQueuedMessage(ctx, chatID) m.queryLatencies.WithLabelValues("PopNextQueuedMessage").Observe(time.Since(start).Seconds()) @@ -6728,3 +6728,67 @@ func (m queryMetricsStore) GetAuthorizedChatsByChatFileID(ctx context.Context, f m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAuthorizedChatsByChatFileID").Inc() return r0, r1 } + +func (m queryMetricsStore) GetChatMessageByID(ctx context.Context, id int64) (database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatMessageByID(ctx, id) + m.queryLatencies.WithLabelValues("GetChatMessageByID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessageByID").Inc() + return r0, r1 +} + +func (m queryMetricsStore) GetChatMessagesByChatID(ctx context.Context, chatID database.GetChatMessagesByChatIDParams) ([]database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatMessagesByChatID(ctx, chatID) + m.queryLatencies.WithLabelValues("GetChatMessagesByChatID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatID").Inc() + return r0, r1 +} + +func (m queryMetricsStore) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatMessagesByRevisionForStream(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByRevisionForStream").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByRevisionForStream").Inc() + return r0, r1 +} + +func (m queryMetricsStore) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDAscPaginatedParams) ([]database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatMessagesByChatIDAscPaginated(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByChatIDAscPaginated").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatIDAscPaginated").Inc() + return r0, r1 +} + +func (m queryMetricsStore) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedParams) ([]database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatMessagesByChatIDDescPaginated(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByChatIDDescPaginated").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByChatIDDescPaginated").Inc() + return r0, r1 +} + +func (m queryMetricsStore) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatMessagesForPromptByChatID(ctx, chatID) + m.queryLatencies.WithLabelValues("GetChatMessagesForPromptByChatID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesForPromptByChatID").Inc() + return r0, r1 +} + +func (m queryMetricsStore) GetLastChatMessageByRole(ctx context.Context, arg database.GetLastChatMessageByRoleParams) (database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetLastChatMessageByRole(ctx, arg) + m.queryLatencies.WithLabelValues("GetLastChatMessageByRole").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetLastChatMessageByRole").Inc() + return r0, r1 +} + +func (m queryMetricsStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.InsertChatMessages(ctx, arg) + m.queryLatencies.WithLabelValues("InsertChatMessages").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatMessages").Inc() + return r0, r1 +} diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 9f43514310b..21053fcba92 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -2983,6 +2983,21 @@ func (mr *MockStoreMockRecorder) GetChatMessageByID(ctx, id any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessageByID", reflect.TypeOf((*MockStore)(nil).GetChatMessageByID), ctx, id) } +// GetChatMessageByIDRaw mocks base method. +func (m *MockStore) GetChatMessageByIDRaw(ctx context.Context, id int64) (database.GetChatMessageByIDRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatMessageByIDRaw", ctx, id) + ret0, _ := ret[0].(database.GetChatMessageByIDRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatMessageByIDRaw indicates an expected call of GetChatMessageByIDRaw. +func (mr *MockStoreMockRecorder) GetChatMessageByIDRaw(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessageByIDRaw", reflect.TypeOf((*MockStore)(nil).GetChatMessageByIDRaw), ctx, id) +} + // GetChatMessageSummariesPerChat mocks base method. func (m *MockStore) GetChatMessageSummariesPerChat(ctx context.Context, createdAfter time.Time) ([]database.GetChatMessageSummariesPerChatRow, error) { m.ctrl.T.Helper() @@ -3028,6 +3043,21 @@ func (mr *MockStoreMockRecorder) GetChatMessagesByChatIDAscPaginated(ctx, arg an return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByChatIDAscPaginated", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByChatIDAscPaginated), ctx, arg) } +// GetChatMessagesByChatIDAscPaginatedRaw mocks base method. +func (m *MockStore) GetChatMessagesByChatIDAscPaginatedRaw(ctx context.Context, arg database.GetChatMessagesByChatIDAscPaginatedRawParams) ([]database.GetChatMessagesByChatIDAscPaginatedRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatMessagesByChatIDAscPaginatedRaw", ctx, arg) + ret0, _ := ret[0].([]database.GetChatMessagesByChatIDAscPaginatedRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatMessagesByChatIDAscPaginatedRaw indicates an expected call of GetChatMessagesByChatIDAscPaginatedRaw. +func (mr *MockStoreMockRecorder) GetChatMessagesByChatIDAscPaginatedRaw(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByChatIDAscPaginatedRaw", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByChatIDAscPaginatedRaw), ctx, arg) +} + // GetChatMessagesByChatIDDescPaginated mocks base method. func (m *MockStore) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedParams) ([]database.ChatMessage, error) { m.ctrl.T.Helper() @@ -3043,6 +3073,36 @@ func (mr *MockStoreMockRecorder) GetChatMessagesByChatIDDescPaginated(ctx, arg a return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByChatIDDescPaginated", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByChatIDDescPaginated), ctx, arg) } +// GetChatMessagesByChatIDDescPaginatedRaw mocks base method. +func (m *MockStore) GetChatMessagesByChatIDDescPaginatedRaw(ctx context.Context, arg database.GetChatMessagesByChatIDDescPaginatedRawParams) ([]database.GetChatMessagesByChatIDDescPaginatedRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatMessagesByChatIDDescPaginatedRaw", ctx, arg) + ret0, _ := ret[0].([]database.GetChatMessagesByChatIDDescPaginatedRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatMessagesByChatIDDescPaginatedRaw indicates an expected call of GetChatMessagesByChatIDDescPaginatedRaw. +func (mr *MockStoreMockRecorder) GetChatMessagesByChatIDDescPaginatedRaw(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByChatIDDescPaginatedRaw", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByChatIDDescPaginatedRaw), ctx, arg) +} + +// GetChatMessagesByChatIDRaw mocks base method. +func (m *MockStore) GetChatMessagesByChatIDRaw(ctx context.Context, arg database.GetChatMessagesByChatIDRawParams) ([]database.GetChatMessagesByChatIDRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatMessagesByChatIDRaw", ctx, arg) + ret0, _ := ret[0].([]database.GetChatMessagesByChatIDRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatMessagesByChatIDRaw indicates an expected call of GetChatMessagesByChatIDRaw. +func (mr *MockStoreMockRecorder) GetChatMessagesByChatIDRaw(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByChatIDRaw", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByChatIDRaw), ctx, arg) +} + // GetChatMessagesByRevisionForStream mocks base method. func (m *MockStore) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { m.ctrl.T.Helper() @@ -3058,6 +3118,21 @@ func (mr *MockStoreMockRecorder) GetChatMessagesByRevisionForStream(ctx, arg any return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByRevisionForStream", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByRevisionForStream), ctx, arg) } +// GetChatMessagesByRevisionForStreamRaw mocks base method. +func (m *MockStore) GetChatMessagesByRevisionForStreamRaw(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamRawParams) ([]database.GetChatMessagesByRevisionForStreamRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatMessagesByRevisionForStreamRaw", ctx, arg) + ret0, _ := ret[0].([]database.GetChatMessagesByRevisionForStreamRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatMessagesByRevisionForStreamRaw indicates an expected call of GetChatMessagesByRevisionForStreamRaw. +func (mr *MockStoreMockRecorder) GetChatMessagesByRevisionForStreamRaw(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByRevisionForStreamRaw", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByRevisionForStreamRaw), ctx, arg) +} + // GetChatMessagesForPromptByChatID mocks base method. func (m *MockStore) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { m.ctrl.T.Helper() @@ -3073,6 +3148,21 @@ func (mr *MockStoreMockRecorder) GetChatMessagesForPromptByChatID(ctx, chatID an return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesForPromptByChatID", reflect.TypeOf((*MockStore)(nil).GetChatMessagesForPromptByChatID), ctx, chatID) } +// GetChatMessagesForPromptByChatIDRaw mocks base method. +func (m *MockStore) GetChatMessagesForPromptByChatIDRaw(ctx context.Context, chatID uuid.UUID) ([]database.GetChatMessagesForPromptByChatIDRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatMessagesForPromptByChatIDRaw", ctx, chatID) + ret0, _ := ret[0].([]database.GetChatMessagesForPromptByChatIDRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatMessagesForPromptByChatIDRaw indicates an expected call of GetChatMessagesForPromptByChatIDRaw. +func (mr *MockStoreMockRecorder) GetChatMessagesForPromptByChatIDRaw(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesForPromptByChatIDRaw", reflect.TypeOf((*MockStore)(nil).GetChatMessagesForPromptByChatIDRaw), ctx, chatID) +} + // GetChatModelConfigByID mocks base method. func (m *MockStore) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) { m.ctrl.T.Helper() @@ -3149,10 +3239,10 @@ func (mr *MockStoreMockRecorder) GetChatPlanModeInstructions(ctx any) *gomock.Ca } // GetChatQueuedMessageByID mocks base method. -func (m *MockStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { +func (m *MockStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.GetChatQueuedMessageByIDRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GetChatQueuedMessageByID", ctx, arg) - ret0, _ := ret[0].(database.ChatQueuedMessage) + ret0, _ := ret[0].(database.GetChatQueuedMessageByIDRow) ret1, _ := ret[1].(error) return ret0, ret1 } @@ -3164,10 +3254,10 @@ func (mr *MockStoreMockRecorder) GetChatQueuedMessageByID(ctx, arg any) *gomock. } // GetChatQueuedMessageHead mocks base method. -func (m *MockStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { +func (m *MockStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.GetChatQueuedMessageHeadRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GetChatQueuedMessageHead", ctx, chatID) - ret0, _ := ret[0].(database.ChatQueuedMessage) + ret0, _ := ret[0].(database.GetChatQueuedMessageHeadRow) ret1, _ := ret[1].(error) return ret0, ret1 } @@ -3179,10 +3269,10 @@ func (mr *MockStoreMockRecorder) GetChatQueuedMessageHead(ctx, chatID any) *gomo } // GetChatQueuedMessages mocks base method. -func (m *MockStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { +func (m *MockStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.GetChatQueuedMessagesRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GetChatQueuedMessages", ctx, chatID) - ret0, _ := ret[0].([]database.ChatQueuedMessage) + ret0, _ := ret[0].([]database.GetChatQueuedMessagesRow) ret1, _ := ret[1].(error) return ret0, ret1 } @@ -3194,10 +3284,10 @@ func (mr *MockStoreMockRecorder) GetChatQueuedMessages(ctx, chatID any) *gomock. } // GetChatQueuedMessagesByPosition mocks base method. -func (m *MockStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { +func (m *MockStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.GetChatQueuedMessagesByPositionRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GetChatQueuedMessagesByPosition", ctx, chatID) - ret0, _ := ret[0].([]database.ChatQueuedMessage) + ret0, _ := ret[0].([]database.GetChatQueuedMessagesByPositionRow) ret1, _ := ret[1].(error) return ret0, ret1 } @@ -4123,6 +4213,21 @@ func (mr *MockStoreMockRecorder) GetLastChatMessageByRole(ctx, arg any) *gomock. return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLastChatMessageByRole", reflect.TypeOf((*MockStore)(nil).GetLastChatMessageByRole), ctx, arg) } +// GetLastChatMessageByRoleRaw mocks base method. +func (m *MockStore) GetLastChatMessageByRoleRaw(ctx context.Context, arg database.GetLastChatMessageByRoleRawParams) (database.GetLastChatMessageByRoleRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetLastChatMessageByRoleRaw", ctx, arg) + ret0, _ := ret[0].(database.GetLastChatMessageByRoleRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetLastChatMessageByRoleRaw indicates an expected call of GetLastChatMessageByRoleRaw. +func (mr *MockStoreMockRecorder) GetLastChatMessageByRoleRaw(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLastChatMessageByRoleRaw", reflect.TypeOf((*MockStore)(nil).GetLastChatMessageByRoleRaw), ctx, arg) +} + // GetLastUpdateCheck mocks base method. func (m *MockStore) GetLastUpdateCheck(ctx context.Context) (string, error) { m.ctrl.T.Helper() @@ -7660,6 +7765,21 @@ func (mr *MockStoreMockRecorder) InsertChatMessages(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatMessages", reflect.TypeOf((*MockStore)(nil).InsertChatMessages), ctx, arg) } +// InsertChatMessagesRaw mocks base method. +func (m *MockStore) InsertChatMessagesRaw(ctx context.Context, arg database.InsertChatMessagesRawParams) ([]database.InsertChatMessagesRawRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InsertChatMessagesRaw", ctx, arg) + ret0, _ := ret[0].([]database.InsertChatMessagesRawRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// InsertChatMessagesRaw indicates an expected call of InsertChatMessagesRaw. +func (mr *MockStoreMockRecorder) InsertChatMessagesRaw(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatMessagesRaw", reflect.TypeOf((*MockStore)(nil).InsertChatMessagesRaw), ctx, arg) +} + // InsertChatModelConfig mocks base method. func (m *MockStore) InsertChatModelConfig(ctx context.Context, arg database.InsertChatModelConfigParams) (database.ChatModelConfig, error) { m.ctrl.T.Helper() @@ -7676,10 +7796,10 @@ func (mr *MockStoreMockRecorder) InsertChatModelConfig(ctx, arg any) *gomock.Cal } // InsertChatQueuedMessage mocks base method. -func (m *MockStore) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.ChatQueuedMessage, error) { +func (m *MockStore) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.InsertChatQueuedMessageRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "InsertChatQueuedMessage", ctx, arg) - ret0, _ := ret[0].(database.ChatQueuedMessage) + ret0, _ := ret[0].(database.InsertChatQueuedMessageRow) ret1, _ := ret[1].(error) return ret0, ret1 } @@ -7691,10 +7811,10 @@ func (mr *MockStoreMockRecorder) InsertChatQueuedMessage(ctx, arg any) *gomock.C } // InsertChatQueuedMessageWithCreator mocks base method. -func (m *MockStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { +func (m *MockStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.InsertChatQueuedMessageWithCreatorRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "InsertChatQueuedMessageWithCreator", ctx, arg) - ret0, _ := ret[0].(database.ChatQueuedMessage) + ret0, _ := ret[0].(database.InsertChatQueuedMessageWithCreatorRow) ret1, _ := ret[1].(error) return ret0, ret1 } @@ -9249,10 +9369,10 @@ func (mr *MockStoreMockRecorder) Ping(ctx any) *gomock.Call { } // PopNextQueuedMessage mocks base method. -func (m *MockStore) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { +func (m *MockStore) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.PopNextQueuedMessageRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "PopNextQueuedMessage", ctx, chatID) - ret0, _ := ret[0].(database.ChatQueuedMessage) + ret0, _ := ret[0].(database.PopNextQueuedMessageRow) ret1, _ := ret[1].(error) return ret0, ret1 } diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go index 8bc15192637..49eed7d49a4 100644 --- a/coderd/database/modelmethods.go +++ b/coderd/database/modelmethods.go @@ -3,6 +3,7 @@ package database import ( "database/sql" "encoding/hex" + "encoding/json" "fmt" "slices" "sort" @@ -20,6 +21,88 @@ import ( "github.com/coder/coder/v2/coderd/rbac/policy" ) +// ChatMessageRow is the persisted message projection used by chat queries. +type ChatMessageRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +// ChatMessage converts the query projection to the schema model. +func (r ChatMessageRow) ChatMessage() ChatMessage { + return ChatMessage{ + ID: r.ID, + ChatID: r.ChatID, + ModelConfigID: r.ModelConfigID, + CreatedAt: r.CreatedAt, + Role: r.Role, + Content: r.Content, + Visibility: r.Visibility, + InputTokens: r.InputTokens, + OutputTokens: r.OutputTokens, + TotalTokens: r.TotalTokens, + ReasoningTokens: r.ReasoningTokens, + CacheCreationTokens: r.CacheCreationTokens, + CacheReadTokens: r.CacheReadTokens, + ContextLimit: r.ContextLimit, + Compressed: r.Compressed, + CreatedBy: r.CreatedBy, + ContentVersion: r.ContentVersion, + TotalCostMicros: r.TotalCostMicros, + RuntimeMs: r.RuntimeMs, + Deleted: r.Deleted, + ProviderResponseID: r.ProviderResponseID, + Revision: r.Revision, + ReasoningEffort: r.ReasoningEffort, + } +} + +// ChatQueuedMessageRow is the persisted queue projection used by chat queries. +type ChatQueuedMessageRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +// ChatQueuedMessage converts the query projection to the schema model. +func (r ChatQueuedMessageRow) ChatQueuedMessage() ChatQueuedMessage { + return ChatQueuedMessage{ + ID: r.ID, + ChatID: r.ChatID, + Content: r.Content, + CreatedAt: r.CreatedAt, + ModelConfigID: r.ModelConfigID, + Position: r.Position, + CreatedBy: r.CreatedBy, + ReasoningEffort: r.ReasoningEffort, + } +} + type WorkspaceStatus string const ( diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index b5c9d0b84e1..efb8f5e5a39 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -53,6 +53,90 @@ type customQuerier interface { connectionLogQuerier aibridgeQuerier chatQuerier + chatMessageQuerier +} + +type chatMessageQuerier interface { + GetChatMessageByID(ctx context.Context, id int64) (ChatMessage, error) + GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error) + GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) + GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedParams) ([]ChatMessage, error) + GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error) + GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) + GetLastChatMessageByRole(ctx context.Context, arg GetLastChatMessageByRoleParams) (ChatMessage, error) + InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) +} + +type ( + GetChatMessagesByChatIDParams = GetChatMessagesByChatIDRawParams + GetChatMessagesByRevisionForStreamParams = GetChatMessagesByRevisionForStreamRawParams + GetChatMessagesByChatIDAscPaginatedParams = GetChatMessagesByChatIDAscPaginatedRawParams + GetChatMessagesByChatIDDescPaginatedParams = GetChatMessagesByChatIDDescPaginatedRawParams + GetLastChatMessageByRoleParams = GetLastChatMessageByRoleRawParams + InsertChatMessagesParams = InsertChatMessagesRawParams +) + +func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMessage, error) { + row, err := q.GetChatMessageByIDRaw(ctx, id) + return ChatMessageRow(row).ChatMessage(), err +} + +func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error) { + rows, err := q.GetChatMessagesByChatIDRaw(ctx, arg) + return convertChatMessageRows(rows, err, func(row GetChatMessagesByChatIDRawRow) ChatMessage { + return ChatMessageRow(row).ChatMessage() + }) +} + +func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) { + rows, err := q.GetChatMessagesByRevisionForStreamRaw(ctx, arg) + return convertChatMessageRows(rows, err, func(row GetChatMessagesByRevisionForStreamRawRow) ChatMessage { + return ChatMessageRow(row).ChatMessage() + }) +} + +func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedParams) ([]ChatMessage, error) { + rows, err := q.GetChatMessagesByChatIDAscPaginatedRaw(ctx, arg) + return convertChatMessageRows(rows, err, func(row GetChatMessagesByChatIDAscPaginatedRawRow) ChatMessage { + return ChatMessageRow(row).ChatMessage() + }) +} + +func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error) { + rows, err := q.GetChatMessagesByChatIDDescPaginatedRaw(ctx, arg) + return convertChatMessageRows(rows, err, func(row GetChatMessagesByChatIDDescPaginatedRawRow) ChatMessage { + return ChatMessageRow(row).ChatMessage() + }) +} + +func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) { + rows, err := q.GetChatMessagesForPromptByChatIDRaw(ctx, chatID) + return convertChatMessageRows(rows, err, func(row GetChatMessagesForPromptByChatIDRawRow) ChatMessage { + return ChatMessageRow(row).ChatMessage() + }) +} + +func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastChatMessageByRoleParams) (ChatMessage, error) { + row, err := q.GetLastChatMessageByRoleRaw(ctx, arg) + return ChatMessageRow(row).ChatMessage(), err +} + +func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) { + rows, err := q.InsertChatMessagesRaw(ctx, arg) + return convertChatMessageRows(rows, err, func(row InsertChatMessagesRawRow) ChatMessage { + return ChatMessageRow(row).ChatMessage() + }) +} + +func convertChatMessageRows[T any](rows []T, err error, convert func(T) ChatMessage) ([]ChatMessage, error) { + if err != nil { + return nil, err + } + messages := make([]ChatMessage, len(rows)) + for i, row := range rows { + messages[i] = convert(row) + } + return messages, nil } type templateQuerier interface { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 31293a9d327..34a4f3b999b 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -433,16 +433,16 @@ type sqlcQuerier interface { // When the toggle is unset, a non-empty custom prompt implies false; // otherwise the setting defaults to true. GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) - GetChatMessageByID(ctx context.Context, id int64) (ChatMessage, error) + GetChatMessageByIDRaw(ctx context.Context, id int64) (GetChatMessageByIDRawRow, error) // Aggregates message-level metrics per chat for messages created // after the given timestamp. Uses message created_at so that // ongoing activity in long-running chats is captured each window. GetChatMessageSummariesPerChat(ctx context.Context, createdAfter time.Time) ([]GetChatMessageSummariesPerChatRow, error) - GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error) - GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedParams) ([]ChatMessage, error) - GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error) - GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) - GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) + GetChatMessagesByChatIDAscPaginatedRaw(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedRawParams) ([]GetChatMessagesByChatIDAscPaginatedRawRow, error) + GetChatMessagesByChatIDDescPaginatedRaw(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedRawParams) ([]GetChatMessagesByChatIDDescPaginatedRawRow, error) + GetChatMessagesByChatIDRaw(ctx context.Context, arg GetChatMessagesByChatIDRawParams) ([]GetChatMessagesByChatIDRawRow, error) + GetChatMessagesByRevisionForStreamRaw(ctx context.Context, arg GetChatMessagesByRevisionForStreamRawParams) ([]GetChatMessagesByRevisionForStreamRawRow, error) + GetChatMessagesForPromptByChatIDRaw(ctx context.Context, chatID uuid.UUID) ([]GetChatMessagesForPromptByChatIDRawRow, error) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error) GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error) // Returns all model configurations for telemetry snapshot collection. @@ -452,12 +452,12 @@ type sqlcQuerier interface { // personal chat model overrides. It defaults to false when unset. GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error) GetChatPlanModeInstructions(ctx context.Context) (string, error) - GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (ChatQueuedMessage, error) + GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (GetChatQueuedMessageByIDRow, error) // Returns the queue head (lowest position, then lowest id). - GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) - GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) + GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (GetChatQueuedMessageHeadRow, error) + GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]GetChatQueuedMessagesRow, error) // Returns queued messages in state-machine order (position ASC, id ASC). - GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) + GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]GetChatQueuedMessagesByPositionRow, error) // Returns the chat retention period in days. Chats archived longer // than this and orphaned chat files older than this are purged by // dbpurge. Returns 30 (days) when no value has been configured. @@ -598,7 +598,7 @@ type sqlcQuerier interface { // param created_at_opt: The created_at timestamp to filter by. This parameter is usd for pagination - it fetches notifications created before the specified timestamp if it is not the zero value // param limit_opt: The limit of notifications to fetch. If the limit is not specified, it defaults to 25 GetInboxNotificationsByUserID(ctx context.Context, arg GetInboxNotificationsByUserIDParams) ([]InboxNotification, error) - GetLastChatMessageByRole(ctx context.Context, arg GetLastChatMessageByRoleParams) (ChatMessage, error) + GetLastChatMessageByRoleRaw(ctx context.Context, arg GetLastChatMessageByRoleRawParams) (GetLastChatMessageByRoleRawRow, error) GetLastUpdateCheck(ctx context.Context) (string, error) GetLatestCryptoKeyByFeature(ctx context.Context, feature CryptoKeyFeature) (CryptoKey, error) GetLatestWorkspaceAgentContextSnapshot(ctx context.Context, workspaceAgentID uuid.UUID) (WorkspaceAgentContextSnapshot, error) @@ -1043,16 +1043,16 @@ type sqlcQuerier interface { // with concurrent FinalizeStale under READ COMMITTED isolation. InsertChatDebugStep(ctx context.Context, arg InsertChatDebugStepParams) (ChatDebugStep, error) InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error) - InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) + InsertChatMessagesRaw(ctx context.Context, arg InsertChatMessagesRawParams) ([]InsertChatMessagesRawRow, error) InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error) // Legacy queue insertion path. When no caller-supplied creator exists, // preserve the created_by invariant by attributing the queued row to the // chat owner. - InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error) + InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (InsertChatQueuedMessageRow, error) // Inserts a queued message that carries a position (from the default // sequence) and an explicit created_by reference. Use this when the // queued-message creator differs from the chat owner. - InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (ChatQueuedMessage, error) + InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (InsertChatQueuedMessageWithCreatorRow, error) InsertChatSyntheticAPIKey(ctx context.Context, arg InsertChatSyntheticAPIKeyParams) (int64, error) InsertCryptoKey(ctx context.Context, arg InsertCryptoKeyParams) (CryptoKey, error) InsertCustomRole(ctx context.Context, arg InsertCustomRoleParams) (CustomRole, error) @@ -1221,7 +1221,7 @@ type sqlcQuerier interface { // pin/unpin/reorder operation's ROW_NUMBER() self-heals the // sequence, so this is acceptable. PinChatByID(ctx context.Context, id uuid.UUID) error - PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) + PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (PopNextQueuedMessageRow, error) ReduceWorkspaceAgentShareLevelToAuthenticatedByTemplate(ctx context.Context, templateID uuid.UUID) error RegisterWorkspaceProxy(ctx context.Context, arg RegisterWorkspaceProxyParams) (WorkspaceProxy, error) RemoveUserFromGroups(ctx context.Context, arg RemoveUserFromGroupsParams) ([]uuid.UUID, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 008b6ef4dbb..01d71ff41fd 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11935,12 +11935,9 @@ func TestInsertChatMessages(t *testing.T) { insertMessage := func(t *testing.T, store database.Store, ctx context.Context, chatID, userID, modelConfigID uuid.UUID, content string) { t.Helper() - apiKey, _ := dbgen.APIKey(t, store, database.APIKey{ID: uuid.NewString(), UserID: userID}) - _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chatID, CreatedBy: []uuid.UUID{userID}, - APIKeyID: []string{apiKey.ID}, ModelConfigID: []uuid.UUID{modelConfigID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, ContentVersion: []int16{chatprompt.CurrentContentVersion}, @@ -11999,12 +11996,9 @@ func TestInsertChatMessages(t *testing.T) { t.Parallel() store, ctx, user, chat, _, modelConfigA := setupChat(t) - apiKey, _ := dbgen.APIKey(t, store, database.APIKey{ID: uuid.NewString(), UserID: user.ID}) - msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, CreatedBy: []uuid.UUID{user.ID, uuid.Nil, uuid.Nil}, - APIKeyID: []string{apiKey.ID, "", ""}, ModelConfigID: []uuid.UUID{modelConfigA.ID, modelConfigA.ID, modelConfigA.ID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleAssistant, database.ChatMessageRoleTool}, ContentVersion: []int16{chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion}, diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index ef4d7b437c8..c2ae8744e53 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -7560,9 +7560,31 @@ func (q *sqlQuerier) GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatP return i, err } -const getChatMessageByID = `-- name: GetChatMessageByID :one +const getChatMessageByIDRaw = `-- name: GetChatMessageByIDRaw :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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -7570,9 +7592,35 @@ WHERE AND deleted = false ` -func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMessage, error) { - row := q.db.QueryRowContext(ctx, getChatMessageByID, id) - var i ChatMessage +type GetChatMessageByIDRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatMessageByIDRaw(ctx context.Context, id int64) (GetChatMessageByIDRawRow, error) { + row := q.db.QueryRowContext(ctx, getChatMessageByIDRaw, id) + var i GetChatMessageByIDRawRow err := row.Scan( &i.ID, &i.ChatID, @@ -7595,7 +7643,6 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ) @@ -7685,9 +7732,31 @@ func (q *sqlQuerier) GetChatMessageSummariesPerChat(ctx context.Context, created return items, nil } -const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many +const getChatMessagesByChatIDAscPaginatedRaw = `-- name: GetChatMessagesByChatIDAscPaginatedRaw :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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -7696,23 +7765,52 @@ WHERE AND visibility IN ('user', 'both') AND deleted = false ORDER BY - created_at ASC + id ASC +LIMIT + COALESCE(NULLIF($3::int, 0), 50) ` -type GetChatMessagesByChatIDParams struct { - ChatID uuid.UUID `db:"chat_id" json:"chat_id"` - AfterID int64 `db:"after_id" json:"after_id"` +type GetChatMessagesByChatIDAscPaginatedRawParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + AfterID int64 `db:"after_id" json:"after_id"` + LimitVal int32 `db:"limit_val" json:"limit_val"` } -func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error) { - rows, err := q.db.QueryContext(ctx, getChatMessagesByChatID, arg.ChatID, arg.AfterID) +type GetChatMessagesByChatIDAscPaginatedRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginatedRaw(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedRawParams) ([]GetChatMessagesByChatIDAscPaginatedRawRow, error) { + rows, err := q.db.QueryContext(ctx, getChatMessagesByChatIDAscPaginatedRaw, arg.ChatID, arg.AfterID, arg.LimitVal) if err != nil { return nil, err } defer rows.Close() - var items []ChatMessage + var items []GetChatMessagesByChatIDAscPaginatedRawRow for rows.Next() { - var i ChatMessage + var i GetChatMessagesByChatIDAscPaginatedRawRow if err := rows.Scan( &i.ID, &i.ChatID, @@ -7735,7 +7833,6 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ); err != nil { @@ -7752,37 +7849,98 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes return items, nil } -const getChatMessagesByChatIDAscPaginated = `-- name: GetChatMessagesByChatIDAscPaginated :many +const getChatMessagesByChatIDDescPaginatedRaw = `-- name: GetChatMessagesByChatIDDescPaginatedRaw :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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE chat_id = $1::uuid - AND id > $2::bigint + AND CASE + WHEN $2::bigint > 0 THEN id < $2::bigint + ELSE true + END + AND CASE + WHEN $3::bigint > 0 THEN id > $3::bigint + ELSE true + END AND visibility IN ('user', 'both') AND deleted = false ORDER BY - id ASC + id DESC LIMIT - COALESCE(NULLIF($3::int, 0), 50) + COALESCE(NULLIF($4::int, 0), 50) ` -type GetChatMessagesByChatIDAscPaginatedParams struct { +type GetChatMessagesByChatIDDescPaginatedRawParams struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + BeforeID int64 `db:"before_id" json:"before_id"` AfterID int64 `db:"after_id" json:"after_id"` LimitVal int32 `db:"limit_val" json:"limit_val"` } -func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedParams) ([]ChatMessage, error) { - rows, err := q.db.QueryContext(ctx, getChatMessagesByChatIDAscPaginated, arg.ChatID, arg.AfterID, arg.LimitVal) +type GetChatMessagesByChatIDDescPaginatedRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginatedRaw(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedRawParams) ([]GetChatMessagesByChatIDDescPaginatedRawRow, error) { + rows, err := q.db.QueryContext(ctx, getChatMessagesByChatIDDescPaginatedRaw, + arg.ChatID, + arg.BeforeID, + arg.AfterID, + arg.LimitVal, + ) if err != nil { return nil, err } defer rows.Close() - var items []ChatMessage + var items []GetChatMessagesByChatIDDescPaginatedRawRow for rows.Next() { - var i ChatMessage + var i GetChatMessagesByChatIDDescPaginatedRawRow if err := rows.Scan( &i.ID, &i.ChatID, @@ -7805,7 +7963,6 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ); err != nil { @@ -7822,50 +7979,82 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar return items, nil } -const getChatMessagesByChatIDDescPaginated = `-- name: GetChatMessagesByChatIDDescPaginated :many +const getChatMessagesByChatIDRaw = `-- name: GetChatMessagesByChatIDRaw :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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE chat_id = $1::uuid - AND CASE - WHEN $2::bigint > 0 THEN id < $2::bigint - ELSE true - END - AND CASE - WHEN $3::bigint > 0 THEN id > $3::bigint - ELSE true - END + AND id > $2::bigint AND visibility IN ('user', 'both') AND deleted = false ORDER BY - id DESC -LIMIT - COALESCE(NULLIF($4::int, 0), 50) + created_at ASC ` -type GetChatMessagesByChatIDDescPaginatedParams struct { - ChatID uuid.UUID `db:"chat_id" json:"chat_id"` - BeforeID int64 `db:"before_id" json:"before_id"` - AfterID int64 `db:"after_id" json:"after_id"` - LimitVal int32 `db:"limit_val" json:"limit_val"` +type GetChatMessagesByChatIDRawParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + AfterID int64 `db:"after_id" json:"after_id"` } -func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error) { - rows, err := q.db.QueryContext(ctx, getChatMessagesByChatIDDescPaginated, - arg.ChatID, - arg.BeforeID, - arg.AfterID, - arg.LimitVal, - ) +type GetChatMessagesByChatIDRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatMessagesByChatIDRaw(ctx context.Context, arg GetChatMessagesByChatIDRawParams) ([]GetChatMessagesByChatIDRawRow, error) { + rows, err := q.db.QueryContext(ctx, getChatMessagesByChatIDRaw, arg.ChatID, arg.AfterID) if err != nil { return nil, err } defer rows.Close() - var items []ChatMessage + var items []GetChatMessagesByChatIDRawRow for rows.Next() { - var i ChatMessage + var i GetChatMessagesByChatIDRawRow if err := rows.Scan( &i.ID, &i.ChatID, @@ -7888,7 +8077,6 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ); err != nil { @@ -7905,9 +8093,31 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a return items, nil } -const getChatMessagesByRevisionForStream = `-- name: GetChatMessagesByRevisionForStream :many +const getChatMessagesByRevisionForStreamRaw = `-- name: GetChatMessagesByRevisionForStreamRaw :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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -7918,20 +8128,46 @@ ORDER BY created_at ASC, id ASC ` -type GetChatMessagesByRevisionForStreamParams struct { +type GetChatMessagesByRevisionForStreamRawParams struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` AfterRevision int64 `db:"after_revision" json:"after_revision"` } -func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) { - rows, err := q.db.QueryContext(ctx, getChatMessagesByRevisionForStream, arg.ChatID, arg.AfterRevision) +type GetChatMessagesByRevisionForStreamRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatMessagesByRevisionForStreamRaw(ctx context.Context, arg GetChatMessagesByRevisionForStreamRawParams) ([]GetChatMessagesByRevisionForStreamRawRow, error) { + rows, err := q.db.QueryContext(ctx, getChatMessagesByRevisionForStreamRaw, arg.ChatID, arg.AfterRevision) if err != nil { return nil, err } defer rows.Close() - var items []ChatMessage + var items []GetChatMessagesByRevisionForStreamRawRow for rows.Next() { - var i ChatMessage + var i GetChatMessagesByRevisionForStreamRawRow if err := rows.Scan( &i.ID, &i.ChatID, @@ -7954,7 +8190,6 @@ func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ); err != nil { @@ -7971,7 +8206,7 @@ func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg return items, nil } -const getChatMessagesForPromptByChatID = `-- name: GetChatMessagesForPromptByChatID :many +const getChatMessagesForPromptByChatIDRaw = `-- name: GetChatMessagesForPromptByChatIDRaw :many WITH latest_compressed_summary AS ( SELECT id @@ -7989,7 +8224,29 @@ 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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -8030,15 +8287,41 @@ ORDER BY id ASC ` -func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) { - rows, err := q.db.QueryContext(ctx, getChatMessagesForPromptByChatID, chatID) +type GetChatMessagesForPromptByChatIDRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatMessagesForPromptByChatIDRaw(ctx context.Context, chatID uuid.UUID) ([]GetChatMessagesForPromptByChatIDRawRow, error) { + rows, err := q.db.QueryContext(ctx, getChatMessagesForPromptByChatIDRaw, chatID) if err != nil { return nil, err } defer rows.Close() - var items []ChatMessage + var items []GetChatMessagesForPromptByChatIDRawRow for rows.Next() { - var i ChatMessage + var i GetChatMessagesForPromptByChatIDRawRow if err := rows.Scan( &i.ID, &i.ChatID, @@ -8061,7 +8344,6 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ); err != nil { @@ -8127,7 +8409,8 @@ func (q *sqlQuerier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]Get } const getChatQueuedMessageByID = `-- name: GetChatQueuedMessageByID :one -SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE id = $1::bigint AND chat_id = $2::uuid ` @@ -8136,16 +8419,26 @@ type GetChatQueuedMessageByIDParams struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` } -func (q *sqlQuerier) GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (ChatQueuedMessage, error) { +type GetChatQueuedMessageByIDRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (GetChatQueuedMessageByIDRow, error) { row := q.db.QueryRowContext(ctx, getChatQueuedMessageByID, arg.ID, arg.ChatID) - var i ChatQueuedMessage + var i GetChatQueuedMessageByIDRow err := row.Scan( &i.ID, &i.ChatID, &i.Content, &i.CreatedAt, &i.ModelConfigID, - &i.APIKeyID, &i.Position, &i.CreatedBy, &i.ReasoningEffort, @@ -8154,23 +8447,34 @@ func (q *sqlQuerier) GetChatQueuedMessageByID(ctx context.Context, arg GetChatQu } const getChatQueuedMessageHead = `-- name: GetChatQueuedMessageHead :one -SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE chat_id = $1::uuid ORDER BY position ASC, id ASC LIMIT 1 ` +type GetChatQueuedMessageHeadRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + // Returns the queue head (lowest position, then lowest id). -func (q *sqlQuerier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) { +func (q *sqlQuerier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (GetChatQueuedMessageHeadRow, error) { row := q.db.QueryRowContext(ctx, getChatQueuedMessageHead, chatID) - var i ChatQueuedMessage + var i GetChatQueuedMessageHeadRow err := row.Scan( &i.ID, &i.ChatID, &i.Content, &i.CreatedAt, &i.ModelConfigID, - &i.APIKeyID, &i.Position, &i.CreatedBy, &i.ReasoningEffort, @@ -8179,27 +8483,38 @@ func (q *sqlQuerier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.U } const getChatQueuedMessages = `-- name: GetChatQueuedMessages :many -SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE chat_id = $1 ORDER BY created_at ASC, id ASC ` -func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) { +type GetChatQueuedMessagesRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]GetChatQueuedMessagesRow, error) { rows, err := q.db.QueryContext(ctx, getChatQueuedMessages, chatID) if err != nil { return nil, err } defer rows.Close() - var items []ChatQueuedMessage + var items []GetChatQueuedMessagesRow for rows.Next() { - var i ChatQueuedMessage + var i GetChatQueuedMessagesRow if err := rows.Scan( &i.ID, &i.ChatID, &i.Content, &i.CreatedAt, &i.ModelConfigID, - &i.APIKeyID, &i.Position, &i.CreatedBy, &i.ReasoningEffort, @@ -8218,28 +8533,39 @@ func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID } const getChatQueuedMessagesByPosition = `-- name: GetChatQueuedMessagesByPosition :many -SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE chat_id = $1::uuid ORDER BY position ASC, id ASC ` +type GetChatQueuedMessagesByPositionRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + // Returns queued messages in state-machine order (position ASC, id ASC). -func (q *sqlQuerier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) { +func (q *sqlQuerier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]GetChatQueuedMessagesByPositionRow, error) { rows, err := q.db.QueryContext(ctx, getChatQueuedMessagesByPosition, chatID) if err != nil { return nil, err } defer rows.Close() - var items []ChatQueuedMessage + var items []GetChatQueuedMessagesByPositionRow for rows.Next() { - var i ChatQueuedMessage + var i GetChatQueuedMessagesByPositionRow if err := rows.Scan( &i.ID, &i.ChatID, &i.Content, &i.CreatedAt, &i.ModelConfigID, - &i.APIKeyID, &i.Position, &i.CreatedBy, &i.ReasoningEffort, @@ -9325,9 +9651,31 @@ func (q *sqlQuerier) GetDatabaseNow(ctx context.Context) (time.Time, error) { return now, err } -const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one +const getLastChatMessageByRoleRaw = `-- name: GetLastChatMessageByRoleRaw :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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -9340,14 +9688,40 @@ LIMIT 1 ` -type GetLastChatMessageByRoleParams struct { +type GetLastChatMessageByRoleRawParams struct { 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) { - row := q.db.QueryRowContext(ctx, getLastChatMessageByRole, arg.ChatID, arg.Role) - var i ChatMessage +type GetLastChatMessageByRoleRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) GetLastChatMessageByRoleRaw(ctx context.Context, arg GetLastChatMessageByRoleRawParams) (GetLastChatMessageByRoleRawRow, error) { + row := q.db.QueryRowContext(ctx, getLastChatMessageByRoleRaw, arg.ChatID, arg.Role) + var i GetLastChatMessageByRoleRawRow err := row.Scan( &i.ID, &i.ChatID, @@ -9370,7 +9744,6 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ) @@ -9806,12 +10179,12 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat return i, err } -const insertChatMessages = `-- name: InsertChatMessages :many +const insertChatMessagesRaw = `-- name: InsertChatMessagesRaw :many WITH batch AS ( SELECT ( SELECT val - FROM UNNEST($4::uuid[]) + FROM UNNEST($3::uuid[]) WITH ORDINALITY AS t(val, ord) WHERE val != '00000000-0000-0000-0000-000000000000'::uuid ORDER BY ord DESC @@ -9819,7 +10192,7 @@ WITH batch AS ( ) AS last_model_config_id, ( SELECT NULLIF(val, '')::chat_reasoning_effort - FROM UNNEST($5::text[]) + FROM UNNEST($4::text[]) WITH ORDINALITY AS t(val, ord) WHERE val != '' ORDER BY ord DESC @@ -9843,7 +10216,6 @@ updated_chat AS ( INSERT INTO chat_messages ( chat_id, created_by, - api_key_id, model_config_id, reasoning_effort, role, @@ -9864,31 +10236,51 @@ INSERT INTO chat_messages ( SELECT $1::uuid, NULLIF(UNNEST($2::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST($3::text[]), ''), - NULLIF(UNNEST($4::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST($5::text[]), '')::chat_reasoning_effort, - UNNEST($6::chat_message_role[]), - UNNEST($7::text[])::jsonb, - UNNEST($8::smallint[]), - UNNEST($9::chat_message_visibility[]), + NULLIF(UNNEST($3::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF(UNNEST($4::text[]), '')::chat_reasoning_effort, + UNNEST($5::chat_message_role[]), + UNNEST($6::text[])::jsonb, + UNNEST($7::smallint[]), + UNNEST($8::chat_message_visibility[]), + NULLIF(UNNEST($9::bigint[]), 0), NULLIF(UNNEST($10::bigint[]), 0), NULLIF(UNNEST($11::bigint[]), 0), NULLIF(UNNEST($12::bigint[]), 0), NULLIF(UNNEST($13::bigint[]), 0), NULLIF(UNNEST($14::bigint[]), 0), NULLIF(UNNEST($15::bigint[]), 0), - NULLIF(UNNEST($16::bigint[]), 0), - UNNEST($17::boolean[]), - NULLIF(UNNEST($18::bigint[]), 0), - NULLIF(UNNEST($19::bigint[]), 0) + UNNEST($16::boolean[]), + NULLIF(UNNEST($17::bigint[]), 0), + NULLIF(UNNEST($18::bigint[]), 0) 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, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort + 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, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort ` -type InsertChatMessagesParams struct { +type InsertChatMessagesRawParams struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` CreatedBy []uuid.UUID `db:"created_by" json:"created_by"` - APIKeyID []string `db:"api_key_id" json:"api_key_id"` ModelConfigID []uuid.UUID `db:"model_config_id" json:"model_config_id"` ReasoningEffort []string `db:"reasoning_effort" json:"reasoning_effort"` Role []ChatMessageRole `db:"role" json:"role"` @@ -9907,11 +10299,36 @@ type InsertChatMessagesParams struct { RuntimeMs []int64 `db:"runtime_ms" json:"runtime_ms"` } -func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) { - rows, err := q.db.QueryContext(ctx, insertChatMessages, +type InsertChatMessagesRawRow struct { + ID int64 `db:"id" json:"id"` + 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 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"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + 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"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) InsertChatMessagesRaw(ctx context.Context, arg InsertChatMessagesRawParams) ([]InsertChatMessagesRawRow, error) { + rows, err := q.db.QueryContext(ctx, insertChatMessagesRaw, arg.ChatID, pq.Array(arg.CreatedBy), - pq.Array(arg.APIKeyID), pq.Array(arg.ModelConfigID), pq.Array(arg.ReasoningEffort), pq.Array(arg.Role), @@ -9933,9 +10350,9 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa return nil, err } defer rows.Close() - var items []ChatMessage + var items []InsertChatMessagesRawRow for rows.Next() { - var i ChatMessage + var i InsertChatMessagesRawRow if err := rows.Scan( &i.ID, &i.ChatID, @@ -9958,7 +10375,6 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa &i.RuntimeMs, &i.Deleted, &i.ProviderResponseID, - &i.APIKeyID, &i.Revision, &i.ReasoningEffort, ); err != nil { @@ -9976,17 +10392,16 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa } const insertChatQueuedMessage = `-- name: InsertChatQueuedMessage :one -INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, api_key_id, created_by) +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, created_by) SELECT $1::uuid, $2::jsonb, $3::uuid, $4::chat_reasoning_effort, - $5::text, chats.owner_id FROM chats WHERE chats.id = $1::uuid -RETURNING id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by, reasoning_effort +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort ` type InsertChatQueuedMessageParams struct { @@ -9994,28 +10409,36 @@ type InsertChatQueuedMessageParams struct { Content json.RawMessage `db:"content" json:"content"` ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` - APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"` +} + +type InsertChatQueuedMessageRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` } // Legacy queue insertion path. When no caller-supplied creator exists, // preserve the created_by invariant by attributing the queued row to the // chat owner. -func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error) { +func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (InsertChatQueuedMessageRow, error) { row := q.db.QueryRowContext(ctx, insertChatQueuedMessage, arg.ChatID, arg.Content, arg.ModelConfigID, arg.ReasoningEffort, - arg.APIKeyID, ) - var i ChatQueuedMessage + var i InsertChatQueuedMessageRow err := row.Scan( &i.ID, &i.ChatID, &i.Content, &i.CreatedAt, &i.ModelConfigID, - &i.APIKeyID, &i.Position, &i.CreatedBy, &i.ReasoningEffort, @@ -10024,16 +10447,15 @@ func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChat } const insertChatQueuedMessageWithCreator = `-- name: InsertChatQueuedMessageWithCreator :one -INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, api_key_id, created_by) +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, created_by) VALUES ( $1::uuid, $2::jsonb, $3::uuid, $4::chat_reasoning_effort, - $5::text, - $6::uuid + $5::uuid ) -RETURNING id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by, reasoning_effort +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort ` type InsertChatQueuedMessageWithCreatorParams struct { @@ -10041,30 +10463,38 @@ type InsertChatQueuedMessageWithCreatorParams struct { Content json.RawMessage `db:"content" json:"content"` ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` - APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"` CreatedBy uuid.UUID `db:"created_by" json:"created_by"` } +type InsertChatQueuedMessageWithCreatorRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + // Inserts a queued message that carries a position (from the default // sequence) and an explicit created_by reference. Use this when the // queued-message creator differs from the chat owner. -func (q *sqlQuerier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (ChatQueuedMessage, error) { +func (q *sqlQuerier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (InsertChatQueuedMessageWithCreatorRow, error) { row := q.db.QueryRowContext(ctx, insertChatQueuedMessageWithCreator, arg.ChatID, arg.Content, arg.ModelConfigID, arg.ReasoningEffort, - arg.APIKeyID, arg.CreatedBy, ) - var i ChatQueuedMessage + var i InsertChatQueuedMessageWithCreatorRow err := row.Scan( &i.ID, &i.ChatID, &i.Content, &i.CreatedAt, &i.ModelConfigID, - &i.APIKeyID, &i.Position, &i.CreatedBy, &i.ReasoningEffort, @@ -10552,19 +10982,29 @@ WHERE id = ( ORDER BY cqm.created_at ASC, cqm.id ASC LIMIT 1 ) -RETURNING id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by, reasoning_effort +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort ` -func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) { +type PopNextQueuedMessageRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` +} + +func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (PopNextQueuedMessageRow, error) { row := q.db.QueryRowContext(ctx, popNextQueuedMessage, chatID) - var i ChatQueuedMessage + var i PopNextQueuedMessageRow err := row.Scan( &i.ID, &i.ChatID, &i.Content, &i.CreatedAt, &i.ModelConfigID, - &i.APIKeyID, &i.Position, &i.CreatedBy, &i.ReasoningEffort, diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index a4ee5fab1a5..80e507ed9f1 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -344,18 +344,62 @@ SET WHERE id = @id::uuid; --- name: GetChatMessageByID :one +-- name: GetChatMessageByIDRaw :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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE id = @id::bigint AND deleted = false; --- name: GetChatMessagesByChatID :many +-- name: GetChatMessagesByChatIDRaw :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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -366,9 +410,31 @@ WHERE ORDER BY created_at ASC; --- name: GetChatMessagesByRevisionForStream :many +-- name: GetChatMessagesByRevisionForStreamRaw :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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -378,9 +444,31 @@ WHERE ORDER BY created_at ASC, id ASC; --- name: GetChatMessagesByChatIDAscPaginated :many +-- name: GetChatMessagesByChatIDAscPaginatedRaw :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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -393,9 +481,31 @@ ORDER BY LIMIT COALESCE(NULLIF(@limit_val::int, 0), 50); --- name: GetChatMessagesByChatIDDescPaginated :many +-- name: GetChatMessagesByChatIDDescPaginatedRaw :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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -447,7 +557,7 @@ ORDER BY LIMIT COALESCE(NULLIF(@limit_val::int, 0), 500); --- name: GetChatMessagesForPromptByChatID :many +-- name: GetChatMessagesForPromptByChatIDRaw :many WITH latest_compressed_summary AS ( SELECT id @@ -465,7 +575,29 @@ 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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -792,7 +924,7 @@ chats_expanded AS ( SELECT * FROM chats_expanded; --- name: InsertChatMessages :many +-- name: InsertChatMessagesRaw :many WITH batch AS ( SELECT ( @@ -829,7 +961,6 @@ updated_chat AS ( INSERT INTO chat_messages ( chat_id, created_by, - api_key_id, model_config_id, reasoning_effort, role, @@ -850,7 +981,6 @@ INSERT INTO chat_messages ( SELECT @chat_id::uuid, NULLIF(UNNEST(@created_by::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST(@api_key_id::text[]), ''), NULLIF(UNNEST(@model_config_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), NULLIF(UNNEST(@reasoning_effort::text[]), '')::chat_reasoning_effort, UNNEST(@role::chat_message_role[]), @@ -868,7 +998,29 @@ SELECT NULLIF(UNNEST(@total_cost_micros::bigint[]), 0), NULLIF(UNNEST(@runtime_ms::bigint[]), 0) 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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort; -- name: UpdateChatByID :one WITH updated_chat AS ( @@ -1851,20 +2003,20 @@ RETURNING -- Legacy queue insertion path. When no caller-supplied creator exists, -- preserve the created_by invariant by attributing the queued row to the -- chat owner. -INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, api_key_id, created_by) +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, created_by) SELECT @chat_id::uuid, @content::jsonb, sqlc.narg('model_config_id')::uuid, sqlc.narg('reasoning_effort')::chat_reasoning_effort, - sqlc.narg('api_key_id')::text, chats.owner_id FROM chats WHERE chats.id = @chat_id::uuid -RETURNING *; +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort; -- name: GetChatQueuedMessages :many -SELECT * FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE chat_id = @chat_id ORDER BY created_at ASC, id ASC; @@ -1882,7 +2034,7 @@ WHERE id = ( ORDER BY cqm.created_at ASC, cqm.id ASC LIMIT 1 ) -RETURNING *; +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort; -- name: ReorderChatQueuedMessageToFront :execrows -- Mutates only created_at on the target row; ids are unchanged so @@ -1895,9 +2047,31 @@ SET created_at = ( ) WHERE target.id = @target_id AND target.chat_id = @chat_id; --- name: GetLastChatMessageByRole :one +-- name: GetLastChatMessageByRoleRaw :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, + content_version, + total_cost_micros, + runtime_ms, + deleted, + provider_response_id, + revision, + reasoning_effort FROM chat_messages WHERE @@ -2860,20 +3034,20 @@ SELECT NOW()::timestamptz AS now; -- Inserts a queued message that carries a position (from the default -- sequence) and an explicit created_by reference. Use this when the -- queued-message creator differs from the chat owner. -INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, api_key_id, created_by) +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, created_by) VALUES ( @chat_id::uuid, @content::jsonb, sqlc.narg('model_config_id')::uuid, sqlc.narg('reasoning_effort')::chat_reasoning_effort, - sqlc.narg('api_key_id')::text, @created_by::uuid ) -RETURNING *; +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort; -- name: GetChatQueuedMessagesByPosition :many -- Returns queued messages in state-machine order (position ASC, id ASC). -SELECT * FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE chat_id = @chat_id::uuid ORDER BY position ASC, id ASC; @@ -2886,13 +3060,15 @@ WHERE chat_id = @chat_id::uuid; -- name: GetChatQueuedMessageHead :one -- Returns the queue head (lowest position, then lowest id). -SELECT * FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE chat_id = @chat_id::uuid ORDER BY position ASC, id ASC LIMIT 1; -- name: GetChatQueuedMessageByID :one -SELECT * FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +FROM chat_queued_messages WHERE id = @id::bigint AND chat_id = @chat_id::uuid; -- name: DeleteChatQueuedMessageReturningCount :execrows diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index b2fe7da0d07..30336058f37 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -2283,7 +2283,7 @@ func (api *API) getChatMessages(rw http.ResponseWriter, r *http.Request) { // Queued messages are only meaningful for the initial top-of-history // load. Suppress them whenever any cursor is set so polling callers do // not receive the snapshot on every page fetch. - var queuedMessages []database.ChatQueuedMessage + var queuedMessages []database.GetChatQueuedMessagesRow if beforeID == 0 && afterID == 0 { queuedMessages, err = api.Database.GetChatQueuedMessages(ctx, chatID) if err != nil { @@ -6705,10 +6705,10 @@ func convertChatQueuedMessagePtr(m database.ChatQueuedMessage) *codersdk.ChatQue return &qm } -func convertChatQueuedMessages(msgs []database.ChatQueuedMessage) []codersdk.ChatQueuedMessage { +func convertChatQueuedMessages(msgs []database.GetChatQueuedMessagesRow) []codersdk.ChatQueuedMessage { result := make([]codersdk.ChatQueuedMessage, 0, len(msgs)) for _, m := range msgs { - result = append(result, convertChatQueuedMessage(m)) + result = append(result, db2sdk.ChatQueuedMessage(database.ChatQueuedMessageRow(m).ChatQueuedMessage())) } return result } diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index a4818d172cf..1096ba7c720 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -138,15 +138,6 @@ func newChatClientWithAPIAndDatabase(t testing.TB, overrides ...func(*coderdtest return codersdk.NewExperimentalClient(client), api.Database, api } -func currentTestAPIKeyID(t testing.TB, client *codersdk.ExperimentalClient) string { - t.Helper() - - apiKeyID, _, ok := strings.Cut(client.SessionToken(), "-") - require.True(t, ok) - require.NotEmpty(t, apiKeyID) - return apiKeyID -} - func insertTestChatQueuedMessage( ctx context.Context, t testing.TB, @@ -154,10 +145,9 @@ func insertTestChatQueuedMessage( chatID uuid.UUID, content json.RawMessage, modelConfigID uuid.UUID, - apiKeyID string, ) database.ChatQueuedMessage { t.Helper() - return insertTestChatQueuedMessageWithReasoningEffort(ctx, t, db, chatID, content, modelConfigID, apiKeyID, "") + return insertTestChatQueuedMessageWithReasoningEffort(ctx, t, db, chatID, content, modelConfigID, "") } func insertTestChatQueuedMessageWithReasoningEffort( @@ -167,7 +157,6 @@ func insertTestChatQueuedMessageWithReasoningEffort( chatID uuid.UUID, content json.RawMessage, modelConfigID uuid.UUID, - apiKeyID string, reasoningEffort string, ) database.ChatQueuedMessage { t.Helper() @@ -179,11 +168,10 @@ func insertTestChatQueuedMessageWithReasoningEffort( Content: content, ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, ReasoningEffort: database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffort(reasoningEffort), Valid: reasoningEffort != ""}, - APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, }, ) require.NoError(t, err) - return queued + return database.ChatQueuedMessageRow(queued).ChatQueuedMessage() } // findUserMessage returns the first user-role message from a slice of chat @@ -5007,11 +4995,9 @@ func TestGetChatUserPrompts(t *testing.T) { t.Helper() content, err := chatprompt.MarshalParts(parts) require.NoError(t, err) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: userID}) msgs, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ ChatID: chatID, CreatedBy: []uuid.UUID{userID}, - APIKeyID: []string{apiKey.ID}, ModelConfigID: []uuid.UUID{modelConfigID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, ContentVersion: []int16{chatprompt.CurrentContentVersion}, @@ -5118,11 +5104,9 @@ func TestGetChatUserPrompts(t *testing.T) { // without the guard, jsonb_array_elements would raise // "cannot extract elements from a scalar" and the request // would 500. - legacyAPIKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.UserID}) _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ ChatID: chat.ID, CreatedBy: []uuid.UUID{user.UserID}, - APIKeyID: []string{legacyAPIKey.ID}, ModelConfigID: []uuid.UUID{modelConfig.ID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, ContentVersion: []int16{chatprompt.ContentVersionV0}, @@ -9785,7 +9769,7 @@ func TestDeleteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued message for delete route"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, deleteContent, modelConfig.ID, currentTestAPIKeyID(t, client)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, deleteContent, modelConfig.ID) res, err := client.Request( ctx, @@ -9866,7 +9850,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := client.Request( ctx, @@ -9931,7 +9915,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) @@ -10026,7 +10010,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued message no agents access"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := memberClient.Request( ctx, @@ -10058,7 +10042,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) // Archive the chat. _, err = db.ArchiveChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) @@ -10148,7 +10132,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := client.Request( ctx, @@ -10241,7 +10225,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("running-promote"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := client.Request( ctx, @@ -15382,7 +15366,6 @@ func TestGetChatMessages_Pagination(t *testing.T) { db database.Store, chatID uuid.UUID, modelConfigID uuid.UUID, - apiKeyID string, ) { t.Helper() @@ -15390,7 +15373,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - _ = insertTestChatQueuedMessage(ctx, t, db, chatID, content, modelConfigID, apiKeyID) + _ = insertTestChatQueuedMessage(ctx, t, db, chatID, content, modelConfigID) } t.Run("NoCursorReturnsAllDESCPlusQueued", func(t *testing.T) { @@ -15402,7 +15385,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID) resp, err := client.GetChatMessages(ctx, chat.ID, nil) require.NoError(t, err) @@ -15427,7 +15410,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ BeforeID: ids[2], @@ -15453,7 +15436,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ AfterID: ids[1], @@ -15481,7 +15464,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ AfterID: ids[0], @@ -15510,7 +15493,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) // Seed a queued message so the Empty assertion below verifies // the cursor suppresses queued rows, not just that none exist. - seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ AfterID: ids[0], @@ -15604,7 +15587,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 3) // Seed a queued message to prove the cursor path suppresses // it even when nothing else comes back. - seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID) // The steady-state polling case: the caller already has every // message, so after_id equals the largest seen id. The server @@ -15841,12 +15824,12 @@ func TestChatReadOnlySharedWriteHandlers(t *testing.T) { t.Run("PromoteChatQueuedMessage", func(t *testing.T) { t.Parallel() - ctx, ownerClient, sharedClient, chat, db := setup(t) + ctx, _, sharedClient, chat, db := setup(t) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, ownerClient)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) res, err := sharedClient.Request( ctx, @@ -15876,12 +15859,12 @@ func TestChatReadOnlySharedWriteHandlers(t *testing.T) { t.Run("DeleteChatQueuedMessage", func(t *testing.T) { t.Parallel() - ctx, ownerClient, sharedClient, chat, db := setup(t) + ctx, _, sharedClient, chat, db := setup(t) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, ownerClient)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) res, err := sharedClient.Request( ctx, @@ -16026,14 +16009,14 @@ func TestChatOwnerOnlyWriteHandlers(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - ownerClient, adminClient, chat, db := setupOrgAdminAndOwnerChat(t) + _, adminClient, chat, db := setupOrgAdminAndOwnerChat(t) // Insert a queued message directly in the DB. queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, ownerClient)) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) // Org admin tries to promote. promoteRes, err := adminClient.Request( diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 40fe057486b..95cd60b7ffd 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -13,7 +13,7 @@ Chatd attributes AI Gateway requests with a synthetic API key owned by the chat Synthetic keys expire after 30 days and renew when less than 24 hours remain. Renewal replaces the mapping but does not delete the previous key, because an in-flight request may still use it. Concurrent mints use conditional insert/update operations, and losing keys are deleted. The generated token is discarded, so the stored key cannot be used as a bearer credential. -The legacy `api_key_id` columns on messages and queued messages are still stamped with the synthetic key for rolling deployment compatibility, but they no longer have foreign keys to `api_keys`. They are not the source of gateway routing. Stale IDs are harmless because chatd resolves attribution from `chats.owner_id`. +The legacy `api_key_id` columns on messages and queued messages have no foreign keys and are not read or written. Any IDs stamped by older replicas remain inert because chatd resolves attribution from `chats.owner_id`. Deleting a synthetic key removes its mapping without updating chat messages, queued messages, or their version fields. Deleting all user API keys or resetting a password therefore causes chatd to mint a replacement on the next request without mutating history. User suspension and deletion still block delegated gateway authorization. diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index ac0439f7f5d..067d4b126c8 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1259,11 +1259,6 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C return database.Chat{}, limitErr } - apiKeyID, err := p.ensureSyntheticAPIKeyID(ctx, opts.OwnerID) - if err != nil { - return database.Chat{}, xerrors.Errorf("ensure synthetic API key: %w", err) - } - labelsJSON, err := json.Marshal(opts.Labels) if err != nil { return database.Chat{}, xerrors.Errorf("marshal labels: %w", err) @@ -1305,7 +1300,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C initialMessages = append(initialMessages, systemMessage(userPromptContent, opts.ModelConfigID)) } initialMessages = append(initialMessages, systemMessage(workspaceAwarenessContent, opts.ModelConfigID)) - initialMessages = append(initialMessages, userMessageWithAPIKeyID(userContent, opts.ModelConfigID, opts.OwnerID, apiKeyID, opts.ReasoningEffort)) + initialMessages = append(initialMessages, userMessage(userContent, opts.ModelConfigID, opts.OwnerID, opts.ReasoningEffort)) result, err := chatstate.CreateChat(ctx, p.db, p.pubsub, chatstate.CreateChatInput{ OrganizationID: opts.OrganizationID, @@ -1385,15 +1380,6 @@ func (p *Server) SendMessage( requestedPlanMode := opts.PlanMode requestedMCPServerIDs := opts.MCPServerIDs - chat, err := p.db.GetChatByID(ctx, opts.ChatID) - if err != nil { - return SendMessageResult{}, xerrors.Errorf("load chat: %w", err) - } - apiKeyID, err := p.ensureSyntheticAPIKeyID(ctx, chat.OwnerID) - if err != nil { - return SendMessageResult{}, xerrors.Errorf("ensure synthetic API key: %w", err) - } - var result SendMessageResult machine := p.newChatMachine(opts.ChatID) updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { @@ -1458,7 +1444,7 @@ func (p *Server) SendMessage( // Queue capacity is enforced inside tx.SendMessage; this // wrapper only propagates the typed error. sendResult, err := tx.SendMessage(chatstate.SendMessageInput{ - Message: userMessageWithAPIKeyID(content, modelConfigID, messageCreatedBy, apiKeyID, opts.ReasoningEffort), + Message: userMessage(content, modelConfigID, messageCreatedBy, opts.ReasoningEffort), BusyBehavior: busyBehaviorToChatState(busyBehavior), }) if err != nil { @@ -1604,15 +1590,6 @@ func (p *Server) EditMessage( if err != nil { return EditMessageResult{}, xerrors.Errorf("marshal message content: %w", err) } - chat, err := p.db.GetChatByID(ctx, opts.ChatID) - if err != nil { - return EditMessageResult{}, xerrors.Errorf("load chat: %w", err) - } - apiKeyID, err := p.ensureSyntheticAPIKeyID(ctx, chat.OwnerID) - if err != nil { - return EditMessageResult{}, xerrors.Errorf("ensure synthetic API key: %w", err) - } - var ( result EditMessageResult editedMsg database.ChatMessage @@ -1682,7 +1659,6 @@ func (p *Server) EditMessage( Content: content, ModelConfigIDOverride: modelOverride, ReasoningEffortOverride: reasoningEffortOverride, - APIKeyID: sql.NullString{String: apiKeyID, Valid: true}, }) if err != nil { if errors.Is(err, chatstate.ErrEditedMessageNotUser) { @@ -2275,7 +2251,6 @@ func (p *Server) recordManualTitleGenerationFailure( chat, generationErr.modelConfig, generationErr.usage, - generationErr.activeAPIKeyID, "", ); recordErr != nil { return errors.Join( @@ -2406,7 +2381,6 @@ func (p *Server) proposeChatTitleWithStore( chat, result.modelConfig, result.usage, - result.activeAPIKeyID, "", ); recordErr != nil { return "", xerrors.Errorf("record manual title usage: %w", recordErr) @@ -2436,7 +2410,6 @@ func (p *Server) regenerateChatTitleWithStore( chat, result.modelConfig, result.usage, - result.activeAPIKeyID, result.title, ) if recordErr != nil { @@ -2785,7 +2758,6 @@ func recordManualTitleUsage( chat database.Chat, modelConfig database.ChatModelConfig, usage fantasy.Usage, - activeAPIKeyID string, newTitle string, ) (database.Chat, bool, error) { hasUsage := usage != (fantasy.Usage{}) @@ -2826,7 +2798,6 @@ func recordManualTitleUsage( messages, err := tx.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, CreatedBy: []uuid.UUID{chat.OwnerID}, - APIKeyID: []string{activeAPIKeyID}, ModelConfigID: []uuid.UUID{modelConfig.ID}, ReasoningEffort: []string{""}, Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, @@ -2899,16 +2870,6 @@ type chatMessage struct { runtimeMs int64 } -type userChatMessage struct { - chatMessage - apiKeyID string -} - -func (m userChatMessage) withCreatedBy(id uuid.UUID) userChatMessage { - m.chatMessage = m.chatMessage.withCreatedBy(id) - return m -} - func newChatMessage( role database.ChatMessageRole, content pqtype.NullRawMessage, @@ -2925,25 +2886,6 @@ func newChatMessage( } } -func newUserChatMessage( - apiKeyID string, - content pqtype.NullRawMessage, - visibility database.ChatMessageVisibility, - modelConfigID uuid.UUID, - contentVersion int16, -) userChatMessage { - return userChatMessage{ - chatMessage: newChatMessage( - database.ChatMessageRoleUser, - content, - visibility, - modelConfigID, - contentVersion, - ), - apiKeyID: apiKeyID, - } -} - func (m chatMessage) withCreatedBy(id uuid.UUID) chatMessage { m.createdBy = id return m @@ -2952,10 +2894,8 @@ func (m chatMessage) withCreatedBy(id uuid.UUID) chatMessage { func appendMessageFields( params *database.InsertChatMessagesParams, msg chatMessage, - apiKeyID string, ) { params.CreatedBy = append(params.CreatedBy, msg.createdBy) - params.APIKeyID = append(params.APIKeyID, apiKeyID) params.ModelConfigID = append(params.ModelConfigID, msg.modelConfigID) params.ReasoningEffort = append(params.ReasoningEffort, "") params.Role = append(params.Role, msg.role) @@ -2975,20 +2915,10 @@ func appendMessageFields( } func appendChatMessage(params *database.InsertChatMessagesParams, msg chatMessage) { - if msg.role == database.ChatMessageRoleUser { - panic("developer error: use appendUserChatMessage for user-role messages") - } - appendMessageFields(params, msg, "") -} - -func appendUserChatMessage(params *database.InsertChatMessagesParams, msg userChatMessage) { - appendMessageFields(params, msg.chatMessage, msg.apiKeyID) + appendMessageFields(params, msg) } -// BuildSingleUserChatMessageInsertParams creates batch insert params for -// one user message, requiring an apiKeyID for AI Gateway attribution. -// BuildSingleChatMessageInsertParams creates batch insert params for one -// non-user message using the shared chat message builder. +// BuildSingleChatMessageInsertParams builds insert parameters for one chat message. func BuildSingleChatMessageInsertParams( chatID uuid.UUID, role database.ChatMessageRole, @@ -3005,31 +2935,7 @@ func BuildSingleChatMessageInsertParams( if createdBy != uuid.Nil { msg = msg.withCreatedBy(createdBy) } - if role == database.ChatMessageRoleUser { - appendMessageFields(¶ms, msg, "") - } else { - appendChatMessage(¶ms, msg) - } - return params -} - -func BuildSingleUserChatMessageInsertParams( - chatID uuid.UUID, - apiKeyID string, - content pqtype.NullRawMessage, - visibility database.ChatMessageVisibility, - modelConfigID uuid.UUID, - contentVersion int16, - createdBy uuid.UUID, -) database.InsertChatMessagesParams { - params := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. - ChatID: chatID, - } - msg := newUserChatMessage(apiKeyID, content, visibility, modelConfigID, contentVersion) - if createdBy != uuid.Nil { - msg = msg.withCreatedBy(createdBy) - } - appendUserChatMessage(¶ms, msg) + appendChatMessage(¶ms, msg) return params } diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 9537547c5dc..a702d13e96d 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -740,11 +740,6 @@ func TestRenameChatTitle(t *testing.T) { }) } -func withChatMessageAPIKeyID(message database.ChatMessage, apiKeyID string) database.ChatMessage { - message.APIKeyID = sqlNullString(apiKeyID) - return message -} - // requireOutgoingRequestModel asserts that the outgoing request body // requests wantModel. This is so that mock transports can still // verify the outgoing request asked for the expected model. @@ -865,12 +860,12 @@ func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) { LimitVal: manualTitleMessageWindowLimit, }, ).Return([]database.ChatMessage{ - withChatMessageAPIKeyID(mustChatMessage( + mustChatMessage( t, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, codersdk.ChatMessageText(userPrompt), - ), activeAPIKeyID), + ), mustChatMessage( t, database.ChatMessageRoleAssistant, @@ -1027,12 +1022,12 @@ func TestRegenerateChatTitle_SkipsPersistWhenTitleChangedConcurrently(t *testing LimitVal: manualTitleMessageWindowLimit, }, ).Return([]database.ChatMessage{ - withChatMessageAPIKeyID(mustChatMessage( + mustChatMessage( t, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, codersdk.ChatMessageText(userPrompt), - ), activeAPIKeyID), + ), }, nil) db.EXPECT().GetChatMessagesByChatIDDescPaginated( gomock.Any(), diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 7c2609ed725..ed51a39f3fe 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -72,12 +72,6 @@ type recordedOpenAIRequest struct { ContentLength int64 } -func testAPIKeyID(t testing.TB, db database.Store, userID uuid.UUID) string { - t.Helper() - key, _ := dbgen.APIKey(t, db, database.APIKey{ID: uuid.NewString(), UserID: userID}) - return key.ID -} - func chatAIGatewayTransportFactoryPointer(factory aibridge.TransportFactory) *atomic.Pointer[aibridge.TransportFactory] { var factoryPtr atomic.Pointer[aibridge.TransportFactory] factoryPtr.Store(&factory) @@ -768,7 +762,6 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { codersdk.ChatMessageText("inspect the codebase"), }) require.NoError(t, err) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) createdExplore, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, OwnerID: user.ID, @@ -792,7 +785,6 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: webSearchModel.ID, Valid: true}, - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, }, }, }) @@ -1601,163 +1593,6 @@ func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) { require.Equal(t, chat.ID, ids[0]) } -func TestCreateChatPersistsSyntheticAPIKeyIDOnInitialUserMessage(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "create-chat-synthetic-api-key-id", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 1) - require.Equal(t, database.ChatMessageRoleUser, messages[0].Role) - require.True(t, messages[0].APIKeyID.Valid) - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) - require.Equal(t, mapping.APIKeyID, messages[0].APIKeyID.String) -} - -func TestSendMessagePersistsSyntheticAPIKeyIDOnUserMessage(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: model.ID, - Title: "send-message-synthetic-api-key-id", - }) - - result, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - CreatedBy: user.ID, - Content: []codersdk.ChatMessagePart{ - codersdk.ChatMessageText("message with synthetic api key id"), - }, - }) - require.NoError(t, err) - require.False(t, result.Queued) - require.True(t, result.Message.APIKeyID.Valid) - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) - require.Equal(t, mapping.APIKeyID, result.Message.APIKeyID.String) - - stored, err := db.GetChatMessageByID(ctx, result.Message.ID) - require.NoError(t, err) - require.True(t, stored.APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, stored.APIKeyID.String) -} - -func TestSendMessagePersistsSyntheticAPIKeyIDOnQueuedUserMessage(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "queue-synthetic-api-key-id", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - result, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - CreatedBy: user.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, result.Queued) - require.NotNil(t, result.QueuedMessage) - - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) - require.True(t, result.QueuedMessage.APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, result.QueuedMessage.APIKeyID.String) - - queued, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Len(t, queued, 1) - require.True(t, queued[0].APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, queued[0].APIKeyID.String) -} - -func TestEditMessagePersistsSyntheticAPIKeyIDOnReplacement(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "edit-synthetic-api-key-id", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, - }) - require.NoError(t, err) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 1) - - result, err := replica.EditMessage(ctx, chatd.EditMessageOptions{ - ChatID: chat.ID, - EditedMessageID: messages[0].ID, - CreatedBy: user.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, - }) - require.NoError(t, err) - - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) - require.True(t, result.Message.APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, result.Message.APIKeyID.String) - - stored, err := db.GetChatMessageByID(ctx, result.Message.ID) - require.NoError(t, err) - require.True(t, stored.APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, stored.APIKeyID.String) -} - func TestSendMessageQueueBehaviorQueuesWhenBusy(t *testing.T) { t.Parallel() @@ -2759,7 +2594,6 @@ func TestRecoverStaleRequiresActionChat(t *testing.T) { codersdk.ChatMessageText("hello"), }) require.NoError(t, err) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, OwnerID: user.ID, @@ -2775,7 +2609,6 @@ func TestRecoverStaleRequiresActionChat(t *testing.T) { ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, }, }, }) @@ -2853,7 +2686,6 @@ func TestNewReplicaRecoversStaleChatFromDeadReplica(t *testing.T) { codersdk.ChatMessageText("hello"), }) require.NoError(t, err) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, OwnerID: user.ID, @@ -2868,7 +2700,6 @@ func TestNewReplicaRecoversStaleChatFromDeadReplica(t *testing.T) { ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, }, }, }) @@ -5451,8 +5282,6 @@ func TestActiveServer_RoutingPreservesAPIKeyAfterCompaction(t *testing.T) { }, }) require.NoError(t, err) - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) contextContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ Type: codersdk.ChatMessagePartTypeContextFile, ContextFileAgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, @@ -5462,9 +5291,9 @@ func TestActiveServer_RoutingPreservesAPIKeyAfterCompaction(t *testing.T) { ContextFileDirectory: "/home/coder/project", }}) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, chatd.BuildSingleUserChatMessageInsertParams( + _, err = db.InsertChatMessages(ctx, chatd.BuildSingleChatMessageInsertParams( chat.ID, - mapping.APIKeyID, + database.ChatMessageRoleUser, contextContent, database.ChatMessageVisibilityBoth, model.ID, @@ -5486,14 +5315,14 @@ func TestActiveServer_RoutingPreservesAPIKeyAfterCompaction(t *testing.T) { chatResult := waitForTerminalChat(ctx, t, db, chat.ID) require.Equal(t, database.ChatStatusWaiting, chatResult.Status) require.False(t, chatResult.LastError.Valid) + mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) + require.NoError(t, err) messages := chatMessages(ctx, t, db, chat.ID) promptMessages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) compressed := compressedChatSummarizedMessages(t, append(promptMessages, messages...)) require.Len(t, compressed.summaries, 1) - require.True(t, compressed.summaries[0].APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, compressed.summaries[0].APIKeyID.String) requests := factory.RequestsSnapshot() require.NotEmpty(t, requests) @@ -6442,7 +6271,6 @@ func userMessageForTest( ContentVersion: chatprompt.CurrentContentVersion, ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, - APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, } } @@ -7964,29 +7792,15 @@ func insertChatMessageParts( t.Helper() content, err := chatprompt.MarshalParts(parts) require.NoError(t, err) - var params database.InsertChatMessagesParams - if role == database.ChatMessageRoleUser { - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: createdBy}) - params = chatd.BuildSingleUserChatMessageInsertParams( - chatID, - apiKey.ID, - content, - database.ChatMessageVisibilityBoth, - modelID, - chatprompt.CurrentContentVersion, - createdBy, - ) - } else { - params = chatd.BuildSingleChatMessageInsertParams( - chatID, - role, - content, - database.ChatMessageVisibilityBoth, - modelID, - chatprompt.CurrentContentVersion, - createdBy, - ) - } + params := chatd.BuildSingleChatMessageInsertParams( + chatID, + role, + content, + database.ChatMessageVisibilityBoth, + modelID, + chatprompt.CurrentContentVersion, + createdBy, + ) messages, err := db.InsertChatMessages(ctx, params) require.NoError(t, err) require.Len(t, messages, 1) @@ -9816,12 +9630,11 @@ func seedAIGatewayOpenAITestDependencies( t *testing.T, db database.Store, openAIURL string, -) (database.User, database.Organization, database.AIProvider, database.ChatModelConfig, database.APIKey) { +) (database.User, database.Organization, database.AIProvider, database.ChatModelConfig) { t.Helper() user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, @@ -9844,7 +9657,7 @@ func seedAIGatewayOpenAITestDependencies( }) require.NoError(t, err) - return user, org, provider, model, apiKey + return user, org, provider, model } func TestProcessChat_RoutingUsesDelegatedAPIKey(t *testing.T) { @@ -9863,7 +9676,7 @@ func TestProcessChat_RoutingUsesDelegatedAPIKey(t *testing.T) { }) factory := chattest.NewMockAIBridgeTransport(t, openAIURL) - user, org, provider, model, _ := seedAIGatewayOpenAITestDependencies(t, db, openAIURL) + user, org, provider, model := seedAIGatewayOpenAITestDependencies(t, db, openAIURL) creator := newTestServer(t, db, ps, uuid.New()) chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ @@ -9876,8 +9689,6 @@ func TestProcessChat_RoutingUsesDelegatedAPIKey(t *testing.T) { }, }) require.NoError(t, err) - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) _, events, cancel, ok := creator.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -9894,6 +9705,8 @@ func TestProcessChat_RoutingUsesDelegatedAPIKey(t *testing.T) { chatResult := waitForTerminalChat(ctx, t, db, chat.ID) require.Equal(t, database.ChatStatusWaiting, chatResult.Status) require.False(t, chatResult.LastError.Valid) + mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) + require.NoError(t, err) requests := factory.RequestsSnapshot() require.NotEmpty(t, requests) @@ -9926,7 +9739,7 @@ func TestProcessChat_RoutingPreservesAPIKeyAfterWorkspaceContext(t *testing.T) { return chattest.OpenAINonStreamingResponse(`{"title":"AI Gateway Workspace"}`) }) factory := chattest.NewMockAIBridgeTransport(t, openAIURL) - user, org, provider, model, _ := seedAIGatewayOpenAITestDependencies(t, db, openAIURL) + user, org, provider, model := seedAIGatewayOpenAITestDependencies(t, db, openAIURL) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) creator := newTestServer(t, db, ps, uuid.New()) @@ -9941,8 +9754,6 @@ func TestProcessChat_RoutingPreservesAPIKeyAfterWorkspaceContext(t *testing.T) { }, }) require.NoError(t, err) - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) const contextText = "# Project instructions\nAlways keep routing metadata." // Workspace context is sourced from the agent's pinned snapshot. Seed it so @@ -9973,6 +9784,8 @@ func TestProcessChat_RoutingPreservesAPIKeyAfterWorkspaceContext(t *testing.T) { pinned, err := db.ListChatContextResourcesByChatID(ctx, chat.ID) require.NoError(t, err) require.NotEmpty(t, pinned, "workspace context should be pinned to the chat") + mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) + require.NoError(t, err) requests := factory.RequestsSnapshot() require.NotEmpty(t, requests) @@ -11788,7 +11601,6 @@ func TestPromoteQueuedPreservesReasoningEffort(t *testing.T) { Content: content.RawMessage, ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, ReasoningEffort: database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortHigh, Valid: true}, - APIKeyID: sql.NullString{String: testAPIKeyID(t, db, user.ID), Valid: true}, CreatedBy: user.ID, }) require.NoError(t, err) diff --git a/coderd/x/chatd/chatstate/machine_test.go b/coderd/x/chatd/chatstate/machine_test.go index c4def781742..e89effc9cef 100644 --- a/coderd/x/chatd/chatstate/machine_test.go +++ b/coderd/x/chatd/chatstate/machine_test.go @@ -2,7 +2,6 @@ package chatstate_test import ( "context" - "database/sql" "encoding/json" "slices" "sync" @@ -34,13 +33,6 @@ type testFixture struct { User database.User Org database.Organization Model database.ChatModelConfig - APIKey database.APIKey -} - -// apiKeyID returns the fixture API key wrapped for the chatstate -// inputs that require a non-null api_key_id (for example EditMessage). -func (f *testFixture) apiKeyID() sql.NullString { - return sql.NullString{String: f.APIKey.ID, Valid: true} } func newTestFixture(t *testing.T) *testFixture { @@ -60,7 +52,6 @@ func newTestFixture(t *testing.T) *testFixture { model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ IsDefault: true, }) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) pub := newRecordingPubsub() return &testFixture{ DB: db, @@ -69,7 +60,6 @@ func newTestFixture(t *testing.T) *testFixture { User: user, Org: org, Model: model, - APIKey: apiKey, } } diff --git a/coderd/x/chatd/chatstate/messages.go b/coderd/x/chatd/chatstate/messages.go index 50d84563f22..b867c1f05ea 100644 --- a/coderd/x/chatd/chatstate/messages.go +++ b/coderd/x/chatd/chatstate/messages.go @@ -35,7 +35,6 @@ type Message struct { ContextLimit sql.NullInt64 TotalCostMicros sql.NullInt64 RuntimeMs sql.NullInt64 - APIKeyID sql.NullString } // toInsertParams converts a batch of Messages into the parallel-array @@ -51,7 +50,6 @@ func toInsertParams(chatID uuid.UUID, messages []Message) database.InsertChatMes CreatedBy: make([]uuid.UUID, n), ModelConfigID: make([]uuid.UUID, n), ReasoningEffort: make([]string, n), - APIKeyID: make([]string, n), Role: make([]database.ChatMessageRole, n), Content: make([]string, n), ContentVersion: make([]int16, n), @@ -73,9 +71,6 @@ func toInsertParams(chatID uuid.UUID, messages []Message) database.InsertChatMes if m.ReasoningEffort.Valid { params.ReasoningEffort[i] = string(m.ReasoningEffort.ChatReasoningEffort) } - if m.APIKeyID.Valid { - params.APIKeyID[i] = m.APIKeyID.String - } params.Role[i] = m.Role if m.Content.Valid { params.Content[i] = string(m.Content.RawMessage) diff --git a/coderd/x/chatd/chatstate/synthetic_cancellation_test.go b/coderd/x/chatd/chatstate/synthetic_cancellation_test.go index 75880aa2f7f..d4056c34929 100644 --- a/coderd/x/chatd/chatstate/synthetic_cancellation_test.go +++ b/coderd/x/chatd/chatstate/synthetic_cancellation_test.go @@ -249,7 +249,6 @@ func testEditMessageSynthesizesToolCancellationsBeforeReplacement(t *testing.T) MessageID: secondUserID, CreatedBy: f.User.ID, Content: editedContent, - APIKeyID: f.apiKeyID(), }) return err })) diff --git a/coderd/x/chatd/chatstate/transitions.go b/coderd/x/chatd/chatstate/transitions.go index 3b2f23479c1..dd1ee357aea 100644 --- a/coderd/x/chatd/chatstate/transitions.go +++ b/coderd/x/chatd/chatstate/transitions.go @@ -230,14 +230,17 @@ func (tx *Tx) insertQueuedMessage(ownerFallback uuid.UUID, m Message) (database. if err := tx.requireQueueCapacity(); err != nil { return database.ChatQueuedMessage{}, err } - return tx.store.InsertChatQueuedMessageWithCreator(tx.ctx, database.InsertChatQueuedMessageWithCreatorParams{ + row, err := tx.store.InsertChatQueuedMessageWithCreator(tx.ctx, database.InsertChatQueuedMessageWithCreatorParams{ ChatID: tx.chatID, Content: rawContent, ModelConfigID: m.ModelConfigID, ReasoningEffort: m.ReasoningEffort, CreatedBy: createdBy, - APIKeyID: m.APIKeyID, }) + if err != nil { + return database.ChatQueuedMessage{}, err + } + return database.ChatQueuedMessageRow(row).ChatQueuedMessage(), nil } // messageFromQueuedRow synthesizes a Message from a stored queued row, @@ -251,7 +254,6 @@ func messageFromQueuedRow(q database.ChatQueuedMessage) Message { ReasoningEffort: q.ReasoningEffort, CreatedBy: uuid.NullUUID{UUID: q.CreatedBy, Valid: true}, ContentVersion: chatprompt.CurrentContentVersion, - APIKeyID: q.APIKeyID, } } @@ -416,11 +418,12 @@ func (tx *Tx) sendMessageE1(chat database.Chat, m Message) (SendMessageResult, e if err != nil { return SendMessageResult{}, xerrors.Errorf("get queue head: %w", err) } + headMessage := database.ChatQueuedMessageRow(head).ChatQueuedMessage() cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by queued message promotion", false) if err != nil { return SendMessageResult{}, err } - promoted := messageFromQueuedRow(head) + promoted := messageFromQueuedRow(headMessage) inserted, err := tx.insertMessages(append(cancels, promoted)) if err != nil { return SendMessageResult{}, xerrors.Errorf("insert promoted queued head: %w", err) @@ -491,7 +494,6 @@ type EditMessageInput struct { Content pqtype.NullRawMessage ModelConfigIDOverride uuid.NullUUID ReasoningEffortOverride database.NullChatReasoningEffort - APIKeyID sql.NullString } // EditMessageResult is returned by [Tx.EditMessage]. @@ -571,10 +573,6 @@ func (tx *Tx) EditMessage(input EditMessageInput) (EditMessageResult, error) { if input.ReasoningEffortOverride.Valid { reasoningEffort = input.ReasoningEffortOverride } - apiKeyID := input.APIKeyID - if !apiKeyID.Valid { - return EditMessageResult{}, xerrors.Errorf("api_key_id is required") - } replacement := Message{ Role: database.ChatMessageRoleUser, Content: input.Content, @@ -583,7 +581,6 @@ func (tx *Tx) EditMessage(input EditMessageInput) (EditMessageResult, error) { ReasoningEffort: reasoningEffort, CreatedBy: uuid.NullUUID{UUID: input.CreatedBy, Valid: true}, ContentVersion: chatprompt.CurrentContentVersion, - APIKeyID: apiKeyID, } insertedReplacement, err := tx.insertMessages([]Message{replacement}) if err != nil { @@ -643,6 +640,7 @@ func (tx *Tx) DeleteQueuedMessage(input DeleteQueuedMessageInput) (DeleteQueuedM if err != nil { return DeleteQueuedMessageResult{}, xerrors.Errorf("get queued: %w", err) } + targetMessage := database.ChatQueuedMessageRow(target).ChatQueuedMessage() rows, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ ID: input.QueuedMessageID, ChatID: tx.chatID, @@ -654,7 +652,7 @@ func (tx *Tx) DeleteQueuedMessage(input DeleteQueuedMessageInput) (DeleteQueuedM return DeleteQueuedMessageResult{}, ErrQueuedMessageNotFound } return DeleteQueuedMessageResult{ - DeletedQueuedMessage: target, + DeletedQueuedMessage: targetMessage, }, nil } @@ -688,6 +686,7 @@ func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueu if err != nil { return PromoteQueuedMessageResult{}, xerrors.Errorf("get queued: %w", err) } + targetMessage := database.ChatQueuedMessageRow(target).ChatQueuedMessage() rows, err := tx.store.ReorderChatQueuedMessageToHead(tx.ctx, database.ReorderChatQueuedMessageToHeadParams{ ID: input.QueuedMessageID, ChatID: tx.chatID, @@ -713,7 +712,7 @@ func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueu return PromoteQueuedMessageResult{}, xerrors.Errorf("set interrupting: %w", err) } return PromoteQueuedMessageResult{ - QueuedMessage: target, + QueuedMessage: targetMessage, ReorderedQueueOnly: reorderOnly, }, nil } @@ -726,7 +725,7 @@ func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueu if err != nil { return PromoteQueuedMessageResult{}, err } - promotedMsg := messageFromQueuedRow(target) + promotedMsg := messageFromQueuedRow(targetMessage) inserted, err := tx.insertMessages(append(cancels, promotedMsg)) if err != nil { return PromoteQueuedMessageResult{}, xerrors.Errorf("insert promoted queued message: %w", err) @@ -756,7 +755,7 @@ func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueu cancellations := inserted[:len(inserted)-1] insertedUserMsg := inserted[len(inserted)-1] return PromoteQueuedMessageResult{ - QueuedMessage: target, + QueuedMessage: targetMessage, InsertedMessage: &insertedUserMsg, CancellationMessages: cancellations, ReorderedQueueOnly: reorderOnly, @@ -1211,7 +1210,8 @@ func (tx *Tx) FinishInterruption(input FinishInterruptionInput) (FinishInterrupt if err != nil { return FinishInterruptionResult{}, xerrors.Errorf("get queue head: %w", err) } - promotedMsg := messageFromQueuedRow(head) + headMessage := database.ChatQueuedMessageRow(head).ChatQueuedMessage() + promotedMsg := messageFromQueuedRow(headMessage) insertedHead, err := tx.insertMessages([]Message{promotedMsg}) if err != nil { return FinishInterruptionResult{}, xerrors.Errorf("insert promoted queue head: %w", err) @@ -1277,11 +1277,12 @@ func (tx *Tx) FinishTurn(_ FinishTurnInput) (FinishTurnResult, error) { if err != nil { return FinishTurnResult{}, xerrors.Errorf("get queue head: %w", err) } + headMessage := database.ChatQueuedMessageRow(head).ChatQueuedMessage() cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by queued message promotion", false) if err != nil { return FinishTurnResult{}, err } - promotedMsg := messageFromQueuedRow(head) + promotedMsg := messageFromQueuedRow(headMessage) inserted, err := tx.insertMessages(append(cancels, promotedMsg)) if err != nil { return FinishTurnResult{}, xerrors.Errorf("insert promoted queue head: %w", err) diff --git a/coderd/x/chatd/chatstate/transitions_helpers_test.go b/coderd/x/chatd/chatstate/transitions_helpers_test.go index 91cc2419f83..2eb54189dfe 100644 --- a/coderd/x/chatd/chatstate/transitions_helpers_test.go +++ b/coderd/x/chatd/chatstate/transitions_helpers_test.go @@ -700,7 +700,7 @@ func requireQueuedMessageByID(ctx context.Context, t *testing.T, f *testFixture, ChatID: chatID, }) require.NoError(t, err) - return msg + return database.ChatQueuedMessageRow(msg).ChatQueuedMessage() } func requireQueuedMessageDeleted(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID, id int64) { @@ -794,9 +794,14 @@ func assertChatMessageText(t *testing.T, msg database.ChatMessage, want string) // matrix cases that need to verify the body inserted into // chat_queued_messages via SendMessage. func assertQueuedMessageText(t *testing.T, queued database.ChatQueuedMessage, want string) { + t.Helper() + assertQueuedMessageContent(t, queued.Content, want) +} + +func assertQueuedMessageContent(t *testing.T, content json.RawMessage, want string) { t.Helper() var parts []codersdk.ChatMessagePart - require.NoError(t, json.Unmarshal(queued.Content, &parts), "unmarshal queued content") + require.NoError(t, json.Unmarshal(content, &parts), "unmarshal queued content") require.Len(t, parts, 1, "expected exactly one queued content part") require.Equal(t, codersdk.ChatMessagePartTypeText, parts[0].Type, "expected a text content part") @@ -814,7 +819,7 @@ func assertQueueBodiesInOrder(ctx context.Context, t *testing.T, f *testFixture, require.NoError(t, err) require.Len(t, rows, len(want), "queue length must match expected bodies") for i, r := range rows { - assertQueuedMessageText(t, r, want[i]) + assertQueuedMessageContent(t, r.Content, want[i]) } } diff --git a/coderd/x/chatd/chatstate/transitions_matrix_test.go b/coderd/x/chatd/chatstate/transitions_matrix_test.go index ef35090f7da..cc2e5793821 100644 --- a/coderd/x/chatd/chatstate/transitions_matrix_test.go +++ b/coderd/x/chatd/chatstate/transitions_matrix_test.go @@ -136,7 +136,6 @@ func applyEditMessage(t *testing.T, f *testFixture, tx *chatstate.Tx, seeded see MessageID: seeded.initialUserMessageID, CreatedBy: f.User.ID, Content: content, - APIKeyID: f.apiKeyID(), }) return err } diff --git a/coderd/x/chatd/chatstate_bridge.go b/coderd/x/chatd/chatstate_bridge.go index 04aae64c073..ae3cf59e114 100644 --- a/coderd/x/chatd/chatstate_bridge.go +++ b/coderd/x/chatd/chatstate_bridge.go @@ -1,8 +1,6 @@ package chatd import ( - "database/sql" - "github.com/google/uuid" "github.com/sqlc-dev/pqtype" @@ -29,7 +27,7 @@ func systemMessage(rawContent pqtype.NullRawMessage, modelConfigID uuid.UUID) ch } } -func userMessageWithAPIKeyID(rawContent pqtype.NullRawMessage, modelConfigID, createdBy uuid.UUID, apiKeyID string, reasoningEffort *string) chatstate.Message { +func userMessage(rawContent pqtype.NullRawMessage, modelConfigID, createdBy uuid.UUID, reasoningEffort *string) chatstate.Message { var effort database.NullChatReasoningEffort if reasoningEffort != nil && *reasoningEffort != "" { effort = database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffort(*reasoningEffort), Valid: true} @@ -42,7 +40,6 @@ func userMessageWithAPIKeyID(rawContent pqtype.NullRawMessage, modelConfigID, cr ReasoningEffort: effort, CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: createdBy != uuid.Nil}, ContentVersion: chatprompt.CurrentContentVersion, - APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, } } diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index bdd7583f084..2ffd24a0027 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -691,7 +691,6 @@ func (s *taskStarter) generateCompaction( } messages, err := buildCompactionMessages(buildCompactionMessagesInput{ modelConfigID: prepared.ModelConfigID, - activeAPIKeyID: prepared.ModelBuildOptions.ActiveAPIKeyID, toolCallID: compactionOpts.ToolCallID, toolName: compactionOpts.ToolName, compaction: compactionOutcome(outcome), diff --git a/coderd/x/chatd/generation_preparer_internal_test.go b/coderd/x/chatd/generation_preparer_internal_test.go index 13cbb4ac13b..44d860de5f8 100644 --- a/coderd/x/chatd/generation_preparer_internal_test.go +++ b/coderd/x/chatd/generation_preparer_internal_test.go @@ -1,7 +1,6 @@ package chatd //nolint:testpackage // Exercises unexported re-derivation helpers. import ( - "database/sql" "encoding/json" "testing" @@ -93,7 +92,6 @@ func TestPrepareGenerationClampsRequestedReasoningEffortToMax(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := chatdTestContext(t) user := dbgen.User(t, db, database.User{}) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) org := dbgen.Organization(t, db, database.Organization{}) dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -135,7 +133,6 @@ func TestPrepareGenerationClampsRequestedReasoningEffortToMax(t *testing.T) { }, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, ContentVersion: chatprompt.CurrentContentVersion, - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, }, }, }) @@ -161,6 +158,71 @@ func TestPrepareGenerationClampsRequestedReasoningEffortToMax(t *testing.T) { require.Equal(t, fantasyopenai.ReasoningEffortMedium, *providerOptions.ReasoningEffort) } +func TestPrepareGenerationSubagentUsesOwnerSyntheticAPIKey(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := chatdTestContext(t) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + provider := dbgen.AIProviderWithOptionalKey(t, db, database.AIProvider{ + Type: database.AIProviderTypeOpenai, + }, "test-key") + modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + }, func(p *database.InsertChatModelConfigParams) { + p.Enabled = true + }) + parent := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: modelConfig.ID, + }) + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, + LastModelConfigID: modelConfig.ID, + Title: "subagent attribution", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: mustMarshalText(t, "inspect the workspace"), + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + }, + }, + }) + require.NoError(t, err) + + server := newInternalTestServer( + t, + db, + ps, + chatprovider.ProviderAPIKeys{}, + withInternalTestServerTransportFactory(&aibridgeTestFactory{}), + ) + prepared, err := server.prepareGeneration(ctx, generationPrepareInput{ + Chat: created.Chat, + Messages: created.InitialMessages, + }) + require.NoError(t, err) + t.Cleanup(prepared.Cleanup) + + mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) + require.NoError(t, err) + require.Equal(t, mapping.APIKeyID, prepared.ModelBuildOptions.ActiveAPIKeyID) +} + // TestDeriveFinalTurnRunResult exercises the re-derivation path that replaces // the old in-memory generationSideEffects stash. The server here never ran // prepareGeneration, so a passing test proves the finish-turn inputs are @@ -196,7 +258,6 @@ func TestDeriveFinalTurnRunResult(t *testing.T) { p.Enabled = true p.IsDefault = true }) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, @@ -212,7 +273,6 @@ func TestDeriveFinalTurnRunResult(t *testing.T) { ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, }, }, }) @@ -314,7 +374,6 @@ func TestDeriveFinalTurnRunResult(t *testing.T) { DisplayName: "gpt-4o-mini", AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, }) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, @@ -330,7 +389,6 @@ func TestDeriveFinalTurnRunResult(t *testing.T) { ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, }, }, }) diff --git a/coderd/x/chatd/helpers_test.go b/coderd/x/chatd/helpers_test.go index 178c51b9461..18e4ec3f0ec 100644 --- a/coderd/x/chatd/helpers_test.go +++ b/coderd/x/chatd/helpers_test.go @@ -202,7 +202,6 @@ func userTextMessage(t *testing.T, text string, createdBy uuid.UUID, modelConfig ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, } } diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index c9d4cc348e6..2197e0aaab4 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -255,7 +255,6 @@ func textFromParts(parts []codersdk.ChatMessagePart) string { type buildCompactionMessagesInput struct { modelConfigID uuid.UUID - activeAPIKeyID string toolCallID string toolName string compaction compactionOutcome @@ -319,7 +318,6 @@ func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMess Visibility: database.ChatMessageVisibilityModel, ModelConfigID: uuid.NullUUID{UUID: input.modelConfigID, Valid: input.modelConfigID != uuid.Nil}, ContentVersion: contentVersion, - APIKeyID: sql.NullString{String: input.activeAPIKeyID, Valid: input.activeAPIKeyID != ""}, }, baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, input.modelConfigID, contentVersion, assistantContent), baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, toolContent), diff --git a/coderd/x/chatd/model_routing_internal_test.go b/coderd/x/chatd/model_routing_internal_test.go index f03b94c0c54..55a01931aec 100644 --- a/coderd/x/chatd/model_routing_internal_test.go +++ b/coderd/x/chatd/model_routing_internal_test.go @@ -365,10 +365,6 @@ func TestAIGatewayModelForwardsProviderAuth(t *testing.T) { }) } -func sqlNullString(value string) sql.NullString { - return sql.NullString{String: value, Valid: value != ""} -} - func TestAIBridgeRoutingFailClosed(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/stream_loop.go b/coderd/x/chatd/stream_loop.go index 5004e8a490e..987c5af3617 100644 --- a/coderd/x/chatd/stream_loop.go +++ b/coderd/x/chatd/stream_loop.go @@ -176,10 +176,14 @@ func (l *streamLoop) loadDBSnapshot(ctx context.Context) (streamDBSnapshot, erro if chat.QueueVersion > l.state.queueVersion { snapshot.queueChanged = true - snapshot.queue, err = tx.GetChatQueuedMessages(ctx, l.chatID) + queued, err := tx.GetChatQueuedMessages(ctx, l.chatID) if err != nil { return xerrors.Errorf("get chat queue: %w", err) } + snapshot.queue = make([]database.ChatQueuedMessage, len(queued)) + for i, row := range queued { + snapshot.queue[i] = database.ChatQueuedMessageRow(row).ChatQueuedMessage() + } } if chat.Status == database.ChatStatusRequiresAction { diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index b046863fb66..5e9fea275b6 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -1056,11 +1056,6 @@ func (p *Server) createChildSubagentChatWithOptions( if modelConfigID == uuid.Nil { return database.Chat{}, xerrors.New("model config is required") } - childAPIKeyID, err := p.ensureSyntheticAPIKeyID(ctx, parent.OwnerID) - if err != nil { - return database.Chat{}, xerrors.Errorf("ensure synthetic API key: %w", err) - } - childPlanMode := parent.PlanMode if opts.planModeOverride != nil { childPlanMode = *opts.planModeOverride @@ -1131,7 +1126,7 @@ func (p *Server) createChildSubagentChatWithOptions( // workspace context the same way a top-level chat does: pinned from the // agent's latest snapshot (see hydrateChatContextOnCreate below). The // parent's context is not copied into child history. - initialMessages = append(initialMessages, userMessageWithAPIKeyID(userContent, modelConfigID, parent.OwnerID, childAPIKeyID, opts.reasoningEffortOverride)) + initialMessages = append(initialMessages, userMessage(userContent, modelConfigID, parent.OwnerID, opts.reasoningEffortOverride)) publisher := p.pubsub if publisher == nil { diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index c50e33647fd..4b582632f64 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -651,83 +651,6 @@ func upsertInternalUserChatPersonalModelOverride( ) } -func TestCreateChildSubagentChatPersistsOwnerSyntheticAPIKeyID(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) - - ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(t, db) - parent := createInternalParentChat( - ctx, t, server, db, org.ID, user.ID, model.ID, "parent-child-key", - ) - - child, err := server.createChildSubagentChatWithOptions( - ctx, - parent, - "inspect the workspace", - "", - childSubagentChatOptions{}, - ) - require.NoError(t, err) - - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: child.ID, - AfterID: 0, - }) - require.NoError(t, err) - for _, message := range messages { - if message.Role != database.ChatMessageRoleUser { - continue - } - require.True(t, message.APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, message.APIKeyID.String) - return - } - require.Fail(t, "child user message not found") -} - -func TestSendSubagentMessagePersistsOwnerSyntheticAPIKeyID(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) - - ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(t, db) - parent, child := createParentChildChats(ctx, t, server, user, org, model) - setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") - - _, err := server.sendSubagentMessage( - ctx, - parent.ID, - child.ID, - "follow up", - SendMessageBusyBehaviorInterrupt, - ) - require.NoError(t, err) - - mapping, err := db.GetChatSyntheticAPIKeyByUserID(ctx, user.ID) - require.NoError(t, err) - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: child.ID, - AfterID: 0, - }) - require.NoError(t, err) - var latestUserMessage database.ChatMessage - for _, message := range messages { - if message.Role == database.ChatMessageRoleUser && message.ID > latestUserMessage.ID { - latestUserMessage = message - } - } - require.NotZero(t, latestUserMessage.ID) - require.True(t, latestUserMessage.APIKeyID.Valid) - require.Equal(t, mapping.APIKeyID, latestUserMessage.APIKeyID.String) -} - func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/synthetickey_internal_test.go b/coderd/x/chatd/synthetickey_internal_test.go index 986479952cf..f76b876da2f 100644 --- a/coderd/x/chatd/synthetickey_internal_test.go +++ b/coderd/x/chatd/synthetickey_internal_test.go @@ -128,18 +128,16 @@ func TestSyntheticAPIKeyDeletionDoesNotMutateChatState(t *testing.T) { OwnerID: user.ID, LastModelConfigID: model.ID, }) - message := dbgen.ChatMessage(t, db, database.ChatMessage{ + 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.ChatMessageRoleUser, - APIKeyID: sql.NullString{String: syntheticID, Valid: true}, }) - queued, err := db.InsertChatQueuedMessage(t.Context(), database.InsertChatQueuedMessageParams{ + _, err = db.InsertChatQueuedMessage(t.Context(), database.InsertChatQueuedMessageParams{ ChatID: chat.ID, Content: json.RawMessage(`[]`), ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - APIKeyID: sql.NullString{String: syntheticID, Valid: true}, }) require.NoError(t, err) @@ -158,16 +156,6 @@ func TestSyntheticAPIKeyDeletionDoesNotMutateChatState(t *testing.T) { require.Equal(t, before.QueueVersion, after.QueueVersion) require.Equal(t, before.GenerationAttempt, after.GenerationAttempt) - stored, err := db.GetChatMessageByID(t.Context(), message.ID) - require.NoError(t, err) - require.Equal(t, sql.NullString{String: syntheticID, Valid: true}, stored.APIKeyID) - storedQueued, err := db.GetChatQueuedMessageByID(t.Context(), database.GetChatQueuedMessageByIDParams{ - ID: queued.ID, - ChatID: chat.ID, - }) - require.NoError(t, err) - require.Equal(t, sql.NullString{String: syntheticID, Valid: true}, storedQueued.APIKeyID) - remintedID, err := server.ensureSyntheticAPIKeyID(t.Context(), user.ID) require.NoError(t, err) require.NotEqual(t, syntheticID, remintedID) diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 31882f59919..8f95da4e193 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -1044,7 +1044,6 @@ func taskUserTextMessage(t *testing.T, text string, createdBy uuid.UUID, modelCo ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, } } diff --git a/coderd/x/chatd/turn_summary_internal_test.go b/coderd/x/chatd/turn_summary_internal_test.go index d5d545b2ed6..1ef5f52e4c4 100644 --- a/coderd/x/chatd/turn_summary_internal_test.go +++ b/coderd/x/chatd/turn_summary_internal_test.go @@ -59,7 +59,6 @@ func TestUpdateLastTurnSummaryRejectsStaleWrites(t *testing.T) { codersdk.ChatMessageText("hello"), }) require.NoError(t, err) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: owner.ID}) created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, OwnerID: owner.ID, @@ -74,7 +73,6 @@ func TestUpdateLastTurnSummaryRejectsStaleWrites(t *testing.T) { ContentVersion: chatprompt.CurrentContentVersion, CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, }, }, })