Use independent view sandbox for tree-agent edit isolation

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

@@ -21,14 +21,13 @@ import type {
 	SharedTreeChatModel,
 	EditResult,
 	SemanticAgentOptions,
-	Logger,
 	AsynchronousEditor,
 	Context,
 	SynchronousEditor,
 	ViewOrTree,
 } from "./api.js";
 import { getPrompt, stringifyTree } from "./prompt.js";
-import { Subtree } from "./subtree.js";
+import { copyIndependentSubtree, Subtree } from "./subtree.js";
 import {
 	llmDefault,
 	findSchemas,
@@ -114,13 +113,17 @@ export class SharedTreeSemanticAgent<TSchema extends ImplicitFieldSchema> {
 			);
 		}
 
-		// Fork a branch that will live for the lifetime of this query (which can be multiple LLM calls if the there are errors or the LLM decides to take multiple steps to accomplish a task).
+		// Fork a branch that will live for the lifetime of this query.
 		// The branch will be merged back into the outer branch if and only if the query succeeds.
-		const queryTree = this.outerTree.fork();
+		const snapshotTree = this.outerTree.fork();
+		// Create an independent view as a "sandbox" for LLM edits.
+		// This provides isolation - errors cannot corrupt the query branch or outer tree.
+		let sandbox = copyIndependentSubtree(snapshotTree.viewOrTree);
 		const maxEditCount = this.options?.maximumSequentialEdits ?? defaultMaxSequentialEdits;
 		let active = true;
 		let editCount = 0;
 		let rollbackEdits = false;
