@fluidframework/tree-agent 2.74.0 → 2.81.0

package/src/methodBinding.ts

@@ -7,7 +7,9 @@
 import type { TreeNodeSchema, TreeNodeSchemaClass } from "@fluidframework/tree";
 import { NodeKind } from "@fluidframework/tree";
 import type { z } from "zod";
+
 import { instanceOf } from "./renderZodTypeScript.js";
+import type { TypeFactoryType } from "./treeAgentTypes.js";
 
 /**
  * A utility type that extracts the method keys from a given type.
@@ -61,7 +63,8 @@ export function getExposedMethods(schemaClass: BindableSchema): {
  * A type that represents a function argument.
  * @alpha
  */
-export type Arg<T extends z.ZodTypeAny = z.ZodTypeAny> = readonly [name: string, type: T];
+export type Arg<T extends z.ZodTypeAny | TypeFactoryType = z.ZodTypeAny | TypeFactoryType> =
+	readonly [name: string, type: T];
 
 /**
  * A function definition interface that describes the structure of a function.
@@ -69,12 +72,24 @@ export type Arg<T extends z.ZodTypeAny = z.ZodTypeAny> = readonly [name: string,
  */
 export interface FunctionDef<
 	Args extends readonly Arg[],
-	Return extends z.ZodTypeAny,
-	Rest extends z.ZodTypeAny | null = null,
+	Return extends z.ZodTypeAny | TypeFactoryType,
+	Rest extends z.ZodTypeAny | TypeFactoryType | null = null,
 > {
+	/**
+	 * Optional description of the function.
+	 */
 	description?: string;
+	/**
+	 * The function's parameters.
+	 */
 	args: Args;
+	/**
+	 * Optional rest parameter type.
+	 */
 	rest?: Rest;
+	/**
+	 * The function's return type.
+	 */
 	returns: Return;
 }
 
@@ -82,15 +97,20 @@ export interface FunctionDef<
  * A class that implements the FunctionDef interface.
  */
 export class FunctionWrapper
-	implements FunctionDef<readonly Arg[], z.ZodTypeAny, z.ZodTypeAny | null>
+	implements
+		FunctionDef<
+			readonly Arg[],
+			z.ZodTypeAny | TypeFactoryType,
+			z.ZodTypeAny | TypeFactoryType | null
+		>
 {
 	public constructor(
 		public readonly name: string,
 		public readonly description: string | undefined,
 		public readonly args: readonly Arg[],
 		// eslint-disable-next-line @rushstack/no-new-null
-		public readonly rest: z.ZodTypeAny | null,
-		public readonly returns: z.ZodTypeAny,
+		public readonly rest: z.ZodTypeAny | TypeFactoryType | null,
+		public readonly returns: z.ZodTypeAny | TypeFactoryType,
 	) {}
 }
 
@@ -109,9 +129,9 @@ export type ArgsTuple<T extends readonly Arg[]> = T extends readonly [infer Sing
  * @alpha
  */
 export function buildFunc<
-	const Return extends z.ZodTypeAny,
+	const Return extends z.ZodTypeAny | TypeFactoryType,
 	const Args extends readonly Arg[],
-	const Rest extends z.ZodTypeAny | null = null,
+	const Rest extends z.ZodTypeAny | TypeFactoryType | null = null,
 >(
 	def: { description?: string; returns: Return; rest?: Rest },
 	...args: Args
@@ -124,12 +144,48 @@ export function buildFunc<
 	};
 }
 
+/**
+ * A utility type that extracts inferred parameter types from Zod args.
+ * @alpha
+ */
+export type InferArgsZod<Args extends readonly Arg<z.ZodTypeAny>[]> = Args extends readonly [
+	infer Head extends Arg<z.ZodTypeAny>,
+	...infer Tail extends readonly Arg<z.ZodTypeAny>[],
+]
+	? [z.infer<Head[1]>, ...InferArgsZod<Tail>]
+	: [];
+
+/**
+ * A utility type that infers the function signature from a Zod function definition with strict type checking.
+ * @alpha
+ */
+export type InferZod<T> = T extends FunctionDef<
+	infer Args extends readonly Arg<z.ZodTypeAny>[],
+	infer Return extends z.ZodTypeAny,
+	any
+>
+	? (...args: InferArgsZod<Args>) => z.infer<Return>
+	: never;
+
+/**
+ * A utility type that infers the function signature from a type factory function definition with relaxed type checking.
+ * @alpha
+ */
+export type InferTypeFactory<T> = T extends FunctionDef<readonly Arg[], infer Return, any>
+	? (...args: any[]) => any
+	: never;
+
 /**
  * A utility type that infers the return type of a function definition.
  * @alpha
+ * @remarks
+ * For Zod types, provides strict compile-time type checking. For type factory types, returns `any`.
+ * @deprecated Use InferZod or InferTypeFactory directly for better type safety.
  */
