diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 9c2df36cab..ff743ba902 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1216,6 +1216,17 @@ BEGIN END; $$; +CREATE FUNCTION remove_mcp_server_config_id_from_chats() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + UPDATE chats + SET mcp_server_ids = array_remove(mcp_server_ids, OLD.id) + WHERE OLD.id = ANY(mcp_server_ids); + RETURN OLD; +END; +$$; + CREATE FUNCTION remove_organization_member_role() RETURNS trigger LANGUAGE plpgsql AS $$ @@ -4435,6 +4446,10 @@ CREATE TRIGGER inhibit_enqueue_if_disabled BEFORE INSERT ON notification_message CREATE TRIGGER protect_deleting_organizations BEFORE UPDATE ON organizations FOR EACH ROW WHEN (((new.deleted = true) AND (old.deleted = false))) EXECUTE FUNCTION protect_deleting_organizations(); +CREATE TRIGGER remove_chat_mcp_server_config_id BEFORE DELETE ON mcp_server_configs FOR EACH ROW EXECUTE FUNCTION remove_mcp_server_config_id_from_chats(); + +COMMENT ON TRIGGER remove_chat_mcp_server_config_id ON mcp_server_configs IS 'When an MCP server config is deleted, this trigger removes its ID from all chats.'; + CREATE TRIGGER remove_organization_member_custom_role BEFORE DELETE ON custom_roles FOR EACH ROW EXECUTE FUNCTION remove_organization_member_role(); COMMENT ON TRIGGER remove_organization_member_custom_role ON custom_roles IS 'When a custom_role is deleted, this trigger removes the role from all organization members.'; diff --git a/coderd/database/migrations/000510_cleanup_chats_mcp_server_ids_on_delete.down.sql b/coderd/database/migrations/000510_cleanup_chats_mcp_server_ids_on_delete.down.sql new file mode 100644 index 0000000000..15c10e19e6 --- /dev/null +++ b/coderd/database/migrations/000510_cleanup_chats_mcp_server_ids_on_delete.down.sql @@ -0,0 +1,2 @@ +DROP TRIGGER IF EXISTS remove_chat_mcp_server_config_id ON mcp_server_configs; +DROP FUNCTION IF EXISTS remove_mcp_server_config_id_from_chats; diff --git a/coderd/database/migrations/000510_cleanup_chats_mcp_server_ids_on_delete.up.sql b/coderd/database/migrations/000510_cleanup_chats_mcp_server_ids_on_delete.up.sql new file mode 100644 index 0000000000..5366328b3c --- /dev/null +++ b/coderd/database/migrations/000510_cleanup_chats_mcp_server_ids_on_delete.up.sql @@ -0,0 +1,41 @@ +-- Remove already-stale MCP server references before future deletes are +-- handled by the trigger below. +UPDATE chats +SET mcp_server_ids = ( + SELECT COALESCE(array_agg(ids.mcp_server_id ORDER BY ids.position), '{}'::uuid[]) + FROM unnest(chats.mcp_server_ids) WITH ORDINALITY AS ids(mcp_server_id, position) + WHERE EXISTS ( + SELECT 1 + FROM mcp_server_configs + WHERE mcp_server_configs.id = ids.mcp_server_id + ) +) +WHERE EXISTS ( + SELECT 1 + FROM unnest(chats.mcp_server_ids) AS ids(mcp_server_id) + WHERE NOT EXISTS ( + SELECT 1 + FROM mcp_server_configs + WHERE mcp_server_configs.id = ids.mcp_server_id + ) +); + +CREATE OR REPLACE FUNCTION remove_mcp_server_config_id_from_chats() + RETURNS TRIGGER AS +$$ +BEGIN + UPDATE chats + SET mcp_server_ids = array_remove(mcp_server_ids, OLD.id) + WHERE OLD.id = ANY(mcp_server_ids); + RETURN OLD; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER remove_chat_mcp_server_config_id + BEFORE DELETE ON mcp_server_configs FOR EACH ROW + EXECUTE PROCEDURE remove_mcp_server_config_id_from_chats(); + +COMMENT ON TRIGGER + remove_chat_mcp_server_config_id + ON mcp_server_configs IS + 'When an MCP server config is deleted, this trigger removes its ID from all chats.'; diff --git a/coderd/mcp_test.go b/coderd/mcp_test.go index add730960f..dde85f12e7 100644 --- a/coderd/mcp_test.go +++ b/coderd/mcp_test.go @@ -1396,10 +1396,11 @@ func TestChatWithMCPServerIDs(t *testing.T) { // Create the chat model config required for creating a chat. _ = createChatModelConfigForMCP(t, expClient) - // Create an enabled MCP server config. - mcpConfig := createMCPServerConfig(t, client, "chat-mcp-server", true) + // Create enabled MCP server configs. + mcpConfigA := createMCPServerConfig(t, client, "chat-mcp-server-a", true) + mcpConfigB := createMCPServerConfig(t, client, "chat-mcp-server-b", true) - // Create a chat referencing the MCP server. + // Create a chat referencing the MCP servers. chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ @@ -1408,16 +1409,24 @@ func TestChatWithMCPServerIDs(t *testing.T) { Text: "hello with mcp server", }, }, - MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + MCPServerIDs: []uuid.UUID{mcpConfigA.ID, mcpConfigB.ID}, }) require.NoError(t, err) require.NotEqual(t, uuid.Nil, chat.ID) - require.Contains(t, chat.MCPServerIDs, mcpConfig.ID) + require.ElementsMatch(t, []uuid.UUID{mcpConfigA.ID, mcpConfigB.ID}, chat.MCPServerIDs) // Fetch the chat and verify the MCP server IDs persist. fetched, err := expClient.GetChat(ctx, chat.ID) require.NoError(t, err) - require.Contains(t, fetched.MCPServerIDs, mcpConfig.ID) + require.ElementsMatch(t, []uuid.UUID{mcpConfigA.ID, mcpConfigB.ID}, fetched.MCPServerIDs) + + err = client.DeleteMCPServerConfig(ctx, mcpConfigA.ID) + require.NoError(t, err) + + fetched, err = expClient.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.NotContains(t, fetched.MCPServerIDs, mcpConfigA.ID) + require.Contains(t, fetched.MCPServerIDs, mcpConfigB.ID) } func createChatModelConfigForMCP(t testing.TB, client *codersdk.ExperimentalClient) codersdk.ChatModelConfig {