@fluidframework/tree-agent 2.72.0 → 2.73.0

package/src/propertyBinding.ts

@@ -0,0 +1,181 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import type { TreeNodeSchema, TreeNodeSchemaClass } from "@fluidframework/tree";
+import type { ZodType, ZodTypeAny, ZodTypeDef, infer as ZodInfer } from "zod";
+
+import type { BindableSchema, Ctor } from "./methodBinding.js";
+import { instanceOf } from "./utils.js";
+
+/**
+ * A symbol used to expose properties to the LLM.
+ * @alpha
+ */
+export const exposePropertiesSymbol: unique symbol = Symbol.for(
+	"@fluidframework/tree-agent/exposeProperties",
+);
+
+/**
+ * Set of property keys from `T` that are not methods.
+ * @alpha
+ */
+export type ExposableKeys<T> = {
+	[K in keyof T]?: T[K] extends (...args: any[]) => any ? never : K;
+}[keyof T];
+
+/**
+ * Type-level equality test used as a helper to evaluate readonly keys.
+ * - If X and Y are the same type, it evaluates to A.
+ * - If X and Y are different, it evaluates to B.
+ * @alpha
+ */
+export type IfEquals<X, Y, A = true, B = false> = (<T>() => T extends X ? 1 : 2) extends <
+	T,
+>() => T extends Y ? 1 : 2
+	? A
+	: B;
+
+/**
+ * Produces a union of keys of `T` which are readonly.
+ * @alpha
+ */
+export type ReadonlyKeys<T> = {
+	[P in keyof T]-?: IfEquals<{ [Q in P]: T[P] }, { -readonly [Q in P]: T[P] }, never, P>;
+}[keyof T];
+
+/**
+ * Type to enforce `readOnly: true` for readonly properties.
+ * @alpha
+ */
+export type ReadOnlyRequirement<TObj, K extends keyof TObj> = {
+	[P in K]-?: P extends ReadonlyKeys<TObj> ? { readOnly: true } : { readOnly?: false };
+}[K];
+
+/**
+ * Emits compile-time error when there is a type mismatch.
+ * @alpha
+ */
+export type TypeMatchOrError<Expected, Received> = [Received] extends [Expected]
+	? unknown
+	: {
+			__error__: "Zod schema value type does not match the property's declared type";
+			expected: Expected;
+			received: Received;
+		};
+
+/**
+ * A property definition class that describes the structure of the property
+ * @alpha
+ */
+export class PropertyDef {
+	public constructor(
+		public readonly name: string,
+		public readonly description: string | undefined,
+		public readonly schema: ZodTypeAny,
+		public readonly readOnly: boolean,
+	) {}
+}
+
+/**
+ * An interface for exposing properties of schema classes to an agent.
+ * @alpha
+ */
+export interface ExposedProperties {
+	exposeProperty<
+		S extends BindableSchema & Ctor,
+		K extends string & ExposableKeys<InstanceType<S>>,
+		TZ extends ZodTypeAny,
+	>(
+		schema: S,
+		name: K,
+		def: { schema: TZ; description?: string } & ReadOnlyRequirement<InstanceType<S>, K> &
+			TypeMatchOrError<InstanceType<S>[K], ZodInfer<TZ>>,
+	): void;
+
+	instanceOf<T extends TreeNodeSchemaClass>(
+		schema: T,
+	): ZodType<InstanceType<T>, ZodTypeDef, InstanceType<T>>;
+}
+
+/**
+ * An interface that SharedTree schema classes should implement to expose their properties to the LLM.
+ *
+ * @remarks
+ * The `getExposedProperties` free function will cause the method here to be called on the class passed to it.
+ *
+ * @privateremarks
+ * Implementing this interface correctly seems tricky?
+ * To actually implement it in a way that satisfies TypeScript,
+ * classes need to declare both a static version and an instance version of the method
+ * (the instance one can just delegate to the static one).
+ *
+ * @alpha
+ */
+export interface IExposedProperties {
+	[exposePropertiesSymbol]?(properties: ExposedProperties): void;
+}
+
+class ExposedPropertiesI implements ExposedProperties {
+	private readonly properties: Record<string, PropertyDef> = {};
+	private readonly referencedTypes = new Set<TreeNodeSchema>();
+
+	public constructor(private readonly schemaClass: BindableSchema) {}
+
+	public exposeProperty<
+		S extends BindableSchema & Ctor,
+		K extends string & ExposableKeys<InstanceType<S>>,
+		TZ extends ZodTypeAny,
+	>(
+		schema: S,
+		name: K,
+		def: { schema: TZ; description?: string } & ReadOnlyRequirement<InstanceType<S>, K> &
+			TypeMatchOrError<InstanceType<S>[K], ZodInfer<TZ>>,
+	): void {
+		if (schema !== this.schemaClass) {
+			throw new Error('Must expose properties on the "this" schema class');
+		}
+		this.properties[name] = new PropertyDef(
+			name,
+			def.description,
+			def.schema,
+			def.readOnly === true,
+		);
+	}
+
+	public instanceOf<T extends TreeNodeSchemaClass>(
+		schema: T,
+	): ZodType<InstanceType<T>, ZodTypeDef, InstanceType<T>> {
+		this.referencedTypes.add(schema);
+		return instanceOf(schema);
+	}
+
+	public static getExposedProperties(schemaClass: BindableSchema): {
+		properties: Record<string, PropertyDef>;
+		referencedTypes: Set<TreeNodeSchema>;
+	} {
+		const exposed = new ExposedPropertiesI(schemaClass);
+		const extractable = schemaClass as unknown as IExposedProperties;
+		if (extractable[exposePropertiesSymbol] !== undefined) {
+			extractable[exposePropertiesSymbol](exposed);
+		}
+		return {
+			properties: exposed.properties,
+			referencedTypes: exposed.referencedTypes,
+		};
+	}
+}
+
+/**
+ * Get the exposed properties of a schema class.
+ * @param schemaClass - The schema class to extract properties from.
+ * @returns A record of property names and their corresponding Zod types.
+ */
+export function getExposedProperties(schemaClass: BindableSchema): {
+	properties: Record<string, PropertyDef>;
+	referencedTypes: Set<TreeNodeSchema>;
+} {
+	return ExposedPropertiesI.getExposedProperties(schemaClass);
+}

