diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/property-editor/uis/tiny-mce/plugins/tiny-mce-mediapicker.plugin.ts b/src/Umbraco.Web.UI.Client/src/packages/core/property-editor/uis/tiny-mce/plugins/tiny-mce-mediapicker.plugin.ts
index a1960177a8..7487b73377 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/property-editor/uis/tiny-mce/plugins/tiny-mce-mediapicker.plugin.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/property-editor/uis/tiny-mce/plugins/tiny-mce-mediapicker.plugin.ts
@@ -51,11 +51,16 @@ export default class UmbTinyMceMediaPickerPlugin extends UmbTinyMcePluginBase {
 		// 	this.#observeCurrentUser();
 		// });
 
-		this.editor.ui.registry.addButton('umbmediapicker', {
+		this.editor.ui.registry.addToggleButton('umbmediapicker', {
 			icon: 'image',
 			tooltip: 'Media Picker',
-			//stateSelector: 'img[data-udi]', TODO => Investigate where stateselector has gone, or if it is still needed
 			onAction: () => this.#onAction(),
+			onSetup: (api) => {
+				const changed = this.editor.selection.selectorChangedWithUnbind('img[data-udi]', (state) =>
+					api.setActive(state),
+				);
+				return () => changed;
+			},
 		});
 
 		// Register global options for the editor