-export type Infer<T> = T extends FunctionDef<infer Args, infer Return, infer Rest>
-	? z.infer<z.ZodFunction<z.ZodTuple<ArgsTuple<Args>, Rest>, Return>>
+export type Infer<T> = T extends FunctionDef<readonly Arg[], infer Return, any>
+	? Return extends z.ZodTypeAny
+		? InferZod<T>
+		: InferTypeFactory<T>
 	: never;
 
 /**
@@ -137,16 +193,32 @@ export type Infer<T> = T extends FunctionDef<infer Args, infer Return, infer Res
  * @alpha
  */
 export interface ExposedMethods {
+	/**
+	 * Expose a method with Zod types (strict compile-time type checking).
+	 */
 	expose<
 		const K extends string & keyof MethodKeys<InstanceType<S>>,
-		S extends BindableSchema & Ctor<Record<K, Infer<Z>>> & IExposedMethods,
-		Z extends FunctionDef<any, any, any>,
+		S extends BindableSchema & Ctor<Record<K, InferZod<Z>>> & IExposedMethods,
+		Z extends FunctionDef<readonly Arg<z.ZodTypeAny>[], z.ZodTypeAny, z.ZodTypeAny | null>,
 	>(schema: S, methodName: K, zodFunction: Z): void;
 
+	/**
+	 * Expose a method with type factory types (relaxed compile-time type checking).
+	 */
+	expose<
+		const K extends string & keyof MethodKeys<InstanceType<S>>,
+		S extends BindableSchema & Ctor & IExposedMethods,
+		Z extends FunctionDef<
+			readonly Arg<TypeFactoryType>[],
+			TypeFactoryType,
+			TypeFactoryType | null
+		>,
+	>(schema: S, methodName: K, tfFunction: Z): void;
+
 	/**
 	 * Create a Zod schema for a SharedTree schema class.
 	 * @remarks
-	 * Use it to "wrap" schema types that are referenced as arguments or return types when exposing methods (with {@link ExposedMethods.expose}).
+	 * Use it to "wrap" schema types that are referenced as arguments or return types when exposing methods with {@link ExposedMethods}.
 	 */
 	instanceOf<T extends TreeNodeSchemaClass>(
 		schema: T,
@@ -174,6 +246,9 @@ export const exposeMethodsSymbol: unique symbol = Symbol("run");
  * @alpha
  */
 export interface IExposedMethods {
+	/**
+	 * Static method that exposes methods of this schema class to an agent.
+	 */
 	[exposeMethodsSymbol](methods: ExposedMethods): void;
 }
 
@@ -185,8 +260,12 @@ class ExposedMethodsI implements ExposedMethods {
 
 	public expose<
 		const K extends string & keyof MethodKeys<InstanceType<S>>,
-		S extends BindableSchema & Ctor<Record<K, Infer<Z>>> & IExposedMethods,
-		Z extends FunctionDef<readonly Arg[], z.ZodTypeAny, z.ZodTypeAny | null>,
+		S extends BindableSchema & Ctor & IExposedMethods,
+		Z extends FunctionDef<
+			readonly Arg[],
+			z.ZodTypeAny | TypeFactoryType,
+			z.ZodTypeAny | TypeFactoryType | null
+		>,
 	>(schema: S, methodName: K, functionDef: Z): void {
 		if (schema !== this.schemaClass) {
 			throw new Error('Must expose methods on the "this" object');

package/src/propertyBinding.ts

@@ -9,6 +9,8 @@ import type { ZodType, ZodTypeAny, ZodTypeDef, infer as ZodInfer } from "zod";
 
 import type { BindableSchema, Ctor } from "./methodBinding.js";
 import { instanceOf } from "./renderZodTypeScript.js";
+import type { TypeFactoryType } from "./treeAgentTypes.js";
+import { isTypeFactoryType } from "./treeAgentTypes.js";
 
 /**
  * A symbol used to expose properties to the LLM.
@@ -73,9 +75,21 @@ export type TypeMatchOrError<Expected, Received> = [Received] extends [Expected]
  */
 export class PropertyDef {
 	public constructor(
+		/**
+		 * The name of the property.
+		 */
 		public readonly name: string,
+		/**
+		 * Optional description of the property.
+		 */
 		public readonly description: string | undefined,
-		public readonly schema: ZodTypeAny,
+		/**
+		 * The schema defining the property's type (either Zod or TypeFactory).
+		 */
+		public readonly schema: ZodTypeAny | TypeFactoryType,
+		/**
+		 * Whether the property is readonly.
+		 */
 		public readonly readOnly: boolean,
 	) {}
 }
@@ -85,6 +99,9 @@ export class PropertyDef {
  * @alpha
  */
 export interface ExposedProperties {
+	/**
+	 * Expose a property with Zod type checking.
+	 */
 	exposeProperty<
 		S extends BindableSchema & Ctor,
 		K extends string & ExposableKeys<InstanceType<S>>,
@@ -96,6 +113,29 @@ export interface ExposedProperties {
 			TypeMatchOrError<InstanceType<S>[K], ZodInfer<TZ>>,
 	): void;
 
+	/**
+	 * Expose a property with type factory type and metadata.
+	 */
+	exposeProperty<
+		S extends BindableSchema & Ctor,
+		K extends string & ExposableKeys<InstanceType<S>>,
+	>(
+		schema: S,
+		name: K,
+		def: { schema: TypeFactoryType; description?: string; readOnly?: boolean },
+	): void;
+
+	/**
+	 * Expose a property with type factory type (simple form).
+	 */
+	exposeProperty<
+		S extends BindableSchema & Ctor,
+		K extends string & ExposableKeys<InstanceType<S>>,
+	>(schema: S, name: K, tfType: TypeFactoryType): void;
+
+	/**
+	 * Create a Zod type that references a SharedTree schema class.
+	 */
 	instanceOf<T extends TreeNodeSchemaClass>(
 		schema: T,
 	): ZodType<InstanceType<T>, ZodTypeDef, InstanceType<T>>;
@@ -116,6 +156,9 @@ export interface ExposedProperties {
  * @alpha
  */
 export interface IExposedProperties {
+	/**
+	 * Static method that exposes properties of this schema class to an agent.
+	 */
 	[exposePropertiesSymbol]?(properties: ExposedProperties): void;
 }
 
@@ -132,18 +175,32 @@ class ExposedPropertiesI implements ExposedProperties {
 	>(
 		schema: S,
 		name: K,
-		def: { schema: TZ; description?: string } & ReadOnlyRequirement<InstanceType<S>, K> &
-			TypeMatchOrError<InstanceType<S>[K], ZodInfer<TZ>>,
+		defOrType:
+			| ({ schema: TZ; description?: string } & ReadOnlyRequirement<InstanceType<S>, K> &
+					TypeMatchOrError<InstanceType<S>[K], ZodInfer<TZ>>)
+			| TypeFactoryType,
 	): 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,
-		);
+
+		// Handle TypeFactoryType (simple case - type passed directly)
+		if (isTypeFactoryType(defOrType)) {
+			this.properties[name] = new PropertyDef(name, undefined, defOrType, false);
+		} else {
+			// Handle object with schema property (works for both Zod and TypeFactory)
+			const def = defOrType as {
+				schema: TZ | TypeFactoryType;
+				description?: string;
+				readOnly?: boolean;
+			};
+			this.properties[name] = new PropertyDef(
+				name,
+				def.description,
+				def.schema,
+				def.readOnly === true,
+			);
+		}
 	}
 
 	public instanceOf<T extends TreeNodeSchemaClass>(

package/src/renderSchemaTypeScript.ts

@@ -18,8 +18,14 @@ import { z } from "zod";
 import type { BindableSchema, FunctionWrapper } from "./methodBinding.js";
 import { getExposedMethods } from "./methodBinding.js";
 import { getExposedProperties, type PropertyDef } from "./propertyBinding.js";
-import { getFriendlyName, isNamedSchema, llmDefault, unqualifySchema } from "./utils.js";
+import {
+	instanceOfsTypeFactory,
+	renderTypeFactoryTypeScript,
+} from "./renderTypeFactoryTypeScript.js";
 import { instanceOfs, renderZodTypeScript } from "./renderZodTypeScript.js";
+import type { TypeFactoryOptional, TypeFactoryType } from "./treeAgentTypes.js";
+import { isTypeFactoryType } from "./treeAgentTypes.js";
+import { getFriendlyName, isNamedSchema, llmDefault, unqualifySchema } from "./utils.js";
 
 interface BoundMembers {
 	methods: Record<string, FunctionWrapper>;
@@ -284,7 +290,9 @@ export function renderSchemaTypeScript(
 					lines.push(`// ${note}`);
 				}
 			}
-			lines.push(formatMethod(name, method));
+			const methodString = formatMethod(name, method);
+			const methodLines = methodString.split("\n");
+			lines.push(...methodLines);
 		}
 		if (lines.length > 0) {
 			hasHelperMethods = true;
@@ -434,7 +442,11 @@ function renderPropertyLines(properties: Record<string, PropertyDef>): string[]
 			}
 		}
 		const modifier = property.readOnly ? "readonly " : "";
-		lines.push(`${modifier}${name}: ${renderZodType(property.schema)};`);
+		const typeString = renderType(property.schema, 0);
+		const propertyLine = `${modifier}${name}: ${typeString};`;
+		// Split multi-line type strings and add to lines array
+		const propertyLines = propertyLine.split("\n");
+		lines.push(...propertyLines);
 	}
 	return lines;
 }
@@ -443,13 +455,13 @@ function formatMethod(name: string, method: FunctionWrapper): string {
 	const args: string[] = [];
 	for (const [argName, argType] of method.args) {
 		const { innerType, optional } = unwrapOptional(argType);
-		const renderedType = renderZodType(innerType);
+		const renderedType = renderType(innerType, 0);
 		args.push(`${argName}${optional ? "?" : ""}: ${renderedType}`);
 	}
 	if (method.rest !== null) {
-		args.push(`...rest: ${renderZodType(method.rest)}[]`);
+		args.push(`...rest: ${renderType(method.rest, 0)}[]`);
 	}
-	return `${name}(${args.join(", ")}): ${renderZodType(method.returns)};`;
+	return `${name}(${args.join(", ")}): ${renderType(method.returns, 0)};`;
 }
 
 function renderLeaf(leafKind: ValueSchema): string {
@@ -482,7 +494,15 @@ function formatExpression(
 /**
  * Detects optional zod wrappers so argument lists can keep TypeScript optional markers in sync.
  */
-function unwrapOptional(type: z.ZodTypeAny): { innerType: z.ZodTypeAny; optional: boolean } {
+function unwrapOptional(type: z.ZodTypeAny | TypeFactoryType): {
+	innerType: z.ZodTypeAny | TypeFactoryType;
+	optional: boolean;
+} {
+	// Handle type factory optional type
+	if (isTypeFactoryType(type) && type._kind === "optional") {
+		return { innerType: (type as TypeFactoryOptional).innerType, optional: true };
+	}
+	// Handle Zod optional type
 	if (type instanceof z.ZodOptional) {
 		const inner = type.unwrap() as z.ZodTypeAny;
 		return { innerType: inner, optional: true };
@@ -516,8 +536,10 @@ function ensureNoMemberConflicts(
 }
 
 /**
- * Converts schema metadata into TypeScript declarations suitable for prompt inclusion.
+ * Dispatches to the correct renderer based on whether the type is Zod or type factory.
  */
-function renderZodType(type: z.ZodTypeAny): string {
-	return renderZodTypeScript(type, getFriendlyName, instanceOfs);
+function renderType(type: z.ZodTypeAny | TypeFactoryType, indentLevel: number = 0): string {
+	return isTypeFactoryType(type)
+		? renderTypeFactoryTypeScript(type, getFriendlyName, instanceOfsTypeFactory, indentLevel)
+		: renderZodTypeScript(type, getFriendlyName, instanceOfs);
 }