+		const successfulEditCodes: string[] = [];
 		const { editToolName } = this.client;
 		const edit = async (editCode: string): Promise<EditResult> => {
 			if (editToolName === undefined) {
@@ -144,15 +147,50 @@ export class SharedTreeSemanticAgent<TSchema extends ImplicitFieldSchema> {
 				};
 			}
 
-			const editResult = await applyTreeFunction(
-				queryTree,
-				editCode,
-				this.editor,
-				this.options?.logger,
+			this.options?.logger?.log(`### Editing Tool Invoked\n\n`);
+			this.options?.logger?.log(
+				`#### Generated Code\n\n\`\`\`javascript\n${editCode}\n\`\`\`\n\n`,
 			);
 
-			rollbackEdits = editResult.type !== "success";
-			return editResult;
+			try {
+				await this.editor(sandbox.viewOrTree, editCode);
+			} catch (error: unknown) {
+				rollbackEdits = true;
+				this.options?.logger?.log(`#### Error\n\n`);
+				this.options?.logger?.log(`\`\`\`JSON\n${toErrorString(error)}\n\`\`\`\n\n`);
+
+				// Rebuild the sandbox from scratch and replay successful edits.
+				// eslint-disable-next-line require-atomic-updates -- edits are called sequentially, not concurrently
+				sandbox = copyIndependentSubtree(snapshotTree.viewOrTree);
+				try {
+					for (const successfulEdit of successfulEditCodes) {
+						await this.editor(sandbox.viewOrTree, successfulEdit);
+					}
+				} catch (replayError: unknown) {
+					successfulEditCodes.length = 0;
+					return {
+						type: "editingError",
+						message: `An internal error occurred. All edits for this query have been discarded and the state of the tree will be reset to its state as it was before the query began. Error: ${toErrorString(replayError)}`,
+					};
+				}
+
+				return {
+					type: "editingError",
+					message: `Running the generated code produced an error. The state of the tree will be reset to its previous state as it was before the code ran. Please try again. Here is the error: ${toErrorString(error)}`,
+				};
+			}
+
+			successfulEditCodes.push(editCode);
+			rollbackEdits = false;
+
+			this.options?.logger?.log(`#### New Tree State\n\n`);
+			this.options?.logger?.log(
+				`${`\`\`\`JSON\n${stringifyTree(sandbox.field)}\n\`\`\``}\n\n`,
+			);
+			return {
+				type: "success",
+				message: `After running the code, the new state of the tree is:\n\n\`\`\`JSON\n${stringifyTree(sandbox.field)}\n\`\`\``,
+			};
 		};
 
 		const responseMessage = await this.client.query({
@@ -161,9 +199,19 @@ export class SharedTreeSemanticAgent<TSchema extends ImplicitFieldSchema> {
 		});
 		active = false;
 
-		if (!rollbackEdits) {
-			this.outerTree.branch.merge(queryTree.branch);
-			this.outerTreeIsDirty = false;
+		try {
+			if (!rollbackEdits && successfulEditCodes.length > 0) {
+				// Replay all successful edits on the query branch, then merge into the outer tree.
+				await snapshotTree.branch.runTransactionAsync(async () => {
+					for (const code of successfulEditCodes) {
+						await this.editor(snapshotTree.viewOrTree, code);
+					}
+				});
+				this.outerTree.branch.merge(snapshotTree.branch, false);
+				this.outerTreeIsDirty = false;
+			}
+		} finally {
+			snapshotTree.branch.dispose();
 		}
 		this.options?.logger?.log(`## Response\n\n`);
 		this.options?.logger?.log(`${responseMessage}\n\n`);
@@ -206,41 +254,6 @@ function constructTreeNode(schema: TreeNodeSchema, content: FactoryContentObject
 	return TreeAlpha.tagContentSchema(schema, toInsert as never);
 }
 
-/**
- * Applies the given function (as a string of JavaScript code or an actual function) to the given tree.
- */
-async function applyTreeFunction<TSchema extends ImplicitFieldSchema>(
-	tree: Subtree<TSchema>,
-	editCode: string,
-	editor: SynchronousEditor<TSchema> | AsynchronousEditor<TSchema>,
-	logger: Logger | undefined,
-): Promise<EditResult> {
-	logger?.log(`### Editing Tool Invoked\n\n`);
-	logger?.log(`#### Generated Code\n\n\`\`\`javascript\n${editCode}\n\`\`\`\n\n`);
-
-	// Fork a branch to edit. If the edit fails or produces an error, we discard this branch, otherwise we merge it.
-	const editTree = tree.fork();
-	try {
-		await editor(editTree.viewOrTree, editCode);
-	} catch (error: unknown) {
-		logger?.log(`#### Error\n\n`);
-		logger?.log(`\`\`\`JSON\n${toErrorString(error)}\n\`\`\`\n\n`);
-		editTree.branch.dispose();
-		return {
-			type: "editingError",
-			message: `Running the generated code produced an error. The state of the tree will be reset to its previous state as it was before the code ran. Please try again. Here is the error: ${toErrorString(error)}`,
-		};
-	}
-
-	tree.branch.merge(editTree.branch);
-	logger?.log(`#### New Tree State\n\n`);
-	logger?.log(`${`\`\`\`JSON\n${stringifyTree(tree.field)}\n\`\`\``}\n\n`);
-	return {
-		type: "success",
-		message: `After running the code, the new state of the tree is:\n\n\`\`\`JSON\n${stringifyTree(tree.field)}\n\`\`\``,
-	};
-}
-
 function createDefaultEditor<
 	TSchema extends ImplicitFieldSchema = ImplicitFieldSchema,
 >(): AsynchronousEditor<TSchema> {

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

@@ -5,18 +5,19 @@
 
 import { fail } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
-import { TreeNode, Tree, NodeKind } from "@fluidframework/tree";
+import { TreeNode, Tree, NodeKind, TreeViewConfiguration } from "@fluidframework/tree";
 import type {
 	ImplicitFieldSchema,
 	TreeFieldFromImplicitField,
 	TreeMapNode,
 	TreeArrayNode,
 } from "@fluidframework/tree";
-import { TreeAlpha } from "@fluidframework/tree/alpha";
+import { independentView, TreeAlpha } from "@fluidframework/tree/alpha";
 import type {
 	ReadableField,
 	TreeRecordNode,
 	TreeBranchAlpha,
+	InsertableField,
 } from "@fluidframework/tree/alpha";
 
 import type { TreeView, ViewOrTree } from "./api.js";
@@ -120,3 +121,20 @@ export class Subtree<TRoot extends ImplicitFieldSchema> {
 		}
 	}
 }
+
+/**
+ * Initializes a {@link Subtree} backed by an independent view.
+ */
+export function copyIndependentSubtree<TSchema extends ImplicitFieldSchema>(
+	view: ViewOrTree<TSchema>,
+): Subtree<TSchema> {
+	const subtree = new Subtree(view);
+	const config = new TreeViewConfiguration({ schema: subtree.schema });
+	const sandboxView = independentView(config);
+	const exported =
+		subtree.field === undefined ? undefined : TreeAlpha.exportVerbose(subtree.field);
+
+	const imported = TreeAlpha.importVerbose(sandboxView.schema, exported);
+	sandboxView.initialize(imported as InsertableField<TSchema>);
+	return new Subtree(sandboxView);
+}

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

@@ -248,6 +248,117 @@ describe("Semantic Agent", () => {
 		assert.equal(view.root, "Initial", "Tree should not have changed");
 	});
 
