From b86647e90ccfc4f7adbf7124a8e73c087e08b93e Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Fri, 10 Jul 2026 20:25:41 +0000 Subject: [PATCH] fix: render chat stream deterministically via history versions --- coderd/apidoc/docs.go | 12 + coderd/apidoc/swagger.json | 12 + coderd/database/db2sdk/db2sdk.go | 2 + coderd/database/db2sdk/db2sdk_test.go | 2 + coderd/x/chatd/ARCHITECTURE.md | 11 +- coderd/x/chatd/stream_loop.go | 6 +- coderd/x/chatd/stream_loop_internal_test.go | 93 +++++++ codersdk/chats.go | 6 +- docs/reference/api/chats.md | 32 +++ docs/reference/api/schemas.md | 20 +- site/src/api/typesGenerated.ts | 4 + .../pages/AgentsPage/AgentChatPage.test.ts | 86 ++++-- site/src/pages/AgentsPage/AgentChatPage.tsx | 60 ++-- .../AgentsPage/components/AgentChatInput.tsx | 3 + .../chatStore.createStore.test.ts | 218 ++++++++++----- .../ChatConversation/chatStore.test.tsx | 263 +++++++++++++----- .../components/ChatConversation/chatStore.ts | 217 ++++++++++----- .../ChatConversation/useChatStore.ts | 86 +++--- .../AgentsPage/components/ChatPageContent.tsx | 3 + .../components/QueuedMessagesList.stories.tsx | 12 + .../components/QueuedMessagesList.tsx | 21 +- 21 files changed, 822 insertions(+), 347 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index d92c1cc16aa..4d4c38706c3 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -16729,10 +16729,16 @@ const docTemplate = `{ "$ref": "#/definitions/codersdk.ChatFileMetadata" } }, + "generation_attempt": { + "type": "integer" + }, "has_unread": { "description": "HasUnread is true when assistant messages exist beyond\nthe owner's read cursor, which updates on stream\nconnect and disconnect.", "type": "boolean" }, + "history_version": { + "type": "integer" + }, "id": { "type": "string", "format": "uuid" @@ -17841,6 +17847,12 @@ const docTemplate = `{ "codersdk.ChatStreamStatus": { "type": "object", "properties": { + "generation_attempt": { + "type": "integer" + }, + "history_version": { + "type": "integer" + }, "status": { "$ref": "#/definitions/codersdk.ChatStatus" } diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index b507c2db77a..0fdf8c47000 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -15026,10 +15026,16 @@ "$ref": "#/definitions/codersdk.ChatFileMetadata" } }, + "generation_attempt": { + "type": "integer" + }, "has_unread": { "description": "HasUnread is true when assistant messages exist beyond\nthe owner's read cursor, which updates on stream\nconnect and disconnect.", "type": "boolean" }, + "history_version": { + "type": "integer" + }, "id": { "type": "string", "format": "uuid" @@ -16090,6 +16096,12 @@ "codersdk.ChatStreamStatus": { "type": "object", "properties": { + "generation_attempt": { + "type": "integer" + }, + "history_version": { + "type": "integer" + }, "status": { "$ref": "#/definitions/codersdk.ChatStatus" } diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 4754bbbe950..ed25011e7a7 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1700,6 +1700,8 @@ func Chat(c database.Chat, diffStatus *database.ChatDiffStatus, files []database LastModelConfigID: c.LastModelConfigID, Title: c.Title, Status: codersdk.ChatStatus(c.Status), + HistoryVersion: c.HistoryVersion, + GenerationAttempt: c.GenerationAttempt, Archived: c.Archived, Shared: len(c.UserACL) > 0 || len(c.GroupACL) > 0, PinOrder: c.PinOrder, diff --git a/coderd/database/db2sdk/db2sdk_test.go b/coderd/database/db2sdk/db2sdk_test.go index 44d43442d20..2631b843f53 100644 --- a/coderd/database/db2sdk/db2sdk_test.go +++ b/coderd/database/db2sdk/db2sdk_test.go @@ -710,6 +710,8 @@ func TestChat_AllFieldsPopulated(t *testing.T) { LastReasoningEffort: database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortHigh, Valid: true}, Title: "all-fields-test", Status: database.ChatStatusRunning, + HistoryVersion: 7, + GenerationAttempt: 2, ClientType: database.ChatClientTypeUi, LastError: pqtype.NullRawMessage{RawMessage: lastErrorRaw, Valid: true}, LastTurnSummary: sql.NullString{String: "turn completed", Valid: true}, diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index d1519231ff6..4ddb6d02867 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -869,10 +869,9 @@ The stream loop powers the `GET /api/experimental/chats/{chat}/stream` endpoint. The following chat stream events, delivered to the client over WebSocket, are supported: -- `message_part`: a streaming message part emitted by the chat worker. - - compared to the current implementation, these should additionally include the `history_version` and `generation_attempt` fields, so a client knows which episode a message part comes from +- `message_part`: a streaming message part emitted by the chat worker. Its payload includes `history_version`, `generation_attempt`, and `seq`, which identify the part's episode and position. - `message`: a committed chat message present in the database. -- `status`: the chat's status. +- `status`: the chat's status plus the current `history_version` and `generation_attempt`. Clients can use this tuple as a lower bound when rejecting stale parts. - `error`: the chat's persisted error payload. - `queue_update`: the full current queued-message list. - `action_required`: a dynamic tool call was issued by the chat worker, the client must execute it and submit the result. @@ -880,6 +879,10 @@ The following chat stream events, delivered to the client over WebSocket, are su - `preview_reset`: a reset of the stream's preview state (message parts), emitted when the worker's LLM call fails mid-way for whatever reason. - `history_reset`: a reset of the stream's history state (committed messages), emitted when the message history is edited and some messages are removed from the history. +Database-derived events from one `Sync` are emitted in the application order described below, including `message` before `queue_update` before `status`. Relayed `message_part` events are multiplexed independently, so there is no total ordering between parts and database-derived events. + +The single-chat REST response exposes the same `history_version` and `generation_attempt` tuple as status events. A pure status transition advances `snapshot_version` but keeps this tuple stable, so it must not fence out parts from the active episode. A history change advances `history_version` and resets `generation_attempt` to `0`; the next generation starts at the same history version with a higher attempt. + ## Endpoint lifecycle When a client connects, the endpoint: @@ -1079,7 +1082,7 @@ Status sync happens inside `Sync`. Flow: 1. Compare database status to local status. -2. If they differ, emit `status`. +2. If they differ, emit `status` with `db.status`, `db.history_version`, and `db.generation_attempt`. ### Error synchronization diff --git a/coderd/x/chatd/stream_loop.go b/coderd/x/chatd/stream_loop.go index 5004e8a490e..fabe9f5285e 100644 --- a/coderd/x/chatd/stream_loop.go +++ b/coderd/x/chatd/stream_loop.go @@ -252,7 +252,11 @@ func (l *streamLoop) applyDBSnapshot(snapshot streamDBSnapshot) []codersdk.ChatS events = append(events, codersdk.ChatStreamEvent{ Type: codersdk.ChatStreamEventTypeStatus, ChatID: l.chatID, - Status: &codersdk.ChatStreamStatus{Status: codersdk.ChatStatus(chat.Status)}, + Status: &codersdk.ChatStreamStatus{ + Status: codersdk.ChatStatus(chat.Status), + HistoryVersion: chat.HistoryVersion, + GenerationAttempt: chat.GenerationAttempt, + }, }) } diff --git a/coderd/x/chatd/stream_loop_internal_test.go b/coderd/x/chatd/stream_loop_internal_test.go index eebd6d0c978..5110a2eb15f 100644 --- a/coderd/x/chatd/stream_loop_internal_test.go +++ b/coderd/x/chatd/stream_loop_internal_test.go @@ -213,6 +213,8 @@ func TestStreamLoopQueueStatusRetryErrorActionRequiredAndPreviewReset(t *testing codersdk.ChatStreamEventTypeRetry, codersdk.ChatStreamEventTypePreviewReset, ) + require.Equal(t, int64(2), events[1].Status.HistoryVersion) + require.Equal(t, int64(2), events[1].Status.GenerationAttempt) require.Equal(t, chatError.Message, events[2].Error.Message) require.Equal(t, retry.Attempt, events[3].Retry.Attempt) @@ -231,9 +233,100 @@ func TestStreamLoopQueueStatusRetryErrorActionRequiredAndPreviewReset(t *testing codersdk.ChatStreamEventTypeActionRequired, codersdk.ChatStreamEventTypePreviewReset, ) + require.Equal(t, int64(1), actionEvents[0].Status.HistoryVersion) + require.Zero(t, actionEvents[0].Status.GenerationAttempt) require.Equal(t, "call-1", actionEvents[1].ActionRequired.ToolCalls[0].ToolCallID) } +func TestStreamLoopStatusCarriesPartConsistentEpisode(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + loop := newStreamLoop(database.Chat{ID: chatID}, nil, slogtest.Make(t, nil), 0) + loop.state.snapshotVersion = 10 + loop.state.historyVersion = 7 + loop.state.generationAttempt = 1 + loop.state.status = database.ChatStatusRunning + + // A pure status transition must keep the active part episode even though + // snapshot_version advances. + events := loop.applyDBSnapshot(streamDBSnapshot{chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusInterrupting, + SnapshotVersion: 11, + HistoryVersion: 7, + GenerationAttempt: 1, + }}) + requireEventTypes(t, events, codersdk.ChatStreamEventTypeStatus) + require.Equal(t, int64(7), events[0].Status.HistoryVersion) + require.Equal(t, int64(1), events[0].Status.GenerationAttempt) + _, accepted, err := loop.part(StreamPart{ + HistoryVersion: 7, + GenerationAttempt: 1, + Seq: 1, + Part: codersdk.ChatMessageText("draining"), + }) + require.NoError(t, err) + require.True(t, accepted) + + // History changes reset the attempt; retries advance it within that history. + events = loop.applyDBSnapshot(streamDBSnapshot{chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusRunning, + SnapshotVersion: 12, + HistoryVersion: 12, + GenerationAttempt: 0, + }}) + requireEventTypes(t, events, + codersdk.ChatStreamEventTypeStatus, + codersdk.ChatStreamEventTypePreviewReset, + ) + require.Equal(t, int64(12), events[0].Status.HistoryVersion) + require.Zero(t, events[0].Status.GenerationAttempt) + + events = loop.applyDBSnapshot(streamDBSnapshot{chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusRunning, + SnapshotVersion: 13, + HistoryVersion: 12, + GenerationAttempt: 1, + }}) + requireEventTypes(t, events, codersdk.ChatStreamEventTypePreviewReset) + _, accepted, err = loop.part(StreamPart{ + HistoryVersion: 12, + GenerationAttempt: 1, + Seq: 1, + Part: codersdk.ChatMessageText("new turn"), + }) + require.NoError(t, err) + require.True(t, accepted) + + events = loop.applyDBSnapshot(streamDBSnapshot{chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusRunning, + SnapshotVersion: 14, + HistoryVersion: 12, + GenerationAttempt: 2, + }}) + requireEventTypes(t, events, codersdk.ChatStreamEventTypePreviewReset) + _, accepted, err = loop.part(StreamPart{ + HistoryVersion: 12, + GenerationAttempt: 1, + Seq: 2, + Part: codersdk.ChatMessageText("stale retry"), + }) + require.NoError(t, err) + require.False(t, accepted) + _, accepted, err = loop.part(StreamPart{ + HistoryVersion: 12, + GenerationAttempt: 2, + Seq: 1, + Part: codersdk.ChatMessageText("retry"), + }) + require.NoError(t, err) + require.True(t, accepted) +} + func TestStreamLoopActionRequiredFromHistory(t *testing.T) { t.Parallel() diff --git a/codersdk/chats.go b/codersdk/chats.go index be2cd99769a..fc9a728b072 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -119,6 +119,8 @@ type Chat struct { LastReasoningEffort *string `json:"last_reasoning_effort,omitempty"` Title string `json:"title"` Status ChatStatus `json:"status"` + HistoryVersion int64 `json:"history_version,omitempty"` + GenerationAttempt int64 `json:"generation_attempt,omitempty"` PlanMode ChatPlanMode `json:"plan_mode,omitempty"` LastError *ChatError `json:"last_error,omitempty"` LastTurnSummary *string `json:"last_turn_summary"` @@ -1649,7 +1651,9 @@ type ChatStreamMessagePart struct { // ChatStreamStatus represents an updated chat status. type ChatStreamStatus struct { - Status ChatStatus `json:"status"` + Status ChatStatus `json:"status"` + HistoryVersion int64 `json:"history_version,omitempty"` + GenerationAttempt int64 `json:"generation_attempt,omitempty"` } // ChatErrorKind classifies chat errors for consistent client rendering. diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index a9bd9bd0426..8c1eae7110e 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -92,7 +92,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -193,7 +195,9 @@ Status Code **200** | `»» name` | string | false | | | | `»» organization_id` | string(uuid) | false | | | | `»» owner_id` | string(uuid) | false | | | +| `» generation_attempt` | integer | false | | | | `» has_unread` | boolean | false | | Has unread is true when assistant messages exist beyond the owner's read cursor, which updates on stream connect and disconnect. | +| `» history_version` | integer | false | | | | `» id` | string(uuid) | false | | | | `» labels` | object | false | | | | `»» [any property]` | string | false | | | @@ -367,7 +371,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -460,7 +466,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -710,7 +718,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -857,7 +867,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -950,7 +962,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -1134,7 +1148,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -1227,7 +1243,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -1409,7 +1427,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -1502,7 +1522,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -2253,7 +2275,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -2346,7 +2370,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -2669,6 +2695,8 @@ Experimental: this endpoint is subject to change. "status_code": 0 }, "status": { + "generation_attempt": 0, + "history_version": 0, "status": "waiting" }, "type": "message_part" @@ -2853,7 +2881,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -2946,7 +2976,9 @@ Experimental: this endpoint is subject to change. "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 7d96c4221bb..0ae685000ca 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -2090,7 +2090,9 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -2183,7 +2185,9 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", @@ -2235,7 +2239,9 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in | `created_at` | string | false | | | | `diff_status` | [codersdk.ChatDiffStatus](#codersdkchatdiffstatus) | false | | | | `files` | array of [codersdk.ChatFileMetadata](#codersdkchatfilemetadata) | false | | | +| `generation_attempt` | integer | false | | | | `has_unread` | boolean | false | | Has unread is true when assistant messages exist beyond the owner's read cursor, which updates on stream connect and disconnect. | +| `history_version` | integer | false | | | | `id` | string | false | | | | `labels` | object | false | | | | » `[any property]` | string | false | | | @@ -3683,6 +3689,8 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "status_code": 0 }, "status": { + "generation_attempt": 0, + "history_version": 0, "status": "waiting" }, "type": "message_part" @@ -3830,15 +3838,19 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in ```json { + "generation_attempt": 0, + "history_version": 0, "status": "waiting" } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|----------|--------------------------------------------|----------|--------------|-------------| -| `status` | [codersdk.ChatStatus](#codersdkchatstatus) | false | | | +| Name | Type | Required | Restrictions | Description | +|----------------------|--------------------------------------------|----------|--------------|-------------| +| `generation_attempt` | integer | false | | | +| `history_version` | integer | false | | | +| `status` | [codersdk.ChatStatus](#codersdkchatstatus) | false | | | ## codersdk.ChatStreamToolCall @@ -3968,7 +3980,9 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" } ], + "generation_attempt": 0, "has_unread": true, + "history_version": 0, "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", "labels": { "property1": "string", diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index fe7c731a48c..1ad6c51c4e1 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1610,6 +1610,8 @@ export interface Chat { readonly last_reasoning_effort?: string; readonly title: string; readonly status: ChatStatus; + readonly history_version?: number; + readonly generation_attempt?: number; readonly plan_mode?: ChatPlanMode; readonly last_error?: ChatError; readonly last_turn_summary: string | null; @@ -3126,6 +3128,8 @@ export interface ChatStreamRetry { */ export interface ChatStreamStatus { readonly status: ChatStatus; + readonly history_version?: number; + readonly generation_attempt?: number; } // From codersdk/chats.go diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 945a5724e6e..51c46e02f5d 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -186,7 +186,9 @@ describe("restoreOptimisticRequestSnapshot", () => { }, ]); store.setChatStatus("running"); - store.applyMessagePart({ type: "text", text: "partial response" }); + store.applyMessagePart({ + part: { type: "text", text: "partial response" }, + }); store.setStreamError({ kind: "generic", message: "old error" }); const previousSnapshot = store.getSnapshot(); @@ -221,66 +223,92 @@ describe("runPromoteQueuedMessage", () => { content: [{ type: "text", text }], }); - it("suppresses the promoted ID and removes it optimistically", async () => { + it("keeps server-derived state and marks the row until the queue removes it", async () => { const store = createChatStore(); const a = buildQueuedMessage(1, "A"); const b = buildQueuedMessage(2, "B"); - const c = buildQueuedMessage(3, "C"); - store.setQueuedMessages([a, b, c]); - store.setChatStatus("running"); - - const promote = vi.fn(async (_id: number) => undefined); - const clearChatErrorReason = vi.fn(); + store.setQueuedMessages([a, b]); + store.setChatStatus("interrupting"); + store.applyMessagePart({ + part: { type: "text", text: "partial response" }, + history_version: 4, + generation_attempt: 1, + seq: 1, + }); + store.setStreamError({ kind: "generic", message: "old error" }); + const before = store.getSnapshot(); + const deferred = createDeferred(); + const promote = vi.fn(() => deferred.promise); const handleUsageLimitError = vi.fn(); - await runPromoteQueuedMessage({ + const promotion = 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("running"); + expect(snapshot.queuedMessages).toBe(before.queuedMessages); + expect(snapshot.chatStatus).toBe(before.chatStatus); + expect(snapshot.streamState).toBe(before.streamState); + expect(snapshot.streamError).toBeNull(); + expect(snapshot.promoteInFlightIDs.has(b.id)).toBe(true); + + deferred.resolve(); + await promotion; + expect(store.getSnapshot().promoteInFlightIDs.has(b.id)).toBe(true); + + store.setQueuedMessages([a]); + expect(store.getSnapshot().promoteInFlightIDs.has(b.id)).toBe(false); }); - it("rolls back queue and status, clears suppression, and rethrows on API error", async () => { + it("clears the marker, reports the error, and preserves chat state", async () => { const store = createChatStore(); - const a = buildQueuedMessage(1, "A"); - const b = buildQueuedMessage(2, "B"); - store.setQueuedMessages([a, b]); + const queued = buildQueuedMessage(2, "B"); + store.setQueuedMessages([queued]); store.setChatStatus("waiting"); - + const before = store.getSnapshot(); const apiError = new Error("boom"); - const promote = vi.fn(async (_id: number) => { + const promote = vi.fn(async () => { throw apiError; }); - const clearChatErrorReason = vi.fn(); const handleUsageLimitError = vi.fn(); await expect( runPromoteQueuedMessage({ - id: b.id, + id: queued.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); + expect(snapshot.queuedMessages).toBe(before.queuedMessages); + expect(snapshot.chatStatus).toBe(before.chatStatus); + expect(snapshot.promoteInFlightIDs.has(queued.id)).toBe(false); + }); + + it("ignores duplicate requests while the same promotion is in flight", async () => { + const store = createChatStore(); + const deferred = createDeferred(); + const promote = vi.fn(() => deferred.promise); + const params = { + id: 7, + store, + promoteQueuedMessage: promote, + handleUsageLimitError: vi.fn(), + }; + + const first = runPromoteQueuedMessage(params); + await runPromoteQueuedMessage(params); + expect(promote).toHaveBeenCalledTimes(1); + + deferred.resolve(); + await first; }); }); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 2bcd8b52938..ce889303ea7 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -154,13 +154,8 @@ 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. + * Requests queued-message promotion without fabricating server state. + * The row remains visible until an authoritative queue update removes it. * * @internal Exported for testing. */ @@ -168,49 +163,26 @@ export const runPromoteQueuedMessage = async (params: { id: number; store: Pick< ChatStore, - | "batch" + | "clearPromoteInFlight" | "clearStreamError" - | "clearStreamState" | "getSnapshot" - | "setChatStatus" - | "setQueuedMessages" - | "setStreamError" - | "setStreamState" - | "suppressQueuedMessageID" - | "unsuppressQueuedMessageID" + | "markPromoteInFlight" >; promoteQueuedMessage: (id: number) => Promise; - agentId: string | undefined; - clearChatErrorReason: (chatID: string) => void; handleUsageLimitError: (error: unknown) => void; }): Promise => { - 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("running"); - }); - if (agentId) { - clearChatErrorReason(agentId); + const { id, store, promoteQueuedMessage, handleUsageLimitError } = params; + if (store.getSnapshot().promoteInFlightIDs.has(id)) { + return; } + store.markPromoteInFlight(id); + store.clearStreamError(); try { await promoteQueuedMessage(id); } catch (error) { - store.unsuppressQueuedMessageID(id); - restoreOptimisticRequestSnapshot(store, previousSnapshot); + store.clearPromoteInFlight(id); handleUsageLimitError(error); + toast.error(getErrorMessage(error, "Failed to send queued message.")); throw error; } }; @@ -1254,15 +1226,17 @@ const AgentChatPage: FC = () => { } }; - const handlePromoteQueuedMessage = (id: number) => - runPromoteQueuedMessage({ + const handlePromoteQueuedMessage = (id: number) => { + if (agentId) { + clearChatErrorReason(agentId); + } + return runPromoteQueuedMessage({ id, store, promoteQueuedMessage, - agentId, - clearChatErrorReason, handleUsageLimitError, }); + }; const editing = useConversationEditingState({ chatID: agentId, diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 98675297337..28048e6df0e 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -143,6 +143,7 @@ interface AgentChatInputProps { isWorkspaceLoading?: boolean; // Queued user messages rendered above the textarea. queuedMessages?: readonly ChatQueuedMessage[]; + promoteInFlightIDs?: ReadonlySet; onDeleteQueuedMessage?: (id: number) => Promise | void; onPromoteQueuedMessage?: (id: number) => Promise | void; // Queue editing state, owned by the parent. @@ -370,6 +371,7 @@ export const AgentChatInput: FC = ({ chatOrganizationId, isWorkspaceLoading, queuedMessages = [], + promoteInFlightIDs, onDeleteQueuedMessage, onPromoteQueuedMessage, editingQueuedMessageID = null, @@ -1064,6 +1066,7 @@ export const AgentChatInput: FC = ({ {queuedMessages.length > 0 && ( { if (id === editingQueuedMessageID) { onCancelQueueEdit?.(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts index 02ea1467792..53e08fab53b 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts @@ -34,6 +34,20 @@ const makeQueuedMessage = ( content: [{ type: "text", text }], }) as TypesGen.ChatQueuedMessage; +const makeStreamPart = ( + text: string, + historyVersion?: number, + generationAttempt?: number, + seq?: number, +): TypesGen.ChatStreamMessagePart => ({ + part: { type: "text", text }, + ...(historyVersion === undefined ? {} : { history_version: historyVersion }), + ...(generationAttempt === undefined + ? {} + : { generation_attempt: generationAttempt }), + ...(seq === undefined ? {} : { seq }), +}); + // --------------------------------------------------------------------------- // replaceMessages // --------------------------------------------------------------------------- @@ -197,7 +211,7 @@ describe("setChatStatus", () => { describe("setStreamState", () => { it("does not notify when setting the same stream state reference", () => { const store = createChatStore(); - store.applyMessagePart({ type: "text", text: "hello" }); + store.applyMessagePart(makeStreamPart("hello")); const streamState = store.getSnapshot().streamState; expect(streamState).not.toBeNull(); @@ -425,70 +439,37 @@ describe("setQueuedMessages", () => { }); // --------------------------------------------------------------------------- -// suppressQueuedMessageID / applyAuthoritativeQueuedMessages +// queued-message promotion state // --------------------------------------------------------------------------- -describe("suppressQueuedMessageID / applyAuthoritativeQueuedMessages", () => { - it("filters suppressed IDs from authoritative writes and auto-clears", () => { +describe("queued-message promotion state", () => { + it("keeps a promoting ID through a reordered queue and clears it when removed", () => { 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.markPromoteInFlight(b.id); + store.setQueuedMessages([b, a, c]); - 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]); - }); + ).toEqual([b.id, a.id, c.id]); + expect(store.getSnapshot().promoteInFlightIDs.has(b.id)).toBe(true); - 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.setQueuedMessages([a, 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]); + expect(store.getSnapshot().promoteInFlightIDs.has(b.id)).toBe(false); }); - it("unsuppressQueuedMessageID removes IDs from the suppression set", () => { + it("can clear a promotion marker after an API error", () => { 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.markPromoteInFlight(42); + store.clearPromoteInFlight(42); - 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); + expect(store.getSnapshot().promoteInFlightIDs.size).toBe(0); }); }); @@ -500,7 +481,7 @@ describe("clearStreamState", () => { it("clears stream state to null", () => { const store = createChatStore(); // Build up some stream state via applyMessagePart. - store.applyMessagePart({ type: "text", text: "hello" }); + store.applyMessagePart(makeStreamPart("hello")); expect(store.getSnapshot().streamState).not.toBeNull(); store.clearStreamState(); @@ -508,6 +489,21 @@ describe("clearStreamState", () => { expect(store.getSnapshot().streamState).toBeNull(); }); + it("preserves episode fencing when clearing visual stream state", () => { + const store = createChatStore(); + store.applyMessagePart(makeStreamPart("hello", 2, 1, 1)); + + store.clearStreamState(); + + const state = store.getSnapshot(); + expect(state.streamState).toBeNull(); + expect(state.streamEpisode).toEqual({ + historyVersion: 2, + generationAttempt: 1, + }); + expect(state.lastStreamPartSeq).toBe(1); + }); + it("is a no-op when stream state is already null", () => { const store = createChatStore(); @@ -526,32 +522,94 @@ describe("clearStreamState", () => { // --------------------------------------------------------------------------- describe("applyMessagePart / applyMessageParts", () => { - it("creates stream state from a text part", () => { + it("accepts legacy parts only before versioned context exists", () => { + const store = createChatStore(); + + store.applyMessagePart(makeStreamPart("legacy")); + store.applyMessagePart(makeStreamPart(" versioned", 1, 1, 1)); + store.applyMessagePart(makeStreamPart(" ignored")); + + expect(store.getSnapshot().streamState?.blocks).toEqual([ + { type: "response", text: " versioned" }, + ]); + }); + + it("adopts a newer episode before any status event", () => { + const store = createChatStore(); + store.applyMessagePart(makeStreamPart("old", 1, 1, 1)); + + store.applyMessagePart(makeStreamPart("new", 2, 1, 1)); + + expect(store.getSnapshot().streamEpisode).toEqual({ + historyVersion: 2, + generationAttempt: 1, + }); + expect(store.getSnapshot().streamState?.blocks).toEqual([ + { type: "response", text: "new" }, + ]); + }); + + it("drops parts older than the server floor or rendered episode", () => { const store = createChatStore(); + store.applyMessagePart(makeStreamPart("current", 4, 2, 1)); + store.updateServerEpisodeFloor(5, 0); + + store.applyMessageParts([ + makeStreamPart(" old history", 4, 3, 1), + makeStreamPart(" old attempt", 4, 1, 2), + ]); + + expect(store.getSnapshot().streamState?.blocks).toEqual([ + { type: "response", text: "current" }, + ]); + }); - store.applyMessagePart({ type: "text", text: "hello" }); + it("replaces the stream for a higher generation attempt", () => { + const store = createChatStore(); + store.applyMessageParts([ + makeStreamPart("first", 7, 1, 1), + makeStreamPart(" attempt", 7, 1, 2), + makeStreamPart("retry", 7, 2, 1), + ]); + expect(store.getSnapshot().streamEpisode).toEqual({ + historyVersion: 7, + generationAttempt: 2, + }); expect(store.getSnapshot().streamState?.blocks).toEqual([ - { type: "response", text: "hello" }, + { type: "response", text: "retry" }, ]); }); - it("appends to existing stream state", () => { + it("accepts equal and greater tuples while rejecting parts below the floor", () => { const store = createChatStore(); - store.applyMessagePart({ type: "text", text: "hello" }); - store.applyMessagePart({ type: "text", text: " world" }); + store.updateServerEpisodeFloor(9, 0); + + store.applyMessagePart(makeStreamPart("stale", 8, 4, 1)); + store.applyMessagePart(makeStreamPart("current", 9, 1, 1)); + store.updateServerEpisodeFloor(9, 1); + store.applyMessagePart(makeStreamPart(" equal", 9, 1, 2)); expect(store.getSnapshot().streamState?.blocks).toEqual([ - { type: "response", text: "hello world" }, + { type: "response", text: "current equal" }, ]); }); - it("applies multiple parts in a single batch", () => { + it("deduplicates by sequence and rebuilds the same episode after reconnect", () => { const store = createChatStore(); + store.applyMessageParts([ + makeStreamPart("one", 3, 1, 1), + makeStreamPart(" two", 3, 1, 2), + makeStreamPart(" duplicate", 3, 1, 2), + ]); + expect(store.getSnapshot().streamState?.blocks).toEqual([ + { type: "response", text: "one two" }, + ]); + store.resetTransportReplayState(); store.applyMessageParts([ - { type: "text", text: "one" }, - { type: "text", text: " two" }, + makeStreamPart("one", 3, 1, 1), + makeStreamPart(" two", 3, 1, 2), ]); expect(store.getSnapshot().streamState?.blocks).toEqual([ @@ -559,13 +617,26 @@ describe("applyMessagePart / applyMessageParts", () => { ]); }); - it("is a no-op for an empty parts array", () => { + it("does not advance across a sequence gap", () => { const store = createChatStore(); + store.applyMessageParts([ + makeStreamPart("one", 3, 1, 1), + makeStreamPart(" gap", 3, 1, 3), + makeStreamPart(" two", 3, 1, 2), + ]); + + expect(store.getSnapshot().streamState?.blocks).toEqual([ + { type: "response", text: "one two" }, + ]); + }); + it("is a no-op for an empty parts array", () => { + const store = createChatStore(); let notified = false; store.subscribe(() => { notified = true; }); + store.applyMessageParts([]); expect(notified).toBe(false); @@ -577,9 +648,9 @@ describe("applyMessagePart / applyMessageParts", () => { // --------------------------------------------------------------------------- describe("resetTransientState", () => { - it("clears streamState, streamError, retryState, reconnectState, and subagentOverrides", () => { + it("clears transient stream state and subagent overrides", () => { const store = createChatStore(); - store.applyMessagePart({ type: "text", text: "stream" }); + store.applyMessagePart(makeStreamPart("stream")); store.setStreamError({ kind: "generic", message: "oops", @@ -638,6 +709,29 @@ describe("resetTransientState", () => { }); }); +describe("resetForChatChange", () => { + it("clears state scoped to the previous chat", () => { + const store = createChatStore(); + store.replaceMessages([makeMessage(1, "user", "hello")]); + store.setQueuedMessages([makeQueuedMessage(1, "queued")]); + store.markPromoteInFlight(1); + store.updateServerEpisodeFloor(4, 0); + store.applyMessagePart(makeStreamPart("stream", 4, 1, 1)); + + store.resetForChatChange(); + + const state = store.getSnapshot(); + expect(state.messagesByID.size).toBe(0); + expect(state.orderedMessageIDs).toEqual([]); + expect(state.streamState).toBeNull(); + expect(state.streamEpisode).toBeNull(); + expect(state.serverEpisodeFloor).toBeNull(); + expect(state.lastStreamPartSeq).toBe(0); + expect(state.queuedMessages).toEqual([]); + expect(state.promoteInFlightIDs.size).toBe(0); + }); +}); + // --------------------------------------------------------------------------- // subscribe // --------------------------------------------------------------------------- @@ -702,7 +796,7 @@ describe("selectIsAwaitingFirstStreamChunk", () => { const store = createChatStore(); store.setChatStatus("running"); store.upsertDurableMessage(makeMessage(1, "user", "hello")); - store.applyMessagePart({ type: "text", text: "response" }); + store.applyMessagePart(makeStreamPart("response")); expect(selectIsAwaitingFirstStreamChunk(store.getSnapshot())).toBe(false); }); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index b4a87870b31..4336cbe0f0e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -212,6 +212,24 @@ const buildChat = (chatID: string): TypesGen.Chat => ({ updated_at: "2025-01-01T00:00:00.000Z", }); +const buildPartEvent = ( + chatID: string, + text: string, + historyVersion: number, + generationAttempt: number, + seq: number, +): TypesGen.ChatStreamEvent => ({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text }, + history_version: historyVersion, + generation_attempt: generationAttempt, + seq, + }, +}); + const buildMessage = ( chatID: string, id: number, @@ -1221,18 +1239,15 @@ describe("useChatStore", () => { }); }); - it("ignores message_part updates while chat is waiting", async () => { + it("accepts a newer episode before status catches up from waiting", async () => { immediateAnimationFrame(); - const chatID = "chat-1"; + const chatID = "chat-waiting-episode"; const existingMessage = buildMessage(chatID, 1, "user", "hello"); const mockSocket = createMockSocket(); mockWatchChatReturn(mockSocket); - const queryClient = createTestQueryClient(); const wrapper = createWrapper(queryClient); - const setChatErrorReason = vi.fn(); - const clearChatErrorReason = vi.fn(); const { result } = renderHook( () => { @@ -1246,8 +1261,8 @@ describe("useChatStore", () => { has_more: false, }, chatQueuedMessages: [], - setChatErrorReason, - clearChatErrorReason, + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), }); return { streamState: useChatSelector(store, selectStreamState), @@ -1260,20 +1275,7 @@ describe("useChatStore", () => { expect(watchChat).toHaveBeenCalledWith(chatID, 1); }); - act(() => { - mockSocket.emitData({ - type: "message_part", - chat_id: chatID, - message_part: { - role: "assistant", - part: { - type: "text", - text: "first", - }, - }, - }); - }); - + act(() => mockSocket.emitData(buildPartEvent(chatID, "first", 1, 1, 1))); await waitFor(() => { expect(result.current.streamState?.blocks).toEqual([ { type: "response", text: "first" }, @@ -1284,41 +1286,23 @@ describe("useChatStore", () => { mockSocket.emitData({ type: "status", chat_id: chatID, - status: { status: "waiting" }, - }); - }); - - await waitFor(() => { - // Stream state is preserved after status=waiting (the - // durable message event handles cleanup via - // needsStreamReset). Only new message_parts should be - // blocked by the shouldApplyMessagePart gate. - expect(result.current.streamState).not.toBeNull(); - expect(result.current.streamState?.blocks).toEqual([ - { type: "response", text: "first" }, - ]); - }); - - act(() => { - mockSocket.emitData({ - type: "message_part", - chat_id: chatID, - message_part: { - role: "assistant", - part: { - type: "text", - text: "late", - }, + status: { + status: "waiting", + history_version: 2, + generation_attempt: 0, }, }); }); + // Status raises the episode floor without clearing the active preview. + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "first" }, + ]); + + act(() => mockSocket.emitData(buildPartEvent(chatID, "new", 2, 1, 1))); await waitFor(() => { - // The late message_part should not be applied because - // shouldApplyMessagePart gates on waiting. - // Stream state still shows the original "first". expect(result.current.streamState?.blocks).toEqual([ - { type: "response", text: "first" }, + { type: "response", text: "new" }, ]); }); }); @@ -2147,17 +2131,14 @@ describe("useChatStore", () => { expect(result.current.queuedMessages).toEqual([]); }); - it("does not apply message parts after status changes to waiting", async () => { + it("drops buffered parts below the status episode floor", async () => { immediateAnimationFrame(); - const chatID = "chat-status-guard"; + const chatID = "chat-status-floor"; const mockSocket = createMockSocket(); mockWatchChatReturn(mockSocket); - const queryClient = createTestQueryClient(); const wrapper = createWrapper(queryClient); - const setChatErrorReason = vi.fn(); - const clearChatErrorReason = vi.fn(); const { result } = renderHook( () => { @@ -2171,8 +2152,8 @@ describe("useChatStore", () => { has_more: false, }, chatQueuedMessages: [], - setChatErrorReason, - clearChatErrorReason, + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), }); return { streamState: useChatSelector(store, selectStreamState), @@ -2185,33 +2166,162 @@ describe("useChatStore", () => { expect(watchChat).toHaveBeenCalledWith(chatID, undefined); }); - // Emit a batch with message_parts followed by a status change - // to "waiting". The status handler clears stream state - // synchronously, and the startTransition guard should prevent - // the deferred applyMessageParts from re-populating it. act(() => { mockSocket.emitDataBatch([ + buildPartEvent(chatID, "stale", 1, 1, 1), { - type: "message_part", + type: "status", chat_id: chatID, - message_part: { - role: "assistant", - part: { type: "text", text: "should be discarded" }, + status: { + status: "waiting", + history_version: 2, + generation_attempt: 0, + }, + }, + ]); + }); + + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(result.current.streamState).toBeNull(); + }); + + it("accepts a newer episode before status and drops a late old episode", async () => { + immediateAnimationFrame(); + + const chatID = "chat-episode-ordering"; + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + const queryClient = createTestQueryClient(); + const wrapper = createWrapper(queryClient); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [], + chatRecord: { + ...buildChat(chatID), + status: "interrupting", + history_version: 1, + generation_attempt: 1, + }, + chatMessagesData: { + messages: [], + queued_messages: [], + has_more: false, }, + chatQueuedMessages: [], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }); + return { + store, + streamState: useChatSelector(store, selectStreamState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, undefined); + }); + + act(() => mockSocket.emitData(buildPartEvent(chatID, "old", 1, 1, 1))); + await waitFor(() => { + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "old" }, + ]); + }); + + act(() => mockSocket.emitData(buildPartEvent(chatID, "new", 2, 1, 1))); + await waitFor(() => { + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "new" }, + ]); + }); + + act(() => { + mockSocket.emitData({ + type: "status", + chat_id: chatID, + status: { + status: "running", + history_version: 2, + generation_attempt: 1, }, + }); + mockSocket.emitData(buildPartEvent(chatID, " stale", 1, 1, 2)); + }); + + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "new" }, + ]); + }); + + it("uses status as a floor without rejecting the next generation attempt", async () => { + immediateAnimationFrame(); + + const chatID = "chat-episode-floor"; + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + const queryClient = createTestQueryClient(); + const wrapper = createWrapper(queryClient); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [], + chatRecord: { + ...buildChat(chatID), + status: "interrupting", + }, + chatMessagesData: { + messages: [], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }); + return { + store, + streamState: useChatSelector(store, selectStreamState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, undefined); + }); + + act(() => { + mockSocket.emitDataBatch([ { type: "status", chat_id: chatID, - status: { status: "waiting" }, + status: { + status: "running", + history_version: 5, + generation_attempt: 0, + }, }, + buildPartEvent(chatID, "current", 5, 1, 1), + buildPartEvent(chatID, "stale", 4, 9, 1), ]); }); - // Stream state should be null — the status change cleared it, - // and the deferred applyMessageParts should not have - // re-populated it. await waitFor(() => { - expect(result.current.streamState).toBeNull(); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "current" }, + ]); }); }); @@ -3926,9 +4036,7 @@ describe("thinking indicator event ordering", () => { expect(watchChat).toHaveBeenCalledWith(chatID, 1); }); - // Server sends message_part then immediately transitions to - // waiting. The buffered parts must be discarded (not applied) - // because waiting status clears stream state. + // The status floor must apply before buffered parts flush. act(() => { mockSocket.emitDataBatch([ { @@ -3936,12 +4044,19 @@ describe("thinking indicator event ordering", () => { chat_id: chatID, message_part: { part: { type: "text", text: "partial response" }, + history_version: 1, + generation_attempt: 1, + seq: 1, }, }, { type: "status", chat_id: chatID, - status: { status: "waiting" }, + status: { + status: "waiting", + history_version: 2, + generation_attempt: 0, + }, }, ]); }); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts index 5f82c01c0ac..7f08ae8233a 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts @@ -144,20 +144,24 @@ export const isActiveChatStatus = ( status: TypesGen.ChatStatus | null, ): boolean => status === "running" || status === "interrupting"; +type StreamEpisode = { + historyVersion: number; + generationAttempt: number; +}; + export type ChatStoreState = { messagesByID: Map; orderedMessageIDs: readonly number[]; streamState: StreamState | null; + streamEpisode: StreamEpisode | null; + serverEpisodeFloor: StreamEpisode | null; + lastStreamPartSeq: number; chatStatus: TypesGen.ChatStatus | null; streamError: ChatDetailError | null; 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; + promoteInFlightIDs: ReadonlySet; subagentStatusOverrides: Map; }; @@ -173,21 +177,17 @@ export type ChatStore = { changed: boolean; }; upsertDurableMessages: (messages: readonly TypesGen.ChatMessage[]) => void; - applyMessagePart: (part: TypesGen.ChatMessagePart) => void; - applyMessageParts: (parts: readonly TypesGen.ChatMessagePart[]) => void; + applyMessagePart: (part: TypesGen.ChatStreamMessagePart) => void; + applyMessageParts: (parts: readonly TypesGen.ChatStreamMessagePart[]) => void; 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, + markPromoteInFlight: (id: number) => void; + clearPromoteInFlight: (id: number) => void; + updateServerEpisodeFloor: ( + historyVersion: number | undefined, + generationAttempt: number | 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; @@ -198,6 +198,7 @@ export type ChatStore = { clearReconnectState: () => void; clearStreamState: () => void; resetTransportReplayState: () => void; + resetForChatChange: () => void; setSubagentStatusOverride: ( chatID: string, status: TypesGen.ChatStatus, @@ -209,12 +210,15 @@ const createInitialState = (): ChatStoreState => ({ messagesByID: new Map(), orderedMessageIDs: [], streamState: null, + streamEpisode: null, + serverEpisodeFloor: null, + lastStreamPartSeq: 0, chatStatus: null, streamError: null, retryState: null, reconnectState: null, queuedMessages: [], - suppressedQueuedMessageIDs: new Set(), + promoteInFlightIDs: new Set(), subagentStatusOverrides: new Map(), }); @@ -319,9 +323,7 @@ export const createChatStore = (): ChatStore => { const nextMessagesByID = new Map(current.messagesByID); nextMessagesByID.set(message.id, message); - const curIsDuplicate = current.messagesByID.has(message.id); - const needsReorder = - !curIsDuplicate || nextMessagesByID.size !== current.messagesByID.size; + const needsReorder = !current.messagesByID.has(message.id); const nextOrderedMessageIDs = needsReorder ? buildOrderedMessageIDs(Array.from(nextMessagesByID.values())) : current.orderedMessageIDs; @@ -335,8 +337,7 @@ export const createChatStore = (): ChatStore => { return { isDuplicate, changed: actuallyChanged }; }; - // Bulk variant that applies all messages in a single pass — - // one Map copy and one sort instead of N copies and N sorts. + // Bulk variant applies all messages with one Map copy and one sort. const upsertDurableMessages = ( messages: readonly TypesGen.ChatMessage[], ): void => { @@ -372,22 +373,81 @@ export const createChatStore = (): ChatStore => { }); }; - const applyMessageParts = (parts: readonly TypesGen.ChatMessagePart[]) => { + const compareEpisodes = (left: StreamEpisode, right: StreamEpisode): number => + left.historyVersion - right.historyVersion || + left.generationAttempt - right.generationAttempt; + + const applyMessageParts = ( + parts: readonly TypesGen.ChatStreamMessagePart[], + ) => { if (parts.length === 0) { return; } setState((current) => { - let nextStreamState: StreamState | null = current.streamState; - for (const part of parts) { - nextStreamState = applyMessagePartToStreamState(nextStreamState, part); + let nextStreamState = current.streamState; + let nextStreamEpisode = current.streamEpisode; + let nextLastStreamPartSeq = current.lastStreamPartSeq; + + for (const messagePart of parts) { + const historyVersion = messagePart.history_version ?? 0; + const generationAttempt = messagePart.generation_attempt ?? 0; + const seq = messagePart.seq ?? 0; + if (historyVersion <= 0 || generationAttempt <= 0 || seq <= 0) { + if (nextStreamEpisode || current.serverEpisodeFloor) { + continue; + } + nextStreamState = applyMessagePartToStreamState( + nextStreamState, + messagePart.part, + ); + continue; + } + + const episode = { historyVersion, generationAttempt }; + if ( + current.serverEpisodeFloor && + compareEpisodes(episode, current.serverEpisodeFloor) < 0 + ) { + continue; + } + if (nextStreamEpisode) { + const comparison = compareEpisodes(episode, nextStreamEpisode); + if (comparison < 0) { + continue; + } + if (comparison > 0) { + nextStreamState = null; + nextStreamEpisode = episode; + nextLastStreamPartSeq = 0; + } + } else { + nextStreamState = null; + nextStreamEpisode = episode; + nextLastStreamPartSeq = 0; + } + if (seq !== nextLastStreamPartSeq + 1) { + continue; + } + nextStreamState = applyMessagePartToStreamState( + nextStreamState, + messagePart.part, + ); + nextLastStreamPartSeq = seq; } - if (nextStreamState === current.streamState) { + + if ( + nextStreamState === current.streamState && + nextStreamEpisode === current.streamEpisode && + nextLastStreamPartSeq === current.lastStreamPartSeq + ) { return current; } return { ...current, streamState: nextStreamState, + streamEpisode: nextStreamEpisode, + lastStreamPartSeq: nextLastStreamPartSeq, }; }); }; @@ -409,84 +469,82 @@ export const createChatStore = (): ChatStore => { setQueuedMessages: (queuedMessages) => { const nextQueuedMessages = queuedMessages ?? []; setState((current) => { - if ( - chatQueuedMessagesEqualByID( - current.queuedMessages, - nextQueuedMessages, - ) - ) { - return current; - } - 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 | null = null; - for (const id of current.suppressedQueuedMessageIDs) { - if (!incomingIDs.has(id)) { - if (!copy) { - copy = new Set(current.suppressedQueuedMessageIDs); - } - copy.delete(id); + let nextPromoteInFlightIDs: ReadonlySet = + current.promoteInFlightIDs; + let mutablePromoteInFlightIDs: Set | null = null; + if (current.promoteInFlightIDs.size > 0) { + const queuedIDs = new Set( + nextQueuedMessages.map((message) => message.id), + ); + for (const id of current.promoteInFlightIDs) { + if (queuedIDs.has(id)) { + continue; } - } - if (copy) { - nextSuppressed = copy; + if (!mutablePromoteInFlightIDs) { + mutablePromoteInFlightIDs = new Set(current.promoteInFlightIDs); + nextPromoteInFlightIDs = mutablePromoteInFlightIDs; + } + mutablePromoteInFlightIDs.delete(id); } } - const filtered = - nextSuppressed.size === 0 - ? incoming - : incoming.filter((message) => !nextSuppressed.has(message.id)); const sameQueue = chatQueuedMessagesEqualByID( current.queuedMessages, - filtered, + nextQueuedMessages, ); - const sameSuppressed = - nextSuppressed === current.suppressedQueuedMessageIDs; - if (sameQueue && sameSuppressed) { + if ( + sameQueue && + nextPromoteInFlightIDs === current.promoteInFlightIDs + ) { return current; } return { ...current, - queuedMessages: sameQueue ? current.queuedMessages : filtered, - suppressedQueuedMessageIDs: nextSuppressed, + queuedMessages: sameQueue + ? current.queuedMessages + : nextQueuedMessages, + promoteInFlightIDs: nextPromoteInFlightIDs, }; }); }, - suppressQueuedMessageID: (id) => { + markPromoteInFlight: (id) => { setState((current) => { - if (current.suppressedQueuedMessageIDs.has(id)) { + if (current.promoteInFlightIDs.has(id)) { return current; } - const next = new Set(current.suppressedQueuedMessageIDs); + const next = new Set(current.promoteInFlightIDs); next.add(id); - return { ...current, suppressedQueuedMessageIDs: next }; + return { ...current, promoteInFlightIDs: next }; }); }, - unsuppressQueuedMessageID: (id) => { + clearPromoteInFlight: (id) => { setState((current) => { - if (!current.suppressedQueuedMessageIDs.has(id)) { + if (!current.promoteInFlightIDs.has(id)) { return current; } - const next = new Set(current.suppressedQueuedMessageIDs); + const next = new Set(current.promoteInFlightIDs); next.delete(id); - return { ...current, suppressedQueuedMessageIDs: next }; + return { ...current, promoteInFlightIDs: next }; }); }, - clearSuppressedQueuedMessageIDs: () => { + updateServerEpisodeFloor: (historyVersion, generationAttempt) => { + if (!historyVersion || historyVersion <= 0) { + return; + } + const nextFloor = { + historyVersion, + generationAttempt: Math.max(0, generationAttempt ?? 0), + }; setState((current) => { - if (current.suppressedQueuedMessageIDs.size === 0) { + if ( + current.serverEpisodeFloor && + compareEpisodes(nextFloor, current.serverEpisodeFloor) <= 0 + ) { return current; } - return { ...current, suppressedQueuedMessageIDs: new Set() }; + return { ...current, serverEpisodeFloor: nextFloor }; }); }, + setChatStatus: (status) => { if (state.chatStatus === status) { return; @@ -583,7 +641,8 @@ export const createChatStore = (): ChatStore => { if ( state.reconnectState === null && state.streamState === null && - state.streamError === null + state.streamError === null && + state.lastStreamPartSeq === 0 ) { return; } @@ -592,8 +651,12 @@ export const createChatStore = (): ChatStore => { reconnectState: null, streamState: null, streamError: null, + lastStreamPartSeq: 0, })); }, + resetForChatChange: () => { + setState(() => createInitialState()); + }, setSubagentStatusOverride: (chatID, status) => { if (state.subagentStatusOverrides.get(chatID) === status) { return; @@ -639,6 +702,8 @@ export const selectChatStatus = (state: ChatStoreState) => state.chatStatus; export const selectStreamError = (state: ChatStoreState) => state.streamError; export const selectQueuedMessages = (state: ChatStoreState) => state.queuedMessages; +export const selectPromoteInFlightIDs = (state: ChatStoreState) => + state.promoteInFlightIDs; export const selectSubagentStatusOverrides = (state: ChatStoreState) => state.subagentStatusOverrides; export const selectRetryState = (state: ChatStoreState) => state.retryState; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index e887682d58d..cab3c1cf111 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -85,9 +85,8 @@ export const useChatStore = ( // the current chat. Once true, the WS is the authoritative // source for chatStatus and the REST-fetched chatRecord.status // must not overwrite it. Without this guard, a React Query - // refetch (e.g. on window focus) can regress chatStatus to a - // stale value like "waiting", causing shouldApplyMessagePart() - // to drop all incoming parts. + // refetch (for example, on window focus) can regress chatStatus + // to a stale value like "waiting". const wsStatusReceivedRef = useRef(false); const activeChatIDRef = useRef(null); const prevChatIDRef = useRef(chatID); @@ -217,7 +216,10 @@ export const useChatStore = ( if (prevChatIDRef.current !== chatID) { prevChatIDRef.current = chatID; lastSyncedMessagesRef.current = []; - store.replaceMessages([]); + queuedMessagesHydratedChatIDRef.current = null; + wsQueueUpdateReceivedRef.current = false; + wsStatusReceivedRef.current = false; + store.resetForChatChange(); } // Merge REST-fetched messages into the store, preserving // any messages the WebSocket delivered that haven't @@ -260,28 +262,21 @@ export const useChatStore = ( }, [chatID, chatMessages, store]); useEffect(() => { - // Only hydrate from REST when the WebSocket hasn't delivered - // a status event yet. Once the WS is the authoritative - // source, a stale REST refetch must not overwrite the - // fresher WS-delivered value. + store.updateServerEpisodeFloor( + chatRecord?.history_version, + chatRecord?.generation_attempt, + ); + // Only hydrate status from REST when the WebSocket hasn't delivered + // a status event yet. The episode floor is independently monotonic. if (!wsStatusReceivedRef.current) { store.setChatStatus(chatRecord?.status ?? null); } - }, [chatRecord?.status, store]); - - useEffect(() => { - queuedMessagesHydratedChatIDRef.current = null; - 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; - } - }, [chatID, store]); + }, [ + chatRecord?.generation_attempt, + chatRecord?.history_version, + chatRecord?.status, + store, + ]); useEffect(() => { if (!chatID || !chatMessagesData) { @@ -299,7 +294,7 @@ export const useChatStore = ( return; } queuedMessagesHydratedChatIDRef.current = chatID; - store.applyAuthoritativeQueuedMessages(chatQueuedMessages); + store.setQueuedMessages(chatQueuedMessages); }, [chatMessagesData, chatID, chatQueuedMessages, store]); useEffect(() => { @@ -374,7 +369,7 @@ export const useChatStore = ( // across WebSocket messages. A rAF-based flush coalesces // parts from multiple WS messages into a single render, // capping stream renders to once per animation frame. - const partsBuf: TypesGen.ChatMessagePart[] = []; + const partsBuf: TypesGen.ChatStreamMessagePart[] = []; let partsFlushTimer: ReturnType | null = null; // History replacement state lives at the effect scope because @@ -387,10 +382,6 @@ export const useChatStore = ( let historyResetPending = false; const historyReplacementBuf: TypesGen.ChatMessage[] = []; - const shouldApplyMessagePart = (): boolean => { - return store.getSnapshot().chatStatus !== "waiting"; - }; - const schedulePartsFlush = () => { if (partsFlushTimer !== null || partsBuf.length === 0) { return; @@ -401,10 +392,9 @@ export const useChatStore = ( return; } const parts = partsBuf.splice(0); - if (parts.length === 0 || !shouldApplyMessagePart()) { - return; + if (parts.length > 0) { + store.applyMessageParts(parts); } - store.applyMessageParts(parts); }, 0); }; @@ -420,10 +410,9 @@ export const useChatStore = ( partsFlushTimer = null; } const parts = partsBuf.splice(0); - if (activeChatIDRef.current !== chatID || !shouldApplyMessagePart()) { - return; + if (activeChatIDRef.current === chatID) { + store.applyMessageParts(parts); } - store.applyMessageParts(parts); }; // Discard buffered parts without applying them. Used when @@ -485,13 +474,10 @@ export const useChatStore = ( continue; } commitHistoryReplacement(); - if (!shouldApplyMessagePart()) { - continue; - } - const part = streamEvent.message_part?.part; - if (part) { + const messagePart = streamEvent.message_part; + if (messagePart?.part) { store.clearRetryState(); - partsBuf.push(part); + partsBuf.push(messagePart); } continue; } @@ -569,23 +555,23 @@ export const useChatStore = ( } case "queue_update": wsQueueUpdateReceivedRef.current = true; - store.applyAuthoritativeQueuedMessages( - streamEvent.queued_messages, - ); + store.setQueuedMessages(streamEvent.queued_messages); updateChatQueuedMessages(streamEvent.queued_messages); continue; case "status": { - const nextStatus = streamEvent.status?.status; + const status = streamEvent.status; + const nextStatus = status?.status; if (!nextStatus) { continue; } wsStatusReceivedRef.current = true; store.clearRetryState(); + store.updateServerEpisodeFloor( + status.history_version, + status.generation_attempt, + ); store.setChatStatus(nextStatus); - if (nextStatus === "waiting") { - discardBufferedParts(); - } if (nextStatus !== "error") { clearChatErrorReasonEvent(chatID); } @@ -656,9 +642,7 @@ export const useChatStore = ( partsFlushTimer = null; } const nextParts = partsBuf.splice(0); - if (shouldApplyMessagePart()) { - store.applyMessageParts(nextParts); - } + store.applyMessageParts(nextParts); } } }); diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 02d5590e1ba..750893fddb8 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -27,6 +27,7 @@ import { selectIsAwaitingFirstStreamChunk, selectMessagesByID, selectOrderedMessageIDs, + selectPromoteInFlightIDs, selectQueuedMessages, useChatSelector, type useChatStore, @@ -290,6 +291,7 @@ export const ChatPageInput: FC = ({ const hasStreamState = useChatSelector(store, selectHasStreamState); const chatStatus = useChatSelector(store, selectChatStatus); const queuedMessages = useChatSelector(store, selectQueuedMessages); + const promoteInFlightIDs = useChatSelector(store, selectPromoteInFlightIDs); const messages = orderedMessageIDs .map((messageID) => { @@ -486,6 +488,7 @@ export const ChatPageInput: FC = ({ remountKey={remountKey} onContentChange={onContentChange} queuedMessages={queuedMessages} + promoteInFlightIDs={promoteInFlightIDs} onDeleteQueuedMessage={onDeleteQueuedMessage} onPromoteQueuedMessage={onPromoteQueuedMessage} editingQueuedMessageID={editingQueuedMessageID} diff --git a/site/src/pages/AgentsPage/components/QueuedMessagesList.stories.tsx b/site/src/pages/AgentsPage/components/QueuedMessagesList.stories.tsx index 6142665f4b1..532c6a2006c 100644 --- a/site/src/pages/AgentsPage/components/QueuedMessagesList.stories.tsx +++ b/site/src/pages/AgentsPage/components/QueuedMessagesList.stories.tsx @@ -46,6 +46,18 @@ export const SingleMessage: Story = { }, }; +export const PromotingMessage: Story = { + args: { + messages: [buildMessage(1, textContent("Run the test suite"))], + promoteInFlightIDs: new Set([1]), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.getByText("Promoting...")).toBeVisible(); + expect(canvas.queryByText("Run the test suite")).not.toBeInTheDocument(); + }, +}; + // Several messages queued up at once. export const SeveralMessages: Story = { args: { diff --git a/site/src/pages/AgentsPage/components/QueuedMessagesList.tsx b/site/src/pages/AgentsPage/components/QueuedMessagesList.tsx index 81d6b394b00..2643be04356 100644 --- a/site/src/pages/AgentsPage/components/QueuedMessagesList.tsx +++ b/site/src/pages/AgentsPage/components/QueuedMessagesList.tsx @@ -20,6 +20,7 @@ interface QueuedMessagesListProps { messages: readonly ChatQueuedMessage[]; onDelete: (id: number) => Promise | void; onPromote: (id: number) => Promise | void; + promoteInFlightIDs?: ReadonlySet; onEdit?: ( id: number, text: string, @@ -69,6 +70,7 @@ export const QueuedMessagesList: FC = ({ messages, onDelete, onPromote, + promoteInFlightIDs, onEdit, editingMessageID = null, className, @@ -150,12 +152,11 @@ export const QueuedMessagesList: FC = ({ const handlePromote = async (id: number) => { setBusyItem({ id, action: "promote" }); - hideItemOptimistically(id); try { await onPromote(id); setBusyItem((current) => (current?.id === id ? null : current)); } catch { - restoreHiddenItem(id); + // The caller owns promotion error feedback. setBusyItem((current) => (current?.id === id ? null : current)); } }; @@ -168,7 +169,7 @@ export const QueuedMessagesList: FC = ({ return null; } - const isBusy = busyItem !== null; + const isBusy = busyItem !== null || (promoteInFlightIDs?.size ?? 0) > 0; return (
= ({ const isEditing = item.id === editingMessageID; const isFirst = index === 0; const isItemBusy = busyItem !== null && busyItem.id === item.id; + const isPromoting = + promoteInFlightIDs?.has(item.id) === true || + (isItemBusy && busyItem.action === "promote"); const isHovered = hoveredID === item.id; // Show actions when: first and nothing else hovered, // or this item is hovered, or being edited. const showActions = isEditing || isHovered || (isFirst && hoveredID === null); + if (isPromoting) { + return ( +
+
+ + Promoting... +
+
+ ); + } + return (