From b6d76532311f8010778721fc9fe2448e2fb3e1b3 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 24 Aug 2026 17:39:52 +0200 Subject: [PATCH] fix: apply MCP server selection when editing a chat message (#28471) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem Editing a user message in an Agent chat silently dropped the MCP server selection. The composer renders the MCP picker in edit mode and toggles update local client state, so the picker displayed the new selection, but the edit request omitted `mcp_server_ids` at every layer (frontend request builder, `codersdk.EditChatMessageRequest`, the `PATCH /chats/{chat}/messages/{message}` handler, and `chatd.EditMessage`). The chat's persisted selection never changed and the regenerated turn ran without the newly enabled MCP tools. Two dogfood users hit this within hours; there is no error anywhere and the UI shows the opposite of the server state. ## Changes - `codersdk`: add `MCPServerIDs *[]uuid.UUID` to `EditChatMessageRequest`, mirroring `CreateChatMessageRequest` (nil preserves the current selection). - `chatd`: extract the send path's MCP update block into one shared `applyRequestedMCPServerIDs` helper (explore-subagent snapshot immutability guard plus Force On enforcement, Cure53 CDM-02-010) and call it from both `SendMessage` and `EditMessage`, so enforcement cannot drift between the two paths. - `coderd`: extract the send handler's request validation (dedupe, enabled-in-organization check, persisted-ID exemption) into `normalizeRequestedChatMCPServerIDs` and wire it into the edit handler, which now threads the selection into `EditMessageOptions`. - Frontend: the edit request now includes `mcp_server_ids: [...effectiveMCPServerIds]`, exactly like the send path, making the picker's displayed state real. - `make gen` artifacts (swagger, API docs, `typesGenerated.ts`). ## Tests Each layer is covered and was proven with independent red toggles (removing one layer's wiring fails only that layer's tests): - chatd (`TestEditMessage_MCPServerIDs`): edit applies a provided selection, nil preserves it, an emptied list cannot remove a `force_on` server, and explore subagent chats keep the spawn-time snapshot. - API (`TestPatchChatMessage/MCPServerIDsApplied`, `MCPServerIDsInvalidRejected`): persistence via the endpoint, omission preserves, unknown IDs get the same 400 as the send path. - Storybook (`EditAppliesMCPServerSelection`): toggling a server on during an edit puts it in the edit request payload. Note: the `AgentChatPage.stories.tsx` story "Queued For Capacity After Polling" fails locally on current main as well (verified against the main baseline with this branch's changes reverted); it is unrelated to this diff. Remote dogfood UAT ran on the exact head and passed, including proof that the regenerated turn actually gains the newly enabled MCP server's tools. > 🤖 Xum acted on Mike's (@ibetitsmike) behalf. • Model: `anthropic:claude-fable-5` --- coderd/apidoc/docs.go | 8 + coderd/apidoc/swagger.json | 8 + coderd/exp_chats.go | 77 ++++++---- coderd/exp_chats_test.go | 92 ++++++++++++ coderd/x/chatd/chatd.go | 66 +++++--- coderd/x/chatd/chatd_test.go | 141 ++++++++++++++++++ codersdk/chats.go | 4 + docs/reference/api/chats.md | 3 + docs/reference/api/schemas.md | 4 + site/src/api/typesGenerated.ts | 6 + .../AgentsPage/AgentChatPage.stories.tsx | 73 +++++++++ site/src/pages/AgentsPage/AgentChatPage.tsx | 1 + 12 files changed, 430 insertions(+), 53 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 4a95c546ddc..434b2367d95 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -22005,6 +22005,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 186c1324f75..4a657b17c14 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -20041,6 +20041,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 7feaa3ccff6..4cbfa58e116 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 5e4c0499d4e..71cbda7ae46 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -10084,6 +10084,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 611f33b9837..38a0ebde20e 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -8680,6 +8680,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" } @@ -8690,6 +8693,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 c8405b14923..7ebd2042386 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -4997,6 +4997,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 8dc1fe418f8..0ea53b6c908 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({