From f1d160c7f44b97fdd721e15cbdddcd49f442ce3d Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 12 May 2026 14:51:55 +0200 Subject: [PATCH] fix: allow changing model when editing earlier chat message (#25084) Editing a previous user message and selecting a different model in the picker silently kept using the original model: the selection was dropped on the frontend, in the SDK, and in the backend, so both the replacement user message and the assistant turn that followed ran against the old model. Plumb the selected model through all three layers (`AgentChatPage`, `codersdk.EditChatMessageRequest`, `chatd.EditMessageOptions` / `Server.EditMessage`), defaulting to the original message's model when the client does not specify one. The existing `InsertChatMessages` CTE already advances `chats.last_model_config_id` when the inserted message's model differs, so the assistant turn picks up the new selection without further changes. The new model is validated inside the transaction, so an unknown ID rolls the edit back and returns a 400 `Invalid model config ID.`, mirroring the `SendMessage` path. Refs: CODAGT-345 This change was generated by a Coder agent.
Implementation plan # CODAGT-345: Editing an earlier message cannot change model ## Problem When editing a previous user message in a chat, the user can change the model in the model picker, but the backend keeps using the original message's model. The model selection is dropped at three layers: 1. **Frontend:** `AgentChatPage.tsx`'s edit branch builds an `EditChatMessageRequest` that omits `model_config_id`. The new-message branch (a few lines below) does include it. 2. **SDK:** `codersdk.EditChatMessageRequest` has no `ModelConfigID` field at all. 3. **Backend:** `chatd.EditMessageOptions` has no model field, and `Server.EditMessage` always copies the original message's `ModelConfigID` into the replacement message. Once the replacement user message is inserted with the original model, the `InsertChatMessages` CTE leaves `chats.last_model_config_id` unchanged, so the assistant turn that follows runs against the old model. ## Fix Plumb the selected model through all three layers, defaulting to the original message's model when the client doesn't override it. This mirrors the `SendMessage` path, which already accepts a `model_config_id` and validates it via `resolveSendMessageModelConfigID`. ### Backend - `codersdk/chats.go`: add `ModelConfigID *uuid.UUID` to `EditChatMessageRequest`. - `coderd/x/chatd/chatd.go`: - Add `ModelConfigID uuid.UUID` to `EditMessageOptions`. - In `EditMessage`, after fetching the edited message, resolve the model: if `opts.ModelConfigID != uuid.Nil`, validate it exists with `tx.GetChatModelConfigByID` (using `chatdModelConfigLookupContext`), otherwise keep `editedMsg.ModelConfigID.UUID`. Pass the resolved ID into `newChatMessage(...)`. - Reuse the existing `ErrInvalidModelConfigID` sentinel. - `coderd/exp_chats.go` (`patchChatMessage`): - Read `req.ModelConfigID` (nil-safe), pass into `chatd.EditMessageOptions`. - Add a `case xerrors.Is(editErr, chatd.ErrInvalidModelConfigID)` arm returning 400 `Invalid model config ID.`, matching the `postChatMessages` handler. ### Frontend - `site/src/pages/AgentsPage/AgentChatPage.tsx`: - In the edit branch, set `model_config_id: effectiveSelectedModel || undefined` on the `EditChatMessageRequest`. - On success, persist the chosen model to `lastModelConfigIDStorageKey` so the next chat from this browser keeps the same default. Mirrors the new-message branch. ### Generated - `make site/src/api/typesGenerated.ts` and `make coderd/apidoc/swagger.json` produce the updated `EditChatMessageRequest` schema in `typesGenerated.ts`, `coderd/apidoc/{docs.go,swagger.json}`, and `docs/reference/api/{chats.md,schemas.md}`. ## Tests - `coderd/x/chatd/chatd_test.go`: - `TestEditMessageWithModelConfigOverride`: edit with a different model -> replacement message and `chats.LastModelConfigID` use the new model. - `TestEditMessagePreservesModelConfigByDefault`: edit without `ModelConfigID` -> original model preserved. - `TestEditMessageRejectsUnknownModelConfig`: passes a random UUID -> `ErrInvalidModelConfigID`, original message still present, `LastModelConfigID` unchanged (rollback). - `coderd/exp_chats_test.go` (under `TestPatchChatMessage`): - `ChangesModel`: end-to-end via SDK; `edited.Message.ModelConfigID` and `chat.LastModelConfigID` both match the new model. - `InvalidModelConfigID`: random UUID -> 400 `Invalid model config ID.`.
--- coderd/apidoc/docs.go | 5 + coderd/apidoc/swagger.json | 5 + coderd/exp_chats.go | 10 ++ coderd/exp_chats_test.go | 95 +++++++++++++ coderd/x/chatd/chatd.go | 37 ++++- coderd/x/chatd/chatd_test.go | 149 ++++++++++++++++++++ codersdk/chats.go | 4 + docs/reference/api/chats.md | 3 +- docs/reference/api/schemas.md | 10 +- site/src/api/typesGenerated.ts | 6 + site/src/pages/AgentsPage/AgentChatPage.tsx | 26 +++- 11 files changed, 342 insertions(+), 8 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index b9364631c4..3286d03eda 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -17983,6 +17983,11 @@ const docTemplate = `{ "items": { "$ref": "#/definitions/codersdk.ChatInputPart" } + }, + "model_config_id": { + "description": "ModelConfigID, when set, overrides the model used for the\nreplacement user message and the assistant turn that follows.\nWhen nil the original message's model is preserved.", + "type": "string", + "format": "uuid" } } }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index f7b890c3a8..4c77d86943 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -16336,6 +16336,11 @@ "items": { "$ref": "#/definitions/codersdk.ChatInputPart" } + }, + "model_config_id": { + "description": "ModelConfigID, when set, overrides the model used for the\nreplacement user message and the assistant turn that follows.\nWhen nil the original message's model is preserved.", + "type": "string", + "format": "uuid" } } }, diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index f7ed4b8498..ffa2c3f8dc 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -3060,11 +3060,17 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { return } + editModelConfigID := uuid.Nil + if req.ModelConfigID != nil { + editModelConfigID = *req.ModelConfigID + } + editResult, editErr := api.chatDaemon.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, CreatedBy: apiKey.UserID, EditedMessageID: messageID, Content: contentBlocks, + ModelConfigID: editModelConfigID, }) if editErr != nil { if maybeWriteLimitErr(ctx, rw, editErr) { @@ -3085,6 +3091,10 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Only user messages can be edited.", }) + case xerrors.Is(editErr, chatd.ErrInvalidModelConfigID): + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Invalid model config ID.", + }) default: httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to edit chat message.", diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 9692f0b024..a41d365524 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -7279,6 +7279,101 @@ func TestPatchChatMessage(t *testing.T) { sdkErr := requireSDKError(t, err, http.StatusBadRequest) require.Contains(t, sdkErr.Message, "archived") }) + + t.Run("ChangesModel", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + defaultModel := createChatModelConfig(t, client) + overrideModel := createAdditionalChatModelConfig( + t, + client, + "openai", + "gpt-4o-mini-edit-override", + ) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "hello before edit", + }}, + }) + require.NoError(t, err) + require.Equal(t, defaultModel.ID, chat.LastModelConfigID, + "chat starts on the default model") + + messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil) + require.NoError(t, err) + var userMessageID int64 + for _, message := range messagesResult.Messages { + if message.Role == codersdk.ChatMessageRoleUser { + userMessageID = message.ID + break + } + } + require.NotZero(t, userMessageID) + + edited, err := client.EditChatMessage(ctx, chat.ID, userMessageID, codersdk.EditChatMessageRequest{ + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "hello after edit with new model", + }}, + ModelConfigID: &overrideModel.ID, + }) + require.NoError(t, err) + require.NotNil(t, edited.Message.ModelConfigID, + "edited message must carry a model config") + require.Equal(t, overrideModel.ID, *edited.Message.ModelConfigID, + "replacement message must use the requested model") + + updatedChat, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, overrideModel.ID, updatedChat.LastModelConfigID, + "chat last_model_config_id must advance so the next assistant turn uses the new model") + }) + + t.Run("InvalidModelConfigID", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "hello", + }}, + }) + require.NoError(t, err) + + messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil) + require.NoError(t, err) + var userMessageID int64 + for _, message := range messagesResult.Messages { + if message.Role == codersdk.ChatMessageRoleUser { + userMessageID = message.ID + break + } + } + require.NotZero(t, userMessageID) + + unknownID := uuid.New() + _, err = client.EditChatMessage(ctx, chat.ID, userMessageID, codersdk.EditChatMessageRequest{ + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "edited", + }}, + ModelConfigID: &unknownID, + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "Invalid model config ID.", sdkErr.Message) + }) } func TestStreamChat(t *testing.T) { diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 00da401d33..bcb493eb64 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1411,6 +1411,10 @@ type EditMessageOptions struct { CreatedBy uuid.UUID EditedMessageID int64 Content []codersdk.ChatMessagePart + // ModelConfigID, when non-zero, overrides the model used for + // the replacement user message. When set to uuid.Nil the + // original message's model is preserved. + ModelConfigID uuid.UUID } // EditMessageResult contains the replacement user message and chat status. @@ -1974,7 +1978,36 @@ func (p *Server) EditMessage( return xerrors.Errorf("soft-delete later chat messages: %w", err) } - // Insert a new message with the updated content. + // Resolve the model for the replacement message. When the + // caller does not specify a model, preserve the original + // message's model so an edit that only changes text keeps + // behaving as before. + messageModelConfigID := editedMsg.ModelConfigID.UUID + if opts.ModelConfigID != uuid.Nil { + if _, err := tx.GetChatModelConfigByID( + chatdModelConfigLookupContext(ctx), + opts.ModelConfigID, + ); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf( + "%w: %s", + ErrInvalidModelConfigID, + opts.ModelConfigID, + ) + } + return xerrors.Errorf( + "get requested model config %s: %w", + opts.ModelConfigID, + err, + ) + } + messageModelConfigID = opts.ModelConfigID + } + + // Insert a new message with the updated content. The + // InsertChatMessages CTE updates chats.last_model_config_id + // when the new message's model differs, so the assistant turn + // that follows picks up the new selection. msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. ChatID: opts.ChatID, } @@ -1982,7 +2015,7 @@ func (p *Server) EditMessage( database.ChatMessageRoleUser, content, editedMsg.Visibility, - editedMsg.ModelConfigID.UUID, + messageModelConfigID, chatprompt.CurrentContentVersion, ).withCreatedBy(opts.CreatedBy)) newMessages, err := insertChatMessageWithStore(ctx, tx, msgParams) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 6074bebd3b..e8430ee232 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -8734,6 +8734,155 @@ func TestEditMessageRejectsArchivedChat(t *testing.T) { require.ErrorIs(t, err, chatd.ErrChatArchived) } +// TestEditMessageWithModelConfigOverride verifies that callers can +// change the model when editing a previous user message. The +// replacement message must persist with the new model and the chat's +// LastModelConfigID must be advanced so the assistant turn that follows +// runs against the new selection. +func TestEditMessageWithModelConfigOverride(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, modelA := seedChatDependencies(t, db) + modelB := insertChatModelConfigWithCallConfig( + t, + db, + user.ID, + "openai", + "gpt-4o-mini-edit-"+uuid.NewString(), + codersdk.ChatModelCallConfig{}, + ) + + chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + OrganizationID: org.ID, + Title: "edit-with-model-override", + ModelConfigID: modelA.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, + }) + require.NoError(t, err) + + initial, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + require.NoError(t, err) + require.Len(t, initial, 1) + require.Equal(t, modelA.ID, initial[0].ModelConfigID.UUID) + + result, err := replica.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + EditedMessageID: initial[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, + ModelConfigID: modelB.ID, + }) + require.NoError(t, err) + require.True(t, result.Message.ModelConfigID.Valid) + require.Equal(t, modelB.ID, result.Message.ModelConfigID.UUID) + + storedChat, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, modelB.ID, storedChat.LastModelConfigID, + "edit must update last_model_config_id so the assistant turn picks up the new model") +} + +// TestEditMessagePreservesModelConfigByDefault verifies that omitting +// ModelConfigID on edit keeps the original message's model. This is the +// existing default for callers that only edit the text. +func TestEditMessagePreservesModelConfigByDefault(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, modelA := seedChatDependencies(t, db) + + chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + OrganizationID: org.ID, + Title: "edit-preserves-model", + ModelConfigID: modelA.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, + }) + require.NoError(t, err) + + initial, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + require.NoError(t, err) + require.Len(t, initial, 1) + + result, err := replica.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + EditedMessageID: initial[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, + }) + require.NoError(t, err) + require.True(t, result.Message.ModelConfigID.Valid) + require.Equal(t, modelA.ID, result.Message.ModelConfigID.UUID) + + storedChat, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, modelA.ID, storedChat.LastModelConfigID, + "edit without model override must not change last_model_config_id") +} + +// TestEditMessageRejectsUnknownModelConfig verifies the edit handler +// returns ErrInvalidModelConfigID when the requested model does not +// exist, mirroring SendMessage's validation. +func TestEditMessageRejectsUnknownModelConfig(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, modelA := seedChatDependencies(t, db) + + chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + OrganizationID: org.ID, + Title: "edit-unknown-model", + ModelConfigID: modelA.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, + }) + require.NoError(t, err) + + initial, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + require.NoError(t, err) + require.Len(t, initial, 1) + + _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + EditedMessageID: initial[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, + ModelConfigID: uuid.New(), + }) + require.ErrorIs(t, err, chatd.ErrInvalidModelConfigID) + + // The edit must roll back: the original message should still be + // present and the chat's LastModelConfigID unchanged. + stillThere, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + require.NoError(t, err) + require.Len(t, stillThere, 1) + require.Equal(t, initial[0].ID, stillThere[0].ID) + + storedChat, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, modelA.ID, storedChat.LastModelConfigID) +} + func TestPromoteQueuedRejectsArchivedChat(t *testing.T) { t.Parallel() diff --git a/codersdk/chats.go b/codersdk/chats.go index fcbba8b5ec..7b0f8c198b 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -521,6 +521,10 @@ type CreateChatMessageRequest struct { // EditChatMessageRequest is the request to edit a user message in a chat. type EditChatMessageRequest struct { Content []ChatInputPart `json:"content"` + // ModelConfigID, when set, overrides the model used for the + // replacement user message and the assistant turn that follows. + // When nil the original message's model is preserved. + ModelConfigID *uuid.UUID `json:"model_config_id,omitempty" format:"uuid"` } // CreateChatMessageResponse is the response from adding a message to a chat. diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index a29c26f1da..15ecdd595d 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -2013,7 +2013,8 @@ Experimental: this endpoint is subject to change. "text": "string", "type": "text" } - ] + ], + "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205" } ``` diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 80f1e39513..729b6946f0 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -6598,15 +6598,17 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "text": "string", "type": "text" } - ] + ], + "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205" } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|-----------|-----------------------------------------------------------|----------|--------------|-------------| -| `content` | array of [codersdk.ChatInputPart](#codersdkchatinputpart) | false | | | +| Name | Type | Required | Restrictions | Description | +|-------------------|-----------------------------------------------------------|----------|--------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `content` | array of [codersdk.ChatInputPart](#codersdkchatinputpart) | false | | | +| `model_config_id` | string | false | | Model config ID when set, overrides the model used for the replacement user message and the assistant turn that follows. When nil the original message's model is preserved. | ## codersdk.EditChatMessageResponse diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 1b5cc35a9c..51d265e8d2 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -3879,6 +3879,12 @@ export interface DynamicToolResponse { */ export interface EditChatMessageRequest { readonly content: readonly ChatInputPart[]; + /** + * ModelConfigID, when set, overrides the model used for the + * replacement user message and the assistant turn that follows. + * When nil the original message's model is preserved. + */ + readonly model_config_id?: string; } // From codersdk/chats.go diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 7ddce57d1f..5840dec74b 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -1374,10 +1374,28 @@ const AgentChatPage: FC = () => { ]); if (editedMessageID !== undefined) { - const request: TypesGen.EditChatMessageRequest = { content }; const originalEditedMessage = chatMessagesList?.find( (existingMessage) => existingMessage.id === editedMessageID, ); + const originalModelConfigID = originalEditedMessage?.model_config_id; + const pickerModelConfigID = effectiveSelectedModel || undefined; + const originalIsSelectable = + originalModelConfigID !== undefined && + modelOptions.some((opt) => opt.id === originalModelConfigID); + // Only override the original model when the user has switched to + // a different selectable option. If the original is no longer + // selectable, the picker is showing a fallback we should not + // silently use; let the backend preserve the original. + const editSelectedModelConfigID = + pickerModelConfigID && + originalIsSelectable && + pickerModelConfigID !== originalModelConfigID + ? pickerModelConfigID + : undefined; + const request: TypesGen.EditChatMessageRequest = { + content, + model_config_id: editSelectedModelConfigID, + }; const optimisticMessage = originalEditedMessage ? buildOptimisticEditedMessage({ requestContent: request.content, @@ -1406,6 +1424,12 @@ const AgentChatPage: FC = () => { handleUsageLimitError(error); }, }); + if (editSelectedModelConfigID) { + localStorage.setItem( + lastModelConfigIDStorageKey, + editSelectedModelConfigID, + ); + } return; }