+	it("recovers from errors by replaying successful edits", async () => {
+		const view = independentView(new TreeViewConfiguration({ schema: sf.string }));
+		view.initialize("Initial");
+		const model: SharedTreeChatModel = {
+			editToolName,
+			async query({ edit }) {
+				// Edit A succeeds
+				const resultA = await edit(`context.root = "A";`);
+				assert.equal(resultA.type, "success", resultA.message);
+				// Edit B fails
+				const resultB = await edit(`throw new Error("boom");`);
+				assert.equal(resultB.type, "editingError", resultB.message);
+				// Edit C succeeds (sandbox should have been rebuilt with A replayed)
+				const resultC = await edit(`context.root = context.root + "+C";`);
+				assert.equal(resultC.type, "success", resultC.message);
+				return resultC.message;
+			},
+		};
+		const agent = new SharedTreeSemanticAgent(model, view);
+		await agent.query("Query");
+		assert.equal(view.root, "A+C");
+	});
+
+	it("recovers from multiple consecutive errors", async () => {
+		const view = independentView(new TreeViewConfiguration({ schema: sf.string }));
+		view.initialize("Initial");
+		const model: SharedTreeChatModel = {
+			editToolName,
+			async query({ edit }) {
+				// Edit A succeeds
+				const resultA = await edit(`context.root = "A";`);
+				assert.equal(resultA.type, "success", resultA.message);
+				// Edit B fails
+				const resultB = await edit(`throw new Error("boom1");`);
+				assert.equal(resultB.type, "editingError", resultB.message);
+				// Edit C also fails
+				const resultC = await edit(`throw new Error("boom2");`);
+				assert.equal(resultC.type, "editingError", resultC.message);
+				// Edit D succeeds (sandbox rebuilt twice, A replayed each time)
+				const resultD = await edit(`context.root = context.root + "+D";`);
+				assert.equal(resultD.type, "success", resultD.message);
+				return resultD.message;
+			},
+		};
+		const agent = new SharedTreeSemanticAgent(model, view);
+		await agent.query("Query");
+		assert.equal(view.root, "A+D");
+	});
+
+	it("rolls back entirely if edit replay fails during recovery", async () => {
+		const view = independentView(new TreeViewConfiguration({ schema: sf.string }));
+		view.initialize("Initial");
+		let editCallCount = 0;
+		const agent = new SharedTreeSemanticAgent(
+			{
+				editToolName,
+				async query({ edit }) {
+					// Edit A succeeds
+					const resultA = await edit(`context.root = "A";`);
+					assert.equal(resultA.type, "success", resultA.message);
+					// Edit B fails, triggering replay of A
+					const resultB = await edit(`throw new Error("boom");`);
+					assert.equal(resultB.type, "editingError", resultB.message);
+					return resultB.message;
+				},
+			} satisfies SharedTreeChatModel,
+			view,
+			{
+				editor: async (tree, code) => {
+					editCallCount++;
+					const context = createContext(tree);
+					if (editCallCount === 3) {
+						// Third call is the replay of A during recovery — make it fail
+						throw new Error("replay failed");
+					}
+					// eslint-disable-next-line no-new-func, @typescript-eslint/no-implied-eval
+					const fn = new Function("context", code);
+					// eslint-disable-next-line @typescript-eslint/no-unsafe-call
+					await fn(context);
+				},
+			},
+		);
+		await agent.query("Query");
+		assert.equal(view.root, "Initial", "Tree should not have changed after replay failure");
+	});
+
+	it("groups multiple edits into a single transaction on the outer branch", async () => {
+		const view = independentView(new TreeViewConfiguration({ schema: sf.string }));
+		view.initialize("Initial");
+		let changeCount = 0;
+		view.events.on("changed", () => {
+			changeCount++;
+		});
+		const model: SharedTreeChatModel = {
+			editToolName,
+			async query({ edit }) {
+				const result1 = await edit(`context.root = "First";`);
+				assert.equal(result1.type, "success", result1.message);
+				const result2 = await edit(`context.root = "Second";`);
+				assert.equal(result2.type, "success", result2.message);
+				const result3 = await edit(`context.root = "Third";`);
+				assert.equal(result3.type, "success", result3.message);
+				return result3.message;
+			},
+		};
+		const agent = new SharedTreeSemanticAgent(model, view);
+		await agent.query("Query");
+		assert.equal(view.root, "Third");
+		assert.equal(changeCount, 1, "Expected exactly one change event on the outer branch");
+	});
+
 	it("supplies the system prompt as context", async () => {
 		const view = independentView(new TreeViewConfiguration({ schema: sf.string }));
 		view.initialize("X");