diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 81bb7709f9e..ca210ce8e51 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -21986,6 +21986,14 @@ const docTemplate = `{ "$ref": "#/definitions/codersdk.ChatInputPart" } }, + "mcp_server_ids": { + "description": "MCPServerIDs, when set, replaces the chat's MCP server selection\nbefore the replacement turn runs. When nil the current selection\nis preserved.", + "type": "array", + "items": { + "type": "string", + "format": "uuid" + } + }, "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", diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 0d53012fdbf..ac5ae7264ba 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -20022,6 +20022,14 @@ "$ref": "#/definitions/codersdk.ChatInputPart" } }, + "mcp_server_ids": { + "description": "MCPServerIDs, when set, replaces the chat's MCP server selection\nbefore the replacement turn runs. When nil the current selection\nis preserved.", + "type": "array", + "items": { + "type": "string", + "format": "uuid" + } + }, "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", diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 362eaf2ff07..1a2f67d660e 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -1102,6 +1102,42 @@ func validateChatMCPServerIDs( return unique, invalid, nil } +// normalizeRequestedChatMCPServerIDs validates a request's MCP server +// selection for an existing chat. When requested is nil there is no +// change to make. IDs already persisted on the chat are exempt from the +// enabled-in-organization check: a server that is disabled or revoked +// after selection must not block sends. The generation path skips +// servers the chat can no longer use, and keeping the ID preserves the +// selection if the server is re-enabled. A non-nil response indicates +// the caller must write it with the returned status and stop. +func (api *API) normalizeRequestedChatMCPServerIDs(ctx context.Context, chat database.Chat, requested *[]uuid.UUID) (*[]uuid.UUID, int, *codersdk.Response) { + if requested == nil { + return nil, 0, nil + } + normalized, invalid, err := validateChatMCPServerIDs(ctx, api.Database, chat.OrganizationID, *requested) + if err != nil { + return nil, http.StatusInternalServerError, &codersdk.Response{ + Message: "Failed to validate MCP server IDs.", + Detail: err.Error(), + } + } + persisted := make(map[uuid.UUID]struct{}, len(chat.MCPServerIDs)) + for _, id := range chat.MCPServerIDs { + persisted[id] = struct{}{} + } + newlyInvalid := make([]uuid.UUID, 0, len(invalid)) + for _, id := range invalid { + if _, ok := persisted[id]; !ok { + newlyInvalid = append(newlyInvalid, id) + } + } + if len(newlyInvalid) > 0 { + resp := invalidChatMCPServerIDsResponse(newlyInvalid) + return nil, http.StatusBadRequest, &resp + } + return &normalized, 0, nil +} + func invalidChatMCPServerIDsResponse(ids []uuid.UUID) codersdk.Response { invalid := make([]string, 0, len(ids)) for _, id := range ids { @@ -2555,36 +2591,12 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { return } - if req.MCPServerIDs != nil { - normalizedMCPServerIDs, invalidMCPServerIDs, err := validateChatMCPServerIDs(ctx, api.Database, chat.OrganizationID, *req.MCPServerIDs) - if err != nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to validate MCP server IDs.", - Detail: err.Error(), - }) - return - } - req.MCPServerIDs = &normalizedMCPServerIDs - // IDs already persisted on the chat are exempt: a server that - // is disabled or revoked after selection must not block sends. - // The generation path skips servers the chat can no longer use, - // and keeping the ID preserves the selection if the server is - // re-enabled. - persisted := make(map[uuid.UUID]struct{}, len(chat.MCPServerIDs)) - for _, id := range chat.MCPServerIDs { - persisted[id] = struct{}{} - } - newlyInvalid := make([]uuid.UUID, 0, len(invalidMCPServerIDs)) - for _, id := range invalidMCPServerIDs { - if _, ok := persisted[id]; !ok { - newlyInvalid = append(newlyInvalid, id) - } - } - if len(newlyInvalid) > 0 { - httpapi.Write(ctx, rw, http.StatusBadRequest, invalidChatMCPServerIDsResponse(newlyInvalid)) - return - } + normalizedMCPServerIDs, status, mcpResp := api.normalizeRequestedChatMCPServerIDs(ctx, chat, req.MCPServerIDs) + if mcpResp != nil { + httpapi.Write(ctx, rw, status, *mcpResp) + return } + req.MCPServerIDs = normalizedMCPServerIDs if req.PlanMode != nil { if !validateChatPlanMode(*req.PlanMode) { @@ -2803,6 +2815,12 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { return } + editMCPServerIDs, status, mcpResp := api.normalizeRequestedChatMCPServerIDs(ctx, chat, req.MCPServerIDs) + if mcpResp != nil { + httpapi.Write(ctx, rw, status, *mcpResp) + return + } + editResult, editErr := api.chatDaemon.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, CreatedBy: apiKey.UserID, @@ -2810,6 +2828,7 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { Content: contentBlocks, ModelConfigID: editModelConfigID, ReasoningEffort: editReasoningEffort, + MCPServerIDs: editMCPServerIDs, }) if editErr != nil { if writeChatHookErr(ctx, rw, editErr, "Chat message denied by lifecycle hook.") { diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 688acc1650d..6d90d60ea8f 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -9942,6 +9942,98 @@ func TestPatchChatMessage(t *testing.T) { require.False(t, foundOriginalInChat) }) + t.Run("MCPServerIDsApplied", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModel(t, client) + + orgConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + OrganizationID: firstUser.OrganizationID, + Enabled: true, + }) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "before mcp edit", + }}, + }) + require.NoError(t, err) + require.Empty(t, chat.MCPServerIDs) + + messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil) + require.NoError(t, err) + userMessageID := messagesResult.Messages[0].ID + + edited, err := client.EditChatMessage(ctx, chat.ID, userMessageID, codersdk.EditChatMessageRequest{ + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "edit enabling the mcp server", + }}, + MCPServerIDs: &[]uuid.UUID{orgConfig.ID}, + }) + require.NoError(t, err) + + storedChat, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Equal(t, []uuid.UUID{orgConfig.ID}, storedChat.MCPServerIDs) + + // Omitting the field preserves the persisted selection. + _, err = client.EditChatMessage(ctx, chat.ID, edited.Message.ID, codersdk.EditChatMessageRequest{ + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "edit without an mcp selection", + }}, + }) + require.NoError(t, err) + + storedChat, err = db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Equal(t, []uuid.UUID{orgConfig.ID}, storedChat.MCPServerIDs) + }) + + t.Run("MCPServerIDsInvalidRejected", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModel(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "before invalid mcp edit", + }}, + }) + require.NoError(t, err) + + messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil) + require.NoError(t, err) + userMessageID := messagesResult.Messages[0].ID + + unknownID := uuid.New() + _, err = client.EditChatMessage(ctx, chat.ID, userMessageID, codersdk.EditChatMessageRequest{ + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "edit selecting an unknown mcp server", + }}, + MCPServerIDs: &[]uuid.UUID{unknownID}, + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "One or more MCP server IDs are invalid or disabled.", sdkErr.Message) + require.Equal(t, "Invalid IDs: "+unknownID.String(), sdkErr.Detail) + + storedChat, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Empty(t, storedChat.MCPServerIDs) + }) + t.Run("CrossOrgModelConfigRejected", func(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 05f71d3ba88..fcf717e61c0 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1173,6 +1173,10 @@ type EditMessageOptions struct { // original message's model is preserved. ModelConfigID uuid.UUID ReasoningEffort *string + // MCPServerIDs, when non-nil, replaces the chat's MCP server + // selection before the replacement turn runs. When nil the + // current selection is preserved. + MCPServerIDs *[]uuid.UUID } // EditMessageResult contains the replacement user message and chat status. @@ -1243,6 +1247,36 @@ func enforceForcedMCPServerIDs(ctx context.Context, store database.Store, organi return merged, nil } +// applyRequestedMCPServerIDs replaces the chat's MCP server selection +// inside the state-machine transaction when a request provides one. +// Explore child chats keep the spawn-time snapshot immutable. Force On +// MCP servers are enforced server-side so a caller cannot remove them +// by tampering with the update (Cure53 CDM-02-010). +func (p *Server) applyRequestedMCPServerIDs(ctx context.Context, store database.Store, lockedChat database.Chat, requested *[]uuid.UUID) (database.Chat, error) { + if requested == nil { + return lockedChat, nil + } + if isExploreSubagentMode(lockedChat.Mode) { + p.logger.Warn(ctx, + "ignoring explore subagent mcp server ids update, snapshot is immutable after spawn", + slog.F("chat_id", lockedChat.ID), + ) + return lockedChat, nil + } + enforcedIDs, err := enforceForcedMCPServerIDs(ctx, store, lockedChat.OrganizationID, lockedChat.OwnerID, *requested) + if err != nil { + return database.Chat{}, err + } + updated, err := store.UpdateChatMCPServerIDs(ctx, database.UpdateChatMCPServerIDsParams{ + ID: lockedChat.ID, + MCPServerIDs: enforcedIDs, + }) + if err != nil { + return database.Chat{}, xerrors.Errorf("update chat mcp server ids: %w", err) + } + return updated, nil +} + // CreateChat creates a chat with its initial history through // chatstate.CreateChat. The new chat starts in `running` status per // the chat execution state model. Ownership hints wake chat workers. @@ -1514,30 +1548,9 @@ func (p *Server) SendMessage( return err } - // Update MCP server IDs on the chat when explicitly provided. - // Explore child chats keep the spawn-time snapshot immutable. - if requestedMCPServerIDs != nil { - if isExploreSubagentMode(lockedChat.Mode) { - p.logger.Warn(ctx, - "ignoring explore subagent mcp server ids update, snapshot is immutable after spawn", - slog.F("chat_id", opts.ChatID), - ) - } else { - // Force On MCP servers are enforced server-side so a - // caller cannot remove them by tampering with the - // update (Cure53 CDM-02-010). - enforcedIDs, enforceErr := enforceForcedMCPServerIDs(ctx, store, lockedChat.OrganizationID, lockedChat.OwnerID, *requestedMCPServerIDs) - if enforceErr != nil { - return enforceErr - } - lockedChat, err = store.UpdateChatMCPServerIDs(ctx, database.UpdateChatMCPServerIDsParams{ - ID: opts.ChatID, - MCPServerIDs: enforcedIDs, - }) - if err != nil { - return xerrors.Errorf("update chat mcp server ids: %w", err) - } - } + lockedChat, err = p.applyRequestedMCPServerIDs(ctx, store, lockedChat, requestedMCPServerIDs) + if err != nil { + return err } messageCreatedBy := opts.CreatedBy @@ -1885,6 +1898,11 @@ func (p *Server) EditMessage( } 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 diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index a60f752a989..7d710597e21 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -2209,6 +2209,147 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { require.True(t, errors.Is(err, chatd.ErrEditedMessageNotUser)) } +func TestEditMessage_MCPServerIDs(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) + + optionalConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + OrganizationID: org.ID, + DisplayName: "Optional MCP", + Slug: "optional-mcp", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + + chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "edit-mcp-server-ids", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{}, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + + latestUserMessageID := func(chatID uuid.UUID) int64 { + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + AfterID: 0, + }) + require.NoError(t, err) + var id int64 + for _, message := range messages { + if message.Role == database.ChatMessageRoleUser && !message.Deleted { + id = message.ID + } + } + require.NotZero(t, id) + return id + } + + // An edit that provides MCP server IDs replaces the selection. + _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: latestUserMessageID(chat.ID), + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edit with mcp")}, + MCPServerIDs: &[]uuid.UUID{optionalConfig.ID}, + }) + require.NoError(t, err) + + dbChat, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.ElementsMatch(t, []uuid.UUID{optionalConfig.ID}, dbChat.MCPServerIDs) + + // A nil selection preserves the persisted one. + _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: latestUserMessageID(chat.ID), + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edit without mcp")}, + }) + require.NoError(t, err) + + dbChat, err = db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.ElementsMatch(t, []uuid.UUID{optionalConfig.ID}, dbChat.MCPServerIDs) + + // An edit that clears the list cannot remove a force_on server + // (Cure53 CDM-02-010 parity with the send path). + forcedConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + OrganizationID: org.ID, + DisplayName: "Forced MCP", + Slug: "forced-mcp", + Availability: "force_on", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: latestUserMessageID(chat.ID), + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edit clearing mcp")}, + MCPServerIDs: &[]uuid.UUID{}, + }) + require.NoError(t, err) + + dbChat, err = db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.ElementsMatch(t, []uuid.UUID{forcedConfig.ID}, dbChat.MCPServerIDs, + "force_on MCP server must survive an emptied mcp_server_ids edit") + + // Explore child chats keep the spawn-time MCP snapshot immutable. + exploreContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("explore"), + }) + require.NoError(t, err) + createdExplore, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + ParentChatID: uuid.NullUUID{UUID: chat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: chat.ID, Valid: true}, + LastModelConfigID: model.ID, + Title: "explore-mcp-immutable", + Mode: database.NullChatMode{ + ChatMode: database.ChatModeExplore, + Valid: true, + }, + MCPServerIDs: []uuid.UUID{optionalConfig.ID}, + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: exploreContent, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + }, + }) + require.NoError(t, err) + exploreChat := createdExplore.Chat + + _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: exploreChat.ID, + CreatedBy: user.ID, + EditedMessageID: latestUserMessageID(exploreChat.ID), + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edit explore")}, + MCPServerIDs: &[]uuid.UUID{}, + }) + require.NoError(t, err) + + dbChat, err = db.GetChatByID(ctx, exploreChat.ID) + require.NoError(t, err) + require.ElementsMatch(t, []uuid.UUID{optionalConfig.ID}, dbChat.MCPServerIDs, + "explore subagent MCP snapshot must be immutable after spawn") +} + // TestEditMessageDebugCleanupDeletesPreEditRuns verifies that // EditMessage schedules the chat debug cleanup goroutine when debug // logging is enabled and that it deletes debug runs tied to the diff --git a/codersdk/chats.go b/codersdk/chats.go index 7925546922a..69bf87009aa 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -645,6 +645,10 @@ type EditChatMessageRequest struct { // When nil the original message's model is preserved. ModelConfigID *uuid.UUID `json:"model_config_id,omitempty" format:"uuid"` ReasoningEffort *string `json:"reasoning_effort,omitempty"` + // MCPServerIDs, when set, replaces the chat's MCP server selection + // before the replacement turn runs. When nil the current selection + // is preserved. + MCPServerIDs *[]uuid.UUID `json:"mcp_server_ids,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 0db7af7b819..420696a4e4f 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -2122,6 +2122,9 @@ Experimental: this endpoint is subject to change. "type": "text" } ], + "mcp_server_ids": [ + "497f6eca-6276-4993-bfeb-53cbbbba6f08" + ], "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205", "reasoning_effort": "string" } diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index a757b639485..d7d2b271ab6 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -8658,6 +8658,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "type": "text" } ], + "mcp_server_ids": [ + "497f6eca-6276-4993-bfeb-53cbbbba6f08" + ], "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205", "reasoning_effort": "string" } @@ -8668,6 +8671,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o | Name | Type | Required | Restrictions | Description | |--------------------|-----------------------------------------------------------|----------|--------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | `content` | array of [codersdk.ChatInputPart](#codersdkchatinputpart) | false | | | +| `mcp_server_ids` | array of string | false | | Mcp server ids when set, replaces the chat's MCP server selection before the replacement turn runs. When nil the current selection is preserved. | | `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. | | `reasoning_effort` | string | false | | | diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 9e73794b016..765337cf7e7 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -4984,6 +4984,12 @@ export interface EditChatMessageRequest { */ readonly model_config_id?: string; readonly reasoning_effort?: string; + /** + * MCPServerIDs, when set, replaces the chat's MCP server selection + * before the replacement turn runs. When nil the current selection + * is preserved. + */ + readonly mcp_server_ids?: string[]; } // From codersdk/chats.go diff --git a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx index dcdc77ced22..1969798f382 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -3735,6 +3735,79 @@ export const RemoveLastMCPServer: Story = { }, }; +const mcpEditableUserMessage: TypesGen.ChatMessage = { + ...MockChatMessage, + id: 5, + chat_id: CHAT_ID, + content: [{ type: "text", text: "Edit this request" }], +}; + +/** + * An MCP server toggled on while editing a message must ride along in + * the edit request. + */ +export const EditAppliesMCPServerSelection: Story = { + parameters: { + queries: buildQueries( + { + id: CHAT_ID, + ...baseChatFields, + title: "Edit applies MCP selection", + status: "waiting", + mcp_server_ids: [], + }, + { + messages: [mcpEditableUserMessage], + queued_messages: [], + has_more: false, + }, + { + diffUrl: undefined, + mcpServers: [MockMCPServerConfig], + }, + ), + }, + beforeEach: () => { + spyOn(API.experimental, "getUserSkills").mockResolvedValue([]); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const body = within(document.body); + const editSpy = spyOn( + API.experimental, + "editChatMessage", + ).mockResolvedValue({ + message: { ...mcpEditableUserMessage, id: 6 }, + }); + + await userEvent.click( + await canvas.findByRole("button", { name: "Edit message" }), + ); + await userEvent.click(canvas.getByRole("button", { name: "More options" })); + await userEvent.click( + await body.findByRole("switch", { + name: `Enable ${MockMCPServerConfig.display_name}`, + }), + ); + // Close the plus menu via its trigger; Escape would exit edit mode. + await userEvent.click(canvas.getByRole("button", { name: "More options" })); + await userEvent.click( + await canvas.findByRole("button", { name: "Save Edit" }), + ); + + await waitFor(() => { + expect(editSpy).toHaveBeenCalledTimes(1); + }); + expect(editSpy).toHaveBeenCalledWith( + CHAT_ID, + 5, + expect.objectContaining({ + mcp_server_ids: [MockMCPServerConfig.id], + }), + ); + }, +}; + /** * The send flow renders the durable user row once the server accepts the * prompt, before the assistant turn produces any output. diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 0adeda7a70a..40ee653b18b 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -1752,6 +1752,7 @@ const AgentChatPage: FC = () => { reasoning_effort: isEditReasoningEffortDirtyRef.current ? effectiveReasoningEffort : undefined, + mcp_server_ids: [...effectiveMCPServerIds], }; const optimisticMessage = originalEditedMessage ? buildOptimisticEditedMessage({