diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 456031add59..731211937ce 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) { @@ -1760,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( @@ -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,53 +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) -} - -// 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. -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 + return (&chatMutator{server: p}).UnarchiveChat(ctx, chat) } // DeleteQueued removes a queued user message through the chatstate @@ -2114,18 +1717,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 @@ -2138,53 +1730,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. @@ -2230,124 +1776,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 -} - -// 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 - } + return (&chatMutator{server: p}).SubmitToolResults(ctx, opts) } // InterruptChat interrupts execution through the chatstate.Interrupt @@ -2361,34 +1790,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 @@ -2409,50 +1811,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 @@ -2466,56 +1825,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 @@ -2533,32 +1843,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 @@ -3305,13 +2590,6 @@ func subscribeWithInitialError(chatID uuid.UUID, message string) ( }}, events, func() {}, true } -// 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 // omit fields that can grow large and that watch consumers already read @@ -3325,30 +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. -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. func (p *Server) ChatQueuedForCapacity(ctx context.Context, chat database.Chat) (bool, error) { diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index c2f022081d2..c208ffc1ca5 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 new file mode 100644 index 00000000000..e00041730f7 --- /dev/null +++ b/coderd/x/chatd/mutator.go @@ -0,0 +1,688 @@ +package chatd + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "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 +} + +func (m *chatMutator) update( + ctx context.Context, + chatID uuid.UUID, + operation string, + transition func(*chatstate.Tx, database.Store, *database.Chat) error, +) (database.Chat, error) { + var updated database.Chat + err := m.server.newChatMachine(chatID).Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if err := transition(tx, store, &updated); err != nil { + return err + } + 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) + } + updated = chat + return nil + }) + if err != nil { + return database.Chat{}, err + } + m.server.publishChatPubsubEvent(updated, codersdk.ChatWatchEventKindStatusChange, nil) + return updated, 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) + } + // 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 + } + // 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) + } + 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, _ *database.Chat) 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 + } + + 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 { + // When the message is not queued, cancellation rows precede it. + result.Message = sendResult.InsertedMessages[len(sendResult.InsertedMessages)-1] + } + // Queued sends may also insert history rows. + 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) + } + // Hooks run before the transaction, so admission is rechecked. + 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 + ) + 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) + } + if lockedChat.Archived { + return ErrChatArchived + } + target, err := validateEditTarget(ctx, store, opts.ChatID, opts.EditedMessageID) + if err != nil { + return err + } + 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 { + // The original model may be disabled, so use the normal fallback path. + 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 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 + } + + 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 + + // 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, + "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 +} + +func (m *chatMutator) ArchiveChat(ctx context.Context, chat database.Chat) error { + 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, chatstate.SetFamilyArchivedInput{}) +} + +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") + } + if chat.ParentChatID.Valid { + return ErrArchiveRequiresRootChat + } + + input.RootID = chat.ID + familyChats, err := chatstate.SetFamilyArchived(ctx, m.server.db, m.server.pubsub, input) + if err != nil { + return err + } + watchKind := codersdk.ChatWatchEventKindCreated + if input.Archived { + m.server.scheduleArchiveDebugCleanup(ctx, familyChats) + watchKind = codersdk.ChatWatchEventKindDeleted + } + m.server.publishChatPubsubEvents(familyChats, 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") + } + + 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 + }) +} + +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, _ *database.Chat) 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...) + } + } + + _, 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) + } + 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) { + return &ToolResultStatusConflictError{ActualStatus: locked.Status} + } + return xerrors.Errorf("complete requires action: %w", err) + } + return nil + }) + if updateErr != nil { + return translateToolResultValidationError(updateErr) + } + + return nil +} + +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, +) (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, _ *database.Chat) 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, updated *database.Chat) 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 first so busy chats report a state conflict before + // no-op validation. + result, err := tx.RequestCompaction(chatstate.RequestCompactionInput{}) + if err != nil { + return err + } + // 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, + }) + if err != nil { + return xerrors.Errorf("load chat messages: %w", err) + } + boundary := latestContextBoundaryIndex(messages) + if _, ok := firstUncompressedAssistantAfter(messages, boundary); !ok { + return ErrNothingToCompact + } + *updated = 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, updated *database.Chat) error { + lockedChat, err := store.GetChatByID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if lockedChat.Archived { + return ErrChatArchived + } + // 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, + }) + 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 + } + // Validate clearable history in the transaction so no-op requests cannot + // add a duplicate boundary. + boundary := latestContextBoundaryIndex(messages) + if !hasClearableMessageAfter(messages, boundary) { + return ErrNothingToClear + } + *updated = 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, _ *database.Chat) 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), + ) + } +} diff --git a/coderd/x/chatd/mutator_internal_test.go b/coderd/x/chatd/mutator_internal_test.go new file mode 100644 index 00000000000..9d5c60c72b9 --- /dev/null +++ b/coderd/x/chatd/mutator_internal_test.go @@ -0,0 +1,139 @@ +package chatd + +import ( + "context" + "encoding/json" + "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" + 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" +) + +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) + 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, _ *database.Chat) error { + if tt.fail { + return wantErr + } + _, err := tx.RequestCompaction(chatstate.RequestCompactionInput{}) + return err + }) + stored, loadErr := fixture.db.GetChatByID(ctx, chat.ID) + require.NoError(t, loadErr) + 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) + 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(fixture.user.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 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 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()}, + } { + 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) + }) + } + }) +}