MUL-7342: index chat session delete lookups (#8375)

* fix(db): index chat session delete lookups

Co-authored-by: multica-agent <[email protected]>

* fix(db): address migration review

Co-authored-by: multica-agent <[email protected]>

---------

Co-authored-by: Sol-Boy <[email protected]>
Co-authored-by: multica-agent <[email protected]>

server/cmd/migrate/README.md

@@ -1,18 +1,18 @@
 # Migration runner operations
 
-## Prebuild the cancelled-chat session guard index
+## Prebuild the chat-session delete lookup index
 
-Migration 465 builds `idx_agent_task_queue_chat_with_session_created_at`, which
-keeps `AdvanceCancelledChatSessionPointer` from scanning the global task queue
-while a cancel or late session pin holds the chat row lock. This is a net-new
-index on `agent_task_queue`, and most historical chat tasks may satisfy its
-predicate. Measure the eligible population before scheduling the build:
+Migration 472 builds `idx_agent_task_queue_chat_session`, which keeps the
+legacy `chat_session` foreign-key delete action from scanning the global task
+queue for every deleted chat. Its ordered keys also continue to cover
+`AdvanceCancelledChatSessionPointer` after migration 473 removes the narrower
+migration 465 index. The replacement includes every task attached to a chat,
+so measure the eligible population before scheduling the build:
 
 ```sql
 SELECT count(*) AS indexed_rows
 FROM agent_task_queue
-WHERE chat_session_id IS NOT NULL
-  AND session_id IS NOT NULL;
+WHERE chat_session_id IS NOT NULL;
 ```
 
 Migrations run during backend startup, whose Helm startup probe allows ten
@@ -21,10 +21,9 @@ prebuild the index in a low-traffic window. Run the statement by itself and
 outside a transaction:
 
 ```sql
-CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_agent_task_queue_chat_with_session_created_at
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_agent_task_queue_chat_session
 ON agent_task_queue (chat_session_id, created_at DESC)
-WHERE chat_session_id IS NOT NULL
-  AND session_id IS NOT NULL;
+WHERE chat_session_id IS NOT NULL;
 ```
 
 Confirm the build is usable before deploying:
@@ -36,17 +35,20 @@ SELECT indexrelid::regclass AS index_name,
        indisready,
        indislive
 FROM pg_index
-WHERE indexrelid = to_regclass('idx_agent_task_queue_chat_with_session_created_at');
+WHERE indexrelid = to_regclass('idx_agent_task_queue_chat_session');
 ```
 
-The migration then becomes a fast no-op. If an interrupted manual build leaves
-the index invalid, drop that invalid index concurrently and retry the standalone
-build before deploying. The migration runner also registers invalid-index
-cleanup so an interrupted startup build can recover on its next attempt.
+Migration 472 then becomes a fast no-op. Migration 473 only drops the narrower
+index concurrently and does not scan `agent_task_queue`; migration 474 builds a
+single-column index on `dingtalk_bot_identity`, not the global task table. If an
+interrupted manual build leaves the replacement index invalid, drop that invalid
+index concurrently and retry the standalone build before deploying. The
+migration runner also registers invalid-index cleanup so an interrupted startup
+build can recover on its next attempt.
 
 Application rollback is compatible with the extra index, so leave it in place.
-Rolling back migration 465 drops the index and is functionally safe, but restores
-the global scan and longer chat-lock hold time.
+A full database rollback restores migration 465's narrower index before dropping
+the replacement, but returns chat-session deletes to the global task-table scan.
 
 ## Issue description search index retirement
 

server/cmd/migrate/main.go

@@ -303,6 +303,8 @@ var concurrentIndexCleanups = map[string]string{
 	"460_agent_task_queue_autopilot_run_created_at_index":       "idx_agent_task_queue_autopilot_run_created_at",
 	"465_agent_task_queue_chat_with_session_index":              "idx_agent_task_queue_chat_with_session_created_at",
 	"466_activity_log_member_assignee_frequency_index":          "idx_activity_log_member_assignee_frequency",
+	"472_agent_task_queue_chat_session_index":                   "idx_agent_task_queue_chat_session",
+	"474_dingtalk_bot_identity_workspace_index":                 "idx_dingtalk_bot_identity_workspace",
 }
 
 // concurrentDownIndexCleanups covers every migration whose down direction
@@ -332,6 +334,7 @@ var concurrentDownIndexCleanups = map[string]string{
 	"455_drop_comment_content_trgm_index":                   "idx_comment_content_trgm",
 	"463_drop_issue_description_bigm_index":                 "idx_issue_description_bigm",
 	"464_drop_issue_description_trgm_index":                 "idx_issue_description_trgm",
+	"473_drop_agent_task_queue_chat_with_session_index":     "idx_agent_task_queue_chat_with_session_created_at",
 }
 
 var preMigrationHooks = func() map[string]preMigrationHook {

server/cmd/migrate/migrate_chat_session_delete_indexes_test.go

@@ -0,0 +1,166 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"math/rand/v2"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func TestChatSessionDeleteIndexMigrationsPreserveCoverageAndRollback(t *testing.T) {
+	adminPool := openTestPool(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+	defer cancel()
+
+	suffix := fmt.Sprintf("%d_%d", time.Now().UnixNano(), rand.Uint32())
+	schema := "migrate_chat_session_delete_idx_" + suffix
+	schemaIdent := pgx.Identifier{schema}.Sanitize()
+	if _, err := adminPool.Exec(ctx, "CREATE SCHEMA "+schemaIdent); err != nil {
+		t.Fatalf("create schema: %v", err)
+	}
+	t.Cleanup(func() {
+		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cleanupCancel()
+		if _, err := adminPool.Exec(cleanupCtx, "DROP SCHEMA IF EXISTS "+schemaIdent+" CASCADE"); err != nil {
+			t.Logf("drop schema %s: %v", schema, err)
+		}
+	})
+
+	pool := openTestPoolWithSearchPath(t, schema)
+	for _, statement := range []string{
+		`CREATE TABLE schema_migrations (
+			version TEXT PRIMARY KEY,
+			applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
+		)`,
+		`CREATE TABLE agent_task_queue (
+			id BIGSERIAL PRIMARY KEY,
+			chat_session_id UUID,
+			created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+			session_id TEXT
+		)`,
+		`CREATE INDEX idx_agent_task_queue_chat_with_session_created_at
+			ON agent_task_queue (chat_session_id, created_at DESC)
+			WHERE chat_session_id IS NOT NULL
+			  AND session_id IS NOT NULL`,
+		`INSERT INTO agent_task_queue (chat_session_id, created_at, session_id)
+			SELECT
+				('00000000-0000-0000-0000-' || lpad(n::text, 12, '0'))::uuid,
+				TIMESTAMPTZ '2026-01-01 00:00:00+00' + make_interval(secs => n),
+				CASE WHEN n % 5 = 0 THEN NULL ELSE 'provider-session-' || n END
+			FROM generate_series(1, 5000) AS n`,
+		`CREATE TABLE dingtalk_bot_identity (
+			workspace_id UUID NOT NULL,
+			installation_id UUID NOT NULL
+		)`,
+	} {
+		if _, err := pool.Exec(ctx, statement); err != nil {
+			t.Fatalf("apply fixture statement %q: %v", statement, err)
+		}
+	}
+
+	versions := []string{
+		"472_agent_task_queue_chat_session_index",
+		"473_drop_agent_task_queue_chat_with_session_index",
+		"474_dingtalk_bot_identity_workspace_index",
+	}
+	if err := runMigrations(ctx, pool, runOptions{
+		Direction:             "up",
+		Files:                 realMigrationFiles(t, versions, "up"),
+		SchemaMigrationsTable: schema + ".schema_migrations",
+		AdvisoryLockKey:       int64(rand.Uint64()&0x7fffffffffffffff) | 1,
+		Hooks:                 hooksForDirection("up"),
+	}); err != nil {
+		t.Fatalf("apply chat-session delete index migrations: %v", err)
+	}
+	if _, err := pool.Exec(ctx, "ANALYZE agent_task_queue"); err != nil {
+		t.Fatalf("analyze agent task fixture: %v", err)
+	}
+
+	assertIndexValidity(t, pool, schema, "idx_agent_task_queue_chat_session", true)
+	assertIndexExists(t, pool, schema, "idx_agent_task_queue_chat_with_session_created_at", false)
+	assertIndexValidity(t, pool, schema, "idx_dingtalk_bot_identity_workspace", true)
+	assertMigrationIndexShape(
+		t, pool, schema, "idx_agent_task_queue_chat_session",
+		"USING btree (chat_session_id, created_at DESC)",
+		"chat_session_idISNOTNULL",
+	)
+	assertMigrationIndexShape(
+		t, pool, schema, "idx_dingtalk_bot_identity_workspace",
+		"USING btree (workspace_id)",
+		"",
+	)
+	// Row 5000 has no session_id, so migration 465's narrower predicate could
+	// not serve this foreign-key-shaped lookup.
+	assertPlanUsesIndex(t, pool, `
+		SELECT id
+		FROM agent_task_queue
+		WHERE chat_session_id = '00000000-0000-0000-0000-000000005000'
+	`, "idx_agent_task_queue_chat_session")
+	// Row 4999 has a provider session and exercises the newer-task guard shape
+	// that migration 465 served before migration 473 removed it.
+	assertPlanUsesIndex(t, pool, `
+		SELECT 1
+		FROM agent_task_queue newer
+		WHERE newer.chat_session_id = '00000000-0000-0000-0000-000000004999'
+		  AND newer.id <> 4999
+		  AND newer.session_id IS NOT NULL
+		  AND newer.created_at > TIMESTAMPTZ '2026-01-01 00:00:00+00'
+	`, "idx_agent_task_queue_chat_session")
+
+	reversedVersions := []string{
+		"474_dingtalk_bot_identity_workspace_index",
+		"473_drop_agent_task_queue_chat_with_session_index",
+		"472_agent_task_queue_chat_session_index",
+	}
+	if err := runMigrations(ctx, pool, runOptions{
+		Direction:             "down",
+		Files:                 realMigrationFiles(t, reversedVersions, "down"),
+		SchemaMigrationsTable: schema + ".schema_migrations",
+		AdvisoryLockKey:       int64(rand.Uint64()&0x7fffffffffffffff) | 1,
+		Hooks:                 hooksForDirection("down"),
+	}); err != nil {
+		t.Fatalf("roll back chat-session delete index migrations: %v", err)
+	}
+
+	assertIndexExists(t, pool, schema, "idx_agent_task_queue_chat_session", false)
+	assertIndexValidity(t, pool, schema, "idx_agent_task_queue_chat_with_session_created_at", true)
+	assertIndexExists(t, pool, schema, "idx_dingtalk_bot_identity_workspace", false)
+	assertMigrationIndexShape(
+		t, pool, schema, "idx_agent_task_queue_chat_with_session_created_at",
+		"USING btree (chat_session_id, created_at DESC)",
+		"chat_session_idISNOTNULLANDsession_idISNOTNULL",
+	)
+}
+
+func assertMigrationIndexShape(
+	t *testing.T,
+	pool *pgxpool.Pool,
+	schema string,
+	indexName string,
+	wantKeys string,
+	wantPredicate string,
+) {
+	t.Helper()
+	var definition string
+	var predicate string
+	if err := pool.QueryRow(context.Background(), `
+		SELECT pg_get_indexdef(indexrelid),
+		       COALESCE(pg_get_expr(indpred, indrelid), '')
+		FROM pg_index
+		WHERE indexrelid = $1::regclass
+	`, pgx.Identifier{schema, indexName}.Sanitize()).Scan(&definition, &predicate); err != nil {
+		t.Fatalf("read index shape for %s.%s: %v", schema, indexName, err)
+	}
+	if !strings.Contains(definition, wantKeys) {
+		t.Fatalf("index %s.%s definition = %q, want keys %q", schema, indexName, definition, wantKeys)
+	}
+	normalizedPredicate := strings.NewReplacer("(", "", ")", "", " ", "").Replace(predicate)
+	if normalizedPredicate != wantPredicate {
+		t.Fatalf("index %s.%s predicate = %q, want %q", schema, indexName, predicate, wantPredicate)
+	}
+}

server/migrations/472_agent_task_queue_chat_session_index.down.sql

@@ -0,0 +1 @@
+DROP INDEX CONCURRENTLY IF EXISTS idx_agent_task_queue_chat_session;

server/migrations/472_agent_task_queue_chat_session_index.up.sql

@@ -0,0 +1,11 @@
+-- Every chat_session delete runs the legacy agent_task_queue foreign-key
+-- action with only chat_session_id in its predicate. The existing chat indexes
+-- add status, session_id, or retired_session_id predicates and cannot serve
+-- that lookup, so each deleted session scans the global task table.
+--
+-- Keep created_at in the key so this broader index also covers migration 465's
+-- cancelled-chat pointer lookup. Migration 473 drops that narrower index only
+-- after this concurrent build has completed successfully.
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_agent_task_queue_chat_session
+ON agent_task_queue (chat_session_id, created_at DESC)
+WHERE chat_session_id IS NOT NULL;

server/migrations/473_drop_agent_task_queue_chat_with_session_index.down.sql

@@ -0,0 +1,6 @@
+-- Restore the narrower index before rollback removes migration 472's broader
+-- replacement.
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_agent_task_queue_chat_with_session_created_at
+ON agent_task_queue (chat_session_id, created_at DESC)
+WHERE chat_session_id IS NOT NULL
+  AND session_id IS NOT NULL;

server/migrations/473_drop_agent_task_queue_chat_with_session_index.up.sql

@@ -0,0 +1,4 @@
+-- Migration 472's index has the same ordered keys and a broader predicate, so
+-- the cancelled-chat pointer lookup remains covered without carrying both
+-- indexes on the hottest write table.
+DROP INDEX CONCURRENTLY IF EXISTS idx_agent_task_queue_chat_with_session_created_at;

server/migrations/474_dingtalk_bot_identity_workspace_index.down.sql

@@ -0,0 +1 @@
+DROP INDEX CONCURRENTLY IF EXISTS idx_dingtalk_bot_identity_workspace;

server/migrations/474_dingtalk_bot_identity_workspace_index.up.sql

@@ -0,0 +1,4 @@
+-- Workspace teardown keeps the denormalized workspace predicate so it also
+-- removes orphaned identities whose installation row is already gone.
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_dingtalk_bot_identity_workspace
+ON dingtalk_bot_identity (workspace_id);

server/pkg/db/generated/chat.sql.go

@@ -87,9 +87,10 @@ WHERE t.id = $1
 // backwards onto the turn the user interrupted.
 //
 // The newer-task lookup depends on
-// idx_agent_task_queue_chat_with_session_created_at (migration 465). Neither
-// chat_pending_v3 nor chat_terminal_resume can replace it: this guard spans
-// both in-flight and terminal tasks and compares created_at.
+// idx_agent_task_queue_chat_session (migration 472). Neither chat_pending_v3
+// nor chat_terminal_resume can replace it: this guard spans both in-flight and
+// terminal tasks and compares created_at; the broader index also serves the
+// chat_session foreign-key delete lookup.
 func (q *Queries) AdvanceCancelledChatSessionPointer(ctx context.Context, taskID pgtype.UUID) error {
 	_, err := q.db.Exec(ctx, advanceCancelledChatSessionPointer, taskID)
 	return err

server/pkg/db/queries/chat.sql

@@ -333,9 +333,10 @@ WHERE id = sqlc.arg('id')
 -- backwards onto the turn the user interrupted.
 --
 -- The newer-task lookup depends on
--- idx_agent_task_queue_chat_with_session_created_at (migration 465). Neither
--- chat_pending_v3 nor chat_terminal_resume can replace it: this guard spans
--- both in-flight and terminal tasks and compares created_at.
+-- idx_agent_task_queue_chat_session (migration 472). Neither chat_pending_v3
+-- nor chat_terminal_resume can replace it: this guard spans both in-flight and
+-- terminal tasks and compares created_at; the broader index also serves the
+-- chat_session foreign-key delete lookup.
 UPDATE chat_session cs
 SET session_id = t.session_id,
     runtime_id = t.runtime_id,