fix: tree agent rendering of fluid handles (#26292)

packages/framework/tree-agent/src/prompt.ts

@@ -8,12 +8,17 @@ import { NodeKind, Tree, TreeNode } from "@fluidframework/tree";
 import type { ImplicitFieldSchema, TreeMapNode } from "@fluidframework/tree";
 import type { ReadableField } from "@fluidframework/tree/alpha";
 import { getSimpleSchema } from "@fluidframework/tree/alpha";
-import { normalizeFieldSchema } from "@fluidframework/tree/internal";
+import { normalizeFieldSchema, ValueSchema } from "@fluidframework/tree/internal";
 
 import type { Subtree } from "./subtree.js";
 import { generateEditTypesForPrompt } from "./typeGeneration.js";
 import { getFriendlyName, communize, findSchemas } from "./utils.js";
 
+/**
+ * The type name used for handles in generated TypeScript.
+ */
+export const fluidHandleTypeName = "_OpaqueHandle";
+
 /**
  * Produces a "system" prompt for the tree agent, based on the provided subtree.
  */
@@ -32,6 +37,7 @@ export function getPrompt(args: {
 	let nodeTypeUnion: string | undefined;
 	let hasArrays = false;
 	let hasMaps = false;
+	let hasFluidHandles = false;
 	let exampleObjectName: string | undefined;
 	for (const s of findSchemas(schema)) {
 		if (s.kind !== NodeKind.Leaf) {
@@ -54,6 +60,10 @@ export function getPrompt(args: {
 				exampleObjectName ??= getFriendlyName(s);
 				break;
 			}
+			case NodeKind.Leaf: {
+				hasFluidHandles ||= s.info === ValueSchema.FluidHandle;
+				break;
+			}
 			// No default
 		}
 	}
@@ -63,6 +73,15 @@ export function getPrompt(args: {
 		schema,
 		getSimpleSchema(schema),
 	);
+	const fluidHandleType = hasFluidHandles
+		? `/**
+ * Opaque handle type representing a reference to a Fluid object.
+ * This type should not be constructed by generated code.
+ */
+type ${fluidHandleTypeName} = unknown;
+
+`
+		: "";
 	const exampleTypeName =
 		nodeTypeUnion === undefined
 			? undefined
@@ -274,7 +293,7 @@ Finally, double check that the edits would accomplish the user's request (if it
 The JSON tree adheres to the following Typescript schema:
 
 \`\`\`typescript
-${typescriptSchemaTypes}
+${fluidHandleType}${typescriptSchemaTypes}
 \`\`\`
 
 If the user asks you a question about the tree, you should inspect the state of the tree and answer the question.

packages/framework/tree-agent/src/renderSchemaTypeScript.ts

@@ -17,6 +17,7 @@ import { z } from "zod";
 
 import type { BindableSchema, FunctionWrapper } from "./methodBinding.js";
 import { getExposedMethods } from "./methodBinding.js";
+import { fluidHandleTypeName } from "./prompt.js";
 import { getExposedProperties, type PropertyDef } from "./propertyBinding.js";
 import {
 	instanceOfsTypeFactory,
@@ -478,8 +479,11 @@ function renderLeaf(leafKind: ValueSchema): string {
 		case ValueSchema.Null: {
 			return "null";
 		}
+		case ValueSchema.FluidHandle: {
+			return fluidHandleTypeName;
+		}
 		default: {
-			throw new Error(`Unsupported leaf kind ${NodeKind[leafKind]}.`);
+			throw new Error(`Unsupported leaf kind.`);
 		}
 	}
 }

packages/framework/tree-agent/src/test/__snapshots__/prompt.md

@@ -2,6 +2,12 @@ You are a helpful assistant collaborating with the user on a document. The docum
 The JSON tree adheres to the following Typescript schema:
 
 ```typescript
+/**
+ * Opaque handle type representing a reference to a Fluid object.
+ * This type should not be constructed by generated code.
+ */
+type _OpaqueHandle = unknown;
+
 // A test map - Note: this map has custom user-defined properties directly on it.
 type TestMap = Map<string, number> & {
     // Readonly map metadata
@@ -23,6 +29,7 @@ type TestArray = TestArrayItem[];
 interface Obj {
     map: TestMap;
     array: TestArray;
+    handle?: _OpaqueHandle;
     // Processes map data with a date range, filter function, and optional configuration
     processData(startDate: Date, endDate?: Date, filter: (value: number) => boolean, options?: {
         mode: ("sync" | "async");

packages/framework/tree-agent/src/test/prompt.spec.ts

@@ -17,7 +17,7 @@ import {
 
 import type { TreeView } from "../api.js";
 import { buildFunc, exposeMethodsSymbol, type ExposedMethods } from "../methodBinding.js";
-import { getPrompt } from "../prompt.js";
+import { fluidHandleTypeName, getPrompt } from "../prompt.js";
 import { exposePropertiesSymbol, type ExposedProperties } from "../propertyBinding.js";
 import { Subtree } from "../subtree.js";
 import { typeFactory as tf } from "../treeAgentTypes.js";
@@ -201,6 +201,33 @@ describe("Prompt generation", () => {
 		}
 	});
 
+	it("includes handle type declaration when handles are present in the schema", () => {
+		// If no handles, then the prompt shouldn't include the handle type declaration
+		{
+			const view = getView(sf.object("Object", {}), {});
+			const prompt = getPrompt({
+				subtree: new Subtree(view),
+				editToolName: "EditTreeTool",
+			});
+			assert.ok(!prompt.includes(`type ${fluidHandleTypeName} = unknown`));
+		}
+		// If there are handles, then the prompt should include the handle type declaration
+		{
+			const view = getView(
+				sf.object("ObjectWithHandle", {
+					handle: sf.optional(sf.handle),
+				}),
+				{ handle: undefined },
+			);
+			const prompt = getPrompt({
+				subtree: new Subtree(view),
+				editToolName: "EditTreeTool",
+			});
+			assert.ok(prompt.includes(`type ${fluidHandleTypeName} = unknown`));
+			assert.ok(prompt.includes(`handle?: ${fluidHandleTypeName}`));
+		}
+	});
+
 	it("sanitizes schema names that contain invalid characters", () => {
 		class InvalidlyNamedObject extends sf.object("Test-Object!", { value: sf.string }) {}
 
@@ -297,6 +324,7 @@ describe("Prompt snapshot", () => {
 		class Obj extends sf.object("Obj", {
 			map: TestMap,
 			array: TestArray,
+			handle: sf.optional(sf.handle),
 		}) {
 			public static [exposeMethodsSymbol](methods: ExposedMethods): void {
 				methods.expose(
@@ -360,6 +388,7 @@ describe("Prompt snapshot", () => {
 				new NumberValue({ value: 2 }),
 				new NumberValue({ value: 3 }),
 			],
+			handle: undefined,
 		});
 
 		const fullPrompt = getPrompt({

packages/framework/tree-agent/src/test/typeGeneration.spec.ts

@@ -16,6 +16,7 @@ import {
 import { z } from "zod";
 
 import { buildFunc, exposeMethodsSymbol, type ExposedMethods } from "../methodBinding.js";
+import { fluidHandleTypeName } from "../prompt.js";
 import { exposePropertiesSymbol, type ExposedProperties } from "../propertyBinding.js";
 import { generateEditTypesForPrompt } from "../typeGeneration.js";
 
@@ -69,6 +70,20 @@ const initialAppState = {
 };
 
 describe("Type generation", () => {
+	it("for handle nodes", () => {
+		class ObjWithHandle extends sf.object("ObjWithHandle", {
+			handle: sf.optional(sf.handle),
+		}) {}
+		const handleSchemaString = getDomainSchemaString(ObjWithHandle, { handle: undefined });
+		assert.deepEqual(
+			handleSchemaString,
+			`interface ObjWithHandle {
+    handle?: ${fluidHandleTypeName};
+}
+`,
+		);
+	});
+
 	describe("for schemas with methods", () => {
 		it("works on object nodes", () => {
 			class ObjWithMethod extends sf.object("ObjWithMethod", {}) {