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
2026-09-24 15:04:27 +08:00 · 2026-05-06 19:11:56 +03:00
parent 2cab1b41ad
commit 6b0518d051
18 changed files with 2431 additions and 101 deletions
@@ -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
@@ -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)
+}