mirror of
https://github.com/coder/coder.git
synced 2026-09-21 20:51:01 +08:00
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:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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}`;
|
||||
|
||||
@@ -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": {
|
||||
|
||||
Reference in New Issue
Block a user