diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/content-type/content-type-structure-manager.class.ts b/src/Umbraco.Web.UI.Client/src/packages/core/content-type/content-type-structure-manager.class.ts
index acc0862b34..7832094e4b 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/content-type/content-type-structure-manager.class.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/content-type/content-type-structure-manager.class.ts
@@ -307,7 +307,7 @@ export class UmbContentTypePropertyStructureManager<R extends UmbDetailRepositor
 		this.#documentTypes.updateOne(documentTypeId, { properties });
 	}
 
-	// TODO: Refactor: These property methods, could be named without structure in their name.
+	// TODO: Refactor: These property methods, should maybe be named without structure in their name.
 	async propertyStructureById(
 		propertyId: string
 	) {
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/dataset-context.token.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/dataset-context.token.ts
index 50fca8b86b..f63ac17653 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/dataset-context.token.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/dataset-context.token.ts
@@ -1,4 +1,4 @@
 import { type UmbDatasetContext } from "./dataset-context.interface.js";
 import { UmbContextToken } from "@umbraco-cms/backoffice/context-api";
 
-export const UMB_DATASET_CONTEXT = new UmbContextToken<UmbDatasetContext>("UmbEntityContext");
+export const UMB_DATASET_CONTEXT = new UmbContextToken<UmbDatasetContext>("UmbDatasetContext");
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/invariant-dataset-context.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/invariant-dataset-context.ts
index c88fe10438..bbf8d820fd 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/invariant-dataset-context.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/invariant-dataset-context.ts
@@ -37,7 +37,6 @@ export class UmbInvariantDatasetContext<WorkspaceType extends UmbInvariantableWo
 		super(host, 'dataSetContext');
 		this._workspace = workspace;
 
-		// TODO: Refactor: use the document dataset context token.
 		this.provideContext(UMB_DATASET_CONTEXT, this);
 	}
 
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/nameable-dataset-context.token.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/nameable-dataset-context.token.ts
index 3bff82bb5f..57f439a0f6 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/nameable-dataset-context.token.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/nameable-dataset-context.token.ts
@@ -5,5 +5,5 @@ import { UmbContextToken } from "@umbraco-cms/backoffice/context-api";
 export const IsNameableDatasetContext = (context: UmbDatasetContext): context is UmbNameableDatasetContext => 'setName' in context;
 
 export const UMB_NAMEABLE_DATASET_CONTEXT = new UmbContextToken<UmbDatasetContext, UmbNameableDatasetContext>(
-	"UmbEntityContext",
+	"UmbDatasetContext",
 	IsNameableDatasetContext);
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/variant-dataset-context.token.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/variant-dataset-context.token.ts
index fd3bb9298d..0643aa530e 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/variant-dataset-context.token.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/dataset-context/variant-dataset-context.token.ts
@@ -5,5 +5,5 @@ import { UmbContextToken } from "@umbraco-cms/backoffice/context-api";
 export const IsVariantDatasetContext = (context: UmbDatasetContext): context is UmbVariantDatasetContext => 'getVariantId' in context;
 
 export const UMB_VARIANT_DATASET_CONTEXT = new UmbContextToken<UmbDatasetContext, UmbVariantDatasetContext>(
-	"UmbEntityContext",
+	"UmbDatasetContext",
 	IsVariantDatasetContext);
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-context.interface.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-context.interface.ts
index d5aaceae55..4ff5c395bb 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-context.interface.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-context.interface.ts
@@ -16,10 +16,4 @@ export interface UmbWorkspaceContextInterface<DataType = unknown> {
 	getIsNew(): boolean | undefined;
 	setIsNew(value: boolean): void;
 
-	/*
-	// TODO: Refactor: This could maybe go away:
-	repository: any; // TODO: add type
-	getData(): DataType | undefined;
-	*/
-
 }
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-variantable-context.interface.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-variantable-context.interface.ts
index ed56862923..30f643854e 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-variantable-context.interface.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-context/workspace-variantable-context.interface.ts
@@ -1,5 +1,5 @@
 import type { UmbWorkspaceSplitViewManager } from '../workspace-split-view-manager.class.js';
-import type { UmbDatasetContext } from '../dataset-context/dataset-context.interface.js';
+import type { UmbVariantDatasetContext } from '../dataset-context/variant-dataset-context.interface.js';
 import type { UmbSaveableWorkspaceContextInterface } from './saveable-workspace-context.interface.js';
 import type { Observable } from '@umbraco-cms/backoffice/external/rxjs';
 import type { UmbVariantId } from '@umbraco-cms/backoffice/variant';
@@ -25,5 +25,5 @@ export interface UmbVariantableWorkspaceContextInterface<EntityType = unknown> e
 	//propertyDataByAlias(alias: string, variantId?: UmbVariantId): Observable<ValueModelBaseModel | undefined>;
 
 	// Dataset methods:
