diff --git a/src/Umbraco.Web.UI.Client/src/backoffice/components/collection/collection.context.ts b/src/Umbraco.Web.UI.Client/src/backoffice/components/collection/collection.context.ts
deleted file mode 100644
index 8b865085f4..0000000000
--- a/src/Umbraco.Web.UI.Client/src/backoffice/components/collection/collection.context.ts
+++ /dev/null
@@ -1,43 +0,0 @@
-import { BehaviorSubject, Observable } from 'rxjs';
-
-export interface UmbCollectionContext {
-	readonly selectable: Observable<boolean>;
-	readonly selection: Observable<Array<string>>;
-	readonly entityKey: string;
-	setSelectable(value: boolean): void;
-	setSelection(value: Array<string>): void;
-	select(key: string): void;
-}
-
-export class UmbCollectionContextBase implements UmbCollectionContext {
-	private _selectable: BehaviorSubject<boolean> = new BehaviorSubject(false);
-	public readonly selectable: Observable<boolean> = this._selectable.asObservable();
-
-	private _selection: BehaviorSubject<Array<string>> = new BehaviorSubject(<Array<string>>[]);
-	public readonly selection: Observable<Array<string>> = this._selection.asObservable();
-
-	public entityKey = '';
-
-	constructor(entityKey: string) {
-		this.entityKey = entityKey;
-	}
-
-	public setSelectable(value: boolean) {
-		this._selectable.next(value);
-	}
-
-	public setSelection(value: Array<string>) {
-		if (!value) return;
-		this._selection.next(value);
-	}
-
-	public select(key: string) {
-		const selection = [...this._selection.getValue(), key];
-		this._selection.next(selection);
-	}
-
-	public deselect(key: string) {
-		const selection = this._selection.getValue();
-		this._selection.next(selection.filter((x) => x !== key));
-	}
-}
diff --git a/src/Umbraco.Web.UI.Client/src/backoffice/components/collection/layouts/collection-layout-media-grid.element.ts b/src/Umbraco.Web.UI.Client/src/backoffice/components/collection/layouts/collection-layout-media-grid.element.ts
index 14f881d533..32a481c32b 100644
--- a/src/Umbraco.Web.UI.Client/src/backoffice/components/collection/layouts/collection-layout-media-grid.element.ts
+++ b/src/Umbraco.Web.UI.Client/src/backoffice/components/collection/layouts/collection-layout-media-grid.element.ts
@@ -7,6 +7,7 @@ import type { MediaDetails } from '@umbraco-cms/models';
 import { UmbMediaStore } from '@umbraco-cms/stores/media/media.store';
 import { UmbContextConsumerMixin } from '@umbraco-cms/context-api';
 import { UmbObserverMixin } from '@umbraco-cms/observable-api';
