feat: Add full text search over chat messages (#27126)

Closes CODAGT-721
Closes CODAGT-722
Closes CODAGT-723
Closes CODAGT-724
Closes CODAGT-725

This PR adds the database and API pieces necessary to support full-text
chat message search.

- Adds required chat schema for full-text search
- Adds dbpurge job to populate search_tsv in the background
- Adds `search` parameter to GetChats query
- Adds `search` filter to `searchquery.Chats`
- Wires chat search filter into chats API

> Implemented by Coder Agents, reviewed and tested by a human.
This commit is contained in:
Cian Johnston
2026-07-16 15:21:57 +01:00
committed by GitHub
parent 9862f10484
commit f7481c5d08
26 changed files with 1756 additions and 35 deletions
+16 -1
View File
@@ -697,7 +697,8 @@ var (
rbac.ResourceApiKey.Type: {policy.ActionDelete},
rbac.ResourceAibridgeInterception.Type: {policy.ActionDelete},
rbac.ResourceWorkspaceBuildOrchestration.Type: {policy.ActionDelete},
// Chat auto-archive sets archived=true on inactive chats.
// Chat auto-archive sets archived=true on inactive chats and computes
// search_tsv tsvector for chat_messages.
rbac.ResourceChat.Type: {policy.ActionRead, policy.ActionUpdate},
// Purge old boundary logs past the retention period.
rbac.ResourceBoundaryLog.Type: {policy.ActionDelete},
@@ -1743,6 +1744,13 @@ func (q *querier) AutoArchiveInactiveChats(ctx context.Context, arg database.Aut
return q.db.AutoArchiveInactiveChats(ctx, arg)
}
func (q *querier) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
return 0, err
}
return q.db.BackfillChatMessagesSearchTsv(ctx, batchSize)
}
func (q *querier) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
// This is a system-level operation used by the gitsync
// background worker to reschedule failed refreshes. Same
@@ -1820,6 +1828,13 @@ func (q *querier) CalculateAIBridgeInterceptionsTelemetrySummary(ctx context.Con
return q.db.CalculateAIBridgeInterceptionsTelemetrySummary(ctx, arg)
}
func (q *querier) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat); err != nil {
return false, err
}
return q.db.ChatSearchQueryIsEmpty(ctx, search)
}
func (q *querier) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
empty := database.ClaimPrebuiltWorkspaceRow{}
+8
View File
@@ -1002,6 +1002,14 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().DeleteOldChats(gomock.Any(), database.DeleteOldChatsParams{}).Return(int64(0), nil).AnyTimes()
check.Args(database.DeleteOldChatsParams{}).Asserts(rbac.ResourceSystem, policy.ActionDelete)
}))
s.Run("BackfillChatMessagesSearchTsv", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), int32(100)).Return(int64(0), nil).AnyTimes()
check.Args(int32(100)).Asserts(rbac.ResourceChat, policy.ActionUpdate)
}))
s.Run("ChatSearchQueryIsEmpty", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().ChatSearchQueryIsEmpty(gomock.Any(), "!!!").Return(true, nil).AnyTimes()
check.Args("!!!").Asserts(rbac.ResourceChat, policy.ActionRead)
}))
s.Run("GetChatRetentionDays", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(30), nil).AnyTimes()
check.Args().Asserts()
+16
View File
@@ -177,6 +177,14 @@ func (m queryMetricsStore) AutoArchiveInactiveChats(ctx context.Context, arg dat
return r0, r1
}
func (m queryMetricsStore) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
start := time.Now()
r0, r1 := m.s.BackfillChatMessagesSearchTsv(ctx, batchSize)
m.queryLatencies.WithLabelValues("BackfillChatMessagesSearchTsv").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BackfillChatMessagesSearchTsv").Inc()
return r0, r1
}
func (m queryMetricsStore) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
start := time.Now()
r0 := m.s.BackoffChatDiffStatus(ctx, arg)
@@ -257,6 +265,14 @@ func (m queryMetricsStore) CalculateAIBridgeInterceptionsTelemetrySummary(ctx co
return r0, r1
}
func (m queryMetricsStore) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
start := time.Now()
r0, r1 := m.s.ChatSearchQueryIsEmpty(ctx, search)
m.queryLatencies.WithLabelValues("ChatSearchQueryIsEmpty").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ChatSearchQueryIsEmpty").Inc()
return r0, r1
}
func (m queryMetricsStore) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
start := time.Now()
r0, r1 := m.s.ClaimPrebuiltWorkspace(ctx, arg)
+30
View File
@@ -178,6 +178,21 @@ func (mr *MockStoreMockRecorder) AutoArchiveInactiveChats(ctx, arg any) *gomock.
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AutoArchiveInactiveChats", reflect.TypeOf((*MockStore)(nil).AutoArchiveInactiveChats), ctx, arg)
}
// BackfillChatMessagesSearchTsv mocks base method.
func (m *MockStore) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "BackfillChatMessagesSearchTsv", ctx, batchSize)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// BackfillChatMessagesSearchTsv indicates an expected call of BackfillChatMessagesSearchTsv.
func (mr *MockStoreMockRecorder) BackfillChatMessagesSearchTsv(ctx, batchSize any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BackfillChatMessagesSearchTsv", reflect.TypeOf((*MockStore)(nil).BackfillChatMessagesSearchTsv), ctx, batchSize)
}
// BackoffChatDiffStatus mocks base method.
func (m *MockStore) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
m.ctrl.T.Helper()
@@ -322,6 +337,21 @@ func (mr *MockStoreMockRecorder) CalculateAIBridgeInterceptionsTelemetrySummary(
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CalculateAIBridgeInterceptionsTelemetrySummary", reflect.TypeOf((*MockStore)(nil).CalculateAIBridgeInterceptionsTelemetrySummary), ctx, arg)
}
// ChatSearchQueryIsEmpty mocks base method.
func (m *MockStore) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ChatSearchQueryIsEmpty", ctx, search)
ret0, _ := ret[0].(bool)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ChatSearchQueryIsEmpty indicates an expected call of ChatSearchQueryIsEmpty.
func (mr *MockStoreMockRecorder) ChatSearchQueryIsEmpty(ctx, search any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ChatSearchQueryIsEmpty", reflect.TypeOf((*MockStore)(nil).ChatSearchQueryIsEmpty), ctx, search)
}
// ClaimPrebuiltWorkspace mocks base method.
func (m *MockStore) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
m.ctrl.T.Helper()
+65 -14
View File
@@ -51,6 +51,12 @@ const (
// Chat debug run deletions can cascade into steps with large JSONB
// payloads, so they use the same conservative batch size.
chatDebugRunsBatchSize = 1000
// Chat search tsvector backfill is capped at 5 batches of 10k
// rows per tick. Benchmarks on a dogfood-class machine (EPYC 9454P)
// with containerized Postgres were measured to take ~800ms per batch.
// This is considered acceptable but may need dialing in later.
chatSearchBackfillBatchSize = 10000
chatSearchBackfillMaxBatches = 5
)
type Option func(*instance)
@@ -61,6 +67,14 @@ func WithClock(clk quartz.Clock) Option {
return func(i *instance) { i.clk = clk }
}
// WithChatSearchBackfillLimits overrides backfill batch size and cap. For tests.
func WithChatSearchBackfillLimits(batchSize int32, maxBatches int) Option {
return func(i *instance) {
i.chatSearchBackfillBatchSize = batchSize
i.chatSearchBackfillMaxBatches = maxBatches
}
}
// New creates a new periodically purging database instance.
// Callers must Close the returned instance.
func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, reg prometheus.Registerer, opts ...Option) io.Closer {
@@ -87,14 +101,25 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
}, []string{"record_type"})
reg.MustRegister(recordsPurged)
chatSearchRowsBackfilled := prometheus.NewCounter(prometheus.CounterOpts{
Namespace: "coderd",
Subsystem: "dbpurge",
Name: "chat_search_rows_backfilled_total",
Help: "Total number of chat message rows whose search_tsv was backfilled.",
})
reg.MustRegister(chatSearchRowsBackfilled)
inst := &instance{
cancel: cancelFunc,
closed: closed,
logger: logger,
vals: vals,
clk: quartz.NewReal(),
iterationDuration: iterationDuration,
recordsPurged: recordsPurged,
cancel: cancelFunc,
closed: closed,
logger: logger,
vals: vals,
clk: quartz.NewReal(),
iterationDuration: iterationDuration,
recordsPurged: recordsPurged,
chatSearchRowsBackfilled: chatSearchRowsBackfilled,
chatSearchBackfillBatchSize: chatSearchBackfillBatchSize,
chatSearchBackfillMaxBatches: chatSearchBackfillMaxBatches,
}
for _, opt := range opts {
opt(inst)
@@ -310,6 +335,25 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
}
}
// Backfill search_tsv tsvector on chat_messages in batches. Doing this here because it's
// potentially too much for a regular migration, especially on larger deployments:
// - Each row with search_tsv = NULL is present in idx_chat_messages_search_tsv_pending.
// - Content of chat_messages is not changed after insert.
// - Rows that are soft-deleted are no longer part of the index.
// NOTE: This should not remain in dbpurge and should be adjusted when the "DBOps" gets
// implemented.
var backfilledChatSearchRows int64
for range i.chatSearchBackfillMaxBatches {
n, err := tx.BackfillChatMessagesSearchTsv(ctx, i.chatSearchBackfillBatchSize)
if err != nil {
return xerrors.Errorf("backfill chat_messages.search_tsv: %w", err)
}
backfilledChatSearchRows += n
if n < int64(i.chatSearchBackfillBatchSize) {
break
}
}
i.logger.Debug(ctx, "purged old database entries",
slog.F("workspace_agent_logs", purgedWorkspaceAgentLogs),
slog.F("expired_api_keys", expiredAPIKeys),
@@ -322,6 +366,7 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
slog.F("chats", purgedChats),
slog.F("chat_files", purgedChatFiles),
slog.F("chat_debug_runs", purgedChatDebugRuns),
slog.F("chat_search_rows_backfilled", backfilledChatSearchRows),
slog.F("duration", i.clk.Since(start)),
)
@@ -338,6 +383,9 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
i.recordsPurged.WithLabelValues("chat_debug_runs").Add(float64(purgedChatDebugRuns))
i.recordsPurged.WithLabelValues("chat_files").Add(float64(purgedChatFiles))
}
if i.chatSearchRowsBackfilled != nil {
i.chatSearchRowsBackfilled.Add(float64(backfilledChatSearchRows))
}
// chatConfigErr is returned after the tx, so do not record this
// iteration as successful when only the deferred config read failed.
@@ -362,13 +410,16 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
}
type instance struct {
cancel context.CancelFunc
closed chan struct{}
logger slog.Logger
vals *codersdk.DeploymentValues
clk quartz.Clock
iterationDuration *prometheus.HistogramVec
recordsPurged *prometheus.CounterVec
cancel context.CancelFunc
closed chan struct{}
logger slog.Logger
vals *codersdk.DeploymentValues
clk quartz.Clock
iterationDuration *prometheus.HistogramVec
recordsPurged *prometheus.CounterVec
chatSearchRowsBackfilled prometheus.Counter
chatSearchBackfillBatchSize int32
chatSearchBackfillMaxBatches int
}
func (i *instance) Close() error {
+375
View File
@@ -14,6 +14,7 @@ import (
"github.com/google/uuid"
"github.com/lib/pq"
"github.com/prometheus/client_golang/prometheus"
"github.com/sqlc-dev/pqtype"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/goleak"
@@ -254,6 +255,7 @@ func TestMetrics(t *testing.T) {
mDB.EXPECT().DeleteOldTelemetryLocks(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
mDB.EXPECT().DeleteOldWorkspaceBuildOrchestrations(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
mDB.EXPECT().DeleteOldAuditLogConnectionEvents(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
mDB.EXPECT().DeleteOldChatDebugRuns(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatDebugRunsParams{})).Return(int64(0), nil).MinTimes(1)
mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).
DoAndReturn(func(f func(database.Store) error, _ *database.TxOptions) error {
@@ -305,6 +307,7 @@ func TestMetrics(t *testing.T) {
mDB.EXPECT().DeleteOldTelemetryLocks(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
mDB.EXPECT().DeleteOldWorkspaceBuildOrchestrations(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
mDB.EXPECT().DeleteOldAuditLogConnectionEvents(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
mDB.EXPECT().DeleteOldChats(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatsParams{})).Return(int64(0), nil).MinTimes(1)
mDB.EXPECT().DeleteOldChatFiles(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatFilesParams{})).Return(int64(0), nil).MinTimes(1)
mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).
@@ -2861,3 +2864,375 @@ func TestDeleteOldChatFiles(t *testing.T) {
})
}
}
func awaitDoTicks(ctx context.Context, t *testing.T, clk *quartz.Mock, n int) func() {
t.Helper()
completed := make(chan struct{})
advance := make(chan struct{})
trapNow := clk.Trap().Now()
trapStop := clk.Trap().TickerStop()
trapReset := clk.Trap().TickerReset()
go func() {
defer close(completed)
defer trapReset.Close()
defer trapStop.Close()
defer trapNow.Close()
trapNow.MustWait(ctx).MustRelease(ctx)
trapReset.MustWait(ctx).MustRelease(ctx)
select {
case completed <- struct{}{}:
case <-ctx.Done():
return
}
for i := 1; i < n; i++ {
select {
case <-advance:
case <-ctx.Done():
return
}
d, w := clk.AdvanceNext()
if !assert.Equal(t, 10*time.Minute, d) {
return
}
w.MustWait(ctx)
trapStop.MustWait(ctx).MustRelease(ctx)
trapReset.MustWait(ctx).MustRelease(ctx)
select {
case completed <- struct{}{}:
case <-ctx.Done():
return
}
}
}()
first := true
return func() {
t.Helper()
if !first {
testutil.RequireSend(ctx, t, advance, struct{}{})
}
first = false
testutil.TryReceive(ctx, t, completed)
}
}
//nolint:paralleltest // It uses LockIDDBPurge.
func TestBackfillChatMessagesSearchTsv(t *testing.T) {
now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC)
type chatSearchDeps struct {
user database.User
modelConfig database.ChatModelConfig
chat database.Chat
}
setupDeps := func(t *testing.T, db database.Store) chatSearchDeps {
t.Helper()
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,
})
_ = dbgen.ChatProvider(t, db, database.ChatProvider{
Provider: "openai",
DisplayName: "OpenAI",
})
modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
Model: "test-model",
ContextLimit: 8192,
})
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: modelConfig.ID,
Title: "search-backfill-test-chat",
})
return chatSearchDeps{user: user, modelConfig: modelConfig, chat: chat}
}
textContent := func(text string) pqtype.NullRawMessage {
return pqtype.NullRawMessage{
RawMessage: json.RawMessage(fmt.Sprintf(`[{"type":"text","text":%q}]`, text)),
Valid: true,
}
}
createMessage := func(t *testing.T, db database.Store, deps chatSearchDeps, role database.ChatMessageRole, visibility database.ChatMessageVisibility, content pqtype.NullRawMessage) database.ChatMessage {
t.Helper()
return dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: deps.chat.ID,
CreatedBy: uuid.NullUUID{UUID: deps.user.ID, Valid: true},
ModelConfigID: uuid.NullUUID{UUID: deps.modelConfig.ID, Valid: true},
Role: role,
Visibility: visibility,
Content: content,
})
}
softDelete := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64) {
t.Helper()
_, err := rawDB.ExecContext(ctx, "UPDATE chat_messages SET deleted = true WHERE id = $1", id)
require.NoError(t, err)
}
// The WHERE clause below must match the predicate of idx_chat_messages_search_tsv_pending.
countPending := func(ctx context.Context, t *testing.T, rawDB *sql.DB) int {
t.Helper()
var count int
err := rawDB.QueryRowContext(ctx, `
SELECT COUNT(*) FROM chat_messages
WHERE search_tsv IS NULL
AND deleted = false
AND visibility IN ('user', 'both')
AND role IN ('user', 'assistant')`).Scan(&count)
require.NoError(t, err)
return count
}
searchTsv := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64) (isNull bool, text string) {
t.Helper()
err := rawDB.QueryRowContext(ctx,
"SELECT search_tsv IS NULL, COALESCE(search_tsv::text, '') FROM chat_messages WHERE id = $1", id).
Scan(&isNull, &text)
require.NoError(t, err)
return isNull, text
}
requireBackfilled := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64, msg string) {
t.Helper()
isNull, _ := searchTsv(ctx, t, rawDB, id)
require.False(t, isNull, msg)
}
// Asserts the row's tsvector matches expectedText, not just non-NULL.
requireTsvFor := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64, expectedText string) {
t.Helper()
var matches bool
err := rawDB.QueryRowContext(ctx,
"SELECT search_tsv = to_tsvector('simple', $2::text) FROM chat_messages WHERE id = $1", id, expectedText).
Scan(&matches)
require.NoError(t, err)
require.True(t, matches, "search_tsv should contain the lexemes of %q", expectedText)
}
requireNotBackfilled := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64, msg string) {
t.Helper()
isNull, _ := searchTsv(ctx, t, rawDB, id)
require.True(t, isNull, msg)
}
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("DrainConverges", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
deps := setupDeps(t, db)
eligibleBoth := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("hello world"))
eligibleUserVis := createMessage(t, db, deps, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, textContent("assistant reply"))
eligibleNoText := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, pqtype.NullRawMessage{RawMessage: json.RawMessage(`[]`), Valid: true})
toolMsg := createMessage(t, db, deps, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, textContent("tool output"))
modelOnlyMsg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, textContent("model only"))
deletedMsg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deleted message"))
softDelete(ctx, t, rawDB, deletedMsg.ID)
tick := awaitDoTicks(ctx, t, clk, 1)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
defer closer.Close()
tick()
require.Zero(t, countPending(ctx, t, rawDB), "queue should be drained")
requireTsvFor(ctx, t, rawDB, eligibleBoth.ID, "hello world")
requireTsvFor(ctx, t, rawDB, eligibleUserVis.ID, "assistant reply")
requireBackfilled(ctx, t, rawDB, eligibleNoText.ID, "eligible message with no text should be backfilled (sentinel)")
requireNotBackfilled(ctx, t, rawDB, toolMsg.ID, "tool message should never be backfilled")
requireNotBackfilled(ctx, t, rawDB, modelOnlyMsg.ID, "model-only message should never be backfilled")
requireNotBackfilled(ctx, t, rawDB, deletedMsg.ID, "deleted message should never be backfilled")
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("BackfillsNewestFirst", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
deps := setupDeps(t, db)
var ids []int64
for i := range 5 {
msg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent(fmt.Sprintf("message %d", i)))
ids = append(ids, msg.ID)
}
tick := awaitDoTicks(ctx, t, clk, 1)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(),
dbpurge.WithClock(clk), dbpurge.WithChatSearchBackfillLimits(2, 1))
defer closer.Close()
tick()
slices.Sort(ids)
requireBackfilled(ctx, t, rawDB, ids[4], "newest message should be backfilled first")
requireBackfilled(ctx, t, rawDB, ids[3], "second-newest message should be backfilled first")
for _, id := range ids[:3] {
requireNotBackfilled(ctx, t, rawDB, id, "older messages should remain pending after one batch")
}
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("NoTextSentinel", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
deps := setupDeps(t, db)
emptyArr := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, pqtype.NullRawMessage{RawMessage: json.RawMessage(`[]`), Valid: true})
noTextParts := createMessage(t, db, deps, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"tool_call","id":"x"}]`), Valid: true})
tick := awaitDoTicks(ctx, t, clk, 1)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
defer closer.Close()
tick()
for _, id := range []int64{emptyArr.ID, noTextParts.ID} {
isNull, text := searchTsv(ctx, t, rawDB, id)
require.False(t, isNull, "no-text row should get the empty-tsvector sentinel, not stay NULL")
require.Empty(t, text, "no-text row should have an empty tsvector")
}
require.Zero(t, countPending(ctx, t, rawDB), "sentinel rows should not reappear as pending")
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("PerTickBound", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
deps := setupDeps(t, db)
for i := range 6 {
createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent(fmt.Sprintf("message %d", i)))
}
tick := awaitDoTicks(ctx, t, clk, 2)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(),
dbpurge.WithClock(clk), dbpurge.WithChatSearchBackfillLimits(2, 2))
defer closer.Close()
tick()
require.Equal(t, 2, countPending(ctx, t, rawDB), "one tick backfills at most maxBatches*batchSize rows")
tick()
require.Zero(t, countPending(ctx, t, rawDB), "next tick continues draining")
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("SkipsDeletedRows", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
deps := setupDeps(t, db)
msg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("soft deleted before backfill"))
softDelete(ctx, t, rawDB, msg.ID)
require.Zero(t, countPending(ctx, t, rawDB), "deleted rows should not appear as pending")
tick := awaitDoTicks(ctx, t, clk, 1)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
defer closer.Close()
tick()
requireNotBackfilled(ctx, t, rawDB, msg.ID, "deleted row should never be backfilled")
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("BackfillsNewMessagesAfterDrain", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
deps := setupDeps(t, db)
initial := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("initial message"))
tick := awaitDoTicks(ctx, t, clk, 2)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
defer closer.Close()
tick()
requireBackfilled(ctx, t, rawDB, initial.ID, "initial message should be backfilled")
require.Zero(t, countPending(ctx, t, rawDB))
fresh := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("post drain message"))
tick()
requireBackfilled(ctx, t, rawDB, fresh.ID, "message inserted after drain should be backfilled on the next tick")
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("SteadyStateNoop", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
_ = setupDeps(t, db)
reg := prometheus.NewRegistry()
tick := awaitDoTicks(ctx, t, clk, 1)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, reg, dbpurge.WithClock(clk))
defer closer.Close()
tick()
require.Zero(t, countPending(ctx, t, rawDB))
backfilled := promhelp.CounterValue(t, reg, "coderd_dbpurge_chat_search_rows_backfilled_total", nil)
require.Zero(t, backfilled, "empty queue should backfill zero rows")
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("MetricsCountsBackfilledRows", func(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
clk := quartz.NewMock(t)
clk.Set(now).MustWait(ctx)
db, _, _ := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
deps := setupDeps(t, db)
reg := prometheus.NewRegistry()
for i := range 3 {
createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent(fmt.Sprintf("message %d", i)))
}
createMessage(t, db, deps, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, textContent("tool output"))
tick := awaitDoTicks(ctx, t, clk, 1)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, reg, dbpurge.WithClock(clk))
defer closer.Close()
tick()
backfilled := promhelp.CounterValue(t, reg, "coderd_dbpurge_chat_search_rows_backfilled_total", nil)
require.Equal(t, 3, backfilled, "counter should count exactly the eligible backfilled rows")
})
//nolint:paralleltest // It uses LockIDDBPurge.
t.Run("SkippedWhenLockHeld", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
defer cancel()
clk := quartz.NewMock(t)
ctrl := gomock.NewController(t)
mDB := dbmock.NewMockStore(ctrl)
mDB.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(0), nil).AnyTimes()
mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays).
Return(int32(0), nil).AnyTimes()
mDB.EXPECT().TryAcquireLock(gomock.Any(), int64(database.LockIDDBPurge)).Return(false, nil).AnyTimes()
mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Times(0)
mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).
DoAndReturn(func(f func(database.Store) error, _ *database.TxOptions) error {
return f(mDB)
}).MinTimes(1)
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
done := awaitDoTick(ctx, t, clk)
closer := dbpurge.New(ctx, logger, mDB, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
defer closer.Close()
testutil.TryReceive(ctx, t, done)
})
}
+37 -3
View File
@@ -783,6 +783,18 @@ BEGIN
END;
$$;
CREATE FUNCTION chat_message_search_text(content jsonb) RETURNS text
LANGUAGE sql IMMUTABLE PARALLEL SAFE
AS $$
SELECT CASE WHEN jsonb_typeof(content) = 'array' THEN (
SELECT string_agg(part->>'text', ' ' ORDER BY ordinality)
FROM jsonb_array_elements(content) WITH ORDINALITY AS t(part, ordinality)
WHERE part->>'type' = 'text'
) END
$$;
COMMENT ON FUNCTION chat_message_search_text(content jsonb) IS 'Extracts searchable content from chat_messages. Returns NULL for scalar JSON strings (content_version=0). Immutable as it is used in indexes.';
CREATE FUNCTION check_workspace_agent_name_unique() RETURNS trigger
LANGUAGE plpgsql
AS $$
@@ -1365,6 +1377,7 @@ CREATE FUNCTION set_chat_message_revision_before() RETURNS trigger
AS $$
DECLARE
chat_snapshot_version bigint;
cmp chat_messages;
BEGIN
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
@@ -1379,7 +1392,9 @@ BEGIN
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
END IF;
IF OLD IS NOT DISTINCT FROM NEW THEN
cmp := NEW;
cmp.search_tsv := OLD.search_tsv;
IF OLD IS NOT DISTINCT FROM cmp THEN
RETURN NEW;
END IF;
END IF;
@@ -1396,6 +1411,8 @@ BEGIN
END;
$$;
COMMENT ON FUNCTION set_chat_message_revision_before() IS 'Component of chatd. Updates chat_snapshot_version when any fields of chat_messages change. Excludes changes to search_tsv as it is not relevant to chatd''s processing loop.';
CREATE FUNCTION sync_chat_retry_state() RETURNS trigger
LANGUAGE plpgsql
AS $$
@@ -1446,7 +1463,7 @@ BEGIN
SELECT DISTINCT n.chat_id
FROM chat_message_history_new_rows n
JOIN chat_message_history_old_rows o ON o.id = n.id
WHERE o IS DISTINCT FROM n
WHERE (to_jsonb(o) - 'search_tsv') IS DISTINCT FROM (to_jsonb(n) - 'search_tsv')
) AS affected
WHERE c.id = affected.chat_id
AND (
@@ -1457,6 +1474,8 @@ BEGIN
END;
$$;
COMMENT ON FUNCTION update_chat_history_after_message_update() IS 'Component of chatd. Updates history_version and generation_attempt on chats when chat_messages is updated. Excludes changes to search_tsv.';
CREATE TABLE ai_gateway_keys (
id uuid NOT NULL,
created_at timestamp with time zone NOT NULL,
@@ -1946,11 +1965,14 @@ CREATE TABLE chat_messages (
provider_response_id text,
api_key_id text,
revision bigint NOT NULL,
reasoning_effort chat_reasoning_effort
reasoning_effort chat_reasoning_effort,
search_tsv tsvector
);
COMMENT ON COLUMN chat_messages.reasoning_effort IS 'Stores the selected effort for the turn triggered by this message.';
COMMENT ON COLUMN chat_messages.search_tsv IS 'Used for full text search. NULL initially, populated async via background job.';
CREATE SEQUENCE chat_messages_id_seq
START WITH 1
INCREMENT BY 1
@@ -4717,6 +4739,8 @@ CREATE UNIQUE INDEX idx_chat_debug_steps_run_step ON chat_debug_steps USING btre
CREATE INDEX idx_chat_debug_steps_stale ON chat_debug_steps USING btree (updated_at) WHERE (finished_at IS NULL);
CREATE INDEX idx_chat_diff_statuses_pr_title_fts ON chat_diff_statuses USING gin (to_tsvector('simple'::regconfig, pull_request_title));
CREATE INDEX idx_chat_diff_statuses_stale_at ON chat_diff_statuses USING btree (stale_at);
CREATE INDEX idx_chat_diff_statuses_url_lower ON chat_diff_statuses USING btree (lower(url)) WHERE ((url IS NOT NULL) AND (url <> ''::text));
@@ -4737,6 +4761,12 @@ CREATE INDEX idx_chat_messages_created_at ON chat_messages USING btree (created_
CREATE INDEX idx_chat_messages_owner_spend ON chat_messages USING btree (chat_id, created_at) WHERE (total_cost_micros IS NOT NULL);
CREATE INDEX idx_chat_messages_search_tsv ON chat_messages USING gin (search_tsv) WHERE ((search_tsv IS NOT NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
COMMENT ON INDEX idx_chat_messages_search_tsv IS 'Partial index over chat_messages used for populating search_tsv in the background. Only defined over ''searchable'' rows of chat_messages where search_tsv is NULL.';
CREATE INDEX idx_chat_messages_search_tsv_pending ON chat_messages USING btree (id DESC) WHERE ((search_tsv IS NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
CREATE INDEX idx_chat_messages_user_prompts ON chat_messages USING btree (chat_id, id DESC) WHERE ((deleted = false) AND (role = 'user'::chat_message_role) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])));
CREATE INDEX idx_chat_model_configs_ai_provider_id ON chat_model_configs USING btree (ai_provider_id);
@@ -4763,6 +4793,10 @@ CREATE INDEX idx_chats_parent_chat_id ON chats USING btree (parent_chat_id);
CREATE INDEX idx_chats_root_chat_id ON chats USING btree (root_chat_id);
CREATE INDEX idx_chats_title_fts ON chats USING gin (to_tsvector('simple'::regconfig, title));
COMMENT ON INDEX idx_chats_title_fts IS 'Used for full text search. Defined over all rows of the chats table.';
CREATE INDEX idx_chats_worker_acquisition_candidates ON chats USING btree (status, updated_at, id) WHERE (archived = false);
CREATE INDEX idx_chats_workspace ON chats USING btree (workspace_id);
@@ -0,0 +1,68 @@
-- Restore the original trigger bodies from 000519.
CREATE OR REPLACE FUNCTION set_chat_message_revision_before()
RETURNS trigger AS $$
DECLARE
chat_snapshot_version bigint;
BEGIN
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
END IF;
IF TG_OP = 'UPDATE' THEN
IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN
RAISE EXCEPTION 'chat_messages.chat_id is immutable';
END IF;
IF OLD.revision IS DISTINCT FROM NEW.revision THEN
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
END IF;
IF OLD IS NOT DISTINCT FROM NEW THEN
RETURN NEW;
END IF;
END IF;
SELECT snapshot_version INTO chat_snapshot_version
FROM chats WHERE id = NEW.chat_id;
IF chat_snapshot_version IS NULL THEN
RAISE EXCEPTION 'chat % does not exist', NEW.chat_id;
END IF;
NEW.revision = chat_snapshot_version;
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
CREATE OR REPLACE FUNCTION update_chat_history_after_message_update()
RETURNS trigger AS $$
BEGIN
UPDATE chats c
SET history_version = c.snapshot_version,
generation_attempt = 0
FROM (
SELECT DISTINCT n.chat_id
FROM chat_message_history_new_rows n
JOIN chat_message_history_old_rows o ON o.id = n.id
WHERE o IS DISTINCT FROM n
) AS affected
WHERE c.id = affected.chat_id
AND (
c.history_version IS DISTINCT FROM c.snapshot_version
OR c.generation_attempt <> 0
);
RETURN NULL;
END;
$$ LANGUAGE plpgsql;
DROP INDEX IF EXISTS idx_chat_diff_statuses_pr_title_fts;
DROP INDEX IF EXISTS idx_chats_title_fts;
DROP INDEX IF EXISTS idx_chat_messages_search_tsv_pending;
DROP INDEX IF EXISTS idx_chat_messages_search_tsv;
ALTER TABLE chat_messages DROP COLUMN IF EXISTS search_tsv;
DROP FUNCTION IF EXISTS chat_message_search_text(jsonb);
@@ -0,0 +1,97 @@
CREATE FUNCTION chat_message_search_text(content jsonb) RETURNS text
LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
SELECT CASE WHEN jsonb_typeof(content) = 'array' THEN (
SELECT string_agg(part->>'text', ' ' ORDER BY ordinality)
FROM jsonb_array_elements(content) WITH ORDINALITY AS t(part, ordinality)
WHERE part->>'type' = 'text'
) END
$$;
COMMENT ON FUNCTION chat_message_search_text IS 'Extracts searchable content from chat_messages. Returns NULL for scalar JSON strings (content_version=0). Immutable as it is used in indexes.';
-- Populated by a background sweep, not at insert time. NULL means pending.
ALTER TABLE chat_messages ADD COLUMN search_tsv tsvector;
COMMENT ON COLUMN chat_messages.search_tsv IS 'Used for full text search. NULL initially, populated async via background job.';
CREATE INDEX idx_chat_messages_search_tsv ON chat_messages
USING GIN (search_tsv)
WHERE ((search_tsv IS NOT NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
COMMENT ON INDEX idx_chat_messages_search_tsv IS 'Partial index over chat_messages used for full text search. Only defined over ''searchable'' rows of chat_messages.';
CREATE INDEX idx_chat_messages_search_tsv_pending ON chat_messages USING btree (id DESC)
WHERE ((search_tsv IS NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
COMMENT ON INDEX idx_chat_messages_search_tsv IS 'Partial index over chat_messages used for populating search_tsv in the background. Only defined over ''searchable'' rows of chat_messages where search_tsv is NULL.';
CREATE INDEX idx_chats_title_fts ON chats USING GIN (to_tsvector('simple', title));
COMMENT ON index idx_chats_title_fts IS 'Used for full text search. Defined over all rows of the chats table.';
CREATE INDEX idx_chat_diff_statuses_pr_title_fts ON chat_diff_statuses USING GIN (to_tsvector('simple', pull_request_title));
COMMENT ON index idx_chats_title_fts IS 'Used for full text search. Defined over all rows of the chats table.';
CREATE OR REPLACE FUNCTION set_chat_message_revision_before()
RETURNS trigger AS $$
DECLARE
chat_snapshot_version bigint;
cmp chat_messages;
BEGIN
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
END IF;
IF TG_OP = 'UPDATE' THEN
IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN
RAISE EXCEPTION 'chat_messages.chat_id is immutable';
END IF;
IF OLD.revision IS DISTINCT FROM NEW.revision THEN
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
END IF;
cmp := NEW;
cmp.search_tsv := OLD.search_tsv;
IF OLD IS NOT DISTINCT FROM cmp THEN
RETURN NEW;
END IF;
END IF;
SELECT snapshot_version INTO chat_snapshot_version
FROM chats WHERE id = NEW.chat_id;
IF chat_snapshot_version IS NULL THEN
RAISE EXCEPTION 'chat % does not exist', NEW.chat_id;
END IF;
NEW.revision = chat_snapshot_version;
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
COMMENT ON FUNCTION set_chat_message_revision_before IS 'Component of chatd. Updates chat_snapshot_version when any fields of chat_messages change. Excludes changes to search_tsv as it is not relevant to chatd''s processing loop.';
CREATE OR REPLACE FUNCTION update_chat_history_after_message_update()
RETURNS trigger AS $$
BEGIN
UPDATE chats c
SET history_version = c.snapshot_version,
generation_attempt = 0
FROM (
SELECT DISTINCT n.chat_id
FROM chat_message_history_new_rows n
JOIN chat_message_history_old_rows o ON o.id = n.id
WHERE (to_jsonb(o) - 'search_tsv') IS DISTINCT FROM (to_jsonb(n) - 'search_tsv')
) AS affected
WHERE c.id = affected.chat_id
AND (
c.history_version IS DISTINCT FROM c.snapshot_version
OR c.generation_attempt <> 0
);
RETURN NULL;
END;
$$ LANGUAGE plpgsql;
COMMENT ON FUNCTION update_chat_history_after_message_update IS 'Component of chatd. Updates history_version and generation_attempt on chats when chat_messages is updated. Excludes changes to search_tsv.';
+243
View File
@@ -19,11 +19,13 @@ import (
"github.com/golang-migrate/migrate/v4/source/stub"
"github.com/google/uuid"
"github.com/lib/pq"
"github.com/sqlc-dev/pqtype"
"github.com/stretchr/testify/require"
"go.uber.org/goleak"
"golang.org/x/sync/errgroup"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/migrations"
"github.com/coder/coder/v2/testutil"
@@ -1866,3 +1868,244 @@ func TestMigration000498SoftDeleteStaleWorkspaceAgents(t *testing.T) {
// TestSoftDeleteWorkspaceAgentsByWorkspaceID, plus integration tests
// under coderd/coderd_test.go; not retested here.
}
func TestMigration000543ChatMessageSearchText(t *testing.T) {
t.Parallel()
if testing.Short() {
t.SkipNow()
}
sqlDB := testSQLDB(t)
require.NoError(t, migrations.Up(sqlDB))
cases := []struct {
name string
content sql.NullString
want sql.NullString
}{
{
name: "SingleTextPart",
content: sql.NullString{String: `[{"type":"text","text":"hello world"}]`, Valid: true},
want: sql.NullString{String: "hello world", Valid: true},
},
{
name: "TextInterleavedWithNonText",
content: sql.NullString{String: `[
{"type":"text","text":"first"},
{"type":"reasoning","text":"thinking"},
{"type":"tool-call","toolName":"execute"},
{"type":"text","text":"second"}
]`, Valid: true},
want: sql.NullString{String: "first second", Valid: true},
},
{
name: "OnlyNonTextParts",
content: sql.NullString{String: `[{"type":"reasoning","text":"thinking"}]`, Valid: true},
want: sql.NullString{},
},
{
name: "ScalarContent",
content: sql.NullString{String: `"hello"`, Valid: true},
want: sql.NullString{},
},
{
name: "EmptyArray",
content: sql.NullString{String: `[]`, Valid: true},
want: sql.NullString{},
},
{
name: "NullInput",
content: sql.NullString{},
want: sql.NullString{},
},
{
name: "ElementsMissingTypeOrText",
content: sql.NullString{String: `[{"text":"no type"},{"type":"text"},{"type":"text","text":"kept"}]`, Valid: true},
want: sql.NullString{String: "kept", Valid: true},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitMedium)
var got sql.NullString
err := sqlDB.QueryRowContext(ctx,
`SELECT chat_message_search_text($1::jsonb)`, tc.content,
).Scan(&got)
require.NoError(t, err)
require.Equal(t, tc.want, got)
})
}
}
// Shared eligibility predicate of the two partial chat_messages search
// indexes. Queries must repeat it verbatim.
const eligibilityPredicate = `deleted = false
AND visibility IN ('user', 'both')
AND role IN ('user', 'assistant')`
func TestMigration000543ChatSearchSchemaIndexes(t *testing.T) {
t.Parallel()
if testing.Short() {
t.SkipNow()
}
sqlDB := testSQLDB(t)
require.NoError(t, migrations.Up(sqlDB))
cases := []struct {
name string
table string
partial bool
}{
{name: "idx_chat_messages_search_tsv", table: "chat_messages", partial: true},
{name: "idx_chat_messages_search_tsv_pending", table: "chat_messages", partial: true},
{name: "idx_chats_title_fts", table: "chats", partial: false},
{name: "idx_chat_diff_statuses_pr_title_fts", table: "chat_diff_statuses", partial: false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitMedium)
var table string
var partial bool
err := sqlDB.QueryRowContext(ctx, `
SELECT i.tablename, x.indpred IS NOT NULL
FROM pg_indexes i
JOIN pg_class c ON c.relname = i.indexname
JOIN pg_index x ON x.indexrelid = c.oid
WHERE i.indexname = $1`, tc.name,
).Scan(&table, &partial)
require.NoError(t, err, "index %s should exist", tc.name)
require.Equal(t, tc.table, table, "index %s table", tc.name)
require.Equal(t, tc.partial, partial, "index %s partial", tc.name)
})
}
}
func TestMigration000543ChatSearchSchemaBehavior(t *testing.T) {
t.Parallel()
if testing.Short() {
t.SkipNow()
}
sqlDB := testSQLDB(t)
require.NoError(t, migrations.Up(sqlDB))
db := database.New(sqlDB)
ctx := testutil.Context(t, testutil.WaitLong)
org := dbgen.Organization(t, db, database.Organization{})
owner := dbgen.User(t, db, database.User{})
_ = dbgen.ChatProvider(t, db, database.ChatProvider{Provider: "openai", DisplayName: "OpenAI"})
modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
IsDefault: true,
})
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: owner.ID,
LastModelConfigID: modelCfg.ID,
})
newMsg := func(role database.ChatMessageRole, visibility database.ChatMessageVisibility, content string) database.ChatMessage {
seed := database.ChatMessage{
ChatID: chat.ID,
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
Role: role,
Visibility: visibility,
}
if content != "" {
seed.Content = pqtype.NullRawMessage{RawMessage: []byte(content), Valid: true}
}
return dbgen.ChatMessage(t, db, seed)
}
textContent := func(text string) string {
return `[{"type":"text","text":"` + text + `"}]`
}
pendingIDs := func(ctx context.Context, limit int) []int64 {
rows, err := sqlDB.QueryContext(ctx, `
SELECT id FROM chat_messages
WHERE search_tsv IS NULL AND `+eligibilityPredicate+`
ORDER BY id DESC
LIMIT $1`, limit)
require.NoError(t, err)
defer rows.Close()
var ids []int64
for rows.Next() {
var id int64
require.NoError(t, rows.Scan(&id))
ids = append(ids, id)
}
require.NoError(t, rows.Err())
return ids
}
// Insert regression: RETURNING * must survive the new column, and new
// rows must start with search_tsv NULL so they enter the pending queue.
eligibleText := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deploy the search feature"))
var tsvIsNull bool
err := sqlDB.QueryRowContext(ctx,
`SELECT search_tsv IS NULL FROM chat_messages WHERE id = $1`, eligibleText.ID,
).Scan(&tsvIsNull)
require.NoError(t, err)
require.True(t, tsvIsNull, "new rows must have search_tsv NULL")
eligibleNoText := newMsg(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, `[{"type":"reasoning","text":"thinking"}]`)
toolMsg := newMsg(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, textContent("tool output about deploy"))
modelOnly := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, textContent("model-only deploy note"))
deletedMsg := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deleted deploy message"))
_, err = sqlDB.ExecContext(ctx, `UPDATE chat_messages SET deleted = true WHERE id = $1`, deletedMsg.ID)
require.NoError(t, err)
// Only eligible rows appear in the queue, newest first. The tool-role,
// model-only, and soft-deleted rows are excluded even though their
// search_tsv is NULL.
require.Equal(t, []int64{eligibleNoText.ID, eligibleText.ID}, pendingIDs(ctx, 10))
// Sweep-style UPDATE. The '' sentinel (not NULL) marks no-text rows as
// swept; NULL means pending, so COALESCE is what drains them from the
// queue.
_, err = sqlDB.ExecContext(ctx, `
UPDATE chat_messages
SET search_tsv = COALESCE(to_tsvector('simple', chat_message_search_text(content)), ''::tsvector)
WHERE id = ANY($1)`, pq.Array([]int64{eligibleText.ID, eligibleNoText.ID}))
require.NoError(t, err)
require.Empty(t, pendingIDs(ctx, 10), "swept rows must leave the queue, including no-text rows")
// Soft-deleting an unswept row removes it from the queue without a sweep.
unswept := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("unswept deploy row"))
require.Equal(t, []int64{unswept.ID}, pendingIDs(ctx, 10))
_, err = sqlDB.ExecContext(ctx, `UPDATE chat_messages SET deleted = true WHERE id = $1`, unswept.ID)
require.NoError(t, err)
require.Empty(t, pendingIDs(ctx, 10))
// Search contract: populate search_tsv on every row (including
// ineligible ones) and assert the search-index predicate filters them.
_, err = sqlDB.ExecContext(ctx, `
UPDATE chat_messages
SET search_tsv = COALESCE(to_tsvector('simple', chat_message_search_text(content)), ''::tsvector)
WHERE chat_id = $1`, chat.ID)
require.NoError(t, err)
rows, err := sqlDB.QueryContext(ctx, `
SELECT id FROM chat_messages
WHERE search_tsv @@ websearch_to_tsquery('simple', $1)
AND search_tsv IS NOT NULL
AND `+eligibilityPredicate+`
ORDER BY id`, "deploy")
require.NoError(t, err)
defer rows.Close()
var matched []int64
for rows.Next() {
var id int64
require.NoError(t, rows.Scan(&id))
matched = append(matched, id)
}
require.NoError(t, rows.Err())
require.Equal(t, []int64{eligibleText.ID}, matched,
"search must exclude deleted, model-only, and tool-role rows (%d %d %d)",
toolMsg.ID, modelOnly.ID, deletedMsg.ID)
}
+1
View File
@@ -786,6 +786,7 @@ func (q *sqlQuerier) GetAuthorizedChats(ctx context.Context, arg GetChatsParams,
arg.PrNumber,
arg.RepoQuery,
arg.PrTitleQuery,
arg.Search,
arg.OffsetOpt,
arg.LimitOpt,
)
+2
View File
@@ -5123,6 +5123,8 @@ type ChatMessage struct {
Revision int64 `db:"revision" json:"revision"`
// Stores the selected effort for the turn triggered by this message.
ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"`
// Used for full text search. NULL initially, populated async via background job.
SearchTsv interface{} `db:"search_tsv" json:"search_tsv"`
}
type ChatModelConfig struct {
+9
View File
@@ -67,6 +67,12 @@ type sqlcQuerier interface {
// created_at ASC flows through to dbpurge's digest truncation; see
// buildDigestData in dbpurge.go for the tradeoff rationale.
AutoArchiveInactiveChats(ctx context.Context, arg AutoArchiveInactiveChatsParams) ([]AutoArchiveInactiveChatsRow, error)
// Backfills chat_messages.search_tsv for pending rows, newest first.
// The WHERE clause must match the predicate of
// idx_chat_messages_search_tsv_pending exactly so the partial index
// serves this query.
// NULL means "pending", empty tsvector means "backfilled, no text".
BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error)
BackoffChatDiffStatus(ctx context.Context, arg BackoffChatDiffStatusParams) error
// Deletes heartbeat rows for the supplied (chat_id, runner_id) pairs.
BatchDeleteChatHeartbeats(ctx context.Context, arg BatchDeleteChatHeartbeatsParams) (int64, error)
@@ -80,6 +86,9 @@ type sqlcQuerier interface {
// Calculates the telemetry summary for a given provider, model, and client
// combination for telemetry reporting.
CalculateAIBridgeInterceptionsTelemetrySummary(ctx context.Context, arg CalculateAIBridgeInterceptionsTelemetrySummaryParams) (CalculateAIBridgeInterceptionsTelemetrySummaryRow, error)
// Reports whether search text tokenizes to an empty tsquery (e.g. '!!!').
// Used to reject input that would silently match nothing.
ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error)
ClaimPrebuiltWorkspace(ctx context.Context, arg ClaimPrebuiltWorkspaceParams) (ClaimPrebuiltWorkspaceRow, error)
CleanTailnetCoordinators(ctx context.Context) error
CleanTailnetLostPeers(ctx context.Context) error
+232
View File
@@ -10,6 +10,7 @@ import (
"net"
"slices"
"sort"
"strconv"
"strings"
"testing"
"time"
@@ -15344,6 +15345,237 @@ func TestGetChatsFilter(t *testing.T) {
}
}
func TestGetChatsSearch(t *testing.T) {
t.Parallel()
store, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
ctx := context.Background()
org := dbgen.Organization(t, store, database.Organization{})
user := dbgen.User(t, store, database.User{})
dbgen.OrganizationMember(t, store, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
provider := dbgen.AIProviderWithOptionalKey(t, store, database.AIProvider{
Type: database.AIProviderTypeOpenai,
}, "test-key")
modelCfg, err := store.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true},
Model: "test-model-" + uuid.NewString(),
DisplayName: "Test Model",
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
Enabled: true,
IsDefault: true,
ContextLimit: 128000,
CompressionThreshold: 80,
Options: json.RawMessage(`{}`),
})
require.NoError(t, err)
createRoot := func(title string) database.Chat {
t.Helper()
chat, err := store.InsertChat(ctx, database.InsertChatParams{
OrganizationID: org.ID,
Status: database.ChatStatusWaiting,
ClientType: database.ChatClientTypeUi,
OwnerID: user.ID,
LastModelConfigID: modelCfg.ID,
Title: title,
})
require.NoError(t, err)
return chat
}
createChild := func(root database.Chat, title string) database.Chat {
t.Helper()
chat, err := store.InsertChat(ctx, database.InsertChatParams{
OrganizationID: org.ID,
Status: database.ChatStatusWaiting,
ClientType: database.ChatClientTypeUi,
OwnerID: user.ID,
LastModelConfigID: modelCfg.ID,
Title: title,
ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true},
RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true},
})
require.NoError(t, err)
return chat
}
insertMsg := func(chatID uuid.UUID, role database.ChatMessageRole, visibility database.ChatMessageVisibility, text string) database.ChatMessage {
t.Helper()
msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{
ChatID: chatID,
CreatedBy: []uuid.UUID{user.ID},
ModelConfigID: []uuid.UUID{modelCfg.ID},
Role: []database.ChatMessageRole{role},
Content: []string{`[{"type":"text","text":` + strconv.Quote(text) + `}]`},
ContentVersion: []int16{1},
Visibility: []database.ChatMessageVisibility{visibility},
InputTokens: []int64{0},
OutputTokens: []int64{0},
TotalTokens: []int64{0},
ReasoningTokens: []int64{0},
CacheCreationTokens: []int64{0},
CacheReadTokens: []int64{0},
ContextLimit: []int64{0},
Compressed: []bool{false},
TotalCostMicros: []int64{0},
RuntimeMs: []int64{0},
})
require.NoError(t, err)
require.Len(t, msgs, 1)
return msgs[0]
}
linkPR := func(chatID uuid.UUID, url, state, prTitle string, prNumber int32, gitRemoteOrigin string) {
t.Helper()
now := time.Now()
_, err := store.UpsertChatDiffStatusReference(ctx, database.UpsertChatDiffStatusReferenceParams{
ChatID: chatID,
Url: sql.NullString{String: url, Valid: true},
GitBranch: "main",
GitRemoteOrigin: gitRemoteOrigin,
StaleAt: now.Add(time.Hour),
})
require.NoError(t, err)
_, err = store.UpsertChatDiffStatus(ctx, database.UpsertChatDiffStatusParams{
ChatID: chatID,
Url: sql.NullString{String: url, Valid: true},
PullRequestState: sql.NullString{String: state, Valid: true},
PullRequestTitle: prTitle,
PrNumber: sql.NullInt32{Int32: prNumber, Valid: prNumber > 0},
Additions: 1,
Deletions: 1,
ChangedFiles: 1,
RefreshedAt: now,
StaleAt: now.Add(time.Hour),
})
require.NoError(t, err)
}
titleChat := createRoot("deploy pipeline alpha")
archivedChat := createRoot("deploy pipeline beta")
prTitleChat := createRoot("widget work")
linkPR(prTitleChat.ID, "https://github.com/acme/widget/pull/42", "open", "Fix authentication bug", 42, "https://github.com/acme/widget.git")
mergedChat := createRoot("other work")
linkPR(mergedChat.ID, "https://github.com/acme/other-repo/pull/7", "merged", "Fix authentication flow", 7, "https://github.com/acme/other-repo.git")
msgChat := createRoot("plain one")
insertMsg(msgChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, "kubernetes cluster restart")
assistantMsgChat := createRoot("plain assistant")
insertMsg(assistantMsgChat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, "grafana dashboard tuning")
userVisMsgChat := createRoot("plain uservis")
insertMsg(userVisMsgChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityUser, "vault token rotation")
assistantUserVisMsgChat := createRoot("plain assistant uservis")
insertMsg(assistantUserVisMsgChat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, "redis eviction policy")
deletedMsgChat := createRoot("plain two")
deletedMsg := insertMsg(deletedMsgChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, "terraform apply failure")
childParent := createRoot("plain parent")
childChat := createChild(childParent, "plain child")
insertMsg(childChat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, "orchestrator saga")
ineligibleChat := createRoot("plain three")
toolMsg := insertMsg(ineligibleChat.ID, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, "forbidden secret token")
modelOnlyMsg := insertMsg(ineligibleChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, "forbidden secret token")
// Ineligible rows keep search_tsv NULL after backfill.
_, err = store.BackfillChatMessagesSearchTsv(ctx, 1000)
require.NoError(t, err)
// Soft-deleted rows stay excluded even though search_tsv remains
// populated.
err = store.SoftDeleteChatMessageByID(ctx, deletedMsg.ID)
require.NoError(t, err)
// Inserted after backfill: search_tsv IS NULL, must match nothing.
pendingChat := createRoot("plain four")
insertMsg(pendingChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, "elasticsearch indexing")
// Prove role/visibility predicates exclude rows even when search_tsv
// is set.
_, err = sqlDB.ExecContext(ctx,
`UPDATE chat_messages SET search_tsv = to_tsvector('simple', 'forbidden secret token') WHERE id = ANY($1)`,
pq.Array([]int64{toolMsg.ID, modelOnlyMsg.ID}))
require.NoError(t, err)
_, err = store.ArchiveChatByID(ctx, archivedChat.ID)
require.NoError(t, err)
allRootIDs := []uuid.UUID{
titleChat.ID, archivedChat.ID, prTitleChat.ID, mergedChat.ID,
msgChat.ID, assistantMsgChat.ID, userVisMsgChat.ID,
assistantUserVisMsgChat.ID, deletedMsgChat.ID, childParent.ID,
ineligibleChat.ID, pendingChat.ID,
}
tests := []struct {
name string
params database.GetChatsParams
want []uuid.UUID
}{
{"Title/Match", database.GetChatsParams{Search: "pipeline alpha"}, []uuid.UUID{titleChat.ID}},
{"Title/CaseInsensitiveMultiWord", database.GetChatsParams{Search: "ALPHA DEPLOY"}, []uuid.UUID{titleChat.ID}},
{"Title/AndSemantics", database.GetChatsParams{Search: "deploy nonexistent"}, nil},
{"PRTitle/Match", database.GetChatsParams{Search: "authentication"}, []uuid.UUID{prTitleChat.ID, mergedChat.ID}},
{"Message/Match", database.GetChatsParams{Search: "kubernetes restart"}, []uuid.UUID{msgChat.ID}},
{"Message/AssistantRoleMatch", database.GetChatsParams{Search: "grafana tuning"}, []uuid.UUID{assistantMsgChat.ID}},
{"Message/UserVisibilityMatch", database.GetChatsParams{Search: "vault rotation"}, []uuid.UUID{userVisMsgChat.ID}},
{"Message/AssistantUserVisibilityMatch", database.GetChatsParams{Search: "redis eviction"}, []uuid.UUID{assistantUserVisMsgChat.ID}},
{"PRNumber/Match", database.GetChatsParams{Search: "42"}, []uuid.UUID{prTitleChat.ID}},
{"PRNumber/NonNumericNoMatch", database.GetChatsParams{Search: "42abc"}, nil},
{"PRNumber/OversizedDigitsNoError", database.GetChatsParams{Search: "1111111111111111111111111"}, nil},
{"NoMatch", database.GetChatsParams{Search: "zzzqqq"}, nil},
{"Message/PendingBackfillNoMatch", database.GetChatsParams{Search: "elasticsearch"}, nil},
{"Message/DeletedNoMatch", database.GetChatsParams{Search: "terraform"}, nil},
// Parent also excluded: EXISTS is per-chat, not per-tree.
{"Message/ChildNotSurfaced", database.GetChatsParams{Search: "orchestrator saga"}, nil},
{"Message/IneligibleMessagesNoMatch", database.GetChatsParams{Search: "forbidden secret"}, nil},
{"Composed/ArchivedDefaultIncludesAll", database.GetChatsParams{Search: "deploy pipeline"}, []uuid.UUID{titleChat.ID, archivedChat.ID}},
{"Composed/ArchivedFalseExcludes", database.GetChatsParams{Search: "deploy pipeline", Archived: sql.NullBool{Bool: false, Valid: true}}, []uuid.UUID{titleChat.ID}},
{"Composed/ArchivedTrueOnly", database.GetChatsParams{Search: "deploy pipeline", Archived: sql.NullBool{Bool: true, Valid: true}}, []uuid.UUID{archivedChat.ID}},
{"Composed/SearchAndRepo", database.GetChatsParams{Search: "authentication", RepoQuery: "acme/widget"}, []uuid.UUID{prTitleChat.ID}},
{"Composed/SearchAndPRStatus", database.GetChatsParams{Search: "authentication", PullRequestStatuses: []string{"merged"}}, []uuid.UUID{mergedChat.ID}},
{"EmptySearch/ReturnsAll", database.GetChatsParams{Search: ""}, allRootIDs},
{"WhitespaceSearch/ReturnsNothing", database.GetChatsParams{Search: " "}, nil},
{"TabOnlySearch/ReturnsNothing", database.GetChatsParams{Search: "\t\t"}, nil},
{"EmptySearch/TitleQueryStillWorks", database.GetChatsParams{Search: "", TitleQuery: "pipeline alpha"}, []uuid.UUID{titleChat.ID}},
{"EmptySearch/PRTitleQueryStillWorks", database.GetChatsParams{Search: "", PrTitleQuery: "authentication bug"}, []uuid.UUID{prTitleChat.ID}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
params := tt.params
params.OwnedOnly = true
params.ViewerID = user.ID
rows, err := store.GetChats(ctx, params)
require.NoError(t, err)
got := make([]uuid.UUID, 0, len(rows))
for _, row := range rows {
got = append(got, row.Chat.ID)
}
if tt.want == nil {
require.Empty(t, got)
} else {
require.ElementsMatch(t, tt.want, got)
}
})
}
}
func TestChatHasUnread(t *testing.T) {
t.Parallel()
+103 -10
View File
@@ -6032,6 +6032,36 @@ func (q *sqlQuerier) AutoArchiveInactiveChats(ctx context.Context, arg AutoArchi
return items, nil
}
const backfillChatMessagesSearchTsv = `-- name: BackfillChatMessagesSearchTsv :execrows
WITH batch AS (
SELECT id FROM chat_messages
WHERE search_tsv IS NULL
AND deleted = false
AND visibility IN ('user', 'both')
AND role IN ('user', 'assistant')
ORDER BY id DESC
LIMIT $1::int
)
UPDATE chat_messages cm
SET search_tsv = COALESCE(
to_tsvector('simple', chat_message_search_text(cm.content)),
''::tsvector)
FROM batch WHERE cm.id = batch.id
`
// Backfills chat_messages.search_tsv for pending rows, newest first.
// The WHERE clause must match the predicate of
// idx_chat_messages_search_tsv_pending exactly so the partial index
// serves this query.
// NULL means "pending", empty tsvector means "backfilled, no text".
func (q *sqlQuerier) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
result, err := q.db.ExecContext(ctx, backfillChatMessagesSearchTsv, batchSize)
if err != nil {
return 0, err
}
return result.RowsAffected()
}
const backoffChatDiffStatus = `-- name: BackoffChatDiffStatus :exec
UPDATE
chat_diff_statuses
@@ -6096,6 +6126,19 @@ func (q *sqlQuerier) BatchUpsertChatHeartbeats(ctx context.Context, arg BatchUps
return err
}
const chatSearchQueryIsEmpty = `-- name: ChatSearchQueryIsEmpty :one
SELECT numnode(websearch_to_tsquery('simple', $1::text)) = 0 AS is_empty
`
// Reports whether search text tokenizes to an empty tsquery (e.g. '!!!').
// Used to reject input that would silently match nothing.
func (q *sqlQuerier) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
row := q.db.QueryRowContext(ctx, chatSearchQueryIsEmpty, search)
var is_empty bool
err := row.Scan(&is_empty)
return is_empty, err
}
const countChatQueuedMessages = `-- name: CountChatQueuedMessages :one
SELECT COUNT(*)::bigint AS count
FROM chat_queued_messages
@@ -7400,7 +7443,7 @@ func (q *sqlQuerier) GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatP
const getChatMessageByID = `-- name: GetChatMessageByID :one
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, 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, api_key_id, revision, reasoning_effort, search_tsv
FROM
chat_messages
WHERE
@@ -7436,6 +7479,7 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
)
return i, err
}
@@ -7525,7 +7569,7 @@ func (q *sqlQuerier) GetChatMessageSummariesPerChat(ctx context.Context, created
const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, 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, api_key_id, revision, reasoning_effort, search_tsv
FROM
chat_messages
WHERE
@@ -7576,6 +7620,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
); err != nil {
return nil, err
}
@@ -7592,7 +7637,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
const getChatMessagesByChatIDAscPaginated = `-- name: GetChatMessagesByChatIDAscPaginated :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, api_key_id, revision, reasoning_effort, search_tsv
FROM
chat_messages
WHERE
@@ -7646,6 +7691,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
); err != nil {
return nil, err
}
@@ -7662,7 +7708,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
const getChatMessagesByChatIDDescPaginated = `-- name: GetChatMessagesByChatIDDescPaginated :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, api_key_id, revision, reasoning_effort, search_tsv
FROM
chat_messages
WHERE
@@ -7729,6 +7775,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
); err != nil {
return nil, err
}
@@ -7745,7 +7792,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
const getChatMessagesByRevisionForStream = `-- name: GetChatMessagesByRevisionForStream :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, api_key_id, revision, reasoning_effort, search_tsv
FROM
chat_messages
WHERE
@@ -7795,6 +7842,7 @@ func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
); err != nil {
return nil, err
}
@@ -7827,7 +7875,7 @@ WITH latest_compressed_summary AS (
1
)
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, 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, api_key_id, revision, reasoning_effort, search_tsv
FROM
chat_messages
WHERE
@@ -7902,6 +7950,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
); err != nil {
return nil, err
}
@@ -8588,6 +8637,46 @@ WHERE
)
ELSE true
END
-- websearch_to_tsquery accepts quoted phrases, OR, and -negation;
-- the 'simple' config folds case and skips stemming.
AND CASE
WHEN $16::text != '' THEN (
-- Served by idx_chats_title_fts.
to_tsvector('simple', chats_expanded.title) @@ websearch_to_tsquery('simple', $16)
-- Served by idx_chat_diff_statuses_pr_title_fts.
OR EXISTS (
SELECT 1
FROM chat_diff_statuses cds
WHERE cds.chat_id = chats_expanded.id
AND to_tsvector('simple', cds.pull_request_title) @@ websearch_to_tsquery('simple', $16)
)
-- The WHERE clause must repeat the predicate of the partial index
-- idx_chat_messages_search_tsv so the planner can use it. Additional
-- filters should still be fine.
OR EXISTS (
SELECT 1
FROM chat_messages cm
WHERE cm.chat_id = chats_expanded.id
AND cm.search_tsv IS NOT NULL
AND cm.deleted = false
AND cm.visibility IN ('user', 'both')
AND cm.role IN ('user', 'assistant')
AND cm.search_tsv @@ websearch_to_tsquery('simple', $16)
)
-- Skip an explicit pr_number lookup unless the search is a valid bigint.
OR CASE
WHEN $16 ~ '^[0-9]{1,18}$' THEN EXISTS (
SELECT 1
FROM chat_diff_statuses cds
WHERE cds.chat_id = chats_expanded.id
AND cds.pr_number IS NOT NULL
AND cds.pr_number = $16::bigint
)
ELSE false
END
)
ELSE true
END
-- Paginate over root chats only. Children are fetched
-- separately via GetChildChatsByParentIDs and embedded under
-- each parent. Other callers that need the full set should
@@ -8604,11 +8693,11 @@ ORDER BY
-chats_expanded.pin_order DESC,
chats_expanded.updated_at DESC,
chats_expanded.id DESC
OFFSET $16
OFFSET $17
LIMIT
-- The chat list is unbounded and expected to grow large.
-- Default to 50 to prevent accidental excessively large queries.
COALESCE(NULLIF($17 :: int, 0), 50)
COALESCE(NULLIF($18 :: int, 0), 50)
`
type GetChatsParams struct {
@@ -8627,6 +8716,7 @@ type GetChatsParams struct {
PrNumber int32 `db:"pr_number" json:"pr_number"`
RepoQuery string `db:"repo_query" json:"repo_query"`
PrTitleQuery string `db:"pr_title_query" json:"pr_title_query"`
Search string `db:"search" json:"search"`
OffsetOpt int32 `db:"offset_opt" json:"offset_opt"`
LimitOpt int32 `db:"limit_opt" json:"limit_opt"`
}
@@ -8653,6 +8743,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha
arg.PrNumber,
arg.RepoQuery,
arg.PrTitleQuery,
arg.Search,
arg.OffsetOpt,
arg.LimitOpt,
)
@@ -9147,7 +9238,7 @@ func (q *sqlQuerier) GetDatabaseNow(ctx context.Context) (time.Time, error) {
const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, 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, api_key_id, revision, reasoning_effort, search_tsv
FROM
chat_messages
WHERE
@@ -9193,6 +9284,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
)
return i, err
}
@@ -9702,7 +9794,7 @@ SELECT
NULLIF(UNNEST($18::bigint[]), 0),
NULLIF(UNNEST($19::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, api_key_id, revision, reasoning_effort, search_tsv
`
type InsertChatMessagesParams struct {
@@ -9781,6 +9873,7 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa
&i.APIKeyID,
&i.Revision,
&i.ReasoningEffort,
&i.SearchTsv,
); err != nil {
return nil, err
}
+66
View File
@@ -312,6 +312,32 @@ SET
WHERE
id = @id::bigint;
-- name: BackfillChatMessagesSearchTsv :execrows
-- Backfills chat_messages.search_tsv for pending rows, newest first.
-- The WHERE clause must match the predicate of
-- idx_chat_messages_search_tsv_pending exactly so the partial index
-- serves this query.
WITH batch AS (
SELECT id FROM chat_messages
WHERE search_tsv IS NULL
AND deleted = false
AND visibility IN ('user', 'both')
AND role IN ('user', 'assistant')
ORDER BY id DESC
LIMIT @batch_size::int
)
UPDATE chat_messages cm
-- NULL means "pending", empty tsvector means "backfilled, no text".
SET search_tsv = COALESCE(
to_tsvector('simple', chat_message_search_text(cm.content)),
''::tsvector)
FROM batch WHERE cm.id = batch.id;
-- name: ChatSearchQueryIsEmpty :one
-- Reports whether search text tokenizes to an empty tsquery (e.g. '!!!').
-- Used to reject input that would silently match nothing.
SELECT numnode(websearch_to_tsquery('simple', @search::text)) = 0 AS is_empty;
-- name: GetChatByID :one
SELECT *
FROM chats_expanded
@@ -651,6 +677,46 @@ WHERE
)
ELSE true
END
-- websearch_to_tsquery accepts quoted phrases, OR, and -negation;
-- the 'simple' config folds case and skips stemming.
AND CASE
WHEN @search::text != '' THEN (
-- Served by idx_chats_title_fts.
to_tsvector('simple', chats_expanded.title) @@ websearch_to_tsquery('simple', @search)
-- Served by idx_chat_diff_statuses_pr_title_fts.
OR EXISTS (
SELECT 1
FROM chat_diff_statuses cds
WHERE cds.chat_id = chats_expanded.id
AND to_tsvector('simple', cds.pull_request_title) @@ websearch_to_tsquery('simple', @search)
)
-- The WHERE clause must repeat the predicate of the partial index
-- idx_chat_messages_search_tsv so the planner can use it. Additional
-- filters should still be fine.
OR EXISTS (
SELECT 1
FROM chat_messages cm
WHERE cm.chat_id = chats_expanded.id
AND cm.search_tsv IS NOT NULL
AND cm.deleted = false
AND cm.visibility IN ('user', 'both')
AND cm.role IN ('user', 'assistant')
AND cm.search_tsv @@ websearch_to_tsquery('simple', @search)
)
-- Skip an explicit pr_number lookup unless the search is a valid bigint.
OR CASE
WHEN @search ~ '^[0-9]{1,18}$' THEN EXISTS (
SELECT 1
FROM chat_diff_statuses cds
WHERE cds.chat_id = chats_expanded.id
AND cds.pr_number IS NOT NULL
AND cds.pr_number = @search::bigint
)
ELSE false
END
)
ELSE true
END
-- Paginate over root chats only. Children are fetched
-- separately via GetChildChatsByParentIDs and embedded under
-- each parent. Other callers that need the full set should