opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: state-aware queued message promotion (#24819)

Browse Source 
PromoteQueued now branches on chat status: synth tool results before
the user message on requires_action, deferred reorder + Waiting on
running so the worker's persist+auto-promote keeps partial output.
Stale heartbeat falls through to the synchronous path; GetStaleChats
picks up Waiting+queue to recover post-cleanup-crash. Endpoint
returns 202.

Closes CODAGT-119
This commit is contained in:
Mathias Fredriksson
2026-05-06 19:11:56 +03:00
committed by GitHub
parent 2cab1b41ad
commit 6b0518d051
18 changed files with 2431 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+11
View File
@@ -6057,6 +6057,17 @@ func (q *querier) RemoveUserFromGroups(ctx context.Context, arg database.RemoveU
return q.db.RemoveUserFromGroups(ctx, arg)
}
func (q *querier) ReorderChatQueuedMessageToFront(ctx context.Context, arg database.ReorderChatQueuedMessageToFrontParams) (int64, error) {
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
if err != nil {
return 0, err
}
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
return 0, err
}
return q.db.ReorderChatQueuedMessageToFront(ctx, arg)
}
func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
return database.ResolveUserChatSpendLimitRow{}, err
coderd/database/dbauthz/dbauthz_test.go
+7
View File
@@ -1042,6 +1042,13 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().PopNextQueuedMessage(gomock.Any(), chat.ID).Return(qm, nil).AnyTimes()
check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(qm)
}))
s.Run("ReorderChatQueuedMessageToFront", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
chat := testutil.Fake(s.T(), faker, database.Chat{})
arg := database.ReorderChatQueuedMessageToFrontParams{ChatID: chat.ID, TargetID: 123}
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
dbm.EXPECT().ReorderChatQueuedMessageToFront(gomock.Any(), arg).Return(int64(1), nil).AnyTimes()
check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(int64(1))
}))
s.Run("UpdateChatByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
chat := testutil.Fake(s.T(), faker, database.Chat{})
arg := database.UpdateChatByIDParams{
coderd/database/dbmetrics/querymetrics.go
+8
View File
@@ -4344,6 +4344,14 @@ func (m queryMetricsStore) RemoveUserFromGroups(ctx context.Context, arg databas
return r0, r1
}
func (m queryMetricsStore) ReorderChatQueuedMessageToFront(ctx context.Context, arg database.ReorderChatQueuedMessageToFrontParams) (int64, error) {
start := time.Now()
r0, r1 := m.s.ReorderChatQueuedMessageToFront(ctx, arg)
m.queryLatencies.WithLabelValues("ReorderChatQueuedMessageToFront").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ReorderChatQueuedMessageToFront").Inc()
return r0, r1
}
func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
start := time.Now()
r0, r1 := m.s.ResolveUserChatSpendLimit(ctx, userID)
coderd/database/dbmock/dbmock.go
+15
View File
@@ -8233,6 +8233,21 @@ func (mr *MockStoreMockRecorder) RemoveUserFromGroups(ctx, arg any) *gomock.Call
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveUserFromGroups", reflect.TypeOf((*MockStore)(nil).RemoveUserFromGroups), ctx, arg)
}
// ReorderChatQueuedMessageToFront mocks base method.
func (m *MockStore) ReorderChatQueuedMessageToFront(ctx context.Context, arg database.ReorderChatQueuedMessageToFrontParams) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ReorderChatQueuedMessageToFront", ctx, arg)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ReorderChatQueuedMessageToFront indicates an expected call of ReorderChatQueuedMessageToFront.
func (mr *MockStoreMockRecorder) ReorderChatQueuedMessageToFront(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReorderChatQueuedMessageToFront", reflect.TypeOf((*MockStore)(nil).ReorderChatQueuedMessageToFront), ctx, arg)
}
// ResolveUserChatSpendLimit mocks base method.
func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
m.ctrl.T.Helper()
coderd/database/querier.go
+9 -3
View File
@@ -591,10 +591,13 @@ type sqlcQuerier interface {
GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]Replica, error)
GetRunningPrebuiltWorkspaces(ctx context.Context) ([]GetRunningPrebuiltWorkspacesRow, error)
GetRuntimeConfig(ctx context.Context, key string) (string, error)
// Find chats that appear stuck and need recovery. This covers:
// Find chats that appear stuck and need recovery:
// 1. Running chats whose heartbeat has expired (worker crash).
// 2. Chats awaiting client action (requires_action) past the
// timeout threshold (client disappeared).
// 2. requires_action chats past the timeout threshold (client
// disappeared).
// 3. Waiting chats with a non-empty queue and stale updated_at
// (deferred-promote stranding when the worker dies before its
// post-cancel cleanup runs).
GetStaleChats(ctx context.Context, staleThreshold time.Time) ([]Chat, error)
GetTailnetPeers(ctx context.Context, id uuid.UUID) ([]TailnetPeer, error)
GetTailnetTunnelPeerBindingsBatch(ctx context.Context, ids []uuid.UUID) ([]GetTailnetTunnelPeerBindingsBatchRow, error)
@@ -1012,6 +1015,9 @@ type sqlcQuerier interface {
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)
// Mutates only created_at on the target row; ids are unchanged so
// consumers can keep tracking queued messages by id.
ReorderChatQueuedMessageToFront(ctx context.Context, arg ReorderChatQueuedMessageToFrontParams) (int64, error)
// Resolves the effective spend limit for a user using the hierarchy:
// 1. Individual user override (highest priority, applies globally across
// all organizations since it lives on the users table)
coderd/database/queries.sql.go
+39 -5
View File
@@ -6808,7 +6808,7 @@ func (q *sqlQuerier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]Get
const getChatQueuedMessages = `-- name: GetChatQueuedMessages :many
SELECT id, chat_id, content, created_at, model_config_id FROM chat_queued_messages
WHERE chat_id = $1
ORDER BY id ASC
ORDER BY created_at ASC, id ASC
`
func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) {
@@ -7311,12 +7311,21 @@ WHERE
AND heartbeat_at < $1::timestamptz)
OR (status = 'requires_action'::chat_status
AND updated_at < $1::timestamptz)
OR (status = 'waiting'::chat_status
AND updated_at < $1::timestamptz
AND EXISTS (
SELECT 1 FROM chat_queued_messages cqm
WHERE cqm.chat_id = chats.id
))
`
// Find chats that appear stuck and need recovery. This covers:
// Find chats that appear stuck and need recovery:
// 1. Running chats whose heartbeat has expired (worker crash).
// 2. Chats awaiting client action (requires_action) past the
// timeout threshold (client disappeared).
// 2. requires_action chats past the timeout threshold (client
// disappeared).
// 3. Waiting chats with a non-empty queue and stale updated_at
// (deferred-promote stranding when the worker dies before its
// post-cancel cleanup runs).
func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time) ([]Chat, error) {
rows, err := q.db.QueryContext(ctx, getStaleChats, staleThreshold)
if err != nil {
@@ -7946,7 +7955,7 @@ DELETE FROM chat_queued_messages
WHERE id = (
SELECT cqm.id FROM chat_queued_messages cqm
WHERE cqm.chat_id = $1
ORDER BY cqm.id ASC
ORDER BY cqm.created_at ASC, cqm.id ASC
LIMIT 1
)
RETURNING id, chat_id, content, created_at, model_config_id
@@ -7965,6 +7974,31 @@ func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID)
return i, err
}
const reorderChatQueuedMessageToFront = `-- name: ReorderChatQueuedMessageToFront :execrows
UPDATE chat_queued_messages AS target
SET created_at = (
SELECT MIN(inner_cqm.created_at) - INTERVAL '1 microsecond'
FROM chat_queued_messages AS inner_cqm
WHERE inner_cqm.chat_id = $1
)
WHERE target.id = $2 AND target.chat_id = $1
`
type ReorderChatQueuedMessageToFrontParams struct {
ChatID uuid.UUID `db:"chat_id" json:"chat_id"`
TargetID int64 `db:"target_id" json:"target_id"`
}
// Mutates only created_at on the target row; ids are unchanged so
// consumers can keep tracking queued messages by id.
func (q *sqlQuerier) ReorderChatQueuedMessageToFront(ctx context.Context, arg ReorderChatQueuedMessageToFrontParams) (int64, error) {
result, err := q.db.ExecContext(ctx, reorderChatQueuedMessageToFront, arg.ChatID, arg.TargetID)
if err != nil {
return 0, err
}
return result.RowsAffected()
}
const resolveUserChatSpendLimit = `-- name: ResolveUserChatSpendLimit :one
SELECT CASE
WHEN NOT cfg.enabled THEN -1
coderd/database/queries/chats.sql
+26 -6
View File
@@ -760,10 +760,13 @@ RETURNING
*;
-- name: GetStaleChats :many
-- Find chats that appear stuck and need recovery. This covers:
-- Find chats that appear stuck and need recovery:
-- 1. Running chats whose heartbeat has expired (worker crash).
-- 2. Chats awaiting client action (requires_action) past the
-- timeout threshold (client disappeared).
-- 2. requires_action chats past the timeout threshold (client
-- disappeared).
-- 3. Waiting chats with a non-empty queue and stale updated_at
-- (deferred-promote stranding when the worker dies before its
-- post-cancel cleanup runs).
SELECT
*
FROM
@@ -772,7 +775,13 @@ WHERE
(status = 'running'::chat_status
AND heartbeat_at < @stale_threshold::timestamptz)
OR (status = 'requires_action'::chat_status
AND updated_at < @stale_threshold::timestamptz);
AND updated_at < @stale_threshold::timestamptz)
OR (status = 'waiting'::chat_status
AND updated_at < @stale_threshold::timestamptz
AND EXISTS (
SELECT 1 FROM chat_queued_messages cqm
WHERE cqm.chat_id = chats.id
));
-- name: UpdateChatHeartbeats :many
-- Bumps the heartbeat timestamp for the given set of chat IDs,
@@ -916,7 +925,7 @@ RETURNING *;
-- name: GetChatQueuedMessages :many
SELECT * FROM chat_queued_messages
WHERE chat_id = @chat_id
ORDER BY id ASC;
ORDER BY created_at ASC, id ASC;
-- name: DeleteChatQueuedMessage :exec
DELETE FROM chat_queued_messages WHERE id = @id AND chat_id = @chat_id;
@@ -929,11 +938,22 @@ DELETE FROM chat_queued_messages
WHERE id = (
SELECT cqm.id FROM chat_queued_messages cqm
WHERE cqm.chat_id = @chat_id
ORDER BY cqm.id ASC
ORDER BY cqm.created_at ASC, cqm.id ASC
LIMIT 1
)
RETURNING *;
-- name: ReorderChatQueuedMessageToFront :execrows
-- Mutates only created_at on the target row; ids are unchanged so
-- consumers can keep tracking queued messages by id.
UPDATE chat_queued_messages AS target
SET created_at = (
SELECT MIN(inner_cqm.created_at) - INTERVAL '1 microsecond'
FROM chat_queued_messages AS inner_cqm
WHERE inner_cqm.chat_id = @chat_id
)
WHERE target.id = @target_id AND target.chat_id = @chat_id;
-- name: GetLastChatMessageByRole :one
SELECT
*