From e3a524d8ce57bade9d9dcff1aae5fd8c139f17f1 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Sat, 29 Aug 2026 09:51:52 +0000 Subject: [PATCH 1/5] refactor(coderd/x/chatd): start chat mutator --- coderd/x/chatd/chatd.go | 13 +------------ coderd/x/chatd/mutator.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 12 deletions(-) create mode 100644 coderd/x/chatd/mutator.go diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 456031add59af..1d8b94452a639 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -2114,18 +2114,7 @@ func (p *Server) DeleteQueued( chatID uuid.UUID, queuedMessageID int64, ) error { - if chatID == uuid.Nil { - return xerrors.New("chat_id is required") - } - - machine := p.newChatMachine(chatID) - err := machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { - _, err := tx.DeleteQueuedMessage(chatstate.DeleteQueuedMessageInput{ - QueuedMessageID: queuedMessageID, - }) - return err - }) - return err + return (&chatMutator{server: p}).DeleteQueued(ctx, chatID, queuedMessageID) } // PromoteQueued promotes a queued message through the chatstate state diff --git a/coderd/x/chatd/mutator.go b/coderd/x/chatd/mutator.go new file mode 100644 index 0000000000000..1cf0930cdb444 --- /dev/null +++ b/coderd/x/chatd/mutator.go @@ -0,0 +1,29 @@ +package chatd + +import ( + "context" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" +) + +type chatMutator struct { + server *Server +} + +func (m *chatMutator) DeleteQueued(ctx context.Context, chatID uuid.UUID, queuedMessageID int64) error { + if chatID == uuid.Nil { + return xerrors.New("chat_id is required") + } + + machine := m.server.newChatMachine(chatID) + return machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.DeleteQueuedMessage(chatstate.DeleteQueuedMessageInput{ + QueuedMessageID: queuedMessageID, + }) + return err + }) +} From 3024921f89cb2945f6cf32ef303e3ce5e989dc1a Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Sat, 29 Aug 2026 09:59:46 +0000 Subject: [PATCH 2/5] refactor(coderd/x/chatd): centralize chat mutations --- coderd/x/chatd/chatd.go | 707 +------------------------------------- coderd/x/chatd/mutator.go | 697 +++++++++++++++++++++++++++++++++++++ 2 files changed, 707 insertions(+), 697 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 1d8b94452a639..f28a66fbb0d74 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1466,158 +1466,7 @@ func (p *Server) SendMessage( ctx context.Context, opts SendMessageOptions, ) (SendMessageResult, error) { - if opts.ChatID == uuid.Nil { - return SendMessageResult{}, xerrors.New("chat_id is required") - } - if len(opts.Content) == 0 { - return SendMessageResult{}, xerrors.New("content is required") - } - - busyBehavior := opts.BusyBehavior - if busyBehavior == "" { - busyBehavior = SendMessageBusyBehaviorQueue - } - switch busyBehavior { - case SendMessageBusyBehaviorQueue, SendMessageBusyBehaviorInterrupt: - default: - return SendMessageResult{}, xerrors.Errorf("invalid busy behavior %q", opts.BusyBehavior) - } - - contentParts := opts.Content - if p.hooks.Enabled() { - turnID := uuid.New() - chat, err := p.db.GetChatByID(ctx, opts.ChatID) - if err != nil { - return SendMessageResult{}, xerrors.Errorf("load chat for user_prompt_submit: %w", err) - } - // Repeat these admission checks under the transaction lock. - if chat.Archived { - return SendMessageResult{}, ErrChatArchived - } - if _, err := resolveSendMessageModelConfigID(ctx, p.db, chat, opts.ModelConfigID); err != nil { - return SendMessageResult{}, err - } - // Check queue capacity before dispatch; the transaction - // rechecks it under lock. - queuedCount, err := p.db.CountChatQueuedMessages(ctx, opts.ChatID) - if err != nil { - return SendMessageResult{}, xerrors.Errorf("count queued messages: %w", err) - } - if queuedCount >= chatstate.MaxQueueSize { - return SendMessageResult{}, &chatstate.MessageQueueFullError{Max: chatstate.MaxQueueSize} - } - promptMessage, err := chathooks.UserPromptMessage(contentParts) - if err != nil { - return SendMessageResult{}, err - } - promptResult, err := p.hooks.Trigger(ctx, chathooks.ChatFor(chat, &turnID), promptMessage, agenthooks.EventUserPromptSubmit, dispatch.CapacityClassAdmission) - if err != nil { - return SendMessageResult{}, p.handleUserPromptDispatchError(ctx, opts.ChatID, chathooks.UserPromptDenial(err)) - } - contentParts, _, err = chathooks.ComposeUserPromptContent(contentParts, promptResult) - if err != nil { - return SendMessageResult{}, err - } - } - - content, err := chatprompt.MarshalParts(contentParts) - if err != nil { - return SendMessageResult{}, xerrors.Errorf("marshal message content: %w", err) - } - - requestedPlanMode := opts.PlanMode - requestedMCPServerIDs := opts.MCPServerIDs - - var result SendMessageResult - machine := p.newChatMachine(opts.ChatID) - updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - lockedChat, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("load chat: %w", err) - } - - if lockedChat.Archived { - return ErrChatArchived - } - - if requestedPlanMode != nil { - lockedChat, err = store.UpdateChatPlanModeByID(ctx, database.UpdateChatPlanModeByIDParams{ - PlanMode: *requestedPlanMode, - ID: opts.ChatID, - }) - if err != nil { - return xerrors.Errorf("update chat plan mode: %w", err) - } - } - - modelConfigID, err := resolveSendMessageModelConfigID( - ctx, - store, - lockedChat, - opts.ModelConfigID, - ) - if err != nil { - return err - } - - lockedChat, err = p.applyRequestedMCPServerIDs(ctx, store, lockedChat, requestedMCPServerIDs) - if err != nil { - return err - } - - messageCreatedBy := opts.CreatedBy - if messageCreatedBy == uuid.Nil { - messageCreatedBy = lockedChat.OwnerID - } - - // Queue capacity is enforced inside tx.SendMessage; this - // wrapper only propagates the typed error. - message := userMessage(content, modelConfigID, messageCreatedBy, opts.ReasoningEffort) - sendResult, err := tx.SendMessage(chatstate.SendMessageInput{ - Message: message, - BusyBehavior: busyBehaviorToChatState(busyBehavior), - }) - if err != nil { - return err - } - - if sendResult.QueuedMessage != nil { - result.Queued = true - result.QueuedMessage = sendResult.QueuedMessage - } else if len(sendResult.InsertedMessages) > 0 { - // The state machine prepends synthetic tool-result - // cancellation messages; the user message is always - // last in the inserted slice. - result.Message = sendResult.InsertedMessages[len(sendResult.InsertedMessages)-1] - } - // A queued send on an errored chat can also promote the - // previous queue head into history; report those inserts so - // clients can update their caches. - result.InsertedMessages = sendResult.InsertedMessages - - // File-link errors must roll back the message. - if err := chatstate.LinkFiles(ctx, store, opts.ChatID, chatprompt.FileIDs(contentParts)); err != nil { - return err - } - // Capture the post-transition chat inside the same - // transaction so the returned chat and the watch event - // reflect the snapshot bump and status change produced by - // the transition itself. - refreshed, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("reload chat after send: %w", err) - } - result.Chat = refreshed - return nil - }) - if updateErr != nil { - return SendMessageResult{}, updateErr - } - - // Sidebar watch event keeps the chat list in sync. Stream side - // effects are handled by chat:update consumers. - p.publishChatPubsubEvent(result.Chat, codersdk.ChatWatchEventKindStatusChange, nil) - return result, nil + return (&chatMutator{server: p}).SendMessage(ctx, opts) } func (p *Server) callerModelConfigContext(ctx context.Context, ownerID uuid.UUID) (context.Context, error) { @@ -1825,201 +1674,7 @@ func (p *Server) EditMessage( ctx context.Context, opts EditMessageOptions, ) (EditMessageResult, error) { - if opts.ChatID == uuid.Nil { - return EditMessageResult{}, xerrors.New("chat_id is required") - } - if opts.EditedMessageID <= 0 { - return EditMessageResult{}, xerrors.New("edited_message_id is required") - } - if len(opts.Content) == 0 { - return EditMessageResult{}, xerrors.New("content is required") - } - - contentParts := opts.Content - var sessionStartHookResult *chathooks.Result - if p.hooks.Enabled() { - turnID := uuid.New() - chat, err := p.db.GetChatByID(ctx, opts.ChatID) - if err != nil { - return EditMessageResult{}, xerrors.Errorf("load chat for edit hooks: %w", err) - } - // Repeat these admission checks under the transaction lock. - if chat.Archived { - return EditMessageResult{}, ErrChatArchived - } - if err := validateEditTarget(ctx, p.db, opts.ChatID, opts.EditedMessageID); err != nil { - return EditMessageResult{}, err - } - if _, err := validateModelConfigOverride(ctx, p.db, chat.OrganizationID, opts.ModelConfigID); err != nil { - return EditMessageResult{}, err - } - sessionStartHookResult, err = p.hooks.Trigger(ctx, chathooks.ChatFor(chat, &turnID), chathooks.Message{Source: chathooks.SessionStartSourceClear}, agenthooks.EventSessionStart, dispatch.CapacityClassAdmission) - if err != nil { - return EditMessageResult{}, p.handleAPIDispatchError(ctx, opts.ChatID, agenthooks.EventSessionStart, err) - } - promptMessage, err := chathooks.UserPromptMessage(contentParts) - if err != nil { - return EditMessageResult{}, err - } - promptResult, err := p.hooks.Trigger(ctx, chathooks.ChatFor(chat, &turnID), promptMessage, agenthooks.EventUserPromptSubmit, dispatch.CapacityClassAdmission) - if err != nil { - return EditMessageResult{}, p.handleUserPromptDispatchError(ctx, opts.ChatID, chathooks.UserPromptDenial(err)) - } - contentParts, _, err = chathooks.ComposeUserPromptContent(contentParts, promptResult) - if err != nil { - return EditMessageResult{}, err - } - } - - content, err := chatprompt.MarshalParts(contentParts) - if err != nil { - return EditMessageResult{}, xerrors.Errorf("marshal message content: %w", err) - } - var ( - result EditMessageResult - editedMsg database.ChatMessage - editedCutoffT time.Time - ) - machine := p.newChatMachine(opts.ChatID) - err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - lockedChat, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("load chat: %w", err) - } - if lockedChat.Archived { - return ErrChatArchived - } - // Capture the target message for the post-commit debug - // cleanup hook below. The transition itself revalidates - // chat ownership and user-message constraints. - target, err := store.GetChatMessageByID(ctx, opts.EditedMessageID) - if err != nil { - if errors.Is(err, sql.ErrNoRows) { - return ErrEditedMessageNotFound - } - return xerrors.Errorf("get edited message: %w", err) - } - if target.ChatID != opts.ChatID { - return ErrEditedMessageNotFound - } - if target.Deleted { - return ErrEditedMessageNotFound - } - if target.Role != database.ChatMessageRoleUser { - return ErrEditedMessageNotUser - } - editedMsg = target - - lockedChat, err = p.applyRequestedMCPServerIDs(ctx, store, lockedChat, opts.MCPServerIDs) - if err != nil { - return err - } - - modelOverride, err := validateModelConfigOverride(ctx, store, lockedChat.OrganizationID, opts.ModelConfigID) - if err != nil { - return err - } - if !modelOverride.Valid { - // Without an explicit override the transition preserves - // the edited message's original model, which may have been - // disabled since; resolve it like a normal message send. - preserved := uuid.Nil - if target.ModelConfigID.Valid { - preserved = target.ModelConfigID.UUID - } - resolved, err := resolveFallbackModelConfigID(ctx, store, lockedChat, preserved) - if err != nil { - return err - } - if resolved != preserved { - modelOverride = uuid.NullUUID{UUID: resolved, Valid: true} - } - } - - modelConfigID := target.ModelConfigID.UUID - if modelOverride.Valid { - modelConfigID = modelOverride.UUID - } - // The prompt response already rides in the replacement content; - // only the session_start(clear) response needs transcript rows. - // They insert after the replacement so a later edit's suffix - // truncation cleans them up. - suffixMessages, err := chathooks.EventMessages(sessionStartHookResult, modelConfigID) - if err != nil { - return err - } - - var reasoningEffortOverride database.NullChatReasoningEffort - if opts.ReasoningEffort != nil && *opts.ReasoningEffort != "" { - reasoningEffortOverride = database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffort(*opts.ReasoningEffort), Valid: true} - } - - editResult, err := tx.EditMessage(chatstate.EditMessageInput{ - MessageID: opts.EditedMessageID, - SuffixMessages: suffixMessages, - CreatedBy: opts.CreatedBy, - Content: content, - ModelConfigIDOverride: modelOverride, - ReasoningEffortOverride: reasoningEffortOverride, - }) - if err != nil { - if errors.Is(err, chatstate.ErrEditedMessageNotUser) { - return ErrEditedMessageNotUser - } - return err - } - result.Message = editResult.ReplacementMessage - inserted := make([]database.ChatMessage, 0, len(editResult.CancellationMessages)+len(editResult.SuffixMessages)+1) - inserted = append(inserted, editResult.CancellationMessages...) - inserted = append(inserted, editResult.ReplacementMessage) - inserted = append(inserted, editResult.SuffixMessages...) - result.InsertedMessages = inserted - result.DeletedMessageIDs = editResult.DeletedMessageIDs - if err := chatstate.LinkFiles(ctx, store, opts.ChatID, chatprompt.FileIDs(contentParts)); err != nil { - return err - } - // Capture the post-edit chat inside the same transaction so - // the returned chat and the debug-cleanup cutoff use the - // snapshot bump and updated_at stamped by the transition. - refreshed, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("reload chat after edit: %w", err) - } - result.Chat = refreshed - editedCutoffT = refreshed.UpdatedAt - return nil - }) - if err != nil { - return EditMessageResult{}, err - } - - // Sidebar watch event keeps the chat list responsive. Stream - // side effects are handled by chat:update consumers. - p.publishChatPubsubEvent(result.Chat, codersdk.ChatWatchEventKindStatusChange, nil) - - // Editing can race with an interrupted worker still flushing its - // final debug writes. Run a short bounded retry loop so we converge - // quickly without relying on the much longer stale-finalization - // sweep. Source editCutoff from the DB-stamped updated_at returned - // by the post-edit chat row so the filter uses the same clock that - // stamps replacement-turn debug rows; subtract - // debugCleanupClockSkew so replica clock drift cannot let the retry - // delete a replacement turn's debug rows. - editCutoff := editedCutoffT.Add(-debugCleanupClockSkew) - p.scheduleDebugCleanup( - ctx, - "failed to delete chat debug rows after edit", - []slog.Field{ - slog.F("chat_id", opts.ChatID), - slog.F("edited_message_id", editedMsg.ID), - }, - func(cleanupCtx context.Context, debugSvc *chatdebug.Service) error { - _, err := debugSvc.DeleteAfterMessageID(cleanupCtx, opts.ChatID, editedMsg.ID-1, editCutoff) - return err - }, - ) - - return result, nil + return (&chatMutator{server: p}).EditMessage(ctx, opts) } // ErrArchiveRequiresRootChat is returned by [Server.ArchiveChat] and @@ -2043,13 +1698,7 @@ var ErrArchiveRequiresRootChat = xerrors.New( // //nolint:staticcheck // Receiver name matches the other Server methods in this file. func (p *Server) ArchiveChat(ctx context.Context, chat database.Chat) error { - if chat.ID == uuid.Nil { - return xerrors.New("chat_id is required") - } - if chat.ParentChatID.Valid { - return ErrArchiveRequiresRootChat - } - return p.setChatFamilyArchived(ctx, chat, true, codersdk.ChatWatchEventKindDeleted) + return (&chatMutator{server: p}).ArchiveChat(ctx, chat) } // UnarchiveChat unarchives a root chat and every child in its family @@ -2057,13 +1706,7 @@ func (p *Server) ArchiveChat(ctx context.Context, chat database.Chat) error { // is atomic; ChildChat unarchive attempts are rejected with // [ErrArchiveRequiresRootChat]. func (p *Server) UnarchiveChat(ctx context.Context, chat database.Chat) error { - if chat.ID == uuid.Nil { - return xerrors.New("chat_id is required") - } - if chat.ParentChatID.Valid { - return ErrArchiveRequiresRootChat - } - return p.setChatFamilyArchived(ctx, chat, false, codersdk.ChatWatchEventKindCreated) + return (&chatMutator{server: p}).UnarchiveChat(ctx, chat) } // setChatFamilyArchived applies SetArchived(archived) to every chat @@ -2072,39 +1715,6 @@ func (p *Server) UnarchiveChat(ctx context.Context, chat database.Chat) error { // events. Callers must only invoke this for root chats. // //nolint:revive // Existing API takes the target archive state as a boolean. -func (p *Server) setChatFamilyArchived( - ctx context.Context, - chat database.Chat, - archived bool, - watchKind codersdk.ChatWatchEventKind, -) error { - if chat.ID == uuid.Nil { - return xerrors.New("chat_id is required") - } - if chat.ParentChatID.Valid { - return ErrArchiveRequiresRootChat - } - - familyChats, err := chatstate.SetFamilyArchived( - ctx, - p.db, - p.pubsub, - chatstate.SetFamilyArchivedInput{ - RootID: chat.ID, - Archived: archived, - }, - ) - if err != nil { - return err - } - - if archived { - p.scheduleArchiveDebugCleanup(ctx, familyChats) - } - - p.publishChatPubsubEvents(familyChats, watchKind) - return nil -} // DeleteQueued removes a queued user message through the chatstate // state machine. Stream side effects are handled by chat:update @@ -2127,53 +1737,7 @@ func (p *Server) PromoteQueued( ctx context.Context, opts PromoteQueuedOptions, ) (PromoteQueuedResult, error) { - if opts.ChatID == uuid.Nil { - return PromoteQueuedResult{}, xerrors.New("chat_id is required") - } - - var ( - result PromoteQueuedResult - refreshChat database.Chat - refreshedOK bool - ) - machine := p.newChatMachine(opts.ChatID) - updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - lockedChat, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("load chat: %w", err) - } - if lockedChat.Archived { - return ErrChatArchived - } - - promoteResult, err := tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ - QueuedMessageID: opts.QueuedMessageID, - }) - if err != nil { - return err - } - if promoteResult.InsertedMessage != nil { - result.PromotedMessage = *promoteResult.InsertedMessage - } - // Capture the chat inside the transaction so the watch event - // published below uses the snapshot bump and status change - // produced by the transition itself. - refreshed, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("reload chat after promote: %w", err) - } - refreshChat = refreshed - refreshedOK = true - return nil - }) - if updateErr != nil { - return PromoteQueuedResult{}, updateErr - } - - if refreshedOK { - p.publishChatPubsubEvent(refreshChat, codersdk.ChatWatchEventKindStatusChange, nil) - } - return result, nil + return (&chatMutator{server: p}).PromoteQueued(ctx, opts) } // SubmitToolResultsOptions controls tool result submission. @@ -2219,88 +1783,7 @@ func (p *Server) SubmitToolResults( ctx context.Context, opts SubmitToolResultsOptions, ) error { - machine := p.newChatMachine(opts.ChatID) - var hookSuffix []chatstate.Message - if p.hooks.Enabled() { - state, err := loadDynamicPostToolUseState(ctx, machine, opts) - if err != nil { - return err - } - for _, result := range opts.Results { - response, err := p.hooks.Trigger(ctx, chathooks.ChatFor(state.chat, nil), chathooks.DynamicPostToolUseMessage(result, state.toolNames[result.ToolCallID]), agenthooks.EventPostToolUse, dispatch.CapacityClassGeneration) - if err != nil { - // Leave pending calls intact so the client can resubmit after recovery. - return chathooks.GenerationDispatchError(agenthooks.EventPostToolUse, err) - } - responseMessages, err := chathooks.EventMessages(response, state.modelConfigID) - if err != nil { - return err - } - hookSuffix = append(hookSuffix, responseMessages...) - } - } - - var ( - statusConflict *ToolResultStatusConflictError - refreshChat database.Chat - refreshedOK bool - ) - updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - locked, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("load chat: %w", err) - } - if locked.Archived { - return ErrChatArchived - } - - toolResults := make([]chatstate.ToolResultInput, 0, len(opts.Results)) - for _, result := range opts.Results { - toolResults = append(toolResults, chatstate.ToolResultInput{ - ToolCallID: result.ToolCallID, - Output: result.Output, - IsError: result.IsError, - }) - } - modelConfigID := opts.ModelConfigID - if modelConfigID == uuid.Nil { - modelConfigID = locked.LastModelConfigID - } - if _, err := tx.CompleteRequiresAction(chatstate.CompleteRequiresActionInput{ - CreatedBy: opts.UserID, - ModelConfigID: modelConfigID, - Results: toolResults, - SuffixMessages: hookSuffix, - }); err != nil { - if !errors.Is(err, chatstate.ErrInvalidState) && - locked.Status != database.ChatStatusRequiresAction && - errors.Is(err, chatstate.ErrTransitionNotAllowed) { - statusConflict = &ToolResultStatusConflictError{ - ActualStatus: locked.Status, - } - return statusConflict - } - return xerrors.Errorf("complete requires action: %w", err) - } - refreshed, err := store.GetChatByID(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("reload chat after tool results: %w", err) - } - refreshChat = refreshed - refreshedOK = true - return nil - }) - if updateErr != nil { - if statusConflict != nil { - return statusConflict - } - return translateToolResultValidationError(updateErr) - } - - if refreshedOK { - p.publishChatPubsubEvent(refreshChat, codersdk.ChatWatchEventKindStatusChange, nil) - } - return nil + return (&chatMutator{server: p}).SubmitToolResults(ctx, opts) } // translateToolResultValidationError converts a chatstate tool-result @@ -2350,34 +1833,7 @@ func (p *Server) InterruptChat( ctx context.Context, chat database.Chat, ) (database.Chat, error) { - if chat.ID == uuid.Nil { - return chat, xerrors.New("chat_id is required") - } - - var refreshed database.Chat - machine := p.newChatMachine(chat.ID) - err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - if _, err := tx.Interrupt(chatstate.InterruptInput{ - Reason: "Tool execution interrupted by user", - }); err != nil { - return err - } - // Capture the post-interrupt chat inside the transaction so - // the returned chat and the watch event reflect the snapshot - // bump and status change produced by the transition itself. - latest, err := store.GetChatByID(ctx, chat.ID) - if err != nil { - return xerrors.Errorf("reload chat after interrupt: %w", err) - } - refreshed = latest - return nil - }) - if err != nil { - return chat, err - } - - p.publishChatPubsubEvent(refreshed, codersdk.ChatWatchEventKindStatusChange, nil) - return refreshed, nil + return (&chatMutator{server: p}).InterruptChat(ctx, chat) } // CompactChat records a manual compaction request through the @@ -2398,50 +1854,7 @@ func (p *Server) CompactChat( ctx context.Context, chat database.Chat, ) (database.Chat, error) { - if chat.ID == uuid.Nil { - return chat, xerrors.New("chat_id is required") - } - - var refreshed database.Chat - machine := p.newChatMachine(chat.ID) - err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - lockedChat, err := store.GetChatByID(ctx, chat.ID) - if err != nil { - return xerrors.Errorf("load chat: %w", err) - } - if lockedChat.Archived { - return ErrChatArchived - } - // Run the transition before content and usage validation so busy - // chats surface the state conflict first. - result, err := tx.RequestCompaction(chatstate.RequestCompactionInput{}) - if err != nil { - return err - } - // Reject requests with nothing to compact inside the same - // transaction (rolling back the transition) so no LLM call - // is ever started for an empty or already-compacted chat. - // This also covers a double-/compact. - messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - if err != nil { - return xerrors.Errorf("load chat messages: %w", err) - } - boundary := latestContextBoundaryIndex(messages) - if _, ok := firstUncompressedAssistantAfter(messages, boundary); !ok { - return ErrNothingToCompact - } - refreshed = result.Chat - return nil - }) - if err != nil { - return chat, err - } - - p.publishChatPubsubEvent(refreshed, codersdk.ChatWatchEventKindStatusChange, nil) - return refreshed, nil + return (&chatMutator{server: p}).CompactChat(ctx, chat) } // ClearChat commits a manual context reset through the @@ -2455,56 +1868,7 @@ func (p *Server) ClearChat( ctx context.Context, chat database.Chat, ) (database.Chat, error) { - if chat.ID == uuid.Nil { - return chat, xerrors.New("chat_id is required") - } - - var refreshed database.Chat - machine := p.newChatMachine(chat.ID) - err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - lockedChat, err := store.GetChatByID(ctx, chat.ID) - if err != nil { - return xerrors.Errorf("load chat: %w", err) - } - if lockedChat.Archived { - return ErrChatArchived - } - // Read the pre-clear history before the transition inserts the - // boundary rows; eligibility is evaluated afterwards so busy - // chats surface the state conflict first. - messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - if err != nil { - return xerrors.Errorf("load chat messages: %w", err) - } - clearMessages, err := buildClearMessages(buildClearMessagesInput{ - modelConfigID: lockedChat.LastModelConfigID, - toolCallID: "chat_cleared_" + uuid.NewString(), - }) - if err != nil { - return xerrors.Errorf("build clear messages: %w", err) - } - result, err := tx.ClearContext(chatstate.ClearContextInput{Messages: clearMessages}) - if err != nil { - return err - } - // Reject no-op clears inside the same transaction so an empty - // or already-cleared chat never gains a duplicate boundary. - boundary := latestContextBoundaryIndex(messages) - if !hasClearableMessageAfter(messages, boundary) { - return ErrNothingToClear - } - refreshed = result.Chat - return nil - }) - if err != nil { - return chat, err - } - - p.publishChatPubsubEvent(refreshed, codersdk.ChatWatchEventKindStatusChange, nil) - return refreshed, nil + return (&chatMutator{server: p}).ClearChat(ctx, chat) } // ReconcileInvalidStateChat recovers a chat stuck in an invalid @@ -2522,32 +1886,7 @@ func (p *Server) ReconcileInvalidStateChat( ctx context.Context, chat database.Chat, ) (database.Chat, error) { - if chat.ID == uuid.Nil { - return chat, xerrors.New("chat_id is required") - } - - var refreshed database.Chat - machine := p.newChatMachine(chat.ID) - err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - if _, err := tx.ReconcileInvalidState(chatstate.ReconcileInvalidStateInput{}); err != nil { - return err - } - // Capture the post-reconcile chat inside the transaction so - // the returned chat and the watch event reflect the snapshot - // bump and status change produced by the transition itself. - latest, err := store.GetChatByID(ctx, chat.ID) - if err != nil { - return xerrors.Errorf("reload chat after reconcile: %w", err) - } - refreshed = latest - return nil - }) - if err != nil { - return chat, err - } - - p.publishChatPubsubEvent(refreshed, codersdk.ChatWatchEventKindStatusChange, nil) - return refreshed, nil + return (&chatMutator{server: p}).ReconcileInvalidStateChat(ctx, chat) } const manualTitleMessageWindowLimit = 50 @@ -3295,11 +2634,6 @@ func subscribeWithInitialError(chatID uuid.UUID, message string) ( } // publishChatPubsubEvents broadcasts a lifecycle event for each affected chat. -func (p *Server) publishChatPubsubEvents(chats []database.Chat, kind codersdk.ChatWatchEventKind) { - for _, chat := range chats { - p.publishChatPubsubEvent(chat, kind, nil) - } -} // chatWatchEventSDKChat builds the chat embedded in ChatWatchEvent // notifications. These payloads travel through PostgreSQL NOTIFY, so @@ -3316,27 +2650,6 @@ func chatWatchEventSDKChat(chat database.Chat, diffStatus *codersdk.ChatDiffStat // publishChatPubsubEvent broadcasts a chat lifecycle event via PostgreSQL // pubsub so that all replicas can push updates to watching clients. -func (p *Server) publishChatPubsubEvent(chat database.Chat, kind codersdk.ChatWatchEventKind, diffStatus *codersdk.ChatDiffStatus) { - event := codersdk.ChatWatchEvent{ - Kind: kind, - Chat: chatWatchEventSDKChat(chat, diffStatus), - } - payload, err := json.Marshal(event) - if err != nil { - p.logger.Error(context.Background(), "failed to marshal chat pubsub event", - slog.F("chat_id", chat.ID), - slog.Error(err), - ) - return - } - if err := p.pubsub.Publish(coderdpubsub.ChatWatchEventChannel(chat.OwnerID), payload); err != nil { - p.logger.Error(context.Background(), "failed to publish chat pubsub event", - slog.F("chat_id", chat.ID), - slog.F("kind", kind), - slog.Error(err), - ) - } -} // ChatQueuedForCapacity reports whether the chat is waiting for a // concurrent-agent capacity slot. Uncapped deployments always return false. diff --git a/coderd/x/chatd/mutator.go b/coderd/x/chatd/mutator.go index 1cf0930cdb444..813676bbe83a5 100644 --- a/coderd/x/chatd/mutator.go +++ b/coderd/x/chatd/mutator.go @@ -2,18 +2,432 @@ package chatd import ( "context" + "database/sql" + "encoding/json" + "errors" + "time" "github.com/google/uuid" "golang.org/x/xerrors" + "cdr.dev/slog/v3" "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/agenthooks/dispatch" + "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" + "github.com/coder/coder/v2/coderd/x/chatd/chathooks" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/x/agenthooks" ) type chatMutator struct { server *Server } +type chatMutation struct { + chat database.Chat +} + +func (m *chatMutator) update( + ctx context.Context, + chatID uuid.UUID, + operation string, + transition func(*chatstate.Tx, database.Store, *chatMutation) error, +) (database.Chat, error) { + mutation := chatMutation{} + err := m.server.newChatMachine(chatID).Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if err := transition(tx, store, &mutation); err != nil { + return err + } + if mutation.chat.ID != uuid.Nil { + return nil + } + chat, err := store.GetChatByID(ctx, chatID) + if err != nil { + return xerrors.Errorf("reload chat after %s: %w", operation, err) + } + mutation.chat = chat + return nil + }) + if err != nil { + return database.Chat{}, err + } + m.server.publishChatPubsubEvent(mutation.chat, codersdk.ChatWatchEventKindStatusChange, nil) + return mutation.chat, nil +} + +func (m *chatMutator) SendMessage( + ctx context.Context, + opts SendMessageOptions, +) (SendMessageResult, error) { + if opts.ChatID == uuid.Nil { + return SendMessageResult{}, xerrors.New("chat_id is required") + } + if len(opts.Content) == 0 { + return SendMessageResult{}, xerrors.New("content is required") + } + + busyBehavior := opts.BusyBehavior + if busyBehavior == "" { + busyBehavior = SendMessageBusyBehaviorQueue + } + switch busyBehavior { + case SendMessageBusyBehaviorQueue, SendMessageBusyBehaviorInterrupt: + default: + return SendMessageResult{}, xerrors.Errorf("invalid busy behavior %q", opts.BusyBehavior) + } + + contentParts := opts.Content + if m.server.hooks.Enabled() { + turnID := uuid.New() + chat, err := m.server.db.GetChatByID(ctx, opts.ChatID) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("load chat for user_prompt_submit: %w", err) + } + // Repeat these admission checks under the transaction lock. + if chat.Archived { + return SendMessageResult{}, ErrChatArchived + } + if _, err := resolveSendMessageModelConfigID(ctx, m.server.db, chat, opts.ModelConfigID); err != nil { + return SendMessageResult{}, err + } + // Check queue capacity before dispatch; the transaction + // rechecks it under lock. + queuedCount, err := m.server.db.CountChatQueuedMessages(ctx, opts.ChatID) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("count queued messages: %w", err) + } + if queuedCount >= chatstate.MaxQueueSize { + return SendMessageResult{}, &chatstate.MessageQueueFullError{Max: chatstate.MaxQueueSize} + } + promptMessage, err := chathooks.UserPromptMessage(contentParts) + if err != nil { + return SendMessageResult{}, err + } + promptResult, err := m.server.hooks.Trigger(ctx, chathooks.ChatFor(chat, &turnID), promptMessage, agenthooks.EventUserPromptSubmit, dispatch.CapacityClassAdmission) + if err != nil { + return SendMessageResult{}, m.server.handleUserPromptDispatchError(ctx, opts.ChatID, chathooks.UserPromptDenial(err)) + } + contentParts, _, err = chathooks.ComposeUserPromptContent(contentParts, promptResult) + if err != nil { + return SendMessageResult{}, err + } + } + + content, err := chatprompt.MarshalParts(contentParts) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("marshal message content: %w", err) + } + + requestedPlanMode := opts.PlanMode + requestedMCPServerIDs := opts.MCPServerIDs + + var result SendMessageResult + refreshed, updateErr := m.update(ctx, opts.ChatID, "send", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + lockedChat, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + + if lockedChat.Archived { + return ErrChatArchived + } + + if requestedPlanMode != nil { + lockedChat, err = store.UpdateChatPlanModeByID(ctx, database.UpdateChatPlanModeByIDParams{ + PlanMode: *requestedPlanMode, + ID: opts.ChatID, + }) + if err != nil { + return xerrors.Errorf("update chat plan mode: %w", err) + } + } + + modelConfigID, err := resolveSendMessageModelConfigID( + ctx, + store, + lockedChat, + opts.ModelConfigID, + ) + if err != nil { + return err + } + + lockedChat, err = m.server.applyRequestedMCPServerIDs(ctx, store, lockedChat, requestedMCPServerIDs) + if err != nil { + return err + } + + messageCreatedBy := opts.CreatedBy + if messageCreatedBy == uuid.Nil { + messageCreatedBy = lockedChat.OwnerID + } + + // Queue capacity is enforced inside tx.SendMessage; this + // wrapper only propagates the typed error. + message := userMessage(content, modelConfigID, messageCreatedBy, opts.ReasoningEffort) + sendResult, err := tx.SendMessage(chatstate.SendMessageInput{ + Message: message, + BusyBehavior: busyBehaviorToChatState(busyBehavior), + }) + if err != nil { + return err + } + + if sendResult.QueuedMessage != nil { + result.Queued = true + result.QueuedMessage = sendResult.QueuedMessage + } else if len(sendResult.InsertedMessages) > 0 { + // The state machine prepends synthetic tool-result + // cancellation messages; the user message is always + // last in the inserted slice. + result.Message = sendResult.InsertedMessages[len(sendResult.InsertedMessages)-1] + } + // A queued send on an errored chat can also promote the + // previous queue head into history; report those inserts so + // clients can update their caches. + result.InsertedMessages = sendResult.InsertedMessages + + // File-link errors must roll back the message. + return chatstate.LinkFiles(ctx, store, opts.ChatID, chatprompt.FileIDs(contentParts)) + }) + if updateErr != nil { + return SendMessageResult{}, updateErr + } + + result.Chat = refreshed + return result, nil +} + +func (m *chatMutator) EditMessage( + ctx context.Context, + opts EditMessageOptions, +) (EditMessageResult, error) { + if opts.ChatID == uuid.Nil { + return EditMessageResult{}, xerrors.New("chat_id is required") + } + if opts.EditedMessageID <= 0 { + return EditMessageResult{}, xerrors.New("edited_message_id is required") + } + if len(opts.Content) == 0 { + return EditMessageResult{}, xerrors.New("content is required") + } + + contentParts := opts.Content + var sessionStartHookResult *chathooks.Result + if m.server.hooks.Enabled() { + turnID := uuid.New() + chat, err := m.server.db.GetChatByID(ctx, opts.ChatID) + if err != nil { + return EditMessageResult{}, xerrors.Errorf("load chat for edit hooks: %w", err) + } + // Repeat these admission checks under the transaction lock. + if chat.Archived { + return EditMessageResult{}, ErrChatArchived + } + if err := validateEditTarget(ctx, m.server.db, opts.ChatID, opts.EditedMessageID); err != nil { + return EditMessageResult{}, err + } + if _, err := validateModelConfigOverride(ctx, m.server.db, chat.OrganizationID, opts.ModelConfigID); err != nil { + return EditMessageResult{}, err + } + sessionStartHookResult, err = m.server.hooks.Trigger(ctx, chathooks.ChatFor(chat, &turnID), chathooks.Message{Source: chathooks.SessionStartSourceClear}, agenthooks.EventSessionStart, dispatch.CapacityClassAdmission) + if err != nil { + return EditMessageResult{}, m.server.handleAPIDispatchError(ctx, opts.ChatID, agenthooks.EventSessionStart, err) + } + promptMessage, err := chathooks.UserPromptMessage(contentParts) + if err != nil { + return EditMessageResult{}, err + } + promptResult, err := m.server.hooks.Trigger(ctx, chathooks.ChatFor(chat, &turnID), promptMessage, agenthooks.EventUserPromptSubmit, dispatch.CapacityClassAdmission) + if err != nil { + return EditMessageResult{}, m.server.handleUserPromptDispatchError(ctx, opts.ChatID, chathooks.UserPromptDenial(err)) + } + contentParts, _, err = chathooks.ComposeUserPromptContent(contentParts, promptResult) + if err != nil { + return EditMessageResult{}, err + } + } + + content, err := chatprompt.MarshalParts(contentParts) + if err != nil { + return EditMessageResult{}, xerrors.Errorf("marshal message content: %w", err) + } + var ( + result EditMessageResult + editedMsg database.ChatMessage + editedCutoffT time.Time + ) + refreshed, err := m.update(ctx, opts.ChatID, "edit", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + lockedChat, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if lockedChat.Archived { + return ErrChatArchived + } + // Capture the target message for the post-commit debug + // cleanup hook below. The transition itself revalidates + // chat ownership and user-message constraints. + target, err := store.GetChatMessageByID(ctx, opts.EditedMessageID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrEditedMessageNotFound + } + return xerrors.Errorf("get edited message: %w", err) + } + if target.ChatID != opts.ChatID { + return ErrEditedMessageNotFound + } + if target.Deleted { + return ErrEditedMessageNotFound + } + if target.Role != database.ChatMessageRoleUser { + return ErrEditedMessageNotUser + } + editedMsg = target + + lockedChat, err = m.server.applyRequestedMCPServerIDs(ctx, store, lockedChat, opts.MCPServerIDs) + if err != nil { + return err + } + + modelOverride, err := validateModelConfigOverride(ctx, store, lockedChat.OrganizationID, opts.ModelConfigID) + if err != nil { + return err + } + if !modelOverride.Valid { + // Without an explicit override the transition preserves + // the edited message's original model, which may have been + // disabled since; resolve it like a normal message send. + preserved := uuid.Nil + if target.ModelConfigID.Valid { + preserved = target.ModelConfigID.UUID + } + resolved, err := resolveFallbackModelConfigID(ctx, store, lockedChat, preserved) + if err != nil { + return err + } + if resolved != preserved { + modelOverride = uuid.NullUUID{UUID: resolved, Valid: true} + } + } + + modelConfigID := target.ModelConfigID.UUID + if modelOverride.Valid { + modelConfigID = modelOverride.UUID + } + // The prompt response already rides in the replacement content; + // only the session_start(clear) response needs transcript rows. + // They insert after the replacement so a later edit's suffix + // truncation cleans them up. + suffixMessages, err := chathooks.EventMessages(sessionStartHookResult, modelConfigID) + if err != nil { + return err + } + + var reasoningEffortOverride database.NullChatReasoningEffort + if opts.ReasoningEffort != nil && *opts.ReasoningEffort != "" { + reasoningEffortOverride = database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffort(*opts.ReasoningEffort), Valid: true} + } + + editResult, err := tx.EditMessage(chatstate.EditMessageInput{ + MessageID: opts.EditedMessageID, + SuffixMessages: suffixMessages, + CreatedBy: opts.CreatedBy, + Content: content, + ModelConfigIDOverride: modelOverride, + ReasoningEffortOverride: reasoningEffortOverride, + }) + if err != nil { + if errors.Is(err, chatstate.ErrEditedMessageNotUser) { + return ErrEditedMessageNotUser + } + return err + } + result.Message = editResult.ReplacementMessage + inserted := make([]database.ChatMessage, 0, len(editResult.CancellationMessages)+len(editResult.SuffixMessages)+1) + inserted = append(inserted, editResult.CancellationMessages...) + inserted = append(inserted, editResult.ReplacementMessage) + inserted = append(inserted, editResult.SuffixMessages...) + result.InsertedMessages = inserted + result.DeletedMessageIDs = editResult.DeletedMessageIDs + return chatstate.LinkFiles(ctx, store, opts.ChatID, chatprompt.FileIDs(contentParts)) + }) + if err != nil { + return EditMessageResult{}, err + } + + result.Chat = refreshed + editedCutoffT = refreshed.UpdatedAt + + // Editing can race with an interrupted worker still flushing its + // final debug writes. Run a short bounded retry loop so we converge + // quickly without relying on the much longer stale-finalization + // sweep. Source editCutoff from the DB-stamped updated_at returned + // by the post-edit chat row so the filter uses the same clock that + // stamps replacement-turn debug rows; subtract + // debugCleanupClockSkew so replica clock drift cannot let the retry + // delete a replacement turn's debug rows. + editCutoff := editedCutoffT.Add(-debugCleanupClockSkew) + m.server.scheduleDebugCleanup( + ctx, + "failed to delete chat debug rows after edit", + []slog.Field{ + slog.F("chat_id", opts.ChatID), + slog.F("edited_message_id", editedMsg.ID), + }, + func(cleanupCtx context.Context, debugSvc *chatdebug.Service) error { + _, err := debugSvc.DeleteAfterMessageID(cleanupCtx, opts.ChatID, editedMsg.ID-1, editCutoff) + return err + }, + ) + + return result, nil +} + +type archiveMutation struct { + archived bool + watchKind codersdk.ChatWatchEventKind +} + +func (m *chatMutator) ArchiveChat(ctx context.Context, chat database.Chat) error { + return m.setChatFamilyArchived(ctx, chat, archiveMutation{ + archived: true, + watchKind: codersdk.ChatWatchEventKindDeleted, + }) +} + +func (m *chatMutator) UnarchiveChat(ctx context.Context, chat database.Chat) error { + return m.setChatFamilyArchived(ctx, chat, archiveMutation{ + watchKind: codersdk.ChatWatchEventKindCreated, + }) +} + +func (m *chatMutator) setChatFamilyArchived(ctx context.Context, chat database.Chat, mutation archiveMutation) error { + if chat.ID == uuid.Nil { + return xerrors.New("chat_id is required") + } + if chat.ParentChatID.Valid { + return ErrArchiveRequiresRootChat + } + + familyChats, err := chatstate.SetFamilyArchived(ctx, m.server.db, m.server.pubsub, chatstate.SetFamilyArchivedInput{ + RootID: chat.ID, + Archived: mutation.archived, + }) + if err != nil { + return err + } + if mutation.archived { + m.server.scheduleArchiveDebugCleanup(ctx, familyChats) + } + m.server.publishChatPubsubEvents(familyChats, mutation.watchKind) + return nil +} + func (m *chatMutator) DeleteQueued(ctx context.Context, chatID uuid.UUID, queuedMessageID int64) error { if chatID == uuid.Nil { return xerrors.New("chat_id is required") @@ -27,3 +441,286 @@ func (m *chatMutator) DeleteQueued(ctx context.Context, chatID uuid.UUID, queued return err }) } + +func (m *chatMutator) PromoteQueued( + ctx context.Context, + opts PromoteQueuedOptions, +) (PromoteQueuedResult, error) { + if opts.ChatID == uuid.Nil { + return PromoteQueuedResult{}, xerrors.New("chat_id is required") + } + + var result PromoteQueuedResult + _, updateErr := m.update(ctx, opts.ChatID, "promote", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + lockedChat, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if lockedChat.Archived { + return ErrChatArchived + } + + promoteResult, err := tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: opts.QueuedMessageID, + }) + if err != nil { + return err + } + if promoteResult.InsertedMessage != nil { + result.PromotedMessage = *promoteResult.InsertedMessage + } + return nil + }) + if updateErr != nil { + return PromoteQueuedResult{}, updateErr + } + + return result, nil +} + +func (m *chatMutator) SubmitToolResults( + ctx context.Context, + opts SubmitToolResultsOptions, +) error { + machine := m.server.newChatMachine(opts.ChatID) + var hookSuffix []chatstate.Message + if m.server.hooks.Enabled() { + state, err := loadDynamicPostToolUseState(ctx, machine, opts) + if err != nil { + return err + } + for _, result := range opts.Results { + response, err := m.server.hooks.Trigger(ctx, chathooks.ChatFor(state.chat, nil), chathooks.DynamicPostToolUseMessage(result, state.toolNames[result.ToolCallID]), agenthooks.EventPostToolUse, dispatch.CapacityClassGeneration) + if err != nil { + // Leave pending calls intact so the client can resubmit after recovery. + return chathooks.GenerationDispatchError(agenthooks.EventPostToolUse, err) + } + responseMessages, err := chathooks.EventMessages(response, state.modelConfigID) + if err != nil { + return err + } + hookSuffix = append(hookSuffix, responseMessages...) + } + } + + var statusConflict *ToolResultStatusConflictError + _, updateErr := m.update(ctx, opts.ChatID, "tool results", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + locked, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if locked.Archived { + return ErrChatArchived + } + + toolResults := make([]chatstate.ToolResultInput, 0, len(opts.Results)) + for _, result := range opts.Results { + toolResults = append(toolResults, chatstate.ToolResultInput{ + ToolCallID: result.ToolCallID, + Output: result.Output, + IsError: result.IsError, + }) + } + modelConfigID := opts.ModelConfigID + if modelConfigID == uuid.Nil { + modelConfigID = locked.LastModelConfigID + } + if _, err := tx.CompleteRequiresAction(chatstate.CompleteRequiresActionInput{ + CreatedBy: opts.UserID, + ModelConfigID: modelConfigID, + Results: toolResults, + SuffixMessages: hookSuffix, + }); err != nil { + if !errors.Is(err, chatstate.ErrInvalidState) && + locked.Status != database.ChatStatusRequiresAction && + errors.Is(err, chatstate.ErrTransitionNotAllowed) { + statusConflict = &ToolResultStatusConflictError{ + ActualStatus: locked.Status, + } + return statusConflict + } + return xerrors.Errorf("complete requires action: %w", err) + } + return nil + }) + if updateErr != nil { + if statusConflict != nil { + return statusConflict + } + return translateToolResultValidationError(updateErr) + } + + return nil +} + +func (m *chatMutator) InterruptChat( + ctx context.Context, + chat database.Chat, +) (database.Chat, error) { + if chat.ID == uuid.Nil { + return chat, xerrors.New("chat_id is required") + } + + refreshed, err := m.update(ctx, chat.ID, "interrupt", func(tx *chatstate.Tx, _ database.Store, _ *chatMutation) error { + if _, err := tx.Interrupt(chatstate.InterruptInput{ + Reason: "Tool execution interrupted by user", + }); err != nil { + return err + } + return nil + }) + if err != nil { + return chat, err + } + + return refreshed, nil +} + +func (m *chatMutator) CompactChat( + ctx context.Context, + chat database.Chat, +) (database.Chat, error) { + if chat.ID == uuid.Nil { + return chat, xerrors.New("chat_id is required") + } + + refreshed, err := m.update(ctx, chat.ID, "compact", func(tx *chatstate.Tx, store database.Store, mutation *chatMutation) error { + lockedChat, err := store.GetChatByID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if lockedChat.Archived { + return ErrChatArchived + } + // Run the transition before content and usage validation so busy + // chats surface the state conflict first. + result, err := tx.RequestCompaction(chatstate.RequestCompactionInput{}) + if err != nil { + return err + } + // Reject requests with nothing to compact inside the same + // transaction (rolling back the transition) so no LLM call + // is ever started for an empty or already-compacted chat. + // This also covers a double-/compact. + messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + if err != nil { + return xerrors.Errorf("load chat messages: %w", err) + } + boundary := latestContextBoundaryIndex(messages) + if _, ok := firstUncompressedAssistantAfter(messages, boundary); !ok { + return ErrNothingToCompact + } + mutation.chat = result.Chat + return nil + }) + if err != nil { + return chat, err + } + + return refreshed, nil +} + +func (m *chatMutator) ClearChat( + ctx context.Context, + chat database.Chat, +) (database.Chat, error) { + if chat.ID == uuid.Nil { + return chat, xerrors.New("chat_id is required") + } + + refreshed, err := m.update(ctx, chat.ID, "clear", func(tx *chatstate.Tx, store database.Store, mutation *chatMutation) error { + lockedChat, err := store.GetChatByID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if lockedChat.Archived { + return ErrChatArchived + } + // Read the pre-clear history before the transition inserts the + // boundary rows; eligibility is evaluated afterwards so busy + // chats surface the state conflict first. + messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + if err != nil { + return xerrors.Errorf("load chat messages: %w", err) + } + clearMessages, err := buildClearMessages(buildClearMessagesInput{ + modelConfigID: lockedChat.LastModelConfigID, + toolCallID: "chat_cleared_" + uuid.NewString(), + }) + if err != nil { + return xerrors.Errorf("build clear messages: %w", err) + } + result, err := tx.ClearContext(chatstate.ClearContextInput{Messages: clearMessages}) + if err != nil { + return err + } + // Reject no-op clears inside the same transaction so an empty + // or already-cleared chat never gains a duplicate boundary. + boundary := latestContextBoundaryIndex(messages) + if !hasClearableMessageAfter(messages, boundary) { + return ErrNothingToClear + } + mutation.chat = result.Chat + return nil + }) + if err != nil { + return chat, err + } + + return refreshed, nil +} + +func (m *chatMutator) ReconcileInvalidStateChat( + ctx context.Context, + chat database.Chat, +) (database.Chat, error) { + if chat.ID == uuid.Nil { + return chat, xerrors.New("chat_id is required") + } + + refreshed, err := m.update(ctx, chat.ID, "reconcile", func(tx *chatstate.Tx, _ database.Store, _ *chatMutation) error { + if _, err := tx.ReconcileInvalidState(chatstate.ReconcileInvalidStateInput{}); err != nil { + return err + } + return nil + }) + if err != nil { + return chat, err + } + + return refreshed, nil +} + +func (p *Server) publishChatPubsubEvents(chats []database.Chat, kind codersdk.ChatWatchEventKind) { + for _, chat := range chats { + p.publishChatPubsubEvent(chat, kind, nil) + } +} + +func (p *Server) publishChatPubsubEvent(chat database.Chat, kind codersdk.ChatWatchEventKind, diffStatus *codersdk.ChatDiffStatus) { + event := codersdk.ChatWatchEvent{ + Kind: kind, + Chat: chatWatchEventSDKChat(chat, diffStatus), + } + payload, err := json.Marshal(event) + if err != nil { + p.logger.Error(context.Background(), "failed to marshal chat pubsub event", + slog.F("chat_id", chat.ID), + slog.Error(err), + ) + return + } + if err := p.pubsub.Publish(coderdpubsub.ChatWatchEventChannel(chat.OwnerID), payload); err != nil { + p.logger.Error(context.Background(), "failed to publish chat pubsub event", + slog.F("chat_id", chat.ID), + slog.F("kind", kind), + slog.Error(err), + ) + } +} From 7b2d852754ef0b4192a4423ebb7aa2b6e88ea02d Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Sat, 29 Aug 2026 10:10:49 +0000 Subject: [PATCH 3/5] test(coderd/x/chatd): consolidate mutation coverage --- coderd/x/chatd/chatd.go | 36 ----- coderd/x/chatd/chatd_test.go | 169 ---------------------- coderd/x/chatd/mutator.go | 25 ++++ coderd/x/chatd/mutator_internal_test.go | 178 ++++++++++++++++++++++++ 4 files changed, 203 insertions(+), 205 deletions(-) create mode 100644 coderd/x/chatd/mutator_internal_test.go diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index f28a66fbb0d74..6919e2469b851 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1786,42 +1786,6 @@ func (p *Server) SubmitToolResults( return (&chatMutator{server: p}).SubmitToolResults(ctx, opts) } -// translateToolResultValidationError converts a chatstate tool-result -// validation error into the legacy chatd.ToolResultValidationError -// shape so HTTP handlers preserve their existing response detail. If -// err is not a tool-result validation error, it is returned -// unchanged. -func translateToolResultValidationError(err error) error { - var v *chatstate.ToolResultValidationError - if !errors.As(err, &v) { - return err - } - switch { - case xerrors.Is(v, chatstate.ErrToolResultDuplicate): - return &ToolResultValidationError{ - Message: "Duplicate tool_call_id in results.", - Detail: fmt.Sprintf("Duplicate tool call ID %q.", v.ToolCallID), - } - case xerrors.Is(v, chatstate.ErrToolResultMissing): - return &ToolResultValidationError{ - Message: "Missing tool result.", - Detail: fmt.Sprintf("Missing result for tool call %q.", v.ToolCallID), - } - case xerrors.Is(v, chatstate.ErrToolResultUnexpected): - return &ToolResultValidationError{ - Message: "Unexpected tool result.", - Detail: fmt.Sprintf("No pending tool call with ID %q.", v.ToolCallID), - } - case xerrors.Is(v, chatstate.ErrToolResultInvalidJSON): - return &ToolResultValidationError{ - Message: "Tool result output must be valid JSON.", - Detail: fmt.Sprintf("Output for tool call %q is not valid JSON.", v.ToolCallID), - } - default: - return err - } -} - // InterruptChat interrupts execution through the chatstate.Interrupt // transition. Active runs land in `interrupting`; requires-action // chats synthesize cancellation messages and return to running. diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index c2f022081d252..c208ffc1ca56e 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -1323,175 +1323,6 @@ func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) { "ask mode should continue exposing workspace MCP tools") } -// TestUnarchiveChildChat covers the deterministic branches of the -// Server.UnarchiveChat child path: every child unarchive attempt is -// rejected with chatd.ErrArchiveRequiresRootChat. -func TestUnarchiveChildChat(t *testing.T) { - t.Parallel() - - t.Run("ChildWithActiveParentRejected", func(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - parent, child := insertParentWithArchivedChild(ctx, t, db, user, org, model) - - err := replica.UnarchiveChat(ctx, child) - require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) - - dbChild, err := db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - require.True(t, dbChild.Archived, "child should remain archived") - - dbParent, err := db.GetChatByID(ctx, parent.ID) - require.NoError(t, err) - require.False(t, dbParent.Archived, "parent should stay active") - }) - - t.Run("ChildWithArchivedParentRejected", func(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - parent, child := insertParentWithArchivedChild(ctx, t, db, user, org, model) - _, err := db.ArchiveChatByID(ctx, parent.ID) - require.NoError(t, err) - - err = replica.UnarchiveChat(ctx, child) - require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) - - dbChild, err := db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - require.True(t, dbChild.Archived, "child should remain archived") - }) - - t.Run("ActiveChildRejected", func(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - _, child := insertParentWithActiveChild(t, db, user, org, model) - - err := replica.UnarchiveChat(ctx, child) - require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) - - dbChild, err := db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - require.False(t, dbChild.Archived, "child should stay active") - }) -} - -// TestArchiveChat_RejectsChildChat verifies that Server.ArchiveChat -// refuses every child chat with chatd.ErrArchiveRequiresRootChat -// regardless of the family's current archive state. Archive state -// changes must always be issued against the root chat so the whole -// family flips together. -func TestArchiveChat_RejectsChildChat(t *testing.T) { - t.Parallel() - - t.Run("ActiveChildRejected", func(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - parent, child := insertParentWithActiveChild(t, db, user, org, model) - - err := replica.ArchiveChat(ctx, child) - require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) - - dbChild, err := db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - require.False(t, dbChild.Archived, "child should stay active after rejected archive") - - dbParent, err := db.GetChatByID(ctx, parent.ID) - require.NoError(t, err) - require.False(t, dbParent.Archived, "parent should stay active after rejected child archive") - }) - - t.Run("AlreadyArchivedChildRejected", func(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - parent, child := insertParentWithArchivedChild(ctx, t, db, user, org, model) - - err := replica.ArchiveChat(ctx, child) - require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat, - "child archive must be rejected even when the child is already archived") - - dbChild, err := db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - require.True(t, dbChild.Archived, "child archived flag should not change") - - dbParent, err := db.GetChatByID(ctx, parent.ID) - require.NoError(t, err) - require.False(t, dbParent.Archived, "parent should stay active") - }) -} - -// insertParentWithActiveChild creates a parent chat and an active -// child chat linked to it. Both are returned in their initial -// (active) state. -func insertParentWithActiveChild( - t *testing.T, - db database.Store, - user database.User, - org database.Organization, - model database.ChatModelConfig, -) (parent database.Chat, child database.Chat) { - t.Helper() - parent = dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: model.ID, - Title: "parent", - }) - child = dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: model.ID, - Title: "child", - ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, - }) - return parent, child -} - -// insertParentWithArchivedChild creates an active parent and an -// individually-archived child. The returned child reflects its -// current (archived) state in the DB. -func insertParentWithArchivedChild( - ctx context.Context, - t *testing.T, - db database.Store, - user database.User, - org database.Organization, - model database.ChatModelConfig, -) (parent database.Chat, child database.Chat) { - t.Helper() - parent, child = insertParentWithActiveChild(t, db, user, org, model) - _, err := db.ArchiveChatByID(ctx, child.ID) - require.NoError(t, err) - child, err = db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - return parent, child -} - func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/mutator.go b/coderd/x/chatd/mutator.go index 813676bbe83a5..cf5589055a326 100644 --- a/coderd/x/chatd/mutator.go +++ b/coderd/x/chatd/mutator.go @@ -5,6 +5,7 @@ import ( "database/sql" "encoding/json" "errors" + "fmt" "time" "github.com/google/uuid" @@ -553,6 +554,30 @@ func (m *chatMutator) SubmitToolResults( return nil } +// translateToolResultValidationError converts a chatstate tool-result +// validation error into the legacy chatd.ToolResultValidationError +// shape so HTTP handlers preserve their existing response detail. If +// err is not a tool-result validation error, it is returned +// unchanged. +func translateToolResultValidationError(err error) error { + var v *chatstate.ToolResultValidationError + if !errors.As(err, &v) { + return err + } + switch { + case xerrors.Is(v, chatstate.ErrToolResultDuplicate): + return &ToolResultValidationError{Message: "Duplicate tool_call_id in results.", Detail: fmt.Sprintf("Duplicate tool call ID %q.", v.ToolCallID)} + case xerrors.Is(v, chatstate.ErrToolResultMissing): + return &ToolResultValidationError{Message: "Missing tool result.", Detail: fmt.Sprintf("Missing result for tool call %q.", v.ToolCallID)} + case xerrors.Is(v, chatstate.ErrToolResultUnexpected): + return &ToolResultValidationError{Message: "Unexpected tool result.", Detail: fmt.Sprintf("No pending tool call with ID %q.", v.ToolCallID)} + case xerrors.Is(v, chatstate.ErrToolResultInvalidJSON): + return &ToolResultValidationError{Message: "Tool result output must be valid JSON.", Detail: fmt.Sprintf("Output for tool call %q is not valid JSON.", v.ToolCallID)} + default: + return err + } +} + func (m *chatMutator) InterruptChat( ctx context.Context, chat database.Chat, diff --git a/coderd/x/chatd/mutator_internal_test.go b/coderd/x/chatd/mutator_internal_test.go new file mode 100644 index 0000000000000..643c2c3bf461e --- /dev/null +++ b/coderd/x/chatd/mutator_internal_test.go @@ -0,0 +1,178 @@ +package chatd + +import ( + "context" + "encoding/json" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/pubsub" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +type mutationRecorder struct { + pubsub.Pubsub + mu sync.Mutex + events []mutationEvent +} + +type mutationEvent struct { + channel string + payload []byte +} + +func (r *mutationRecorder) Publish(channel string, payload []byte) error { + r.mu.Lock() + r.events = append(r.events, mutationEvent{channel: channel, payload: append([]byte(nil), payload...)}) + r.mu.Unlock() + return r.Pubsub.Publish(channel, payload) +} + +func (r *mutationRecorder) snapshot() []mutationEvent { + r.mu.Lock() + defer r.mu.Unlock() + return append([]mutationEvent(nil), r.events...) +} + +func TestChatMutator(t *testing.T) { + t.Parallel() + + t.Run("update", func(t *testing.T) { + t.Parallel() + for _, tt := range []struct { + name string + fail bool + }{{name: "success"}, {name: "failure", fail: true}} { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + owner := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.ChatProvider(t, db, database.ChatProvider{Provider: "openai", DisplayName: "openai", BaseUrl: "http://example.invalid"}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{OrganizationID: org.ID, IsDefault: true}) + chat := dbgen.Chat(t, db, database.Chat{OrganizationID: org.ID, OwnerID: owner.ID, LastModelConfigID: model.ID}) + recorder := &mutationRecorder{Pubsub: ps} + mutator := chatMutator{server: &Server{db: db, pubsub: recorder, logger: slogtest.Make(t, nil)}} + wantErr := xerrors.New("transition failed") + + updated, err := mutator.update(ctx, chat.ID, "test", func(tx *chatstate.Tx, _ database.Store, _ *chatMutation) error { + if tt.fail { + return wantErr + } + _, err := tx.RequestCompaction(chatstate.RequestCompactionInput{}) + return err + }) + stored, loadErr := db.GetChatByID(ctx, chat.ID) + require.NoError(t, loadErr) + events := recorder.snapshot() + if tt.fail { + require.ErrorIs(t, err, wantErr) + require.Equal(t, chat.SnapshotVersion, stored.SnapshotVersion) + require.Empty(t, events) + return + } + + require.NoError(t, err) + require.Equal(t, stored, updated) + require.Greater(t, updated.SnapshotVersion, chat.SnapshotVersion) + require.Equal(t, database.ChatStatusRunning, updated.Status) + stateIndex, watchIndex := -1, -1 + for i, event := range events { + switch event.channel { + case coderdpubsub.ChatStateUpdateChannel(chat.ID): + stateIndex = i + case coderdpubsub.ChatWatchEventChannel(owner.ID): + watchIndex = i + var payload codersdk.ChatWatchEvent + require.NoError(t, json.Unmarshal(event.payload, &payload)) + require.Equal(t, codersdk.ChatWatchEventKindStatusChange, payload.Kind) + require.Equal(t, updated.ID, payload.Chat.ID) + require.Equal(t, codersdk.ChatStatus(updated.Status), payload.Chat.Status) + } + } + require.NotEqual(t, -1, stateIndex) + require.Greater(t, watchIndex, stateIndex) + }) + } + }) + + t.Run("validation", func(t *testing.T) { + t.Parallel() + mutator := chatMutator{server: &Server{}} + child := database.Chat{ID: uuid.New(), ParentChatID: uuid.NullUUID{UUID: uuid.New(), Valid: true}} + for _, tt := range []struct { + name string + call func() error + want string + }{ + {name: "send chat ID", call: func() error { _, err := mutator.SendMessage(context.Background(), SendMessageOptions{}); return err }, want: "chat_id is required"}, + {name: "send content", call: func() error { + _, err := mutator.SendMessage(context.Background(), SendMessageOptions{ChatID: uuid.New()}) + return err + }, want: "content is required"}, + {name: "send busy behavior", call: func() error { + _, err := mutator.SendMessage(context.Background(), SendMessageOptions{ChatID: uuid.New(), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hi")}, BusyBehavior: "invalid"}) + return err + }, want: "invalid busy behavior \"invalid\""}, + {name: "edit chat ID", call: func() error { _, err := mutator.EditMessage(context.Background(), EditMessageOptions{}); return err }, want: "chat_id is required"}, + {name: "edit message ID", call: func() error { + _, err := mutator.EditMessage(context.Background(), EditMessageOptions{ChatID: uuid.New()}) + return err + }, want: "edited_message_id is required"}, + {name: "edit content", call: func() error { + _, err := mutator.EditMessage(context.Background(), EditMessageOptions{ChatID: uuid.New(), EditedMessageID: 1}) + return err + }, want: "content is required"}, + {name: "archive child", call: func() error { return mutator.ArchiveChat(context.Background(), child) }, want: ErrArchiveRequiresRootChat.Error()}, + {name: "unarchive child", call: func() error { return mutator.UnarchiveChat(context.Background(), child) }, want: ErrArchiveRequiresRootChat.Error()}, + {name: "delete queued chat ID", call: func() error { return mutator.DeleteQueued(context.Background(), uuid.Nil, 1) }, want: "chat_id is required"}, + {name: "promote queued chat ID", call: func() error { + _, err := mutator.PromoteQueued(context.Background(), PromoteQueuedOptions{}) + return err + }, want: "chat_id is required"}, + {name: "interrupt chat ID", call: func() error { _, err := mutator.InterruptChat(context.Background(), database.Chat{}); return err }, want: "chat_id is required"}, + {name: "compact chat ID", call: func() error { _, err := mutator.CompactChat(context.Background(), database.Chat{}); return err }, want: "chat_id is required"}, + {name: "clear chat ID", call: func() error { _, err := mutator.ClearChat(context.Background(), database.Chat{}); return err }, want: "chat_id is required"}, + {name: "reconcile chat ID", call: func() error { + _, err := mutator.ReconcileInvalidStateChat(context.Background(), database.Chat{}) + return err + }, want: "chat_id is required"}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.EqualError(t, tt.call(), tt.want) + }) + } + }) + + t.Run("tool result error translation", func(t *testing.T) { + t.Parallel() + for _, tt := range []struct { + cause error + want string + }{ + {cause: chatstate.ErrToolResultDuplicate, want: "Duplicate tool_call_id in results.: Duplicate tool call ID \"call-1\"."}, + {cause: chatstate.ErrToolResultMissing, want: "Missing tool result.: Missing result for tool call \"call-1\"."}, + {cause: chatstate.ErrToolResultUnexpected, want: "Unexpected tool result.: No pending tool call with ID \"call-1\"."}, + {cause: chatstate.ErrToolResultInvalidJSON, want: "Tool result output must be valid JSON.: Output for tool call \"call-1\" is not valid JSON."}, + } { + t.Run(tt.cause.Error(), func(t *testing.T) { + t.Parallel() + err := translateToolResultValidationError(&chatstate.ToolResultValidationError{Cause: tt.cause, ToolCallID: "call-1"}) + require.EqualError(t, err, tt.want) + }) + } + }) +} From 24043e62916247283f26880e68e56ab36e09ff2d Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Sat, 29 Aug 2026 10:25:57 +0000 Subject: [PATCH 4/5] refactor(coderd/x/chatd): simplify mutation module --- coderd/x/chatd/chatd.go | 24 ++---- coderd/x/chatd/mutator.go | 106 ++++++++---------------- coderd/x/chatd/mutator_internal_test.go | 67 ++++----------- 3 files changed, 53 insertions(+), 144 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 6919e2469b851..731211937cec7 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1609,21 +1609,21 @@ func validateModelConfigOverride( return uuid.NullUUID{UUID: requested, Valid: true}, nil } -func validateEditTarget(ctx context.Context, store database.Store, chatID uuid.UUID, messageID int64) error { +func validateEditTarget(ctx context.Context, store database.Store, chatID uuid.UUID, messageID int64) (database.ChatMessage, error) { target, err := store.GetChatMessageByID(ctx, messageID) if err != nil { if errors.Is(err, sql.ErrNoRows) { - return ErrEditedMessageNotFound + return database.ChatMessage{}, ErrEditedMessageNotFound } - return xerrors.Errorf("get edited message: %w", err) + return database.ChatMessage{}, xerrors.Errorf("get edited message: %w", err) } if target.ChatID != chatID || target.Deleted { - return ErrEditedMessageNotFound + return database.ChatMessage{}, ErrEditedMessageNotFound } if target.Role != database.ChatMessageRoleUser { - return ErrEditedMessageNotUser + return database.ChatMessage{}, ErrEditedMessageNotUser } - return nil + return target, nil } func loadEffectiveChatModelConfigs( @@ -1709,13 +1709,6 @@ func (p *Server) UnarchiveChat(ctx context.Context, chat database.Chat) error { return (&chatMutator{server: p}).UnarchiveChat(ctx, chat) } -// setChatFamilyArchived applies SetArchived(archived) to every chat -// in chat's family through chatstate. The transaction-captured -// family rows feed the post-commit debug cleanup and sidebar watch -// events. Callers must only invoke this for root chats. -// -//nolint:revive // Existing API takes the target archive state as a boolean. - // DeleteQueued removes a queued user message through the chatstate // state machine. Stream side effects are handled by chat:update // consumers. @@ -2597,8 +2590,6 @@ func subscribeWithInitialError(chatID uuid.UUID, message string) ( }}, events, func() {}, true } -// publishChatPubsubEvents broadcasts a lifecycle event for each affected chat. - // chatWatchEventSDKChat builds the chat embedded in ChatWatchEvent // notifications. These payloads travel through PostgreSQL NOTIFY, so // omit fields that can grow large and that watch consumers already read @@ -2612,9 +2603,6 @@ func chatWatchEventSDKChat(chat database.Chat, diffStatus *codersdk.ChatDiffStat return sdkChat } -// publishChatPubsubEvent broadcasts a chat lifecycle event via PostgreSQL -// pubsub so that all replicas can push updates to watching clients. - // ChatQueuedForCapacity reports whether the chat is waiting for a // concurrent-agent capacity slot. Uncapped deployments always return false. func (p *Server) ChatQueuedForCapacity(ctx context.Context, chat database.Chat) (bool, error) { diff --git a/coderd/x/chatd/mutator.go b/coderd/x/chatd/mutator.go index cf5589055a326..a5c092ee4bab3 100644 --- a/coderd/x/chatd/mutator.go +++ b/coderd/x/chatd/mutator.go @@ -2,11 +2,9 @@ package chatd import ( "context" - "database/sql" "encoding/json" "errors" "fmt" - "time" "github.com/google/uuid" "golang.org/x/xerrors" @@ -27,36 +25,32 @@ type chatMutator struct { server *Server } -type chatMutation struct { - chat database.Chat -} - func (m *chatMutator) update( ctx context.Context, chatID uuid.UUID, operation string, - transition func(*chatstate.Tx, database.Store, *chatMutation) error, + transition func(*chatstate.Tx, database.Store, *database.Chat) error, ) (database.Chat, error) { - mutation := chatMutation{} + var updated database.Chat err := m.server.newChatMachine(chatID).Update(ctx, func(tx *chatstate.Tx, store database.Store) error { - if err := transition(tx, store, &mutation); err != nil { + if err := transition(tx, store, &updated); err != nil { return err } - if mutation.chat.ID != uuid.Nil { + if updated.ID != uuid.Nil { return nil } chat, err := store.GetChatByID(ctx, chatID) if err != nil { return xerrors.Errorf("reload chat after %s: %w", operation, err) } - mutation.chat = chat + updated = chat return nil }) if err != nil { return database.Chat{}, err } - m.server.publishChatPubsubEvent(mutation.chat, codersdk.ChatWatchEventKindStatusChange, nil) - return mutation.chat, nil + m.server.publishChatPubsubEvent(updated, codersdk.ChatWatchEventKindStatusChange, nil) + return updated, nil } func (m *chatMutator) SendMessage( @@ -126,7 +120,7 @@ func (m *chatMutator) SendMessage( requestedMCPServerIDs := opts.MCPServerIDs var result SendMessageResult - refreshed, updateErr := m.update(ctx, opts.ChatID, "send", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + refreshed, updateErr := m.update(ctx, opts.ChatID, "send", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error { lockedChat, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { return xerrors.Errorf("load chat: %w", err) @@ -228,7 +222,7 @@ func (m *chatMutator) EditMessage( if chat.Archived { return EditMessageResult{}, ErrChatArchived } - if err := validateEditTarget(ctx, m.server.db, opts.ChatID, opts.EditedMessageID); err != nil { + if _, err := validateEditTarget(ctx, m.server.db, opts.ChatID, opts.EditedMessageID); err != nil { return EditMessageResult{}, err } if _, err := validateModelConfigOverride(ctx, m.server.db, chat.OrganizationID, opts.ModelConfigID); err != nil { @@ -257,11 +251,10 @@ func (m *chatMutator) EditMessage( return EditMessageResult{}, xerrors.Errorf("marshal message content: %w", err) } var ( - result EditMessageResult - editedMsg database.ChatMessage - editedCutoffT time.Time + result EditMessageResult + editedMsg database.ChatMessage ) - refreshed, err := m.update(ctx, opts.ChatID, "edit", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + refreshed, err := m.update(ctx, opts.ChatID, "edit", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error { lockedChat, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { return xerrors.Errorf("load chat: %w", err) @@ -269,24 +262,9 @@ func (m *chatMutator) EditMessage( if lockedChat.Archived { return ErrChatArchived } - // Capture the target message for the post-commit debug - // cleanup hook below. The transition itself revalidates - // chat ownership and user-message constraints. - target, err := store.GetChatMessageByID(ctx, opts.EditedMessageID) + target, err := validateEditTarget(ctx, store, opts.ChatID, opts.EditedMessageID) if err != nil { - if errors.Is(err, sql.ErrNoRows) { - return ErrEditedMessageNotFound - } - return xerrors.Errorf("get edited message: %w", err) - } - if target.ChatID != opts.ChatID { - return ErrEditedMessageNotFound - } - if target.Deleted { - return ErrEditedMessageNotFound - } - if target.Role != database.ChatMessageRoleUser { - return ErrEditedMessageNotUser + return err } editedMsg = target @@ -362,7 +340,6 @@ func (m *chatMutator) EditMessage( } result.Chat = refreshed - editedCutoffT = refreshed.UpdatedAt // Editing can race with an interrupted worker still flushing its // final debug writes. Run a short bounded retry loop so we converge @@ -372,7 +349,7 @@ func (m *chatMutator) EditMessage( // stamps replacement-turn debug rows; subtract // debugCleanupClockSkew so replica clock drift cannot let the retry // delete a replacement turn's debug rows. - editCutoff := editedCutoffT.Add(-debugCleanupClockSkew) + editCutoff := refreshed.UpdatedAt.Add(-debugCleanupClockSkew) m.server.scheduleDebugCleanup( ctx, "failed to delete chat debug rows after edit", @@ -389,25 +366,15 @@ func (m *chatMutator) EditMessage( return result, nil } -type archiveMutation struct { - archived bool - watchKind codersdk.ChatWatchEventKind -} - func (m *chatMutator) ArchiveChat(ctx context.Context, chat database.Chat) error { - return m.setChatFamilyArchived(ctx, chat, archiveMutation{ - archived: true, - watchKind: codersdk.ChatWatchEventKindDeleted, - }) + return m.setChatFamilyArchived(ctx, chat, chatstate.SetFamilyArchivedInput{Archived: true}) } func (m *chatMutator) UnarchiveChat(ctx context.Context, chat database.Chat) error { - return m.setChatFamilyArchived(ctx, chat, archiveMutation{ - watchKind: codersdk.ChatWatchEventKindCreated, - }) + return m.setChatFamilyArchived(ctx, chat, chatstate.SetFamilyArchivedInput{}) } -func (m *chatMutator) setChatFamilyArchived(ctx context.Context, chat database.Chat, mutation archiveMutation) error { +func (m *chatMutator) setChatFamilyArchived(ctx context.Context, chat database.Chat, input chatstate.SetFamilyArchivedInput) error { if chat.ID == uuid.Nil { return xerrors.New("chat_id is required") } @@ -415,17 +382,17 @@ func (m *chatMutator) setChatFamilyArchived(ctx context.Context, chat database.C return ErrArchiveRequiresRootChat } - familyChats, err := chatstate.SetFamilyArchived(ctx, m.server.db, m.server.pubsub, chatstate.SetFamilyArchivedInput{ - RootID: chat.ID, - Archived: mutation.archived, - }) + input.RootID = chat.ID + familyChats, err := chatstate.SetFamilyArchived(ctx, m.server.db, m.server.pubsub, input) if err != nil { return err } - if mutation.archived { + watchKind := codersdk.ChatWatchEventKindCreated + if input.Archived { m.server.scheduleArchiveDebugCleanup(ctx, familyChats) + watchKind = codersdk.ChatWatchEventKindDeleted } - m.server.publishChatPubsubEvents(familyChats, mutation.watchKind) + m.server.publishChatPubsubEvents(familyChats, watchKind) return nil } @@ -452,7 +419,7 @@ func (m *chatMutator) PromoteQueued( } var result PromoteQueuedResult - _, updateErr := m.update(ctx, opts.ChatID, "promote", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + _, updateErr := m.update(ctx, opts.ChatID, "promote", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error { lockedChat, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { return xerrors.Errorf("load chat: %w", err) @@ -504,8 +471,7 @@ func (m *chatMutator) SubmitToolResults( } } - var statusConflict *ToolResultStatusConflictError - _, updateErr := m.update(ctx, opts.ChatID, "tool results", func(tx *chatstate.Tx, store database.Store, _ *chatMutation) error { + _, updateErr := m.update(ctx, opts.ChatID, "tool results", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error { locked, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { return xerrors.Errorf("load chat: %w", err) @@ -535,19 +501,13 @@ func (m *chatMutator) SubmitToolResults( if !errors.Is(err, chatstate.ErrInvalidState) && locked.Status != database.ChatStatusRequiresAction && errors.Is(err, chatstate.ErrTransitionNotAllowed) { - statusConflict = &ToolResultStatusConflictError{ - ActualStatus: locked.Status, - } - return statusConflict + return &ToolResultStatusConflictError{ActualStatus: locked.Status} } return xerrors.Errorf("complete requires action: %w", err) } return nil }) if updateErr != nil { - if statusConflict != nil { - return statusConflict - } return translateToolResultValidationError(updateErr) } @@ -586,7 +546,7 @@ func (m *chatMutator) InterruptChat( return chat, xerrors.New("chat_id is required") } - refreshed, err := m.update(ctx, chat.ID, "interrupt", func(tx *chatstate.Tx, _ database.Store, _ *chatMutation) error { + refreshed, err := m.update(ctx, chat.ID, "interrupt", func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error { if _, err := tx.Interrupt(chatstate.InterruptInput{ Reason: "Tool execution interrupted by user", }); err != nil { @@ -609,7 +569,7 @@ func (m *chatMutator) CompactChat( return chat, xerrors.New("chat_id is required") } - refreshed, err := m.update(ctx, chat.ID, "compact", func(tx *chatstate.Tx, store database.Store, mutation *chatMutation) error { + refreshed, err := m.update(ctx, chat.ID, "compact", func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error { lockedChat, err := store.GetChatByID(ctx, chat.ID) if err != nil { return xerrors.Errorf("load chat: %w", err) @@ -638,7 +598,7 @@ func (m *chatMutator) CompactChat( if _, ok := firstUncompressedAssistantAfter(messages, boundary); !ok { return ErrNothingToCompact } - mutation.chat = result.Chat + *updated = result.Chat return nil }) if err != nil { @@ -656,7 +616,7 @@ func (m *chatMutator) ClearChat( return chat, xerrors.New("chat_id is required") } - refreshed, err := m.update(ctx, chat.ID, "clear", func(tx *chatstate.Tx, store database.Store, mutation *chatMutation) error { + refreshed, err := m.update(ctx, chat.ID, "clear", func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error { lockedChat, err := store.GetChatByID(ctx, chat.ID) if err != nil { return xerrors.Errorf("load chat: %w", err) @@ -691,7 +651,7 @@ func (m *chatMutator) ClearChat( if !hasClearableMessageAfter(messages, boundary) { return ErrNothingToClear } - mutation.chat = result.Chat + *updated = result.Chat return nil }) if err != nil { @@ -709,7 +669,7 @@ func (m *chatMutator) ReconcileInvalidStateChat( return chat, xerrors.New("chat_id is required") } - refreshed, err := m.update(ctx, chat.ID, "reconcile", func(tx *chatstate.Tx, _ database.Store, _ *chatMutation) error { + refreshed, err := m.update(ctx, chat.ID, "reconcile", func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error { if _, err := tx.ReconcileInvalidState(chatstate.ReconcileInvalidStateInput{}); err != nil { return err } diff --git a/coderd/x/chatd/mutator_internal_test.go b/coderd/x/chatd/mutator_internal_test.go index 643c2c3bf461e..9d5c60c72b9c4 100644 --- a/coderd/x/chatd/mutator_internal_test.go +++ b/coderd/x/chatd/mutator_internal_test.go @@ -3,7 +3,6 @@ package chatd import ( "context" "encoding/json" - "sync" "testing" "github.com/google/uuid" @@ -13,38 +12,12 @@ import ( "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" - "github.com/coder/coder/v2/coderd/database/dbtestutil" - "github.com/coder/coder/v2/coderd/database/pubsub" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) -type mutationRecorder struct { - pubsub.Pubsub - mu sync.Mutex - events []mutationEvent -} - -type mutationEvent struct { - channel string - payload []byte -} - -func (r *mutationRecorder) Publish(channel string, payload []byte) error { - r.mu.Lock() - r.events = append(r.events, mutationEvent{channel: channel, payload: append([]byte(nil), payload...)}) - r.mu.Unlock() - return r.Pubsub.Publish(channel, payload) -} - -func (r *mutationRecorder) snapshot() []mutationEvent { - r.mu.Lock() - defer r.mu.Unlock() - return append([]mutationEvent(nil), r.events...) -} - func TestChatMutator(t *testing.T) { t.Parallel() @@ -57,26 +30,28 @@ func TestChatMutator(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - db, ps := dbtestutil.NewDB(t) - owner := dbgen.User(t, db, database.User{}) - org := dbgen.Organization(t, db, database.Organization{}) - dbgen.ChatProvider(t, db, database.ChatProvider{Provider: "openai", DisplayName: "openai", BaseUrl: "http://example.invalid"}) - model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{OrganizationID: org.ID, IsDefault: true}) - chat := dbgen.Chat(t, db, database.Chat{OrganizationID: org.ID, OwnerID: owner.ID, LastModelConfigID: model.ID}) - recorder := &mutationRecorder{Pubsub: ps} - mutator := chatMutator{server: &Server{db: db, pubsub: recorder, logger: slogtest.Make(t, nil)}} + fixture := newWorkerTestFixture(t) + chat := dbgen.Chat(t, fixture.db, database.Chat{ + OrganizationID: fixture.org.ID, + OwnerID: fixture.user.ID, + LastModelConfigID: fixture.model.ID, + }) + recorder := newRecordingPubsub(fixture.pubsub) + mutator := chatMutator{server: &Server{db: fixture.db, pubsub: recorder, logger: slogtest.Make(t, nil)}} wantErr := xerrors.New("transition failed") - updated, err := mutator.update(ctx, chat.ID, "test", func(tx *chatstate.Tx, _ database.Store, _ *chatMutation) error { + updated, err := mutator.update(ctx, chat.ID, "test", func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error { if tt.fail { return wantErr } _, err := tx.RequestCompaction(chatstate.RequestCompactionInput{}) return err }) - stored, loadErr := db.GetChatByID(ctx, chat.ID) + stored, loadErr := fixture.db.GetChatByID(ctx, chat.ID) require.NoError(t, loadErr) - events := recorder.snapshot() + recorder.mu.Lock() + events := append([]publishedEvent(nil), recorder.events...) + recorder.mu.Unlock() if tt.fail { require.ErrorIs(t, err, wantErr) require.Equal(t, chat.SnapshotVersion, stored.SnapshotVersion) @@ -93,7 +68,7 @@ func TestChatMutator(t *testing.T) { switch event.channel { case coderdpubsub.ChatStateUpdateChannel(chat.ID): stateIndex = i - case coderdpubsub.ChatWatchEventChannel(owner.ID): + case coderdpubsub.ChatWatchEventChannel(fixture.user.ID): watchIndex = i var payload codersdk.ChatWatchEvent require.NoError(t, json.Unmarshal(event.payload, &payload)) @@ -117,7 +92,6 @@ func TestChatMutator(t *testing.T) { call func() error want string }{ - {name: "send chat ID", call: func() error { _, err := mutator.SendMessage(context.Background(), SendMessageOptions{}); return err }, want: "chat_id is required"}, {name: "send content", call: func() error { _, err := mutator.SendMessage(context.Background(), SendMessageOptions{ChatID: uuid.New()}) return err @@ -126,7 +100,6 @@ func TestChatMutator(t *testing.T) { _, err := mutator.SendMessage(context.Background(), SendMessageOptions{ChatID: uuid.New(), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hi")}, BusyBehavior: "invalid"}) return err }, want: "invalid busy behavior \"invalid\""}, - {name: "edit chat ID", call: func() error { _, err := mutator.EditMessage(context.Background(), EditMessageOptions{}); return err }, want: "chat_id is required"}, {name: "edit message ID", call: func() error { _, err := mutator.EditMessage(context.Background(), EditMessageOptions{ChatID: uuid.New()}) return err @@ -137,18 +110,6 @@ func TestChatMutator(t *testing.T) { }, want: "content is required"}, {name: "archive child", call: func() error { return mutator.ArchiveChat(context.Background(), child) }, want: ErrArchiveRequiresRootChat.Error()}, {name: "unarchive child", call: func() error { return mutator.UnarchiveChat(context.Background(), child) }, want: ErrArchiveRequiresRootChat.Error()}, - {name: "delete queued chat ID", call: func() error { return mutator.DeleteQueued(context.Background(), uuid.Nil, 1) }, want: "chat_id is required"}, - {name: "promote queued chat ID", call: func() error { - _, err := mutator.PromoteQueued(context.Background(), PromoteQueuedOptions{}) - return err - }, want: "chat_id is required"}, - {name: "interrupt chat ID", call: func() error { _, err := mutator.InterruptChat(context.Background(), database.Chat{}); return err }, want: "chat_id is required"}, - {name: "compact chat ID", call: func() error { _, err := mutator.CompactChat(context.Background(), database.Chat{}); return err }, want: "chat_id is required"}, - {name: "clear chat ID", call: func() error { _, err := mutator.ClearChat(context.Background(), database.Chat{}); return err }, want: "chat_id is required"}, - {name: "reconcile chat ID", call: func() error { - _, err := mutator.ReconcileInvalidStateChat(context.Background(), database.Chat{}) - return err - }, want: "chat_id is required"}, } { t.Run(tt.name, func(t *testing.T) { t.Parallel() From 116d0c21c9e8dd844992b9a7056d6d7079c85428 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Sat, 29 Aug 2026 10:39:46 +0000 Subject: [PATCH 5/5] refactor(coderd/x/chatd): polish mutation comments --- coderd/x/chatd/mutator.go | 63 +++++++++++++-------------------------- 1 file changed, 20 insertions(+), 43 deletions(-) diff --git a/coderd/x/chatd/mutator.go b/coderd/x/chatd/mutator.go index a5c092ee4bab3..e00041730f731 100644 --- a/coderd/x/chatd/mutator.go +++ b/coderd/x/chatd/mutator.go @@ -81,15 +81,15 @@ func (m *chatMutator) SendMessage( if err != nil { return SendMessageResult{}, xerrors.Errorf("load chat for user_prompt_submit: %w", err) } - // Repeat these admission checks under the transaction lock. + // Hooks run before the transaction, so admission is rechecked. if chat.Archived { return SendMessageResult{}, ErrChatArchived } if _, err := resolveSendMessageModelConfigID(ctx, m.server.db, chat, opts.ModelConfigID); err != nil { return SendMessageResult{}, err } - // Check queue capacity before dispatch; the transaction - // rechecks it under lock. + // Avoid dispatching hooks for a known-full queue; the transaction + // rechecks it. queuedCount, err := m.server.db.CountChatQueuedMessages(ctx, opts.ChatID) if err != nil { return SendMessageResult{}, xerrors.Errorf("count queued messages: %w", err) @@ -160,8 +160,6 @@ func (m *chatMutator) SendMessage( messageCreatedBy = lockedChat.OwnerID } - // Queue capacity is enforced inside tx.SendMessage; this - // wrapper only propagates the typed error. message := userMessage(content, modelConfigID, messageCreatedBy, opts.ReasoningEffort) sendResult, err := tx.SendMessage(chatstate.SendMessageInput{ Message: message, @@ -175,14 +173,10 @@ func (m *chatMutator) SendMessage( result.Queued = true result.QueuedMessage = sendResult.QueuedMessage } else if len(sendResult.InsertedMessages) > 0 { - // The state machine prepends synthetic tool-result - // cancellation messages; the user message is always - // last in the inserted slice. + // When the message is not queued, cancellation rows precede it. result.Message = sendResult.InsertedMessages[len(sendResult.InsertedMessages)-1] } - // A queued send on an errored chat can also promote the - // previous queue head into history; report those inserts so - // clients can update their caches. + // Queued sends may also insert history rows. result.InsertedMessages = sendResult.InsertedMessages // File-link errors must roll back the message. @@ -218,7 +212,7 @@ func (m *chatMutator) EditMessage( if err != nil { return EditMessageResult{}, xerrors.Errorf("load chat for edit hooks: %w", err) } - // Repeat these admission checks under the transaction lock. + // Hooks run before the transaction, so admission is rechecked. if chat.Archived { return EditMessageResult{}, ErrChatArchived } @@ -278,9 +272,7 @@ func (m *chatMutator) EditMessage( return err } if !modelOverride.Valid { - // Without an explicit override the transition preserves - // the edited message's original model, which may have been - // disabled since; resolve it like a normal message send. + // The original model may be disabled, so use the normal fallback path. preserved := uuid.Nil if target.ModelConfigID.Valid { preserved = target.ModelConfigID.UUID @@ -298,10 +290,8 @@ func (m *chatMutator) EditMessage( if modelOverride.Valid { modelConfigID = modelOverride.UUID } - // The prompt response already rides in the replacement content; - // only the session_start(clear) response needs transcript rows. - // They insert after the replacement so a later edit's suffix - // truncation cleans them up. + // The replacement already contains the prompt response. Append only the + // session-start response so later edits discard it with the suffix. suffixMessages, err := chathooks.EventMessages(sessionStartHookResult, modelConfigID) if err != nil { return err @@ -341,14 +331,9 @@ func (m *chatMutator) EditMessage( result.Chat = refreshed - // Editing can race with an interrupted worker still flushing its - // final debug writes. Run a short bounded retry loop so we converge - // quickly without relying on the much longer stale-finalization - // sweep. Source editCutoff from the DB-stamped updated_at returned - // by the post-edit chat row so the filter uses the same clock that - // stamps replacement-turn debug rows; subtract - // debugCleanupClockSkew so replica clock drift cannot let the retry - // delete a replacement turn's debug rows. + // An interrupted worker may write stale debug rows after the edit. Use the + // database timestamp with a skew allowance so retries do not delete rows + // from the replacement turn. editCutoff := refreshed.UpdatedAt.Add(-debugCleanupClockSkew) m.server.scheduleDebugCleanup( ctx, @@ -514,11 +499,6 @@ func (m *chatMutator) SubmitToolResults( return nil } -// translateToolResultValidationError converts a chatstate tool-result -// validation error into the legacy chatd.ToolResultValidationError -// shape so HTTP handlers preserve their existing response detail. If -// err is not a tool-result validation error, it is returned -// unchanged. func translateToolResultValidationError(err error) error { var v *chatstate.ToolResultValidationError if !errors.As(err, &v) { @@ -577,16 +557,14 @@ func (m *chatMutator) CompactChat( if lockedChat.Archived { return ErrChatArchived } - // Run the transition before content and usage validation so busy - // chats surface the state conflict first. + // Run the transition first so busy chats report a state conflict before + // no-op validation. result, err := tx.RequestCompaction(chatstate.RequestCompactionInput{}) if err != nil { return err } - // Reject requests with nothing to compact inside the same - // transaction (rolling back the transition) so no LLM call - // is ever started for an empty or already-compacted chat. - // This also covers a double-/compact. + // Validate compactable history in the transaction so no-op requests roll + // back before starting an LLM call. messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: chat.ID, AfterID: 0, @@ -624,9 +602,8 @@ func (m *chatMutator) ClearChat( if lockedChat.Archived { return ErrChatArchived } - // Read the pre-clear history before the transition inserts the - // boundary rows; eligibility is evaluated afterwards so busy - // chats surface the state conflict first. + // Read history before adding the boundary, but transition first so busy + // chats report a state conflict before no-op validation. messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: chat.ID, AfterID: 0, @@ -645,8 +622,8 @@ func (m *chatMutator) ClearChat( if err != nil { return err } - // Reject no-op clears inside the same transaction so an empty - // or already-cleared chat never gains a duplicate boundary. + // Validate clearable history in the transaction so no-op requests cannot + // add a duplicate boundary. boundary := latestContextBoundaryIndex(messages) if !hasClearableMessageAfter(messages, boundary) { return ErrNothingToClear