-	createVariableDatasetContext(host: UmbControllerHost, variantId: UmbVariantId): UmbDatasetContext;
+	createVariableDatasetContext(host: UmbControllerHost, variantId: UmbVariantId): UmbVariantDatasetContext;
 }
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-split-view/workspace-split-view.context.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-split-view/workspace-split-view.context.ts
index 64b141ed37..98279c5588 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-split-view/workspace-split-view.context.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-split-view/workspace-split-view.context.ts
@@ -1,5 +1,4 @@
-import { UmbVariantableWorkspaceContextInterface } from '../workspace-context/workspace-variantable-context.interface.js';
-import { UmbDatasetContext } from '../dataset-context/index.js';
+import { UmbVariantDatasetContext } from '../dataset-context/index.js';
 import { UMB_VARIANT_WORKSPACE_CONTEXT_TOKEN } from '../index.js';
 import { UmbVariantId } from '@umbraco-cms/backoffice/variant';
 import {
@@ -18,8 +17,7 @@ export class UmbWorkspaceSplitViewContext extends UmbBaseController {
 		return this.#workspaceContext;
 	}
 
-	// TODO: Refactor: use a variable data set context interface here.
-	#dataSetContext?: UmbDatasetContext;
+	#dataSetContext?: UmbVariantDatasetContext;
 
 	#index = new UmbNumberState(undefined);
 	index = this.#index.asObservable();
diff --git a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-splitview/workspace-split-view.context.ts b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-splitview/workspace-split-view.context.ts
index 7727926401..ed7d40696e 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-splitview/workspace-split-view.context.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/core/workspace/workspace-splitview/workspace-split-view.context.ts
@@ -1,4 +1,4 @@
-import { UmbDatasetContext } from '../dataset-context/index.js';
+import { UmbVariantDatasetContext } from '../dataset-context/index.js';
 import { UMB_VARIANT_WORKSPACE_CONTEXT_TOKEN } from '../index.js';
 import { UmbVariantId } from '@umbraco-cms/backoffice/variant';
 import {
@@ -17,8 +17,7 @@ export class UmbWorkspaceSplitViewContext extends UmbBaseController {
 		return this.#workspaceContext;
 	}
 
-	// TODO: Refactor: use a variable data set context interface here.
-	#dataSetContext?: UmbDatasetContext;
+	#dataSetContext?: UmbVariantDatasetContext;
 
 	#index = new UmbNumberState(undefined);
 	index = this.#index.asObservable();
diff --git a/src/Umbraco.Web.UI.Client/src/packages/documents/documents/dataset-context/document-dataset-context.token.ts b/src/Umbraco.Web.UI.Client/src/packages/documents/documents/dataset-context/document-dataset-context.token.ts
new file mode 100644
index 0000000000..e47014995b
--- /dev/null
+++ b/src/Umbraco.Web.UI.Client/src/packages/documents/documents/dataset-context/document-dataset-context.token.ts
@@ -0,0 +1,9 @@
+import type { UmbDocumentDatasetContext } from "./document-dataset-context.js";
+import { UmbContextToken } from "@umbraco-cms/backoffice/context-api";
+import { UmbDatasetContext } from "@umbraco-cms/backoffice/workspace";
+
+export const IsDocumentDatasetContext = (context: UmbDatasetContext): context is UmbDocumentDatasetContext => context.getType() === 'document';
+
+export const UMB_DOCUMENT_DATASET_CONTEXT = new UmbContextToken<UmbDatasetContext, UmbDocumentDatasetContext>(
+	"UmbDatasetContext",
+	IsDocumentDatasetContext);
diff --git a/src/Umbraco.Web.UI.Client/src/packages/settings/data-types/dataset-context/data-type-dataset-context.token.ts b/src/Umbraco.Web.UI.Client/src/packages/settings/data-types/dataset-context/data-type-dataset-context.token.ts
index 30fdc73f28..c41307f186 100644
--- a/src/Umbraco.Web.UI.Client/src/packages/settings/data-types/dataset-context/data-type-dataset-context.token.ts
+++ b/src/Umbraco.Web.UI.Client/src/packages/settings/data-types/dataset-context/data-type-dataset-context.token.ts
@@ -3,5 +3,5 @@ import { UmbDatasetContext } from "@umbraco-cms/backoffice/workspace";
 import { UmbContextToken } from "@umbraco-cms/backoffice/context-api";
 
 export const UMB_DATA_TYPE_DATASET_CONTEXT = new UmbContextToken<UmbDatasetContext, UmbDataTypeDatasetContext>(
-	"UmbEntityContext",
+	"UmbDatasetContext",
 (context): context is UmbDataTypeDatasetContext => 'properties' in context && context.getType() === 'data-type');