@fluidframework/tree 2.51.0 → 2.52.0

package/src/simple-tree/api/schemaCompatibilityTester.ts

@@ -3,39 +3,25 @@
  * Licensed under the MIT License.
  */
 
-import { assert, unreachableCase } from "@fluidframework/core-utils/internal";
-
 import type { TreeStoredSchema } from "../../core/index.js";
-import {
-	allowsRepoSuperset,
-	FieldKinds,
-	type FullSchemaPolicy,
-	isNeverTree,
-} from "../../feature-libraries/index.js";
-import type { FieldSchema } from "../fieldSchema.js";
+import { allowsRepoSuperset, defaultSchemaPolicy } from "../../feature-libraries/index.js";
 
 import type { SchemaCompatibilityStatus } from "./tree.js";
-import {
-	comparePosetElements,
-	fieldRealizer,
-	getAllowedContentDiscrepancies,
-	PosetComparisonResult,
-	type FieldDiscrepancy,
-} from "../discrepancies.js";
+import { getDiscrepanciesInAllowedContent } from "./discrepancies.js";
 import { toStoredSchema } from "../toStoredSchema.js";
+import type { TreeSchema } from "./configuration.js";
 
 /**
- * A collection of View information for schema, including policy.
+ * A collection of View information for schema
  * @remarks
  * This contains everything needed to determine compatibility with a given stored schema.
  */
 export class SchemaCompatibilityTester {
-	/**
-	 * @param viewSchemaRoot - Schema for the root field.
-	 */
 	public constructor(
-		public readonly policy: FullSchemaPolicy,
-		public readonly viewSchemaRoot: FieldSchema,
+		/**
+		 * Schema for the view
+		 */
+		public readonly viewSchema: TreeSchema,
 	) {}
 
 	/**
@@ -51,147 +37,37 @@ export class SchemaCompatibilityTester {
 	public checkCompatibility(
 		stored: TreeStoredSchema,
 	): Omit<SchemaCompatibilityStatus, "canInitialize"> {
+		// The public API surface assumes defaultSchemaPolicy
+		const policy = defaultSchemaPolicy;
+
 		// View schema allows a subset of documents that stored schema does, and the discrepancies are allowed by policy
 		// determined by the view schema (i.e. objects with extra optional fields in the stored schema have opted into allowing this.
 		// In the future, this would also include things like:
 		// - fields with more allowed types in the stored schema than in the view schema have out-of-schema "unknown content" adapters
 		let canView = true;
-		// View schema allows a superset of documents that stored schema does, hence the document could be upgraded to use a persisted version
-		// of this view schema as its stored schema.
-		let canUpgrade = true;
-
-		const updateCompatibilityFromFieldDiscrepancy = (discrepancy: FieldDiscrepancy): void => {
-			switch (discrepancy.mismatch) {
-				case "allowedTypes": {
-					// Since we only track the symmetric difference between the allowed types in the view and
-					// stored schemas, it's sufficient to check if any extra allowed types still exist in the
-					// stored schema.
-					if (
-						discrepancy.stored.some(
-							(identifier) =>
-								!isNeverTree(this.policy, stored, stored.nodeSchema.get(identifier)),
-						)
-					) {
-						// Stored schema has extra allowed types that the view schema does not.
-						canUpgrade = false;
-						canView = false;
-					}
-
-					if (discrepancy.view.length > 0) {
-						// View schema has extra allowed types that the stored schema does not.
-						canView = false;
-					}
-					break;
-				}
-				case "fieldKind": {
-					const result = comparePosetElements(
-						discrepancy.stored,
-						discrepancy.view,
-						fieldRealizer,
-					);
-
-					if (result === PosetComparisonResult.Greater) {
-						// Stored schema is more relaxed than view schema.
-						canUpgrade = false;
-						if (
-							discrepancy.view === FieldKinds.forbidden.identifier &&
-							discrepancy.identifier !== undefined &&
-							this.policy.allowUnknownOptionalFields(discrepancy.identifier)
-						) {
-							// When the application has opted into it, we allow viewing documents which have additional
-							// optional fields in the stored schema that are not present in the view schema.
-						} else {
-							canView = false;
-						}
-					}
 
-					if (result === PosetComparisonResult.Less) {
-						// View schema is more relaxed than stored schema.
-						canView = false;
-					}
-
-					if (result === PosetComparisonResult.Incomparable) {
-						canUpgrade = false;
-						canView = false;
-					}
-
-					break;
-				}
-				case "valueSchema": {
-					canView = false;
-					canUpgrade = false;
-					break;
-				}
-				default:
-					unreachableCase(discrepancy);
-			}
-		};
-
-		for (const discrepancy of getAllowedContentDiscrepancies(this.viewSchemaRoot, stored)) {
-			if (!canView && !canUpgrade) {
-				break;
-			}
-
-			switch (discrepancy.mismatch) {
-				case "nodeKind": {
-					const storedNodeSchema = stored.nodeSchema.get(discrepancy.identifier);
-					// We conservatively do not allow node types to change.
-					// The only time this might be valid in the sense that the data canonically converts is converting an object node
-					// to a map node over the union of all the object fields' types.
-					if (discrepancy.stored === undefined) {
-						// View schema has added a node type that the stored schema doesn't know about.
-						// Note that all cases which trigger this should also trigger an AllowedTypeDiscrepancy (where the type is used).
-						// This means this case should be redundant and could be removed in the future if there is a reason to do so
-						// (like simplifying enablable type support).
-						// See the TODO in getAllowedContentDiscrepancies.
-						canView = false;
-					} else if (discrepancy.view === undefined) {
-						const storedIsNever =
-							storedNodeSchema !== undefined
-								? isNeverTree(this.policy, stored, storedNodeSchema)
-								: true;
-						if (!storedIsNever) {
-							// Stored schema has a node type that the view schema doesn't know about.
-							// The design of allowUnknownOptionalFields allows adding new optional content to types,
-							// and the new optional content can use new types.
-							// Therefore this case needs to be supported for viewing.
-							// However, it is not supported for upgrade as the stored document must be newer (or at least more general) than the view in this case.
-							// The fact that there might be removed trees with a root of this type which only appears in the stored schema is why this must set canUpgrade to false.
-							canUpgrade = false;
-						}
-					} else {
-						// Node type exists in both schemas but kind has changed. We conservatively never allow this.
-						// See note above about cases where this could be allowed if needed.
-						canView = false;
-						canUpgrade = false;
-					}
-					break;
-				}
-				case "valueSchema":
-				case "allowedTypes":
-				case "fieldKind": {
-					updateCompatibilityFromFieldDiscrepancy(discrepancy);
-					break;
-				}
-				case "fields": {
-					discrepancy.differences.forEach(updateCompatibilityFromFieldDiscrepancy);
-					break;
-				}
-				// No default
-			}
+		for (const _discrepancy of getDiscrepanciesInAllowedContent(this.viewSchema, stored)) {
+			canView = false;
+			break;
 		}
 
-		if (canUpgrade) {
-			assert(
-				allowsRepoSuperset(this.policy, stored, toStoredSchema(this.viewSchemaRoot)),
-				0xbf2 /* View schema must be a superset of the stored schema to allow upgrade */,
-			);
-		}
+		const canUpgrade = allowsRepoSuperset(
+			policy,
+			stored,
+			toStoredSchema(this.viewSchema.root),
+		);
+
+		// If true, then upgrading has no effect on what can be stored in the document.
+		// TODO: This should likely be changed to indicate up a schema upgrade would be a no-op, including stored schema metadata.
+		const isEquivalent =
+			canView &&
+			canUpgrade &&
+			allowsRepoSuperset(policy, toStoredSchema(this.viewSchema.root), stored);
 
 		return {
 			canView,
 			canUpgrade,
-			isEquivalent: canView && canUpgrade,
+			isEquivalent,
 		};
 	}
 }