package/src/subtree.ts

@@ -11,45 +11,40 @@ import type {
 	TreeFieldFromImplicitField,
 	TreeMapNode,
 	TreeArrayNode,
-	TreeNodeSchema,
 } from "@fluidframework/tree";
 import { TreeAlpha } from "@fluidframework/tree/alpha";
 import type {
-	UnsafeUnknownSchema,
 	ReadableField,
 	TreeRecordNode,
-	ReadSchema,
 	TreeBranchAlpha,
 } from "@fluidframework/tree/alpha";
 
 import { getNodeOnBranch } from "./getNodeOnBranch.js";
-import type { TreeView } from "./utils.js";
+import type { TreeView, ViewOrTree } from "./api.js";
 
 /**
  * Wraps either a {@link TreeView} or a {@link TreeNode} and provides a common interface over them.
  */
-export class Subtree<TRoot extends ImplicitFieldSchema | UnsafeUnknownSchema> {
-	public constructor(
-		private readonly viewOrNode: TreeView<TRoot> | (ReadableField<TRoot> & TreeNode),
-	) {
-		if (viewOrNode instanceof TreeNode && TreeAlpha.branch(viewOrNode) === undefined) {
+export class Subtree<TRoot extends ImplicitFieldSchema> {
+	public constructor(public readonly viewOrTree: ViewOrTree<TRoot>) {
+		if (viewOrTree instanceof TreeNode && TreeAlpha.branch(viewOrTree) === undefined) {
 			throw new UsageError("The provided node must belong to a branch.");
 		}
 	}
 
 	public get branch(): TreeBranchAlpha {
-		return this.viewOrNode instanceof TreeNode
-			? (TreeAlpha.branch(this.viewOrNode) ?? fail("Node cannot be raw."))
-			: this.viewOrNode;
+		return this.viewOrTree instanceof TreeNode
+			? (TreeAlpha.branch(this.viewOrTree) ?? fail("Node cannot be raw."))
+			: this.viewOrTree;
 	}
 
 	public get field(): ReadableField<TRoot> {
-		return this.viewOrNode instanceof TreeNode ? this.viewOrNode : this.viewOrNode.root;
+		return this.viewOrTree instanceof TreeNode ? this.viewOrTree : this.viewOrTree.root;
 	}
 
-	public set field(value: TreeFieldFromImplicitField<ReadSchema<TRoot>>) {
-		if (this.viewOrNode instanceof TreeNode) {
-			const parent = Tree.parent(this.viewOrNode);
+	public set field(value: TreeFieldFromImplicitField<TRoot>) {
+		if (this.viewOrTree instanceof TreeNode) {
+			const parent = Tree.parent(this.viewOrTree);
 			if (parent === undefined) {
 				// In general, this is not a correct cast, but we know that the root of the branch at least allows the type of `value`
 				const view = this.branch as TreeView<TRoot>;
@@ -58,12 +53,12 @@ export class Subtree<TRoot extends ImplicitFieldSchema | UnsafeUnknownSchema> {
 				const schema = Tree.schema(parent);
 				switch (schema.kind) {
 					case NodeKind.Object: {
-						const key = Tree.key(this.viewOrNode) as string;
+						const key = Tree.key(this.viewOrTree) as string;
 						(parent as unknown as Record<string, unknown>)[key] = value;
 						break;
 					}
 					case NodeKind.Record: {
-						const key = Tree.key(this.viewOrNode) as string;
+						const key = Tree.key(this.viewOrTree) as string;
 						if (value === undefined) {
 							// eslint-disable-next-line @typescript-eslint/no-dynamic-delete
 							delete (parent as TreeRecordNode)[key];
@@ -73,7 +68,7 @@ export class Subtree<TRoot extends ImplicitFieldSchema | UnsafeUnknownSchema> {
 						break;
 					}
 					case NodeKind.Map: {
-						const key = Tree.key(this.viewOrNode) as string;
+						const key = Tree.key(this.viewOrTree) as string;
 						if (value === undefined) {
 							(parent as TreeMapNode).delete(key);
 						} else {
@@ -82,7 +77,7 @@ export class Subtree<TRoot extends ImplicitFieldSchema | UnsafeUnknownSchema> {
 						break;
 					}
 					case NodeKind.Array: {
-						const index = Tree.key(this.viewOrNode) as number;
+						const index = Tree.key(this.viewOrTree) as number;
 						const arrayNode = parent as TreeArrayNode;
 						if (value === undefined) {
 							arrayNode.removeAt(index);
@@ -99,26 +94,28 @@ export class Subtree<TRoot extends ImplicitFieldSchema | UnsafeUnknownSchema> {
 				}
 			}
 		} else {
-			this.viewOrNode.root = value;
+			this.viewOrTree.root = value;
 		}
 	}
 
-	public get schema(): TreeNodeSchema | ReadSchema<TRoot> {
-		return this.viewOrNode instanceof TreeNode
-			? Tree.schema(this.viewOrNode)
-			: this.viewOrNode.schema;
+	public get schema(): TRoot {
+		return (
+			this.viewOrTree instanceof TreeNode
+				? Tree.schema(this.viewOrTree)
+				: this.viewOrTree.schema
+		) as TRoot;
 	}
 
 	public fork(): Subtree<TRoot> {
-		if (this.viewOrNode instanceof TreeNode) {
-			const branch = TreeAlpha.branch(this.viewOrNode) ?? fail("Node cannot be raw.");
+		if (this.viewOrTree instanceof TreeNode) {
+			const branch = TreeAlpha.branch(this.viewOrTree) ?? fail("Node cannot be raw.");
 			const node =
-				getNodeOnBranch(this.viewOrNode, branch.fork()) ??
+				getNodeOnBranch(this.viewOrTree, branch.fork()) ??
 				fail("Expected node to be on new fork.");
 
 			return new Subtree<TRoot>(node);
 		} else {
-			return new Subtree<TRoot>(this.viewOrNode.fork());
+			return new Subtree<TRoot>(this.viewOrTree.fork());
 		}
 	}
 }

package/src/typeGeneration.ts

@@ -27,6 +27,7 @@ import { z, type ZodTypeAny } from "zod";
 
 import type { BindableSchema } from "./methodBinding.js";
 import { FunctionWrapper, getExposedMethods } from "./methodBinding.js";
+import { getExposedProperties } from "./propertyBinding.js";
 import {
 	fail,
 	getOrCreate,
@@ -83,6 +84,11 @@ export function generateEditTypesForPrompt(
 						allSchemas.add(t);
 						objectTypeSchemas.add(t);
 					}
+					const exposedProperties = getExposedProperties(n);
+					for (const t of exposedProperties.referencedTypes) {
+						allSchemas.add(t);
+						objectTypeSchemas.add(t);
+					}
 				}
 			},
 		});
@@ -131,27 +137,39 @@ function generateEditTypes(
 	};
 }
 
-function getBoundMethodsForBindable(bindableSchema: BindableSchema): {
+function getBoundMembersForBindable(bindableSchema: BindableSchema): {
 	referencedTypes: Set<TreeNodeSchema>;
-	methods: [string, ZodTypeAny][];
+	members: [string, ZodTypeAny][];
 } {
-	const methodTypes: [string, ZodTypeAny][] = [];
+	const memberTypes: [string, ZodTypeAny][] = [];
+
 	const methods = getExposedMethods(bindableSchema);
 	for (const [name, method] of Object.entries(methods.methods)) {
 		const zodFunction = z.instanceof(FunctionWrapper);
 		// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
 		(zodFunction as any).method = method;
-		methodTypes.push([name, zodFunction]);
+		memberTypes.push([name, zodFunction]);
 	}
-	return { methods: methodTypes, referencedTypes: methods.referencedTypes };
+	const props = getExposedProperties(bindableSchema);
+	for (const [name, prop] of Object.entries(props.properties)) {
+		const schema = prop.schema;
+		// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
+		(schema as any).property = prop;
+		memberTypes.push([name, prop.schema]);
+	}
+
+	return {
+		members: memberTypes,
+		referencedTypes: new Set([...props.referencedTypes, ...methods.referencedTypes]),
+	};
 }
 
-function getBoundMethods(
+function getBoundMembers(
 	definition: string,
 	bindableSchemas: Map<string, BindableSchema>,
 ): [string, ZodTypeAny][] {
-	const bindableSchema = bindableSchemas.get(definition) ?? fail("Unknown definition");
-	return getBoundMethodsForBindable(bindableSchema).methods;
+	const bindableSchema = bindableSchemas.get(definition) ?? fail("unknown definition");
+	return getBoundMembersForBindable(bindableSchema).members;
 }
 
 function addBindingIntersectionIfNeeded(
@@ -163,21 +181,68 @@ function addBindingIntersectionIfNeeded(
 ): ZodTypeAny {
 	let zodType = zodTypeBound;
 	let description = simpleNodeSchema.metadata?.description ?? "";
-	const boundMethods = getBoundMethods(definition, bindableSchemas);
-	if (boundMethods.length > 0) {
-		const methods: Record<string, z.ZodTypeAny> = {};
-		for (const [name, zodFunction] of boundMethods) {
-			if (methods[name] !== undefined) {
-				throw new UsageError(
-					`Method ${name} conflicts with field of the same name in schema ${definition}`,
-				);
+
+	const boundMembers = getBoundMembers(definition, bindableSchemas);
+	const methods: Record<string, z.ZodTypeAny> = {};
+	const properties: Record<string, z.ZodTypeAny> = {};
+
+	let hasProperties: boolean = false;
+	let hasMethods: boolean = false;
+
+	if (boundMembers.length > 0) {
+		for (const [name, zodMember] of boundMembers) {
+			const kind =
+				// eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access
+				(zodMember as any).method === undefined
+					? // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access
+						(zodMember as any).property === undefined
+						? "Unknown"
+						: "Property"
+					: "Method";
+
+			if (kind === "Unknown") {
+				throw new UsageError(`Unrecognized bound member ${name} in schema ${definition}`);
+			}
+			if (kind === "Method") {
+				if (methods[name] !== undefined) {
+					throw new UsageError(
+						`${kind} ${name} conflicts with field of the same name in schema ${definition}`,
+					);
+				}
+				methods[name] = zodMember;
 			}
-			methods[name] = zodFunction;
+			if (kind === "Property") {
+				if (properties[name] !== undefined) {
+					throw new UsageError(
+						`${kind} ${name} conflicts with field of the same name in schema ${definition}`,
+					);
+				}
+				properties[name] = zodMember;
+			}
+		}
+		hasMethods = Object.keys(methods).length > 0 ? true : false;
+		hasProperties = Object.keys(properties).length > 0 ? true : false;
+
+		if (hasMethods) {
+			zodType = z.intersection(zodType, z.object(methods));
+		}
+		if (hasProperties) {
+			zodType = z.intersection(zodType, z.object(properties));
 		}
-		zodType = z.intersection(zodType, z.object(methods));
-		const methodNote = `Note: this ${typeString} has custom user-defined methods directly on it.`;
-		description = description === "" ? methodNote : `${description} - ${methodNote}`;
 	}
+
+	let note = "";
+	if (hasMethods && hasProperties) {
+		note = `Note: this ${typeString} has custom user-defined methods and properties directly on it.`;
+	} else if (hasMethods) {
+		note = `Note: this ${typeString} has custom user-defined methods directly on it.`;
+	} else if (hasProperties) {
+		note = `Note: this ${typeString} has custom user-defined properties directly on it.`;
+	}
+	if (note !== "") {
+		description = description === "" ? note : `${description} - ${note}`;
+	}
+
 	return zodType.describe(description);
 }
 
@@ -209,14 +274,15 @@ function getOrCreateType(
 						.filter(([, value]) => value !== undefined),
 				);
 
-				// Unlike arrays/maps/records, object nodes include methods directly on them rather than using an intersection
-				for (const [name, zodFunction] of getBoundMethods(definition, bindableSchemas)) {
+				for (const [name, zodMember] of getBoundMembers(definition, bindableSchemas)) {
+					// eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access
+					const kind = (zodMember as any).method === undefined ? "Property" : "Method";
 					if (properties[name] !== undefined) {
 						throw new UsageError(
-							`Method ${name} conflicts with field of the same name in schema ${definition}`,
+							`${kind} ${name} conflicts with field of the same name in schema ${definition}`,
 						);
 					}
-					properties[name] = zodFunction;
+					properties[name] = zodMember;
 				}
 
 				return (type = z

package/src/utils.ts

@@ -12,10 +12,8 @@ import { UsageError } from "@fluidframework/telemetry-utils/internal";
 import type { ImplicitFieldSchema, TreeNodeSchemaClass } from "@fluidframework/tree";
 import type {
 	InsertableContent,
-	TreeBranchAlpha,
 	TreeNode,
 	TreeNodeSchema,
-	TreeViewAlpha,
 	UnsafeUnknownSchema,
 } from "@fluidframework/tree/alpha";
 import {
@@ -29,6 +27,7 @@ import { NodeKind, normalizeFieldSchema } from "@fluidframework/tree/internal";
 import { z } from "zod";
 
 import { FunctionWrapper } from "./methodBinding.js";
+import { PropertyDef } from "./propertyBinding.js";
 
 /**
  * Subset of Map interface.
@@ -86,17 +85,6 @@ export function getOrCreate<K, V>(
 	return value;
 }
 
-/**
- * TODO
- * @alpha
- * @privateRemarks This is a subset of the TreeViewAlpha functionality because if take it wholesale, it causes problems with invariance of the generic parameters.
- */
-export type TreeView<TRoot extends ImplicitFieldSchema | UnsafeUnknownSchema> = Pick<
-	TreeViewAlpha<TRoot>,
-	"root" | "fork" | "merge" | "rebaseOnto" | "schema" | "events"
-> &
-	TreeBranchAlpha;
-
 /**
  * TODO
  */
@@ -434,6 +422,24 @@ export function getZodSchemaAsTypeScript(
 		}
 	}
 
+	function appendBoundProperties(type: z.ZodType): void {
+		const property = (type as unknown as { property?: PropertyDef }).property;
+
+		if (!(property instanceof PropertyDef)) {
+			if (type.description !== undefined && type.description !== "") {
+				append(` // ${type.description}`);
+			}
+			return;
+		}
+
+		if (property.readOnly === true) {
+			append(" // readonly");
+		}
+		if (property.description !== undefined && property.description !== "") {
+			append(` - ${property.description}`);
+		}
+	}
+
 	function appendArrayType(arrayType: z.ZodType) {
 		appendType((arrayType._def as z.ZodArrayDef).type, TypePrecedence.Object);
 		append("[]");
@@ -446,6 +452,7 @@ export function getZodSchemaAsTypeScript(
 		// eslint-disable-next-line prefer-const
 		for (let [name, type] of Object.entries((objectType._def as z.ZodObjectDef).shape())) {
 			const method = (type as unknown as { method: object | undefined }).method;
+
 			if (method === undefined || !(method instanceof FunctionWrapper)) {
 				append(name);
 				if (getTypeKind(type) === z.ZodFirstPartyTypeKind.ZodOptional) {
@@ -455,8 +462,7 @@ export function getZodSchemaAsTypeScript(
 				append(": ");
 				appendType(type);
 				append(";");
-				const comment = type.description;
-				if (comment !== undefined && comment !== "") append(` // ${comment}`);
+				appendBoundProperties(type);
 				appendNewLine();
 			}
 		}