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

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
+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
+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{
@@ -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)
+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()
+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)
+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
+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
*
+4 -2
View File
@@ -3193,7 +3193,7 @@ func (api *API) promoteChatQueuedMessage(rw http.ResponseWriter, r *http.Request
return
}
promoteResult, txErr := api.chatDaemon.PromoteQueued(ctx, chatd.PromoteQueuedOptions{
_, txErr := api.chatDaemon.PromoteQueued(ctx, chatd.PromoteQueuedOptions{
ChatID: chatID,
CreatedBy: apiKey.UserID,
QueuedMessageID: queuedMessageID,
@@ -3216,7 +3216,9 @@ func (api *API) promoteChatQueuedMessage(rw http.ResponseWriter, r *http.Request
return
}
httpapi.Write(ctx, rw, http.StatusOK, convertChatMessage(promoteResult.PromotedMessage))
httpapi.Write(ctx, rw, http.StatusAccepted, codersdk.Response{
Message: "Queued message promotion accepted.",
})
}
// markChatAsRead updates the last read message ID for a chat to the
+241 -31
View File
@@ -32,6 +32,7 @@ import (
"github.com/coder/coder/v2/coderd/database/dbfake"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/externalauth"
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
"github.com/coder/coder/v2/coderd/rbac"
@@ -6096,7 +6097,7 @@ func TestWatchChatsStatusChangeCarriesUpdatedLastModelConfigID(t *testing.T) {
)
require.NoError(t, err)
defer promoteRes.Body.Close()
require.Equal(t, http.StatusOK, promoteRes.StatusCode)
require.Equal(t, http.StatusAccepted, promoteRes.StatusCode)
event := waitForChatWatchStatusChangeEvent(ctx, t, conn, chat.ID)
require.Equal(t, modelConfigB.ID, event.Chat.LastModelConfigID)
@@ -8163,24 +8164,11 @@ func TestPromoteChatQueuedMessage(t *testing.T) {
)
require.NoError(t, err)
defer promoteRes.Body.Close()
require.Equal(t, http.StatusOK, promoteRes.StatusCode)
require.Equal(t, http.StatusAccepted, promoteRes.StatusCode)
var promoted codersdk.ChatMessage
err = json.NewDecoder(promoteRes.Body).Decode(&promoted)
require.NoError(t, err)
require.NotZero(t, promoted.ID)
require.Equal(t, chat.ID, promoted.ChatID)
require.Equal(t, codersdk.ChatMessageRoleUser, promoted.Role)
foundPromotedText := false
for _, part := range promoted.Content {
if part.Type == codersdk.ChatMessagePartTypeText &&
part.Text == queuedText {
foundPromotedText = true
break
}
}
require.True(t, foundPromotedText)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(promoteRes.Body).Decode(&resp))
require.NotEmpty(t, resp.Message)
messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil)
require.NoError(t, err)
@@ -8188,6 +8176,19 @@ func TestPromoteChatQueuedMessage(t *testing.T) {
require.NotEqual(t, queuedMessage.ID, queued.ID)
}
foundPromoted := false
for _, msg := range messagesResult.Messages {
if msg.Role != codersdk.ChatMessageRoleUser {
continue
}
for _, part := range msg.Content {
if part.Type == codersdk.ChatMessagePartTypeText && part.Text == queuedText {
foundPromoted = true
}
}
}
require.True(t, foundPromoted, "promoted message must appear in chat history")
queuedMessages, err := db.GetChatQueuedMessages(dbauthz.AsSystemRestricted(ctx), chat.ID)
require.NoError(t, err)
for _, queued := range queuedMessages {
@@ -8246,23 +8247,26 @@ func TestPromoteChatQueuedMessage(t *testing.T) {
)
require.NoError(t, err)
defer promoteRes.Body.Close()
require.Equal(t, http.StatusOK, promoteRes.StatusCode)
require.Equal(t, http.StatusAccepted, promoteRes.StatusCode)
var promoted codersdk.ChatMessage
err = json.NewDecoder(promoteRes.Body).Decode(&promoted)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(promoteRes.Body).Decode(&resp))
require.NotEmpty(t, resp.Message)
messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil)
require.NoError(t, err)
require.NotZero(t, promoted.ID)
require.Equal(t, chat.ID, promoted.ChatID)
require.Equal(t, codersdk.ChatMessageRoleUser, promoted.Role)
foundPromotedText := false
for _, part := range promoted.Content {
if part.Type == codersdk.ChatMessagePartTypeText && part.Text == queuedText {
foundPromotedText = true
break
foundPromoted := false
for _, msg := range messagesResult.Messages {
if msg.Role != codersdk.ChatMessageRoleUser {
continue
}
for _, part := range msg.Content {
if part.Type == codersdk.ChatMessagePartTypeText && part.Text == queuedText {
foundPromoted = true
}
}
}
require.True(t, foundPromotedText)
require.True(t, foundPromoted, "promoted message must appear in chat history")
queuedMessages, err := db.GetChatQueuedMessages(dbauthz.AsSystemRestricted(ctx), chat.ID)
require.NoError(t, err)
@@ -8392,6 +8396,212 @@ func TestPromoteChatQueuedMessage(t *testing.T) {
require.ErrorAs(t, promoteErr, &promoteSDKErr)
require.Contains(t, promoteSDKErr.Message, "archived")
})
t.Run("WhileRequiresAction", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
user := coderdtest.CreateFirstUser(t, client.Client)
modelConfig := createChatModelConfig(t, client)
const dynamicToolName = "my_dynamic_tool"
dynamicTools := []mcp.Tool{{
Name: dynamicToolName,
Description: "a test dynamic tool",
InputSchema: mcp.ToolInputSchema{Type: "object"},
}}
dtJSON, err := json.Marshal(dynamicTools)
require.NoError(t, err)
chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{
OrganizationID: user.OrganizationID,
Status: database.ChatStatusWaiting,
ClientType: database.ChatClientTypeUi,
OwnerID: user.UserID,
LastModelConfigID: modelConfig.ID,
Title: "promote queued requires-action route test",
DynamicTools: pqtype.NullRawMessage{RawMessage: dtJSON, Valid: true},
})
require.NoError(t, err)
const pendingToolCallID = "call_pending"
assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{
Type: codersdk.ChatMessagePartTypeToolCall,
ToolCallID: pendingToolCallID,
ToolName: dynamicToolName,
Args: json.RawMessage(`{"x":1}`),
}})
require.NoError(t, err)
_, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{
ChatID: chat.ID,
CreatedBy: []uuid.UUID{uuid.Nil},
ModelConfigID: []uuid.UUID{modelConfig.ID},
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
ContentVersion: []int16{chatprompt.CurrentContentVersion},
Content: []string{string(assistantContent.RawMessage)},
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
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)
_, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{
ID: chat.ID,
Status: database.ChatStatusRequiresAction,
})
require.NoError(t, err)
const queuedText = "queued message for requires-action promote"
queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{
codersdk.ChatMessageText(queuedText),
})
require.NoError(t, err)
queuedMessage, err := db.InsertChatQueuedMessage(
dbauthz.AsSystemRestricted(ctx),
database.InsertChatQueuedMessageParams{
ChatID: chat.ID,
Content: queuedContent,
},
)
require.NoError(t, err)
promoteRes, err := client.Request(
ctx,
http.MethodPost,
fmt.Sprintf("/api/experimental/chats/%s/queue/%d/promote", chat.ID, queuedMessage.ID),
nil,
)
require.NoError(t, err)
defer promoteRes.Body.Close()
require.Equal(t, http.StatusAccepted, promoteRes.StatusCode)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(promoteRes.Body).Decode(&resp))
require.NotEmpty(t, resp.Message)
messages, err := db.GetChatMessagesByChatID(dbauthz.AsSystemRestricted(ctx), database.GetChatMessagesByChatIDParams{
ChatID: chat.ID,
AfterID: 0,
})
require.NoError(t, err)
var (
syntheticID int64
promotedID int64
)
for _, msg := range messages {
parts, parseErr := chatprompt.ParseContent(msg)
require.NoError(t, parseErr)
for _, part := range parts {
if msg.Role == database.ChatMessageRoleTool &&
part.Type == codersdk.ChatMessagePartTypeToolResult &&
part.ToolCallID == pendingToolCallID &&
part.IsError {
syntheticID = msg.ID
}
if msg.Role == database.ChatMessageRoleUser &&
part.Type == codersdk.ChatMessagePartTypeText &&
part.Text == queuedText {
promotedID = msg.ID
}
}
}
require.NotZero(t, syntheticID,
"expected a synthetic error tool result for the pending tool call")
require.NotZero(t, promotedID,
"expected the promoted user message in chat history")
require.Less(t, syntheticID, promotedID,
"synthetic tool result must precede the promoted user message")
queuedRemaining, err := db.GetChatQueuedMessages(dbauthz.AsSystemRestricted(ctx), chat.ID)
require.NoError(t, err)
for _, qm := range queuedRemaining {
require.NotEqual(t, queuedMessage.ID, qm.ID)
}
})
t.Run("WhileRunning", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
user := coderdtest.CreateFirstUser(t, client.Client)
modelConfig := createChatModelConfig(t, client)
chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{
OrganizationID: user.OrganizationID,
Status: database.ChatStatusWaiting,
ClientType: database.ChatClientTypeUi,
OwnerID: user.UserID,
LastModelConfigID: modelConfig.ID,
Title: "promote queued running route test",
})
require.NoError(t, err)
// Simulate an active worker by setting status to running.
// We do not start a real worker; the running-case behavior
// (reorder + set waiting + clear worker) does not depend on
// one. The deferred auto-promote is exercised by the
// chatd-package tests where a real worker is involved.
_, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{
ID: chat.ID,
Status: database.ChatStatusRunning,
WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
StartedAt: sql.NullTime{Time: dbtime.Now(), Valid: true},
HeartbeatAt: sql.NullTime{Time: dbtime.Now(), Valid: true},
})
require.NoError(t, err)
queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{
codersdk.ChatMessageText("running-promote"),
})
require.NoError(t, err)
queuedMessage, err := db.InsertChatQueuedMessage(
dbauthz.AsSystemRestricted(ctx),
database.InsertChatQueuedMessageParams{
ChatID: chat.ID,
Content: queuedContent,
},
)
require.NoError(t, err)
promoteRes, err := client.Request(
ctx,
http.MethodPost,
fmt.Sprintf("/api/experimental/chats/%s/queue/%d/promote", chat.ID, queuedMessage.ID),
nil,
)
require.NoError(t, err)
defer promoteRes.Body.Close()
require.Equal(t, http.StatusAccepted, promoteRes.StatusCode)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(promoteRes.Body).Decode(&resp))
require.NotEmpty(t, resp.Message)
after, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID)
require.NoError(t, err)
require.Equal(t, database.ChatStatusWaiting, after.Status,
"running-case promote must transition chat to waiting")
require.False(t, after.WorkerID.Valid,
"running-case promote must clear WorkerID")
queuedRemaining, err := db.GetChatQueuedMessages(dbauthz.AsSystemRestricted(ctx), chat.ID)
require.NoError(t, err)
require.Len(t, queuedRemaining, 1)
require.Equal(t, queuedMessage.ID, queuedRemaining[0].ID,
"queued message ID must stay stable across reorder")
})
}
func TestChatUsageLimitOverrideRoutes(t *testing.T) {
+199 -25
View File
@@ -1195,6 +1195,9 @@ type PromoteQueuedOptions struct {
// PromoteQueuedResult contains post-promotion message metadata.
type PromoteQueuedResult struct {
// PromotedMessage is the inserted user message. For a chat that
// was running at promote time, the insertion is deferred to the
// worker's auto-promote and PromotedMessage is the zero value.
PromotedMessage database.ChatMessage
}
@@ -2042,7 +2045,10 @@ func (p *Server) DeleteQueued(
return nil
}
// PromoteQueued promotes a queued message into chat history and marks the chat pending.
// PromoteQueued promotes a queued message into chat history. On a
// running chat with a fresh worker heartbeat the promote is deferred
// to the worker's persist+auto-promote so partial assistant output
// is not lost; otherwise it inserts the user message synchronously.
func (p *Server) PromoteQueued(
ctx context.Context,
opts PromoteQueuedOptions,
@@ -2052,10 +2058,12 @@ func (p *Server) PromoteQueued(
}
var (
result PromoteQueuedResult
promoted database.ChatMessage
updatedChat database.Chat
remainingQueue []database.ChatQueuedMessage
result PromoteQueuedResult
promoted database.ChatMessage
updatedChat database.Chat
remainingQueue []database.ChatQueuedMessage
deferred bool
syntheticResults []database.ChatMessage
)
txErr := p.db.InTx(func(tx database.Store) error {
@@ -2087,7 +2095,46 @@ func (p *Server) PromoteQueued(
}
}
if !found {
return xerrors.New("queued message not found")
return xerrors.Errorf("queued message %d not found in chat %s", opts.QueuedMessageID, opts.ChatID)
}
// Setting pending would trip persistStep's ownership guard
// and drop the worker's partial output. Set waiting and
// reorder the queued row so the worker's auto-promote picks
// it up after the persist.
heartbeatFresh := lockedChat.HeartbeatAt.Valid &&
p.clock.Now().Sub(lockedChat.HeartbeatAt.Time) < p.inFlightChatStaleAfter
if lockedChat.Status == database.ChatStatusRunning && heartbeatFresh {
rowsAffected, err := tx.ReorderChatQueuedMessageToFront(ctx, database.ReorderChatQueuedMessageToFrontParams{
ChatID: opts.ChatID,
TargetID: opts.QueuedMessageID,
})
if err != nil {
return xerrors.Errorf("reorder queued message to front: %w", err)
}
// Defensive guard against a future non-chat-locked
// queue mutator. The found check above makes this a
// no-op on the current code path.
if rowsAffected != 1 {
return xerrors.Errorf("reorder queued message to front affected %d rows, want 1", rowsAffected)
}
updatedChat, err = tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
ID: opts.ChatID,
Status: database.ChatStatusWaiting,
WorkerID: uuid.NullUUID{},
StartedAt: sql.NullTime{},
HeartbeatAt: sql.NullTime{},
LastError: pqtype.NullRawMessage{},
})
if err != nil {
return xerrors.Errorf("set chat to waiting for deferred promote: %w", err)
}
remainingQueue, err = tx.GetChatQueuedMessages(ctx, opts.ChatID)
if err != nil {
return xerrors.Errorf("get remaining queue after reorder: %w", err)
}
deferred = true
return nil
}
effectiveModelConfigID, err := resolveQueuedMessageModelConfigID(
@@ -2100,6 +2147,20 @@ func (p *Server) PromoteQueued(
return err
}
// Without synthetic results, the next turn would carry
// unresolved tool_call parts; the LLM API rejects this and the
// chat dead-ends in error.
if lockedChat.Status == database.ChatStatusRequiresAction {
inserted, err := insertSyntheticToolResultsTx(
ctx, tx, lockedChat,
"Tool execution interrupted by queued message promotion",
)
if err != nil {
return xerrors.Errorf("insert synthetic tool results: %w", err)
}
syntheticResults = inserted
}
err = tx.DeleteChatQueuedMessage(ctx, database.DeleteChatQueuedMessageParams{
ID: opts.QueuedMessageID,
ChatID: opts.ChatID,
@@ -2135,6 +2196,22 @@ func (p *Server) PromoteQueued(
return PromoteQueuedResult{}, txErr
}
if deferred {
// Skip publishMessage and signalWake: there is no synchronous
// user message yet, and the active worker's interrupt path
// signals its own auto-promote follow-up.
p.publishEvent(opts.ChatID, codersdk.ChatStreamEvent{
Type: codersdk.ChatStreamEventTypeQueueUpdate,
QueuedMessages: db2sdk.ChatQueuedMessages(remainingQueue),
})
p.publishChatStreamNotify(opts.ChatID, coderdpubsub.ChatStreamNotifyMessage{
QueueUpdate: true,
})
p.publishStatus(opts.ChatID, updatedChat.Status, updatedChat.WorkerID)
p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindStatusChange, nil)
return result, nil
}
p.publishEvent(opts.ChatID, codersdk.ChatStreamEvent{
Type: codersdk.ChatStreamEventTypeQueueUpdate,
QueuedMessages: db2sdk.ChatQueuedMessages(remainingQueue),
@@ -2142,6 +2219,11 @@ func (p *Server) PromoteQueued(
p.publishChatStreamNotify(opts.ChatID, coderdpubsub.ChatStreamNotifyMessage{
QueueUpdate: true,
})
// Publish synth rows before the user message so live viewers
// see the interruption inline.
for _, msg := range syntheticResults {
p.publishMessage(opts.ChatID, msg)
}
p.publishMessage(opts.ChatID, promoted)
p.publishStatus(opts.ChatID, updatedChat.Status, updatedChat.WorkerID)
p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindStatusChange, nil)
@@ -2410,7 +2492,8 @@ func (p *Server) InterruptChat(
if locked.Status != database.ChatStatusRequiresAction {
return nil
}
return insertSyntheticToolResultsTx(ctx, tx, locked, "Tool execution interrupted by user")
_, err := insertSyntheticToolResultsTx(ctx, tx, locked, "Tool execution interrupted by user")
return err
}, nil); txErr != nil {
p.logger.Error(ctx, "failed to insert synthetic tool results during interrupt",
slog.F("chat_id", chat.ID),
@@ -5223,6 +5306,7 @@ func (p *Server) trackWorkspaceUsage(
type finishActiveChatResult struct {
updatedChat database.Chat
promotedMessage *database.ChatMessage
syntheticToolResults []database.ChatMessage
remainingQueuedMessages []database.ChatQueuedMessage
shouldPublishQueueUpdate bool
}
@@ -5259,6 +5343,32 @@ func (p *Server) finishActiveChat(
switch {
case latestChat.Status == database.ChatStatusPending:
status = database.ChatStatusPending
case latestChat.Status == database.ChatStatusWaiting && status != database.ChatStatusWaiting && !latestChat.Archived:
// PromoteQueued's deferred path won the status race.
// Insert synthetic tool results before auto-promoting,
// or a RequiresAction worker outcome reintroduces the
// stops-dead bug this PR exists to fix.
inserted, synthErr := insertSyntheticToolResultsTx(
ctx, tx, latestChat,
"Tool execution interrupted by queued message promotion",
)
if synthErr != nil {
return xerrors.Errorf("insert synthetic tool results during promote-driven cleanup: %w", synthErr)
}
result.syntheticToolResults = inserted
var promoteErr error
result.promotedMessage, result.remainingQueuedMessages, result.shouldPublishQueueUpdate, promoteErr = p.tryAutoPromoteQueuedMessage(ctx, tx, latestChat)
if promoteErr != nil {
logger.Error(ctx, "auto-promote queued message failed during promote-driven cleanup", slog.Error(promoteErr))
return xerrors.Errorf("auto-promote queued message: %w", promoteErr)
}
if result.promotedMessage != nil {
status = database.ChatStatusPending
} else {
// Queue drained between snapshot and lock; honor
// the external Waiting.
status = database.ChatStatusWaiting
}
case status == database.ChatStatusWaiting && !latestChat.Archived:
// Queued messages were already admitted through SendMessage,
// so auto-promotion only preserves FIFO order here. Archived
@@ -5464,6 +5574,10 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
remainingQueuedMessages = finishResult.remainingQueuedMessages
shouldPublishQueueUpdate = finishResult.shouldPublishQueueUpdate
// Publish synth rows before the promoted user message.
for _, msg := range finishResult.syntheticToolResults {
p.publishMessage(chat.ID, msg)
}
if promotedMessage != nil {
p.publishMessage(chat.ID, *promotedMessage)
}
@@ -8032,7 +8146,7 @@ func formatPlanPathBlock(chatPath, home string) string {
}
func (p *Server) recoverStaleChats(ctx context.Context) {
staleAfter := time.Now().Add(-p.inFlightChatStaleAfter)
staleAfter := p.clock.Now().Add(-p.inFlightChatStaleAfter)
staleChats, err := p.db.GetStaleChats(ctx, staleAfter)
if err != nil {
p.logger.Error(ctx, "failed to get stale chats", slog.Error(err))
@@ -8074,6 +8188,14 @@ func (p *Server) recoverStaleChats(ctx context.Context) {
slog.F("chat_id", chat.ID))
return nil
}
case database.ChatStatusWaiting:
// Deferred-promote stranding: worker died before its
// post-cancel cleanup ran. Re-check freshness.
if !locked.UpdatedAt.Before(staleAfter) {
p.logger.Debug(ctx, "chat updated since snapshot, skipping recovery",
slog.F("chat_id", chat.ID))
return nil
}
default:
// Status changed since our snapshot; skip.
p.logger.Debug(ctx, "chat status changed since snapshot, skipping recovery",
@@ -8113,7 +8235,7 @@ func (p *Server) recoverStaleChats(ctx context.Context) {
// so the LLM history remains valid if the user
// retries the chat later.
if locked.Status == database.ChatStatusRequiresAction {
if synthErr := insertSyntheticToolResultsTx(ctx, tx, locked, "Dynamic tool execution timed out"); synthErr != nil {
if _, synthErr := insertSyntheticToolResultsTx(ctx, tx, locked, "Dynamic tool execution timed out"); synthErr != nil {
p.logger.Warn(ctx, "failed to insert synthetic tool results during stale recovery",
slog.F("chat_id", chat.ID),
slog.Error(synthErr),
@@ -8123,6 +8245,25 @@ func (p *Server) recoverStaleChats(ctx context.Context) {
}
}
if locked.Status == database.ChatStatusWaiting {
// Close pending dynamic tool calls; otherwise the
// promoted user message would feed the LLM a turn it
// rejects. Propagate errors so the next recovery
// tick retries instead of promoting incomplete
// history.
if _, synthErr := insertSyntheticToolResultsTx(ctx, tx, locked, "Tool execution interrupted by queued message promotion"); synthErr != nil {
return xerrors.Errorf("insert synthetic tool results during stale recovery: %w", synthErr)
}
promoted, _, _, promoteErr := p.tryAutoPromoteQueuedMessage(ctx, tx, locked)
if promoteErr != nil {
return xerrors.Errorf("auto-promote during stale recovery: %w", promoteErr)
}
if promoted == nil {
// Empty queue means nothing to recover.
return nil
}
}
// Reset so any replica can pick it up (pending) or
// the client sees the failure (error).
_, updateErr := tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
@@ -8150,37 +8291,66 @@ func (p *Server) recoverStaleChats(ctx context.Context) {
}
}
// insertSyntheticToolResultsTx inserts error tool-result messages for
// every pending dynamic tool call in the last assistant message. This
// keeps the LLM message history valid (every tool-call has a matching
// tool-result) when a requires_action chat times out or is interrupted.
// It operates on the provided store, which may be a transaction handle.
// insertSyntheticToolResultsTx inserts IsError tool-result messages
// for unresolved dynamic tool calls in the last assistant message,
// skipping calls already handled (e.g. by chatloop dispatching a
// name-colliding dynamic tool as a built-in). It operates on the
// provided store, which may be a transaction handle.
func insertSyntheticToolResultsTx(
ctx context.Context,
store database.Store,
chat database.Chat,
reason string,
) error {
) ([]database.ChatMessage, error) {
dynamicToolNames, err := parseDynamicToolNames(chat.DynamicTools)
if err != nil {
return xerrors.Errorf("parse dynamic tools: %w", err)
return nil, xerrors.Errorf("parse dynamic tools: %w", err)
}
if len(dynamicToolNames) == 0 {
return nil
return nil, nil
}
// Get the last assistant message to find pending tool calls.
// No assistant means nothing to close: a deferred promote can
// race a worker that fails before any persist, and the cleanup
// TX must still advance.
lastAssistant, err := store.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{
ChatID: chat.ID,
Role: database.ChatMessageRoleAssistant,
})
if errors.Is(err, sql.ErrNoRows) {
return nil, nil
}
if err != nil {
return xerrors.Errorf("get last assistant message: %w", err)
return nil, xerrors.Errorf("get last assistant message: %w", err)
}
parts, err := chatprompt.ParseContent(lastAssistant)
if err != nil {
return xerrors.Errorf("parse assistant message: %w", err)
return nil, xerrors.Errorf("parse assistant message: %w", err)
}
// Mirrors SubmitToolResults.
afterMsgs, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
ChatID: chat.ID,
AfterID: lastAssistant.ID,
})
if err != nil {
return nil, xerrors.Errorf("get messages after assistant: %w", err)
}
handledCallIDs := make(map[string]bool)
for _, msg := range afterMsgs {
if msg.Role != database.ChatMessageRoleTool {
continue
}
msgParts, err := chatprompt.ParseContent(msg)
if err != nil {
continue
}
for _, mp := range msgParts {
if mp.Type == codersdk.ChatMessagePartTypeToolResult {
handledCallIDs[mp.ToolCallID] = true
}
}
}
// Collect dynamic tool calls that need synthetic results.
@@ -8189,6 +8359,9 @@ func insertSyntheticToolResultsTx(
if part.Type != codersdk.ChatMessagePartTypeToolCall || !dynamicToolNames[part.ToolName] {
continue
}
if handledCallIDs[part.ToolCallID] {
continue
}
resultPart := codersdk.ChatMessagePart{
Type: codersdk.ChatMessagePartTypeToolResult,
ToolCallID: part.ToolCallID,
@@ -8198,13 +8371,13 @@ func insertSyntheticToolResultsTx(
}
marshaled, marshalErr := chatprompt.MarshalParts([]codersdk.ChatMessagePart{resultPart})
if marshalErr != nil {
return xerrors.Errorf("marshal synthetic tool result: %w", marshalErr)
return nil, xerrors.Errorf("marshal synthetic tool result: %w", marshalErr)
}
resultContents = append(resultContents, marshaled)
}
if len(resultContents) == 0 {
return nil
return nil, nil
}
// Insert tool-result messages using the same pattern as
@@ -8238,11 +8411,12 @@ func insertSyntheticToolResultsTx(
params.ContentVersion[i] = chatprompt.CurrentContentVersion
params.Visibility[i] = database.ChatMessageVisibilityBoth
}
if _, err := store.InsertChatMessages(ctx, params); err != nil {
return xerrors.Errorf("insert synthetic tool results: %w", err)
inserted, err := store.InsertChatMessages(ctx, params)
if err != nil {
return nil, xerrors.Errorf("insert synthetic tool results: %w", err)
}
return nil
return inserted, nil
}
// parseDynamicToolNames unmarshals the dynamic tools JSON column
File diff suppressed because it is too large Load Diff
+61
View File
@@ -1,5 +1,15 @@
package chatd
import (
"context"
"github.com/sqlc-dev/pqtype"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/codersdk"
)
// WaitUntilIdleForTest waits for background chat work tracked by the server to
// finish without shutting the server down. Tests use this to assert final
// database state only after asynchronous chat processing has completed.
@@ -7,3 +17,54 @@ package chatd
func WaitUntilIdleForTest(server *Server) {
server.drainInflight()
}
// FinishActiveChatForTest exposes the unexported cleanup TX so tests
// can drive the post-run state machine deterministically. Returns the
// resulting chat, the promoted message (if any), the synthetic
// tool-result rows the cleanup TX inserted (if any), and the cleanup
// error. The lastError string is encoded into a structured payload
// the same way runChat does, so callers do not need to know about
// the structured-error wrapper.
func FinishActiveChatForTest(
ctx context.Context,
server *Server,
chat database.Chat,
status database.ChatStatus,
lastError string,
) (database.Chat, *database.ChatMessage, []database.ChatMessage, error) {
logger := server.logger.With(slog.F("chat_id", chat.ID))
var encoded pqtype.NullRawMessage
if lastError != "" {
var err error
encoded, err = encodeChatLastErrorPayload(&codersdk.ChatError{
Message: lastError,
})
if err != nil {
return database.Chat{}, nil, nil, err
}
}
result, err := server.finishActiveChat(ctx, logger, chat, status, encoded)
if err != nil {
return database.Chat{}, nil, nil, err
}
return result.updatedChat, result.promotedMessage, result.syntheticToolResults, nil
}
// RecoverStaleChatsForTest exposes the unexported stale-recovery loop
// so tests can assert the recovery state machine without waiting for
// the periodic ticker.
func RecoverStaleChatsForTest(ctx context.Context, server *Server) {
server.recoverStaleChats(ctx)
}
// InsertSyntheticToolResultsTxForTest exposes the unexported helper
// so tests can verify the dedup path against pre-existing tool
// results.
func InsertSyntheticToolResultsTxForTest(
ctx context.Context,
store database.Store,
chat database.Chat,
reason string,
) ([]database.ChatMessage, error) {
return insertSyntheticToolResultsTx(ctx, store, chat, reason)
}
+2 -3
View File
@@ -3207,11 +3207,10 @@ class ExperimentalApiMethods {
promoteChatQueuedMessage = async (
chatId: string,
queuedMessageId: number,
): Promise<TypesGen.ChatMessage> => {
const response = await this.axios.post<TypesGen.ChatMessage>(
): Promise<void> => {
await this.axios.post(
`/api/experimental/chats/${chatId}/queue/${queuedMessageId}/promote`,
);
return response.data;
};
getChatDiffContents = async (
@@ -1,6 +1,7 @@
import { act, renderHook } from "@testing-library/react";
import { createRef } from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { ChatQueuedMessage } from "#/api/typesGenerated";
import {
clearPersistedSidebarTabId,
draftInputStorageKeyPrefix,
@@ -8,6 +9,7 @@ import {
getPersistedSidebarTabId,
lastActiveSidebarTabStorageKeyPrefix,
restoreOptimisticRequestSnapshot,
runPromoteQueuedMessage,
savePersistedSidebarTabId,
submitEditAndScroll,
useConversationEditingState,
@@ -181,6 +183,78 @@ describe("restoreOptimisticRequestSnapshot", () => {
});
});
describe("runPromoteQueuedMessage", () => {
const makeQueuedMessage = (id: number, text: string, chatID = "chat-1") =>
({
id,
chat_id: chatID,
created_at: "2025-01-01T00:00:00Z",
content: [{ type: "text", text }],
}) as ChatQueuedMessage;
it("suppresses the promoted ID and removes it optimistically", async () => {
const store = createChatStore();
const a = makeQueuedMessage(1, "A");
const b = makeQueuedMessage(2, "B");
const c = makeQueuedMessage(3, "C");
store.setQueuedMessages([a, b, c]);
store.setChatStatus("running");
const promote = vi.fn(async (_id: number) => undefined);
const clearChatErrorReason = vi.fn();
const handleUsageLimitError = vi.fn();
await runPromoteQueuedMessage({
id: b.id,
store,
promoteQueuedMessage: promote,
agentId: "chat-1",
clearChatErrorReason,
handleUsageLimitError,
});
expect(promote).toHaveBeenCalledWith(b.id);
const snapshot = store.getSnapshot();
expect(snapshot.queuedMessages.map((m) => m.id)).toEqual([a.id, c.id]);
expect(snapshot.suppressedQueuedMessageIDs.has(b.id)).toBe(true);
expect(snapshot.chatStatus).toBe("pending");
});
it("rolls back queue and status, clears suppression, and rethrows on API error", async () => {
const store = createChatStore();
const a = makeQueuedMessage(1, "A");
const b = makeQueuedMessage(2, "B");
store.setQueuedMessages([a, b]);
store.setChatStatus("waiting");
const apiError = new Error("boom");
const promote = vi.fn(async (_id: number) => {
throw apiError;
});
const clearChatErrorReason = vi.fn();
const handleUsageLimitError = vi.fn();
await expect(
runPromoteQueuedMessage({
id: b.id,
store,
promoteQueuedMessage: promote,
agentId: "chat-1",
clearChatErrorReason,
handleUsageLimitError,
}),
).rejects.toBe(apiError);
expect(handleUsageLimitError).toHaveBeenCalledWith(apiError);
const snapshot = store.getSnapshot();
expect(snapshot.queuedMessages.map((m) => m.id)).toEqual([a.id, b.id]);
expect(snapshot.chatStatus).toBe("waiting");
expect(snapshot.suppressedQueuedMessageIDs.has(b.id)).toBe(false);
});
});
describe("useConversationEditingState", () => {
const chatID = "chat-abc-123";
const expectedKey = `${draftInputStorageKeyPrefix}${chatID}`;
+71 -24
View File
@@ -191,6 +191,68 @@ export const restoreOptimisticRequestSnapshot = (
});
};
/**
* Runs the optimistic queued-message promotion flow.
*
* The promote endpoint returns 202 Accepted with no message body, so the
* actual user message is delivered via SSE or the messages REST endpoint.
* Suppress the promoted ID so the transient reordered queue published by
* the running-case backend does not flash the message back into the
* visible queue. Roll back queue, status, and suppression on API error.
*
* @internal Exported for testing.
*/
export const runPromoteQueuedMessage = async (params: {
id: number;
store: Pick<
ChatStore,
| "batch"
| "clearStreamError"
| "clearStreamState"
| "getSnapshot"
| "setChatStatus"
| "setQueuedMessages"
| "setStreamError"
| "setStreamState"
| "suppressQueuedMessageID"
| "unsuppressQueuedMessageID"
>;
promoteQueuedMessage: (id: number) => Promise<void>;
agentId: string | undefined;
clearChatErrorReason: (chatID: string) => void;
handleUsageLimitError: (error: unknown) => void;
}): Promise<void> => {
const {
id,
store,
promoteQueuedMessage,
agentId,
clearChatErrorReason,
handleUsageLimitError,
} = params;
const previousSnapshot = store.getSnapshot();
store.batch(() => {
store.suppressQueuedMessageID(id);
store.setQueuedMessages(
previousSnapshot.queuedMessages.filter((message) => message.id !== id),
);
store.clearStreamState();
store.clearStreamError();
store.setChatStatus("pending");
});
if (agentId) {
clearChatErrorReason(agentId);
}
try {
await promoteQueuedMessage(id);
} catch (error) {
store.unsuppressQueuedMessageID(id);
restoreOptimisticRequestSnapshot(store, previousSnapshot);
handleUsageLimitError(error);
throw error;
}
};
export async function submitEditAndScroll({
editMessage,
editArgs,
@@ -1139,30 +1201,15 @@ const AgentChatPage: FC = () => {
}
};
const handlePromoteQueuedMessage = async (id: number) => {
const previousSnapshot = store.getSnapshot();
store.setQueuedMessages(
previousSnapshot.queuedMessages.filter((message) => message.id !== id),
);
store.clearStreamState();
if (agentId) {
clearChatErrorReason(agentId);
}
store.clearStreamError();
store.setChatStatus("pending");
try {
const promotedMessage = await promoteQueuedMessage(id);
// Insert the promoted message into the store and cache
// immediately so it appears in the timeline without
// waiting for the WebSocket to deliver it.
store.upsertDurableMessage(promotedMessage);
upsertCacheMessages([promotedMessage]);
} catch (error) {
restoreOptimisticRequestSnapshot(store, previousSnapshot);
handleUsageLimitError(error);
throw error;
}
};
const handlePromoteQueuedMessage = (id: number) =>
runPromoteQueuedMessage({
id,
store,
promoteQueuedMessage,
agentId,
clearChatErrorReason,
handleUsageLimitError,
});
const editing = useConversationEditingState({
chatID: agentId,
@@ -424,6 +424,74 @@ describe("setQueuedMessages", () => {
});
});
// ---------------------------------------------------------------------------
// suppressQueuedMessageID / applyAuthoritativeQueuedMessages
// ---------------------------------------------------------------------------
describe("suppressQueuedMessageID / applyAuthoritativeQueuedMessages", () => {
it("filters suppressed IDs from authoritative writes and auto-clears", () => {
const store = createChatStore();
const a = makeQueuedMessage(1, "A");
const b = makeQueuedMessage(2, "B");
const c = makeQueuedMessage(3, "C");
store.setQueuedMessages([a, b, c]);
store.suppressQueuedMessageID(b.id);
expect(store.getSnapshot().suppressedQueuedMessageIDs.has(b.id)).toBe(true);
// Transient reordered queue from the running-case backend
// must not surface the suppressed message.
store.applyAuthoritativeQueuedMessages([b, a, c]);
expect(
store.getSnapshot().queuedMessages.map((message) => message.id),
).toEqual([a.id, c.id]);
expect(store.getSnapshot().suppressedQueuedMessageIDs.has(b.id)).toBe(true);
store.applyAuthoritativeQueuedMessages([a, c]);
expect(store.getSnapshot().suppressedQueuedMessageIDs.has(b.id)).toBe(
false,
);
expect(
store.getSnapshot().queuedMessages.map((message) => message.id),
).toEqual([a.id, c.id]);
});
it("filters suppressed IDs from REST hydration via applyAuthoritativeQueuedMessages", () => {
const store = createChatStore();
const a = makeQueuedMessage(1, "A");
const b = makeQueuedMessage(2, "B");
const c = makeQueuedMessage(3, "C");
store.suppressQueuedMessageID(b.id);
// REST hydration delivers the unfiltered queue [B, A, C].
store.applyAuthoritativeQueuedMessages([b, a, c]);
expect(
store.getSnapshot().queuedMessages.map((message) => message.id),
).toEqual([a.id, c.id]);
});
it("unsuppressQueuedMessageID removes IDs from the suppression set", () => {
const store = createChatStore();
store.suppressQueuedMessageID(42);
expect(store.getSnapshot().suppressedQueuedMessageIDs.has(42)).toBe(true);
store.unsuppressQueuedMessageID(42);
expect(store.getSnapshot().suppressedQueuedMessageIDs.has(42)).toBe(false);
});
it("setQueuedMessages does not auto-clear suppression", () => {
const store = createChatStore();
const a = makeQueuedMessage(1, "A");
store.suppressQueuedMessageID(99);
// setQueuedMessages is the optimistic path: it must not
// touch the suppression set, otherwise the optimistic write
// would lift suppression before the authoritative reordered
// queue arrives.
store.setQueuedMessages([a]);
expect(store.getSnapshot().suppressedQueuedMessageIDs.has(99)).toBe(true);
});
});
// ---------------------------------------------------------------------------
// clearStreamState
// ---------------------------------------------------------------------------
@@ -153,6 +153,11 @@ export type ChatStoreState = {
retryState: RetryState | null;
reconnectState: ReconnectState | null;
queuedMessages: readonly TypesGen.ChatQueuedMessage[];
// Hides queued IDs from the visible queue while the backend is
// in a transient state that would briefly include them. Used by
// the running-case promote, where the backend reorders the
// queued message to the front before auto-promoting it.
suppressedQueuedMessageIDs: ReadonlySet<number>;
subagentStatusOverrides: Map<string, TypesGen.ChatStatus>;
};
@@ -173,6 +178,16 @@ export type ChatStore = {
setQueuedMessages: (
queuedMessages: readonly TypesGen.ChatQueuedMessage[] | undefined,
) => void;
// Server-truthful queue snapshot, filtered through the
// suppression set. Use for SSE queue_update and REST hydration;
// optimistic writes go through setQueuedMessages so they don't
// lift suppression.
applyAuthoritativeQueuedMessages: (
queuedMessages: readonly TypesGen.ChatQueuedMessage[] | undefined,
) => void;
suppressQueuedMessageID: (id: number) => void;
unsuppressQueuedMessageID: (id: number) => void;
clearSuppressedQueuedMessageIDs: () => void;
setChatStatus: (status: TypesGen.ChatStatus | null) => void;
setStreamState: (streamState: StreamState | null) => void;
setStreamError: (reason: ChatDetailError | null) => void;
@@ -199,6 +214,7 @@ const createInitialState = (): ChatStoreState => ({
retryState: null,
reconnectState: null,
queuedMessages: [],
suppressedQueuedMessageIDs: new Set(),
subagentStatusOverrides: new Map(),
});
@@ -404,6 +420,73 @@ export const createChatStore = (): ChatStore => {
return { ...current, queuedMessages: nextQueuedMessages };
});
},
applyAuthoritativeQueuedMessages: (queuedMessages) => {
const incoming = queuedMessages ?? [];
setState((current) => {
let nextSuppressed = current.suppressedQueuedMessageIDs;
if (current.suppressedQueuedMessageIDs.size > 0) {
const incomingIDs = new Set(incoming.map((message) => message.id));
let copy: Set<number> | null = null;
for (const id of current.suppressedQueuedMessageIDs) {
if (!incomingIDs.has(id)) {
if (!copy) {
copy = new Set(current.suppressedQueuedMessageIDs);
}
copy.delete(id);
}
}
if (copy) {
nextSuppressed = copy;
}
}
const filtered =
nextSuppressed.size === 0
? incoming
: incoming.filter((message) => !nextSuppressed.has(message.id));
const sameQueue = chatQueuedMessagesEqualByID(
current.queuedMessages,
filtered,
);
const sameSuppressed =
nextSuppressed === current.suppressedQueuedMessageIDs;
if (sameQueue && sameSuppressed) {
return current;
}
return {
...current,
queuedMessages: sameQueue ? current.queuedMessages : filtered,
suppressedQueuedMessageIDs: nextSuppressed,
};
});
},
suppressQueuedMessageID: (id) => {
setState((current) => {
if (current.suppressedQueuedMessageIDs.has(id)) {
return current;
}
const next = new Set(current.suppressedQueuedMessageIDs);
next.add(id);
return { ...current, suppressedQueuedMessageIDs: next };
});
},
unsuppressQueuedMessageID: (id) => {
setState((current) => {
if (!current.suppressedQueuedMessageIDs.has(id)) {
return current;
}
const next = new Set(current.suppressedQueuedMessageIDs);
next.delete(id);
return { ...current, suppressedQueuedMessageIDs: next };
});
},
clearSuppressedQueuedMessageIDs: () => {
setState((current) => {
if (current.suppressedQueuedMessageIDs.size === 0) {
return current;
}
return { ...current, suppressedQueuedMessageIDs: new Set() };
});
},
setChatStatus: (status) => {
if (state.chatStatus === status) {
return;
@@ -237,6 +237,10 @@ export const useChatStore = (
wsQueueUpdateReceivedRef.current = false;
wsStatusReceivedRef.current = false;
store.setQueuedMessages([]);
// Suppression entries are scoped to the current chat; clear
// them on chat change so a stale promote suppression doesn't
// hide queued messages in another chat.
store.clearSuppressedQueuedMessageIDs();
if (!chatID) {
return;
}
@@ -258,7 +262,7 @@ export const useChatStore = (
return;
}
queuedMessagesHydratedChatIDRef.current = chatID;
store.setQueuedMessages(chatQueuedMessages);
store.applyAuthoritativeQueuedMessages(chatQueuedMessages);
}, [chatMessagesData, chatID, chatQueuedMessages, store]);
useEffect(() => {
@@ -473,7 +477,9 @@ export const useChatStore = (
continue;
}
wsQueueUpdateReceivedRef.current = true;
store.setQueuedMessages(streamEvent.queued_messages);
store.applyAuthoritativeQueuedMessages(
streamEvent.queued_messages,
);
updateChatQueuedMessages(streamEvent.queued_messages);
continue;
case "status": {