package/src/simple-tree/api/schemaFactory.ts

@@ -137,12 +137,15 @@ export interface SchemaFactoryObjectOptions<TCustomMetadata = unknown>
 	 * // Then the alleged clone wouldn't actually clone the entire person in either case, it would drop the nickname.
 	 * ```
 	 *
-	 * If an application wants to be particularly careful to preserve all data on a node when editing it, it can use
-	 * {@link (TreeAlpha:interface).importVerbose|import}/{@link (TreeAlpha:interface).exportVerbose|export} APIs with persistent keys.
+	 * The existing import and export APIs have similar problems.
+	 * For example currently the {@link (TreeAlpha:interface).exportVerbose|exportVerbose} API with stored keys preserves unknown optional fields,
+	 * but {@link Unhydrated} nodes produced by {@link TreeNode} constructors, insertable content, and {@link (TreeAlpha:interface).importVerbose|importVerbose} do not.
+	 * {@link (TreeBeta:interface).clone} however can be used to clone a node preserving unknown optional fields.
 	 *
 	 * Note that public API methods which operate on entire nodes (such as `moveTo`, `moveToEnd`, etc. on arrays) do not encounter
-	 * this problem as SharedTree's implementation stores the entire node in its lower layers. It's only when application code
-	 * reaches into a node (either by accessing its fields, spreading it, or some other means) that this problem arises.
+	 * this problem as SharedTree's implementation stores the entire node in its lower layers.
+	 * It's only when application code reaches into a node
+	 * (either by accessing its fields, spreading it, or some other means) that this problem arises.
 	 */
 	allowUnknownOptionalFields?: boolean;
 }

package/src/simple-tree/api/storedSchema.ts

@@ -5,11 +5,7 @@
 
 import type { FluidClientVersion, ICodecOptions } from "../../codec/index.js";
 import { SchemaVersion } from "../../core/index.js";
-import {
-	defaultSchemaPolicy,
-	encodeTreeSchema,
-	makeSchemaCodec,
-} from "../../feature-libraries/index.js";
+import { encodeTreeSchema, makeSchemaCodec } from "../../feature-libraries/index.js";
 import {
 	clientVersionToSchemaVersion,
 	type FormatV1,
@@ -19,6 +15,7 @@ import type { JsonCompatible } from "../../util/index.js";
 import { normalizeFieldSchema, type ImplicitFieldSchema } from "../fieldSchema.js";
 import type { SimpleTreeSchema } from "../simpleSchema.js";
 import { simpleToStoredSchema } from "../toStoredSchema.js";
+import { TreeViewConfigurationAlpha } from "./configuration.js";
 
 import { SchemaCompatibilityTester } from "./schemaCompatibilityTester.js";
 import type { SchemaCompatibilityStatus } from "./tree.js";
@@ -101,9 +98,9 @@ export function comparePersistedSchema(
 	// We only use the decode part, which always dispatches to the correct codec based on the version in the data, not the version passed to `makeSchemaCodec`.
 	const schemaCodec = makeSchemaCodec(options, SchemaVersion.v1);
 	const stored = schemaCodec.decode(persisted as FormatV1);
-	const viewSchema = new SchemaCompatibilityTester(
-		defaultSchemaPolicy,
-		normalizeFieldSchema(view),
-	);
+	const config = new TreeViewConfigurationAlpha({
+		schema: normalizeFieldSchema(view),
+	});
+	const viewSchema = new SchemaCompatibilityTester(config);
 	return viewSchema.checkCompatibility(stored);
 }

package/src/simple-tree/api/treeBeta.ts

@@ -3,14 +3,18 @@
  * Licensed under the MIT License.
  */
 
+import { defaultSchemaPolicy } from "../../feature-libraries/index.js";
 import {
+	Context,
 	getKernel,
 	isTreeNode,
+	UnhydratedContext,
 	type NodeKind,
 	type TreeNode,
 	type Unhydrated,
 	type WithType,
 } from "../core/index.js";
+import { getUnhydratedContext } from "../createContext.js";
 import type { ImplicitFieldSchema, TreeFieldFromImplicitField } from "../fieldSchema.js";
 
 import { createFromCursor } from "./create.js";
@@ -128,8 +132,8 @@ export interface TreeBeta {
 	 *
 	 * - The identifiers in the node's subtree will be preserved, i.e., they are not replaced with new values.
 	 *
-	 * @privateRemarks
-	 * TODO: AB#43548: How this handles unknown optional fields needs to be figured out, tested and documented.
+	 * - If the node (or any node in its subtree) contains {@link SchemaFactoryObjectOptions.allowUnknownOptionalFields|unknown optional fields},
+	 * those fields will be cloned just like the known fields.
 	 */
 	clone<const TSchema extends ImplicitFieldSchema>(
 		node: TreeFieldFromImplicitField<TSchema>,
@@ -137,7 +141,7 @@ export interface TreeBeta {
 
 	// TODO: support more clone options
 	// /**
-	//  * Like {@link TreeBeta.create}, except deeply clones existing nodes.
+	//  * Like {@link (TreeBeta:interface).create}, except deeply clones existing nodes.
 	//  * @remarks
 	//  * This only clones the persisted data associated with a node.
 	//  * Local state, such as properties added to customized schema classes, will not be cloned:
@@ -178,7 +182,16 @@ export const TreeBeta: TreeBeta = {
 
 		const kernel = getKernel(node);
 		const cursor = kernel.getOrCreateInnerNode().borrowCursor();
-		return createFromCursor(kernel.schema, cursor) as Unhydrated<
+
+		// To handle when the node transitively contains unknown optional fields,
+		// derive the context from the source node's stored schema which has stored schema for any such fields and their contents.
+		const flexContext = new UnhydratedContext(
+			defaultSchemaPolicy,
+			kernel.context.flexContext.schema,
+		);
+		const context = new Context(flexContext, getUnhydratedContext(kernel.schema).schema);
+
+		return createFromCursor(kernel.schema, cursor, context) as Unhydrated<
 			TreeFieldFromImplicitField<TSchema>
 		>;
 	},

package/src/simple-tree/api/verboseTree.ts

@@ -12,14 +12,15 @@ import {
 	aboveRootPlaceholder,
 	EmptyKey,
 	keyAsDetachedField,
+	LeafNodeStoredSchema,
 	type FieldKey,
 	type ITreeCursor,
 	type ITreeCursorSynchronous,
+	type TreeNodeStoredSchema,
 } from "../../core/index.js";
 import { brand } from "../../util/index.js";
 import type { ImplicitFieldSchema } from "../fieldSchema.js";
-import { NodeKind } from "../core/index.js";
-import type { TreeNodeSchema, ImplicitAllowedTypes, TreeLeafValue } from "../core/index.js";
+import type { Context, TreeLeafValue, TreeNodeSchema } from "../core/index.js";
 import {
 	isTreeValue,
 	stackTreeFieldCursor,
@@ -279,7 +280,7 @@ function verboseTreeAdapter(options: SchemalessParseOptions): CursorAdapter<Verb
  */
 export function verboseFromCursor(
 	reader: ITreeCursor,
-	rootSchema: ImplicitAllowedTypes,
+	context: Context,
 	options: TreeEncodingOptions,
 ): VerboseTree {
 	const config: Required<TreeEncodingOptions> = {
@@ -287,20 +288,28 @@ export function verboseFromCursor(
 		...options,
 	};
 
-	const schemaMap = getUnhydratedContext(rootSchema).schema;
+	const storedSchemaMap = context.flexContext.schema.nodeSchema;
+	const schemaMap = context.schema;
 
-	return verboseFromCursorInner(reader, config, schemaMap);
+	return verboseFromCursorInner(reader, config, storedSchemaMap, schemaMap);
 }
 
 function verboseFromCursorInner(
 	reader: ITreeCursor,
 	options: Required<TreeEncodingOptions>,
+	storedSchema: ReadonlyMap<string, TreeNodeStoredSchema>,
 	schema: ReadonlyMap<string, TreeNodeSchema>,
 ): VerboseTree {
-	const fields = customFromCursor(reader, options, schema, verboseFromCursorInner);
+	const fields = customFromCursor(
+		reader,
+		options,
+		storedSchema,
+		schema,
+		verboseFromCursorInner,
+	);
 	const nodeSchema =
-		schema.get(reader.type) ?? fail(0xb3c /* missing schema for type in cursor */);
-	if (nodeSchema.kind === NodeKind.Leaf) {
+		storedSchema.get(reader.type) ?? fail(0xb3c /* missing schema for type in cursor */);
+	if (nodeSchema instanceof LeafNodeStoredSchema) {
 		return fields as TreeLeafValue;
 	}
 

package/src/simple-tree/core/allowedTypes.ts

@@ -76,6 +76,7 @@ export interface NormalizedAnnotatedAllowedTypes {
 export function isAnnotatedAllowedTypes(
 	allowedTypes: ImplicitAnnotatedAllowedTypes,
 ): allowedTypes is AnnotatedAllowedTypes {
+	checkForUninitializedSchema(allowedTypes);
 	return (
 		// Class based schema, and lazy schema references report type "function": filtering them out with typeof makes narrowing based on members mostly safe
 		typeof allowedTypes === "object" && "metadata" in allowedTypes && "types" in allowedTypes
@@ -116,7 +117,9 @@ export interface AllowedTypesMetadata {
 export function isAnnotatedAllowedType(
 	allowedType: AnnotatedAllowedType | LazyItem<TreeNodeSchema>,
 ): allowedType is AnnotatedAllowedType {
-	return "metadata" in allowedType && "type" in allowedType;
+	checkForUninitializedSchema(allowedType);
+	// Class based schema, and lazy schema references report type "function": filtering them out with typeof makes narrowing based on members mostly safe
+	return typeof allowedType === "object" && "metadata" in allowedType && "type" in allowedType;
 }
 
 /**
@@ -224,7 +227,7 @@ export type UnannotateAllowedTypes<T extends AnnotatedAllowedTypes> =
  * Removes annotations from an allowed type.
  * @system @alpha
  */
-export type UnannotateAllowedType<T extends AnnotatedAllowedType> =
+export type UnannotateAllowedType<T extends AnnotatedAllowedType | LazyItem<TreeNodeSchema>> =
 	T extends AnnotatedAllowedType<infer X> ? [X] : T;
 
 /**
@@ -237,29 +240,40 @@ export type UnannotateAllowedType<T extends AnnotatedAllowedType> =
  * @internal
  */
 export function normalizeAllowedTypes(
-	types: ImplicitAllowedTypes,
+	types: ImplicitAnnotatedAllowedTypes,
 ): ReadonlySet<TreeNodeSchema> {
+	// remove annotations before normalizing
+	const unannotated = unannotateImplicitAllowedTypes(types);
 	const normalized = new Set<TreeNodeSchema>();
-	if (isReadonlyArray(types)) {
+	if (isReadonlyArray(unannotated)) {
 		// Types array must not be modified after it is normalized since that would result in the user of the normalized data having wrong (out of date) content.
-		Object.freeze(types);
-		for (const lazyType of types) {
+		Object.freeze(unannotated);
+		for (const lazyType of unannotated) {
 			normalized.add(evaluateLazySchema(lazyType));
 		}
 	} else {
-		normalized.add(evaluateLazySchema(types));
+		normalized.add(evaluateLazySchema(unannotated));
 	}
 	return normalized;
 }
 
 /**
- * Normalizes an allowed type to an {@link AnnotatedAllowedType}, by adding empty annotations if they don't already exist.
+ * Normalizes an allowed type to an {@link AnnotatedAllowedType}, by adding empty annotations if they don't already exist
+ * and eagerly evaluating any lazy schema declarations.
+ *
+ * @remarks
+ * Note: this must only be called after all required schemas have been declared, otherwise evaluation of
+ * recursive schemas may fail.
+ * type is frozen and should not be modified after being passed in.
  */
 export function normalizeToAnnotatedAllowedType<T extends TreeNodeSchema>(
 	type: T | AnnotatedAllowedType<T> | AnnotatedAllowedType<LazyItem<T>>,
-): AnnotatedAllowedType<T> | AnnotatedAllowedType<LazyItem<T>> {
+): AnnotatedAllowedType<T> {
 	return isAnnotatedAllowedType(type)
-		? type
+		? {
+				metadata: type.metadata,
+				type: evaluateLazySchema(type.type),
+			}
 		: {
 				metadata: {},
 				type,
@@ -317,19 +331,28 @@ export function unannotateImplicitAllowedTypes<Types extends ImplicitAnnotatedAl
 ): UnannotateImplicitAllowedTypes<Types> {
 	return (
 		isAnnotatedAllowedTypes(types)
-			? types.types.map((allowedType) =>
-					isAnnotatedAllowedType(allowedType) ? allowedType.type : allowedType,
-				)
+			? types.types.map(unannotateAllowedType)
 			: isReadonlyArray(types)
-				? types.map((allowedType) =>
-						isAnnotatedAllowedType(allowedType) ? allowedType.type : allowedType,
-					)
+				? types.map(unannotateAllowedType)
 				: isAnnotatedAllowedType(types)
-					? (types.type as UnannotateImplicitAllowedTypes<Types>)
+					? types.type
 					: types
 	) as UnannotateImplicitAllowedTypes<Types>;
 }
 
+/**
+ * Converts an {@link AnnotatedAllowedType} to an {@link LazyItem} by removing any annotations.
+ * @remarks
+ * This does not evaluate any lazy schemas.
+ */
+export function unannotateAllowedType<
+	Type extends AnnotatedAllowedType | LazyItem<TreeNodeSchema>,
+>(allowedType: Type): UnannotateAllowedType<Type> {
+	return isAnnotatedAllowedType(allowedType)
+		? (allowedType.type as UnannotateAllowedType<Type>)
+		: (allowedType as UnannotateAllowedType<Type>);
+}
+
 const cachedLazyItem = new WeakMap<() => unknown, unknown>();
 
 /**
@@ -342,13 +365,23 @@ export function evaluateLazySchema<T extends TreeNodeSchema>(value: LazyItem<T>)
 	const evaluatedSchema = isLazy(value)
 		? (getOrCreate(cachedLazyItem, value, value) as T)
 		: value;
-	if (evaluatedSchema === undefined) {
+
+	checkForUninitializedSchema(evaluatedSchema);
+	markSchemaMostDerived(evaluatedSchema);
+	return evaluatedSchema;
+}
+
+/**
+ * Throws a UsageError if the provided schema is undefined, most likely due to being used before it was initialized.
+ */
+export function checkForUninitializedSchema(
+	schema: ImplicitAnnotatedAllowedTypes | LazyItem<TreeNodeSchema>,
+): void {
+	if (schema === undefined) {
 		throw new UsageError(
-			`Encountered an undefined schema. This could indicate that some referenced schema has not yet been instantiated.`,
+			`Encountered an undefined schema. This could indicate that some referenced schema has not yet been instantiated. Consider using a lazy schema reference (like "() => schema") or delaying the evaluation of the lazy reference if one is already being used.`,
 		);
 	}
-	markSchemaMostDerived(evaluatedSchema);
-	return evaluatedSchema;
 }
 
 /**

package/src/simple-tree/core/context.ts

@@ -28,34 +28,36 @@ import { walkAllowedTypes } from "./walkSchema.js";
  * This design is the same as {@link FlexTreeContext} with its base type and {@link FlexTreeHydratedContext} extending it.
  */
 export class Context {
-	/**
-	 * All schema which could transitively be used under the associated node.
-	 * @remarks
-	 * While generally {@link TreeNodeSchema} are referenced as objects and thus do not need to be looked up by identifier,
-	 * there are a few cases (mainly constructing new TreeNodes from existing tree data) where such a lookup is useful.
-	 * Having this map in the context addresses this use-case.
-	 * @privateRemarks
-	 * This design mirrors how {@link FlexTreeSchema} are accessed off the {@link FlexTreeContext}, making the migration away from them simpler.
-	 */
-	public readonly schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema>;
-
-	/**
-	 * Builds the context.
-	 * @remarks
-	 * Since this walks the schema, it must not be invoked during schema declaration or schema forward references could fail to be resolved.
-	 */
-	public constructor(
+	public static schemaMapFromRootSchema(
 		rootSchema: NormalizedAnnotatedAllowedTypes,
-		public readonly flexContext: FlexTreeContext,
-	) {
+	): ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema> {
 		const schema: Map<TreeNodeSchemaIdentifier, TreeNodeSchema> = new Map();
 		walkAllowedTypes(rootSchema, {
 			node(nodeSchema) {
 				schema.set(brand(nodeSchema.identifier), nodeSchema);
 			},
 		});
-		this.schema = schema;
+		return schema;
 	}
+
+	/**
+	 * Builds the context.
+	 * @remarks
+	 * Since this walks the schema, it must not be invoked during schema declaration or schema forward references could fail to be resolved.
+	 */
+	public constructor(
+		public readonly flexContext: FlexTreeContext,
+		/**
+		 * All schema which could transitively be used under the associated node.
+		 * @remarks
+		 * While generally {@link TreeNodeSchema} are referenced as objects and thus do not need to be looked up by identifier,
+		 * there are a few cases (mainly constructing new TreeNodes from existing tree data) where such a lookup is useful.
+		 * Having this map in the context addresses this use-case.
+		 * @privateRemarks
+		 * This design mirrors how {@link FlexTreeSchema} are accessed off the {@link FlexTreeContext}, making the migration away from them simpler.
+		 */
+		public readonly schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema>,
+	) {}
 }
 
 /**
@@ -64,9 +66,9 @@ export class Context {
  */
 export class HydratedContext extends Context {
 	public constructor(
-		rootSchema: NormalizedAnnotatedAllowedTypes,
 		public override readonly flexContext: FlexTreeHydratedContext,
+		schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema>,
 	) {
-		super(rootSchema, flexContext);
+		super(flexContext, schema);
 	}
 }

package/src/simple-tree/core/index.ts

@@ -53,6 +53,7 @@ export {
 	isAnnotatedAllowedType,
 	normalizeAllowedTypes,
 	normalizeAnnotatedAllowedTypes,
+	normalizeToAnnotatedAllowedType,
 	unannotateImplicitAllowedTypes,
 	markSchemaMostDerived,
 	evaluateLazySchema,
@@ -76,7 +77,7 @@ export type {
 	AllowedTypeMetadata,
 	AnnotatedAllowedTypes,
 } from "./allowedTypes.js";
-export { walkAllowedTypes, type SchemaVisitor } from "./walkSchema.js";
+export { walkAllowedTypes, walkNodeSchema, type SchemaVisitor } from "./walkSchema.js";
 export { Context, HydratedContext } from "./context.js";
 export {
 	getOrCreateNodeFromInnerNode,

package/src/simple-tree/core/treeNodeKernel.ts

@@ -131,7 +131,7 @@ export class TreeNodeKernel {
 	/**
 	 * Create a TreeNodeKernel which can be looked up with {@link getKernel}.
 	 *
-	 * @param initialContext - context from when this node was originally crated.
+	 * @param initialContext - context from when this node was originally created. Only used when unhydrated.
 	 * @param innerNode - When unhydrated/raw or marinated the MapTreeNode. FlexTreeNode when cooked.
 	 * @remarks
 	 * Exactly one kernel per TreeNode should be created.

package/src/simple-tree/core/treeNodeValid.ts

@@ -10,7 +10,7 @@ import { type FlexTreeNode, isFlexTreeNode } from "../../feature-libraries/index
 
 import { markEager } from "./flexList.js";
 import { inPrototypeChain, privateToken, TreeNode } from "./treeNode.js";
-import type { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
+import { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
 import {
 	NodeKind,
 	type TreeNodeSchema,
@@ -193,10 +193,17 @@ export abstract class TreeNodeValid<TInput> extends TreeNode {
 		);
 
 		const result = schema.prepareInstance(this, node);
+		// For unhydrated nodes, grab the context from the `simpleContext` (which likely is from cache.oneTimeInitialized.context,
+		// but might be customized for better schema evolution support like unknown optional fields like in clone)
+		// For hydrated nodes, this context is unused, so using the default from cache.oneTimeInitialized.context is always fine.
+		const context =
+			node instanceof UnhydratedFlexTreeNode
+				? node.simpleContext
+				: cache.oneTimeInitialized.context;
 		// The TreeNodeKernel associates itself the TreeNode (result here, not node) so it can be looked up later via getKernel.
 		// If desired this could be put in a non-enumerable symbol property for lookup instead, but that gets messy going through proxies,
 		// so just relying on the WeakMap seems like the cleanest approach.
-		new TreeNodeKernel(result, schema, node, cache.oneTimeInitialized.context);
+		new TreeNodeKernel(result, schema, node, context);
 
 		return result;
 	}