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; }