+import UmbDashboardMediaManagementElement from 'src/backoffice/dashboards/media-management/dashboard-media-management.element';
 
 @customElement('umb-collection-layout-media-grid')
 export class UmbCollectionLayoutMediaGridElement extends UmbContextConsumerMixin(UmbObserverMixin(LitElement)) {
@@ -69,9 +70,10 @@ export class UmbCollectionLayoutMediaGridElement extends UmbContextConsumerMixin
 	@state()
 	private _mediaItems: Array<MediaDetails> = [];
 
-	private _collectionContext?: UmbCollectionContext;
+	@state()
+	private _selection: Array<string> = [];
 
-	private _mediaStore?: UmbMediaStore;
+	private _mediaContext?: UmbDashboardMediaManagementElement;
 	constructor() {
 		super();
 		document.addEventListener('dragenter', (e) => {
@@ -84,32 +86,22 @@ export class UmbCollectionLayoutMediaGridElement extends UmbContextConsumerMixin
 			e.preventDefault();
 			this.toggleAttribute('dragging', false);
 		});
-		this.consumeAllContexts(['umbMediaStore', 'umbCollectionContext'], (instance) => {
-			this._mediaStore = instance['umbMediaStore'];
-			this._collectionContext = instance['umbCollectionContext'];
-			this._observeMediaItems();
+		this.consumeAllContexts(['umbMediaContext'], (instance) => {
+			this._mediaContext = instance['umbMediaContext'];
+			this._observeMediaContext();
 		});
 	}
 
-	private _observeMediaItems() {
-		if (!this._mediaStore || !this._collectionContext) return;
+	private _observeMediaContext() {
+		if (!this._mediaContext) return;
 
-		console.log('key', this._collectionContext.entityKey);
+		this.observe<Array<MediaDetails>>(this._mediaContext.mediaItems, (mediaItems) => {
+			this._mediaItems = mediaItems;
+		});
 
-		if (this._collectionContext.entityKey) {
-			this.observe<Array<MediaDetails>>(
-				this._mediaStore?.getTreeItemChildren(this._collectionContext.entityKey),
-				(items) => {
-					this._mediaItems = items;
-					console.log('media store', this._mediaStore);
-				}
-			);
-		} else {
-			this.observe<Array<MediaDetails>>(this._mediaStore?.getTreeRoot(), (items) => {
-				this._mediaItems = items;
-				console.log('media store', this._mediaStore);
-			});
-		}
+		this.observe<Array<string>>(this._mediaContext.selection, (selection) => {
+			this._selection = selection;
+		});
 	}
 
 	private _handleOpenItem(key: string) {
diff --git a/src/Umbraco.Web.UI.Client/src/backoffice/dashboards/media-management/dashboard-media-management.element.ts b/src/Umbraco.Web.UI.Client/src/backoffice/dashboards/media-management/dashboard-media-management.element.ts
index 60044e07a5..24d6b94532 100644
--- a/src/Umbraco.Web.UI.Client/src/backoffice/dashboards/media-management/dashboard-media-management.element.ts
+++ b/src/Umbraco.Web.UI.Client/src/backoffice/dashboards/media-management/dashboard-media-management.element.ts
@@ -4,14 +4,14 @@ import { customElement, property, state } from 'lit/decorators.js';
 import '../../components/collection/collection-toolbar.element';
 import '../../components/collection/collection-selection-actions.element';
 import '../../components/collection/collection-view.element';
-import { map } from 'rxjs';
+import { BehaviorSubject, map, Observable } from 'rxjs';
 import { IRoutingInfo } from 'router-slot';
-import type { ManifestCollectionLayout } from '@umbraco-cms/models';
+import type { ManifestCollectionLayout, MediaDetails } from '@umbraco-cms/models';
 import { UmbContextConsumerMixin, UmbContextProviderMixin } from '@umbraco-cms/context-api';
 import { UmbObserverMixin } from '@umbraco-cms/observable-api';
 import { umbExtensionsRegistry } from '@umbraco-cms/extensions-registry';
 import { createExtensionElement } from '@umbraco-cms/extensions-api';
-import { UmbCollectionContextBase } from '@umbraco-cms/components/collection/collection.context';
+import { UmbMediaStore } from '@umbraco-cms/stores/media/media.store';
 
 @customElement('umb-dashboard-media-management')
 export class UmbDashboardMediaManagementElement extends UmbContextProviderMixin(
@@ -43,14 +43,25 @@ export class UmbDashboardMediaManagementElement extends UmbContextProviderMixin(
 	@state()
 	private _currentLayoutPathName = '';
 
+	private _mediaStore?: UmbMediaStore;
+
+	private _selection: BehaviorSubject<Array<string>> = new BehaviorSubject(<Array<string>>[]);
+	public readonly selection: Observable<Array<string>> = this._selection.asObservable();
+
+	private _mediaItems: BehaviorSubject<Array<MediaDetails>> = new BehaviorSubject(<Array<MediaDetails>>[]);
+	public readonly mediaItems: Observable<Array<MediaDetails>> = this._mediaItems.asObservable();
+
+	private _search: BehaviorSubject<string> = new BehaviorSubject('');
+	public readonly search: Observable<string> = this._search.asObservable();
+
 	constructor() {
 		super();
 		this._observeCollectionLayouts();
-	}
-
-	connectedCallback(): void {
-		super.connectedCallback();
-		this.provideContext('umbCollectionContext', new UmbCollectionContextBase(this.entityKey));
+		this.provideContext('umbMediaContext', this);
+		this.consumeAllContexts(['umbMediaStore'], (instance) => {
+			this._mediaStore = instance['umbMediaStore'];
+			this._observeMediaItems();
+		});
 	}
 
 	private _observeCollectionLayouts() {
@@ -68,6 +79,20 @@ export class UmbDashboardMediaManagementElement extends UmbContextProviderMixin(
 		);
 	}
 
+	private _observeMediaItems() {
+		if (!this._mediaStore) return;
+
+		if (this.entityKey) {
+			this.observe<Array<MediaDetails>>(this._mediaStore?.getTreeItemChildren(this.entityKey), (items) => {
+				this._mediaItems.next(items);
+			});
+		} else {
+			this.observe<Array<MediaDetails>>(this._mediaStore?.getTreeRoot(), (items) => {
+				this._mediaItems.next(items);
+			});
+		}
+	}
+
 	private _createRoutes() {
 		this._routes = [];