diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/repository/data-source/tree-data-source.interface.ts b/src/Umbraco.Web.UI.Client/src/packages/core/repository/data-source/tree-data-source.interface.ts
index b973f626b7..3a09c78491 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/repository/data-source/tree-data-source.interface.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/repository/data-source/tree-data-source.interface.ts
@@ -1,9 +1,16 @@
 import type { UmbPagedData } from './types.js';
 import type { DataSourceResponse } from './data-source-response.interface.js';
+import { UmbControllerHost } from '@umbraco-cms/backoffice/controller-api';
+import { UmbTreeItemModel } from '@umbraco-cms/backoffice/tree';
 
-export interface UmbTreeDataSource<ItemType = any, PagedItemType = UmbPagedData<ItemType>> {
-	getRootItems(): Promise<DataSourceResponse<PagedItemType>>;
-	getChildrenOf(parentUnique: string | null): Promise<DataSourceResponse<PagedItemType>>;
+export interface UmbTreeDataSourceConstructor<TreeItemType extends UmbTreeItemModel = UmbTreeItemModel> {
+	new (host: UmbControllerHost): UmbTreeDataSource<TreeItemType>;
+}
+
+export interface UmbTreeDataSource<TreeItemType extends UmbTreeItemModel = UmbTreeItemModel> {
+	getTreeRoot?(): Promise<DataSourceResponse<TreeItemType>>;
+	getRootItems(): Promise<DataSourceResponse<UmbPagedData<TreeItemType>>>;
+	getChildrenOf(parentUnique: string | null): Promise<DataSourceResponse<UmbPagedData<TreeItemType>>>;
 	// TODO: remove this when all repositories are migrated to the new items interface
 	getItems(unique: Array<string>): Promise<DataSourceResponse<Array<any>>>;
 }
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/repository/tree-repository.interface.ts b/src/Umbraco.Web.UI.Client/src/packages/core/repository/tree-repository.interface.ts
index 46f52b5f06..aaab061e98 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/repository/tree-repository.interface.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/repository/tree-repository.interface.ts
@@ -1,15 +1,12 @@
+import { type UmbTreeItemModel } from '../tree/types.js';
 import { type UmbPagedData } from './data-source/types.js';
 import { type Observable } from '@umbraco-cms/backoffice/external/rxjs';
-import type { UmbTreeRootEntityModel, UmbTreeRootModel } from '@umbraco-cms/backoffice/tree';
-import { ProblemDetails, EntityTreeItemResponseModel } from '@umbraco-cms/backoffice/backend-api';
+import { ProblemDetails } from '@umbraco-cms/backoffice/backend-api';
 import { UmbApi } from '@umbraco-cms/backoffice/extension-api';
 
-export interface UmbTreeRepository<
-	TreeItemType extends EntityTreeItemResponseModel,
-	TreeRootType extends UmbTreeRootModel = UmbTreeRootEntityModel
-> extends UmbApi {
+export interface UmbTreeRepository<TreeItemType extends UmbTreeItemModel = UmbTreeItemModel> extends UmbApi {
 	requestTreeRoot: () => Promise<{
-		data?: TreeRootType;
+		data?: TreeItemType;
 		error?: ProblemDetails;
 	}>;