diff --git a/ts/packages/agents/markdown/src/agent/documentUpdatePersistence.ts b/ts/packages/agents/markdown/src/agent/documentUpdatePersistence.ts new file mode 100644 index 0000000000..83806e9914 --- /dev/null +++ b/ts/packages/agents/markdown/src/agent/documentUpdatePersistence.ts @@ -0,0 +1,123 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import fs from "node:fs"; +import path from "node:path"; +import { createHash } from "node:crypto"; +import { applyDocumentOperations } from "./documentOperations.js"; +import type { DocumentOperation } from "./markdownOperationSchema.js"; +import { + resolveRealDirectory, + resolveWritableFileWithinRoot, +} from "./pathPolicy.js"; + +export interface DocumentBinding { + token: string | undefined; + root: string; + relativePath: string; + filePath: string; +} + +export interface UpdateExpectations { + bindingToken: string | undefined; + root: string | undefined; + relativePath: string | undefined; + revision: string; + updatedRevision: string | undefined; +} + +export function computeContentRevision(content: string): string { + return createHash("sha256").update(content, "utf8").digest("hex"); +} + +function validateIdentity( + binding: DocumentBinding, + expected: UpdateExpectations, +): void { + if ( + expected.bindingToken !== undefined && + expected.bindingToken !== binding.token + ) { + throw new Error("Document binding token changed"); + } + if (expected.root !== undefined && expected.root !== binding.root) { + throw new Error("Document binding root changed"); + } + if ( + expected.relativePath !== undefined && + expected.relativePath !== binding.relativePath + ) { + throw new Error("Document binding path changed"); + } +} + +function resolveBoundFile(binding: DocumentBinding): string { + if (resolveRealDirectory(binding.root) !== binding.root) { + throw new Error("The authorized markdown workspace root changed"); + } + const resolved = resolveWritableFileWithinRoot( + binding.root, + binding.relativePath, + ); + if ( + resolved === undefined || + path.relative(resolved, binding.filePath) !== "" + ) { + throw new Error( + "The markdown document binding changed or is outside its authorized workspace", + ); + } + return resolved; +} + +export function readBoundDocument(binding: DocumentBinding) { + const filePath = resolveBoundFile(binding); + const content = fs.readFileSync(filePath, "utf-8"); + return { content, revision: computeContentRevision(content), filePath }; +} + +export function persistDocumentOperations( + binding: DocumentBinding, + operations: DocumentOperation[], + expected: UpdateExpectations, +) { + validateIdentity(binding, expected); + let filePath = resolveBoundFile(binding); + const currentContent = fs.readFileSync(filePath, "utf-8"); + const currentRevision = computeContentRevision(currentContent); + if (expected.updatedRevision === currentRevision) { + return { + content: currentContent, + revision: currentRevision, + alreadyApplied: true, + filePath, + }; + } + if (currentRevision !== expected.revision) { + throw new Error( + "Document changed between read and apply (revision mismatch)", + ); + } + + const content = applyDocumentOperations(currentContent, operations); + const revision = computeContentRevision(content); + if ( + expected.updatedRevision !== undefined && + expected.updatedRevision !== revision + ) { + throw new Error("Updated document revision does not match operations"); + } + + validateIdentity(binding, expected); + filePath = resolveBoundFile(binding); + if ( + computeContentRevision(fs.readFileSync(filePath, "utf-8")) !== + currentRevision + ) { + throw new Error( + "Document changed between validation and write (revision mismatch)", + ); + } + fs.writeFileSync(filePath, content, "utf-8"); + return { content, revision, alreadyApplied: false, filePath }; +} diff --git a/ts/packages/agents/markdown/src/agent/ipcTypes.ts b/ts/packages/agents/markdown/src/agent/ipcTypes.ts index 34f538ce8b..fd3778960a 100644 --- a/ts/packages/agents/markdown/src/agent/ipcTypes.ts +++ b/ts/packages/agents/markdown/src/agent/ipcTypes.ts @@ -36,28 +36,50 @@ export interface UICommandResult { // Agent → View: Content requests export interface GetDocumentContentMessage { type: "getDocumentContent"; + requestId: string; + expectedBindingToken?: string; + expectedRoot?: string; + expectedRelativePath?: string; } export interface DocumentContentMessage { type: "documentContent"; + requestId: string; content: string; - source?: "client-serializer" | "yjs-fallback" | "error"; + source?: "file" | "error"; error?: string; timestamp: number; + bindingToken: string | null; + boundFilePath: string | null; + boundRoot: string | null; + boundRelativePath: string | null; + revision: string | null; + identityMismatch?: boolean; } // Agent → View: LLM operations export interface LLMOperationsMessage { type: "applyLLMOperations"; + requestId: string; operations: any[]; // DocumentOperation[] timestamp: number; + expectedBindingToken?: string; + expectedRoot?: string; + expectedRelativePath?: string; + expectedRevision: string; + expectedUpdatedRevision?: string; } export interface OperationsAppliedMessage { type: "operationsApplied"; + requestId: string; success: boolean; operationCount?: number; error?: string; + identityMismatch?: boolean; + revisionMismatch?: boolean; + bindingToken?: string | null; + revision?: string | null; } // View → Frontend: Auto-save notifications diff --git a/ts/packages/agents/markdown/src/agent/markdownActionHandler.ts b/ts/packages/agents/markdown/src/agent/markdownActionHandler.ts index 21eab3d795..16855d06e1 100644 --- a/ts/packages/agents/markdown/src/agent/markdownActionHandler.ts +++ b/ts/packages/agents/markdown/src/agent/markdownActionHandler.ts @@ -28,6 +28,13 @@ import { resolveRealDirectory, resolveWritableFileWithinRoot, } from "./pathPolicy.js"; +import { + computeContentRevision, + persistDocumentOperations, + readBoundDocument, + type DocumentBinding, +} from "./documentUpdatePersistence.js"; +import { applyDocumentOperations } from "./documentOperations.js"; const debug = registerDebug("typeagent:markdown:agent"); @@ -56,6 +63,7 @@ type MarkdownActionContext = { currentFileName?: string | undefined; currentFilePath?: string | undefined; currentWorkspaceRoot?: string | undefined; + currentBindingToken?: string | undefined; viewProcess?: ChildProcess | undefined; localHostPort: number; // Handle returned by sessionContext.registerPort for the markdown @@ -314,6 +322,11 @@ async function updateMarkdownContext( if (!context.agentContext.viewProcess) { const fullPath = await getFullMarkdownFilePath(fileName, storage!); if (fullPath) { + const root = fs.realpathSync(path.dirname(fullPath)); + context.agentContext.currentWorkspaceRoot = root; + context.agentContext.currentFileName = path.basename(fullPath); + context.agentContext.currentFilePath = + fs.realpathSync(fullPath); process.env.MARKDOWN_FILE = fullPath; // Fork the express view service in the background instead of // blocking agent enable (and therefore agent-server startup) @@ -336,6 +349,19 @@ async function updateMarkdownContext( context.agentContext.viewPortRegistration?.release(); context.agentContext.viewPortRegistration = context.registerPort("view", result.port); + if ( + context.agentContext.currentWorkspaceRoot && + context.agentContext.currentFileName && + context.agentContext.currentFilePath + ) { + viewProcess.send({ + type: "setFile", + workspaceRoot: + context.agentContext.currentWorkspaceRoot, + relativePath: + context.agentContext.currentFileName, + }); + } // Defensive cleanup if the child crashes mid-session. // The identity guard prevents a late-firing `exit` // event on a previously-replaced process from @@ -391,34 +417,11 @@ async function handleStreamingMarkdownAction( ); const agent = await createMarkdownAgent("GPT_4o"); - const storage = actionContext.sessionContext.sessionStorage; - - // Get current document content - const filePath = `${actionContext.sessionContext.agentContext.currentFileName}`; - let markdownContent = ""; - - if (actionContext.sessionContext.agentContext.viewProcess) { - try { - markdownContent = await getDocumentContentFromView( - actionContext.sessionContext.agentContext.viewProcess, - ); - debug( - `Got content from view process for streaming: ${markdownContent?.length || 0} chars`, - ); - } catch (error) { - console.warn( - "[STREAMING] Failed to get content from view, falling back to storage:", - error, - ); - if (await storage?.exists(filePath)) { - markdownContent = (await storage?.read(filePath, "utf8")) || ""; - } - } - } else { - if (await storage?.exists(filePath)) { - markdownContent = (await storage?.read(filePath, "utf8")) || ""; - } - } + const { + content: markdownContent, + binding, + revision, + } = await readCurrentDocumentContent(actionContext); try { // Call agent with streaming callback @@ -467,6 +470,21 @@ async function handleStreamingMarkdownAction( updateResult.operations || [], actionContext, ); + if (updateResult.operations?.length) { + const operations = + updateResult.operations as DocumentOperation[]; + const updatedContent = applyDocumentOperations( + markdownContent, + operations, + ); + await applyOperationsForCurrentDocument( + actionContext, + operations, + binding, + revision, + computeContentRevision(updatedContent), + ); + } return createActionResult( updateResult.operationSummary || @@ -546,7 +564,9 @@ async function getFullMarkdownFilePath(fileName: string, storage: Storage) { } async function handleCreateDocument( - action: CreateDocumentAction, + action: + | CreateDocumentAction + | Extract, actionContext: ActionContext, ): Promise { const rawName = action.parameters.name; @@ -582,7 +602,10 @@ async function handleCreateDocument( ); } - const initialContent = action.parameters.content ?? ""; + const initialContent = + action.actionName === "createDocument" + ? (action.parameters.content ?? "") + : ""; const documentExisted = fs.existsSync(absoluteFilePath); if (!documentExisted) { fs.writeFileSync(absoluteFilePath, initialContent, { @@ -607,8 +630,8 @@ async function handleCreateDocument( if (agentContext.viewProcess) { agentContext.viewProcess.send({ type: "setFile", - filePath: path.basename(absoluteFilePath), - folderPath: path.dirname(absoluteFilePath), + workspaceRoot: canonicalRoot, + relativePath: relativeName, }); } @@ -631,6 +654,186 @@ async function handleCreateDocument( return result; } +type DocumentUpdateAction = Extract< + MarkdownAction, + { actionName: "updateDocument" | "streamingUpdateDocument" } +>; + +function getCurrentDocumentBinding( + agentContext: MarkdownActionContext, +): DocumentBinding { + const root = agentContext.currentWorkspaceRoot; + const relativePath = agentContext.currentFileName; + const filePath = agentContext.currentFilePath; + if (!root || !relativePath || !filePath) { + throw new Error( + "No markdown document is open. Use createDocument or openDocument first.", + ); + } + return { + token: agentContext.currentBindingToken, + root, + relativePath, + filePath, + }; +} + +async function readCurrentDocumentContent( + actionContext: ActionContext, +): Promise<{ + content: string; + binding: DocumentBinding; + revision: string; +}> { + const agentContext = actionContext.sessionContext.agentContext; + const binding = getCurrentDocumentBinding(agentContext); + if (!agentContext.viewProcess) { + const document = readBoundDocument(binding); + agentContext.currentFilePath = document.filePath; + return { + content: document.content, + binding, + revision: document.revision, + }; + } + + const response = await getDocumentContentFromView( + agentContext.viewProcess, + { + expectedBindingToken: binding.token, + expectedRoot: binding.root, + expectedRelativePath: binding.relativePath, + }, + ); + if (response.identityMismatch) { + throw new Error( + "Document identity changed while reading; refusing to update the wrong file", + ); + } + if (response.error) { + throw new Error(response.error); + } + if (typeof response.bindingToken === "string") { + agentContext.currentBindingToken = response.bindingToken; + } + return { + content: response.content, + binding: { + ...binding, + token: agentContext.currentBindingToken, + }, + revision: response.revision ?? computeContentRevision(response.content), + }; +} + +async function applyOperationsForCurrentDocument( + actionContext: ActionContext, + operations: DocumentOperation[], + binding: DocumentBinding, + revision: string, + expectedUpdatedRevision?: string, +): Promise { + const agentContext = actionContext.sessionContext.agentContext; + const currentBinding = getCurrentDocumentBinding(agentContext); + if ( + currentBinding.token !== binding.token || + currentBinding.root !== binding.root || + currentBinding.relativePath !== binding.relativePath || + currentBinding.filePath !== binding.filePath + ) { + throw new Error("Document binding changed while generating update"); + } + const expectations = { + expectedBindingToken: binding.token, + expectedRoot: binding.root, + expectedRelativePath: binding.relativePath, + expectedRevision: revision, + expectedUpdatedRevision, + }; + if (!agentContext.viewProcess) { + const persisted = persistDocumentOperations(binding, operations, { + bindingToken: binding.token, + root: binding.root, + relativePath: binding.relativePath, + revision, + updatedRevision: expectedUpdatedRevision, + }); + agentContext.currentFilePath = persisted.filePath; + return; + } + + const applied = await sendOperationsToView( + agentContext.viewProcess, + operations, + expectations, + ); + if (applied.success) { + return; + } + if (applied.identityMismatch) { + throw new Error( + "Document identity changed while applying operations; refusing to write to the wrong file", + ); + } + if (applied.revisionMismatch) { + throw new Error( + "Document changed between read and apply; refusing to overwrite (revision mismatch)", + ); + } + throw new Error( + applied.error ?? "Failed to apply operations in view process", + ); +} + +function parseEditorContext(serializedContext: string | undefined): unknown { + if (!serializedContext) { + return undefined; + } + try { + return JSON.parse(serializedContext); + } catch (error) { + debug( + `[AGENT] Failed to parse context JSON: ${ + error instanceof Error ? error.message : String(error) + }, using undefined`, + ); + return undefined; + } +} + +async function updateCurrentDocument( + action: DocumentUpdateAction, + actionContext: ActionContext, + agent: Awaited>, +): Promise { + const { content, binding, revision } = + await readCurrentDocumentContent(actionContext); + const response = await agent.updateDocument( + content, + action.parameters.originalRequest, + action.parameters.cursorPosition, + parseEditorContext(action.parameters.context), + ); + if (!response.success) { + const message = + (response as { message?: string }).message ?? + "Unknown error occurred"; + return createActionResult(`Failed to update document: ${message}`); + } + + if (response.data.operations?.length) { + await applyOperationsForCurrentDocument( + actionContext, + response.data.operations, + binding, + revision, + ); + } + return createActionResult( + response.data.operationSummary ?? "Updated document", + ); +} + async function handleMarkdownAction( action: MarkdownAction, actionContext: ActionContext, @@ -651,298 +854,16 @@ async function handleMarkdownAction( return agent; }; - const storage = actionContext.sessionContext.sessionStorage; - switch (action.actionName) { - case "createDocument": { - result = await handleCreateDocument(action, actionContext); - break; - } + case "createDocument": case "openDocument": { - if (!action.parameters.name) { - result = createActionResult( - "Document could not be created: no name was provided", - ); - } else { - result = createActionResult("Opening document ..."); - - let newFileName = action.parameters.name.trim(); - if (!newFileName.endsWith(".md")) { - newFileName += ".md"; - } - - actionContext.sessionContext.agentContext.currentFileName = - newFileName; - - if (!(await storage?.exists(newFileName))) { - await storage?.write(newFileName, ""); - } - - if (actionContext.sessionContext.agentContext.viewProcess) { - const fullPath = await getFullMarkdownFilePath( - newFileName, - storage!, - ); - - actionContext.sessionContext.agentContext.viewProcess.send({ - type: "setFile", - filePath: path.basename(fullPath!), - folderPath: path.dirname(fullPath!), - }); - } - result = createActionResult("Document opened"); - result.activityContext = { - activityName: "editingMarkdown", - description: "Editing a Markdown document", - state: { - fileName: newFileName, - }, - openLocalView: true, - }; - } - break; - } - case "updateDocument": { - const agent = await createAgent(); - debug("Starting updateDocument action in agent process"); - result = createActionResult("Updating document ..."); - - const filePath = `${actionContext.sessionContext.agentContext.currentFileName}`; - - let markdownContent = ""; - - if (actionContext.sessionContext.agentContext.viewProcess) { - try { - markdownContent = await getDocumentContentFromView( - actionContext.sessionContext.agentContext.viewProcess, - ); - debug( - `Got content from view process: ${markdownContent?.length || 0} chars`, - ); - debug( - `Content preview: ${markdownContent?.substring(0, 200)}...`, - ); - } catch (error) { - console.warn( - "Failed to get content from view, using empty content fallback:", - error, - ); - // Use empty content as fallback to allow agent to continue processing - markdownContent = ""; - debug("Using empty content fallback"); - } - } else { - // Fallback if no view process - if (await storage?.exists(filePath)) { - markdownContent = - (await storage?.read(filePath, "utf8")) || ""; - debug( - "No view process, read content from storage:", - markdownContent?.length, - "chars", - ); - } - } - - // Handle synchronous requests through the agent - const originalRequest = - "originalRequest" in action.parameters - ? action.parameters.originalRequest - : ""; - - const cursorPosition = - "cursorPosition" in action.parameters - ? action.parameters.cursorPosition - : undefined; - - const context = - "context" in action.parameters && action.parameters.context - ? (() => { - try { - return JSON.parse(action.parameters.context); - } catch (error) { - debug( - `[AGENT] Failed to parse context JSON: ${error}, using undefined`, - ); - return undefined; - } - })() - : undefined; - - debug( - `[AGENT] About to call LLM service with request: "${originalRequest}"`, - ); - debug( - `[AGENT] Document content length: ${markdownContent?.length || 0} chars`, - ); - - const response = await agent.updateDocument( - markdownContent, - originalRequest, - cursorPosition, - context, - ); - - debug(`[AGENT] LLM service returned, success: ${response.success}`); - - if (response.success) { - const updateResult = response.data; - debug( - `[AGENT] LLM processing successful, operations count: ${updateResult.operations?.length || 0}`, - ); - - // Apply operations to the document - if ( - updateResult.operations && - updateResult.operations.length > 0 - ) { - // Send operations to view process for application - if (actionContext.sessionContext.agentContext.viewProcess) { - debug( - "Agent sending operations to view process for Yjs application", - ); - - const success = await sendOperationsToView( - actionContext.sessionContext.agentContext - .viewProcess, - updateResult.operations, - ); - - if (!success) { - throw new Error( - "Failed to apply operations in view process", - ); - } - - debug( - "Operations applied successfully via view process", - ); - } else { - console.warn( - "No view process available, operations not applied", - ); - } - } else { - debug("[AGENT] No operations returned from LLM"); - } - - if (updateResult.operationSummary) { - result = createActionResult(updateResult.operationSummary); - } else { - result = createActionResult("Updated document"); - } - - debug(`[AGENT] updateDocument case completed successfully`); - } else { - const errorMessage = - (response as any).message || "Unknown error occurred"; - console.error("Translation failed:", errorMessage); - result = createActionResult( - "Failed to update document: " + errorMessage, - ); - } + result = await handleCreateDocument(action, actionContext); break; } + case "updateDocument": case "streamingUpdateDocument": { const agent = await createAgent(); - // Handle streaming AI commands - now unified with regular updateDocument flow - debug( - "Starting streamingUpdateDocument action - using standard translator flow", - ); - result = createActionResult("Updating document ..."); - - const filePath = `${actionContext.sessionContext.agentContext.currentFileName}`; - - let markdownContent = ""; - - if (actionContext.sessionContext.agentContext.viewProcess) { - try { - markdownContent = await getDocumentContentFromView( - actionContext.sessionContext.agentContext.viewProcess, - ); - debug( - `Got content from view process: ${markdownContent?.length || 0} chars`, - ); - debug( - `Content preview: ${markdownContent?.substring(0, 200)}...`, - ); - } catch (error) { - console.warn( - "Failed to get content from view, using empty content fallback:", - error, - ); - // Use empty content as fallback to allow agent to continue processing - markdownContent = ""; - debug("Using empty content fallback"); - } - } else { - // Fallback if no view process - if (await storage?.exists(filePath)) { - markdownContent = - (await storage?.read(filePath, "utf8")) || ""; - debug( - "No view process, read content from storage:", - markdownContent?.length, - "chars", - ); - } - } - - // Handle streaming requests through the standard agent (same as updateDocument) - const response = await agent.updateDocument( - markdownContent, - action.parameters.originalRequest, - ); - - if (response.success) { - const updateResult = response.data; - - // Apply operations to the document - if ( - updateResult.operations && - updateResult.operations.length > 0 - ) { - // Send operations to view process for application - if (actionContext.sessionContext.agentContext.viewProcess) { - debug( - "Agent sending operations to view process for Yjs application", - ); - - const success = await sendOperationsToView( - actionContext.sessionContext.agentContext - .viewProcess, - updateResult.operations, - ); - - if (!success) { - throw new Error( - "Failed to apply operations in view process", - ); - } - - debug( - "Operations applied successfully via view process", - ); - } else { - console.warn( - "No view process available, operations not applied", - ); - } - } - - if (updateResult.operationSummary) { - result = createActionResult(updateResult.operationSummary); - } else { - result = createActionResult("Updated document"); - } - } else { - const errorMessage = - (response as any).message || "Unknown error occurred"; - console.error("Translation failed:", errorMessage); - result = createActionResult( - "Failed to update document: " + errorMessage, - ); - } + result = await updateCurrentDocument(action, actionContext, agent); break; } } @@ -956,104 +877,144 @@ async function handleMarkdownAction( return result; } -/** - * Send operations to view process for application (Flow 1 implementation) - */ -async function sendOperationsToView( +let applyRequestCounter = 0; + +type ApplyExpectations = { + expectedBindingToken: string | undefined; + expectedRoot: string; + expectedRelativePath: string; + expectedRevision: string; + expectedUpdatedRevision: string | undefined; +}; + +type ApplyResult = { + success: boolean; + identityMismatch: boolean; + revisionMismatch: boolean; + error: string | undefined; +}; + +export async function sendOperationsToView( viewProcess: ChildProcess | undefined, operations: DocumentOperation[], -): Promise { + expectations: ApplyExpectations, +): Promise { if (!viewProcess) { - return false; + return { + success: false, + identityMismatch: false, + revisionMismatch: false, + error: "No view process", + }; } + const requestId = `apply_${++applyRequestCounter}`; return new Promise((resolve) => { const timeout = setTimeout(() => { console.error("[AGENT] View process operation timeout"); - resolve(false); - }, 5000); - - // Listen for response - const responseHandler = (message: any) => { - if (message.type === "operationsApplied") { - clearTimeout(timeout); - viewProcess.off("message", responseHandler); - - if (message.success) { - resolve(true); - } else { - console.error( - "[AGENT] View failed to apply operations:", - message.error, - ); - resolve(false); - } + viewProcess.off("message", responseHandler); + resolve({ + success: false, + identityMismatch: false, + revisionMismatch: false, + error: "View process operation timeout", + }); + }, 15000); + + const responseHandler = (message: Record) => { + if ( + message.type !== "operationsApplied" || + message.requestId !== requestId + ) { + return; } + clearTimeout(timeout); + viewProcess.off("message", responseHandler); + resolve({ + success: message.success === true, + identityMismatch: message.identityMismatch === true, + revisionMismatch: message.revisionMismatch === true, + error: + typeof message.error === "string" + ? message.error + : undefined, + }); }; viewProcess.on("message", responseHandler); - - // Send operations viewProcess.send({ type: "applyLLMOperations", - operations: operations, + requestId, + operations, timestamp: Date.now(), + ...expectations, }); - - debug(`[AGENT] Sent ${operations.length} operations to view process`); }); } -/** - * Get document content from view process (Flow 1 implementation) - */ -async function getDocumentContentFromView( - viewProcess: ChildProcess, -): Promise { - return new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - debug( - "[AGENT] Content request timeout, trying fallback to empty content", - ); - - // Use empty content as fallback when view process fails - // This allows the agent to continue processing even if content retrieval fails - console.warn( - "[AGENT] View process content request timed out, using empty content fallback", - ); - resolve(""); - }, 15000); // 15 second timeout - - const responseHandler = (message: any) => { - if (message.type === "documentContent") { - clearTimeout(timeout); - viewProcess.off("message", responseHandler); +type ViewDocumentContentResponse = { + content: string; + bindingToken: string | null; + revision: string | null; + identityMismatch: boolean; + error: string | undefined; +}; - // Log the source of the content for debugging - const source = message.source || "unknown"; - debug( - `[AGENT] Received document content from ${source}: ${message.content?.length || 0} chars`, - ); +type ReadExpectations = { + expectedBindingToken: string | undefined; + expectedRoot: string | undefined; + expectedRelativePath: string | undefined; +}; - if (message.error) { - debug( - `[AGENT] Content retrieval had error: ${message.error}`, - ); - // Still resolve with content even if there was an error - } +let getContentRequestCounter = 0; - resolve(message.content || ""); +export async function getDocumentContentFromView( + viewProcess: ChildProcess, + expectations: ReadExpectations, +): Promise { + const requestId = `get_${++getContentRequestCounter}`; + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + viewProcess.off("message", responseHandler); + reject(new Error("View process content request timed out")); + }, 15000); + + const responseHandler = (message: Record) => { + if ( + message.type !== "documentContent" || + message.requestId !== requestId + ) { + return; } + clearTimeout(timeout); + viewProcess.off("message", responseHandler); + resolve({ + content: + typeof message.content === "string" ? message.content : "", + bindingToken: + typeof message.bindingToken === "string" + ? message.bindingToken + : null, + revision: + typeof message.revision === "string" + ? message.revision + : null, + identityMismatch: message.identityMismatch === true, + error: + typeof message.error === "string" + ? message.error + : undefined, + }); }; viewProcess.on("message", responseHandler); - - debug("[AGENT] Sending getDocumentContent request to view process"); - viewProcess.send({ type: "getDocumentContent" }); + viewProcess.send({ + type: "getDocumentContent", + requestId, + ...expectations, + }); }); } -// NOTE: Function commented out per Flow 1 consolidation -// Collaboration server now managed by view process async function createViewServiceHost( filePath: string, @@ -1090,7 +1051,8 @@ async function createViewServiceHost( childProcess.send({ type: "setFile", - filePath: path.basename(filePath), + workspaceRoot: folderPath, + relativePath: path.basename(filePath), }); childProcess.on("message", function (message: any) { @@ -1118,6 +1080,7 @@ async function createViewServiceHost( // Global process message handler for UI commands let currentAgentContext: MarkdownActionContext | null = null; +const wiredViewProcesses = new WeakSet(); // Store agent context for UI command processing export function setCurrentAgentContext(context: MarkdownActionContext) { @@ -1125,9 +1088,28 @@ export function setCurrentAgentContext(context: MarkdownActionContext) { const viewProcess = context.viewProcess; - if (typeof viewProcess !== "undefined" && viewProcess.on) { + if ( + typeof viewProcess !== "undefined" && + viewProcess.on && + !wiredViewProcesses.has(viewProcess) + ) { + wiredViewProcesses.add(viewProcess); viewProcess.on("message", async (message: any) => { - if (message.type === "uiCommand" && currentAgentContext) { + if (message.type === "bindingUpdated" && currentAgentContext) { + if ( + message.boundRoot === + currentAgentContext.currentWorkspaceRoot && + message.boundRelativePath === + currentAgentContext.currentFileName && + message.boundFilePath === + currentAgentContext.currentFilePath + ) { + currentAgentContext.currentBindingToken = + typeof message.bindingToken === "string" + ? message.bindingToken + : undefined; + } + } else if (message.type === "uiCommand" && currentAgentContext) { debug( `[AGENT] Received UI command: ${message.command}, requestId: ${message.requestId}, cursorPosition: ${message.parameters?.cursorPosition}, context: ${message.parameters?.context ? "serialized" : "none"}`, ); diff --git a/ts/packages/agents/markdown/src/view/route/service.ts b/ts/packages/agents/markdown/src/view/route/service.ts index 8e8f518e8c..88d3c59492 100644 --- a/ts/packages/agents/markdown/src/view/route/service.ts +++ b/ts/packages/agents/markdown/src/view/route/service.ts @@ -18,12 +18,20 @@ import * as encoding from "lib0/encoding"; import * as decoding from "lib0/decoding"; import registerDebug from "debug"; import sanitizeFilename from "sanitize-filename"; +import { randomUUID } from "node:crypto"; import { isAllowedViewOrigin } from "./originAllowlist.js"; +import { resolvePathWithinRoot } from "./pathPolicy.js"; import { - resolveExistingFileWithinRoot, - resolvePathWithinRoot, + normalizeRelativeDocumentPath, + resolveRealDirectory, resolveWritableFileWithinRoot, -} from "./pathPolicy.js"; +} from "../../agent/pathPolicy.js"; +import { + persistDocumentOperations, + readBoundDocument, + type DocumentBinding, +} from "../../agent/documentUpdatePersistence.js"; +import type { DocumentOperation } from "../../agent/markdownOperationSchema.js"; const debug = registerDebug("typeagent:markdown:service"); @@ -107,7 +115,7 @@ app.post( // Construct and normalize file path const documentPath = resolvePathWithinRoot( - ROOT_DIR, + getValidatedCurrentRoot(), `${sanitizedDocumentName}.md`, ); @@ -121,7 +129,7 @@ app.post( } let safeDocumentPath = resolveWritableFileWithinRoot( - ROOT_DIR, + getValidatedCurrentRoot(), documentPath, ); if (safeDocumentPath === undefined) { @@ -142,6 +150,9 @@ app.post( } filePath = safeDocumentPath; + boundRelativePath = `${sanitizedDocumentName}.md`; + bindingToken = randomUUID(); + notifyBindingToParent(); // Initialize collaboration for new document const documentId = sanitizedDocumentName; @@ -231,7 +242,9 @@ app.post( ); let clients: any[] = []; -let filePath: string | null; +let filePath: string | null = null; +let boundRelativePath: string | null = null; +let bindingToken: string | null = null; let collaborationManager: CollaborationManager; // UI Command routing state @@ -242,8 +255,71 @@ const pendingCommands = new Map(); let markdownRequestCounter = 0; const pendingMarkdownRequests = new Map(); const userHomeDir = os.homedir(); -const ROOT_DIR = +const INITIAL_ROOT_DIR = process.env.TYPEAGENT_MARKDOWN_ROOT || path.join(userHomeDir, "Documents"); +let currentRoot = + resolveRealDirectory(INITIAL_ROOT_DIR) ?? path.resolve(INITIAL_ROOT_DIR); + +function resolveCanonicalRoot(root: string): string | undefined { + const canonicalRoot = resolveRealDirectory(root); + return canonicalRoot !== undefined && + path.relative(path.resolve(root), canonicalRoot) === "" + ? canonicalRoot + : undefined; +} + +function getValidatedCurrentRoot(): string { + if (resolveCanonicalRoot(currentRoot) === undefined) { + throw new Error("The document root is no longer accessible"); + } + return currentRoot; +} + +type BindingSnapshot = { + bindingToken: string | null; + currentRoot: string; + filePath: string | null; + boundRelativePath: string | null; +}; + +function captureBindingSnapshot(): BindingSnapshot { + return { bindingToken, currentRoot, filePath, boundRelativePath }; +} + +function bindingError( + message: Record, + snapshot: BindingSnapshot, +): string | undefined { + if ( + typeof message.expectedBindingToken === "string" && + message.expectedBindingToken !== snapshot.bindingToken + ) { + return "Document binding token changed"; + } + if ( + typeof message.expectedRoot === "string" && + message.expectedRoot !== snapshot.currentRoot + ) { + return "Document binding root changed"; + } + if ( + typeof message.expectedRelativePath === "string" && + message.expectedRelativePath !== snapshot.boundRelativePath + ) { + return "Document binding path changed"; + } + return undefined; +} + +function notifyBindingToParent(): void { + process.send?.({ + type: "bindingUpdated", + bindingToken, + boundFilePath: filePath, + boundRoot: filePath ? currentRoot : null, + boundRelativePath, + }); +} // Streaming state for LLM responses const activeStreamingSessions = new Map< @@ -440,6 +516,8 @@ async function requestMarkdownFromClient(retryCount: number = 0): Promise<{ }); } +void requestMarkdownFromClient; + /** * Determine if a command should use streaming */ @@ -579,7 +657,7 @@ app.post("/document", express.json(), (req: Request, res: Response) => { try { const writableFilePath = resolveWritableFileWithinRoot( - ROOT_DIR, + getValidatedCurrentRoot(), filePath, ); if (writableFilePath === undefined) { @@ -691,7 +769,7 @@ app.post("/autosave", express.json(), (req: Request, res: Response) => { } const resolvedFilePath = resolvePathWithinRoot( - ROOT_DIR, + getValidatedCurrentRoot(), sanitizedFilePath, ); if (resolvedFilePath === undefined) { @@ -849,11 +927,11 @@ app.post("/file/load", express.json(), (req: Request, res: Response) => { return; } - const resolvedPath = resolveExistingFileWithinRoot( - ROOT_DIR, + const resolvedPath = resolveWritableFileWithinRoot( + getValidatedCurrentRoot(), newFilePath, ); - if (resolvedPath === undefined) { + if (resolvedPath === undefined || !fs.existsSync(resolvedPath)) { res.status(403).json({ error: "Access to the file is forbidden or file not found", }); @@ -862,6 +940,12 @@ app.post("/file/load", express.json(), (req: Request, res: Response) => { // Set new file path filePath = resolvedPath; + boundRelativePath = path + .relative(currentRoot, resolvedPath) + .split(path.sep) + .join("/"); + bindingToken = randomUUID(); + notifyBindingToParent(); // Initialize collaboration for new document const documentId = path.basename(resolvedPath, ".md"); @@ -1501,11 +1585,22 @@ process.on("message", async (message: any) => { ); if (message.type == "setFile") { - if (message.filePath) { - // Resolve and validate the file path + if (message.relativePath) { + const nextRoot = + typeof message.workspaceRoot === "string" && + resolveCanonicalRoot(message.workspaceRoot) !== undefined + ? resolveCanonicalRoot(message.workspaceRoot) + : undefined; + const relativePath = normalizeRelativeDocumentPath( + message.relativePath, + ); + if (nextRoot === undefined || relativePath === undefined) { + debug("Invalid document binding provided in message"); + return; + } const resolvedFilePath = resolveWritableFileWithinRoot( - ROOT_DIR, - path.basename(message.filePath), + nextRoot, + relativePath, ); if (resolvedFilePath === undefined) { debug("Invalid file path provided in message"); @@ -1513,10 +1608,14 @@ process.on("message", async (message: any) => { } const oldFilePath = filePath; + currentRoot = nextRoot; filePath = resolvedFilePath; + boundRelativePath = relativePath; + bindingToken = randomUUID(); + notifyBindingToParent(); // Initialize collaboration for this document using authoritative document - const documentId = path.basename(message.filePath, ".md"); + const documentId = path.basename(relativePath, ".md"); // Get or create the authoritative Y.js document const ydoc = getAuthoritativeDocument(documentId); @@ -1531,7 +1630,7 @@ process.on("message", async (message: any) => { ytext.insert(0, content); // Insert file content debug( - `File loaded into authoritative document: ${documentId}, ${content.length} chars from ${message.filePath}`, + `File loaded into authoritative document: ${documentId}, ${content.length} chars from ${relativePath}`, ); } else { debug( @@ -1547,10 +1646,8 @@ process.on("message", async (message: any) => { `data: ${JSON.stringify({ type: "documentChanged", newDocumentId: documentId, - newDocumentName: path.basename( - message.filePath, - ".md", - ), + newDocumentName: path.basename(relativePath, ".md"), + bindingToken, timestamp: Date.now(), })}\n\n`, ); @@ -1559,6 +1656,9 @@ process.on("message", async (message: any) => { } else { // No file mode - initialize with default content using authoritative document filePath = null; + boundRelativePath = null; + bindingToken = null; + notifyBindingToParent(); debug("Running in memory-only mode (no file)"); const documentId = "default"; @@ -1635,198 +1735,149 @@ Start typing to see the editor in action! ); }); } else if (message.type === "applyLLMOperations") { - // PRODUCTION: Send operations to PRIMARY client only via SSE to prevent duplicates + const requestId = + typeof message.requestId === "string" ? message.requestId : ""; + const snapshot = captureBindingSnapshot(); try { - debug( - `[VIEW] Forwarding ${message.operations?.length || 0} operations to primary client via SSE`, - ); - - if (clients.length === 0) { - console.warn( - `[SSE] No clients connected to receive operations`, - ); + if ( + !Array.isArray(message.operations) || + !snapshot.filePath || + !snapshot.boundRelativePath || + typeof message.expectedRevision !== "string" + ) { + throw new Error("Invalid document update request"); + } + const identityError = bindingError(message, snapshot); + if (identityError) { process.send?.({ type: "operationsApplied", + requestId, success: false, - error: "No clients connected", - method: "sse-forwarded", + identityMismatch: true, + error: identityError, + bindingToken: snapshot.bindingToken, }); return; } - // Send operations to ONLY the first client to prevent duplicates - const primaryClient = clients[0]; - const operationsEvent = { - type: "llmOperations", - operations: message.operations, - timestamp: message.timestamp || Date.now(), - source: "agent", - clientRole: "primary", // Mark this client as the primary applier + const binding: DocumentBinding = { + token: snapshot.bindingToken ?? undefined, + root: snapshot.currentRoot, + relativePath: snapshot.boundRelativePath, + filePath: snapshot.filePath, }; + const persisted = persistDocumentOperations( + binding, + message.operations as DocumentOperation[], + { + bindingToken: + typeof message.expectedBindingToken === "string" + ? message.expectedBindingToken + : undefined, + root: + typeof message.expectedRoot === "string" + ? message.expectedRoot + : undefined, + relativePath: + typeof message.expectedRelativePath === "string" + ? message.expectedRelativePath + : undefined, + revision: message.expectedRevision, + updatedRevision: + typeof message.expectedUpdatedRevision === "string" + ? message.expectedUpdatedRevision + : undefined, + }, + ); - try { - primaryClient.write( - `data: ${JSON.stringify(operationsEvent)}\n\n`, - ); - debug( - `[SSE] Sent ${message.operations?.length || 0} operations to PRIMARY client (${clients.indexOf(primaryClient)} of ${clients.length} clients)`, - ); - - debug(`data: ${JSON.stringify(operationsEvent)}\n\n`); - - // Notify other clients that operations are being applied (optional) - if (clients.length > 1) { - const notificationEvent = { - type: "operationsBeingApplied", - timestamp: Date.now(), - operationCount: message.operations?.length || 0, - source: "agent", - }; - - clients.slice(1).forEach((client, index) => { - try { - client.write( - `data: ${JSON.stringify(notificationEvent)}\n\n`, - ); - debug( - `[SSE] Notified secondary client ${index + 1} of pending operations`, - ); - } catch (error) { - console.error( - `[SSE] Failed to notify secondary client ${index + 1}:`, - error, - ); - } - }); - } - } catch (error) { - console.error( - "[SSE] Failed to send operations to primary client:", - error, - ); - throw error; - } - - // Send success confirmation back to agent + const documentId = path.basename(snapshot.boundRelativePath, ".md"); + collaborationManager.setDocumentContent( + documentId, + persisted.content, + ); process.send?.({ type: "operationsApplied", + requestId, success: true, - operationCount: message.operations?.length || 0, - method: "sse-forwarded", - clientsNotified: clients.length, + operationCount: message.operations.length, + bindingToken: snapshot.bindingToken, + revision: persisted.revision, }); - - debug(`[VIEW] Operations forwarded to primary client successfully`); } catch (error) { - console.error( - "[VIEW] Failed to forward operations via SSE:", - error, - ); + const errorMessage = + error instanceof Error ? error.message : "Unknown error"; process.send?.({ type: "operationsApplied", + requestId, success: false, - error: error instanceof Error ? error.message : "Unknown error", - method: "sse-forwarded", + identityMismatch: /binding|workspace root/.test(errorMessage), + revisionMismatch: /revision mismatch/.test(errorMessage), + error: errorMessage, + bindingToken: snapshot.bindingToken, }); } } else if (message.type === "getDocumentContent") { - debug( - `[VIEW] Processing getDocumentContent request at ${new Date().toISOString()}`, - ); - // Handle content requests from agent - try client markdown first, fallback to Y.js - // Process this asynchronously to avoid blocking other messages - (async () => { - try { - let documentId = ""; - - if (!filePath) { - // Use default document ID for memory-only mode - documentId = "default"; - } else { - documentId = path.basename(filePath, ".md"); - } - - debug("Using documentID " + documentId); - - let content = ""; - let source = "unknown"; - - try { - // PRIMARY: Try to get proper markdown from connected client - if (clients.length > 0) { - debug( - `[VIEW] Attempting to get markdown from connected client...`, - ); - const markdownResponse = - await requestMarkdownFromClient(); - content = markdownResponse.markdown; - source = "client-serializer"; - debug( - `[VIEW] Retrieved markdown from client: ${content.length} chars`, - ); - } else { - throw new Error("No clients connected"); - } - } catch (clientError) { - const errorMessage = - clientError instanceof Error - ? clientError.message - : String(clientError); - debug( - `[VIEW] Failed to get markdown from client (${errorMessage}), falling back to Y.js`, - ); - - // FALLBACK: Get content from authoritative Y.js document - const ydoc = getAuthoritativeDocument(documentId); - const yText = ydoc.getText("content"); - content = yText.toString(); - source = "yjs-fallback"; - debug( - `[VIEW] Retrieved content from Y.js fallback: ${content.length} chars`, - ); - - // If Y.js is also empty, try reading from file as last resort - if (!content && filePath && fs.existsSync(filePath)) { - try { - content = fs.readFileSync(filePath, "utf-8"); - source = "file-fallback"; - debug( - `[VIEW] Retrieved content from file fallback: ${content.length} chars`, - ); - } catch (fileError) { - debug( - `[VIEW] File fallback also failed: ${fileError}`, - ); - } - } - } - - debug( - `[VIEW] Sending document content to agent (source: ${source}, ${content.length} chars)`, - ); - - process.send?.({ - type: "documentContent", - content: content, - source: source, - timestamp: Date.now(), - }); - - debug("[SENT] [VIEW] Sent document content to agent process"); - } catch (error) { - console.error("[VIEW] Failed to get document content:", error); + const requestId = + typeof message.requestId === "string" ? message.requestId : ""; + const snapshot = captureBindingSnapshot(); + try { + const identityError = bindingError(message, snapshot); + if (identityError) { process.send?.({ type: "documentContent", + requestId, content: "", source: "error", - error: - error instanceof Error - ? error.message - : "Unknown error", + error: identityError, + identityMismatch: true, + bindingToken: snapshot.bindingToken, + boundFilePath: snapshot.filePath, + boundRoot: snapshot.filePath ? snapshot.currentRoot : null, + boundRelativePath: snapshot.boundRelativePath, + revision: null, timestamp: Date.now(), }); + return; + } + if (!snapshot.filePath || !snapshot.boundRelativePath) { + throw new Error("No markdown document is bound"); } - })(); + const document = readBoundDocument({ + token: snapshot.bindingToken ?? undefined, + root: snapshot.currentRoot, + relativePath: snapshot.boundRelativePath, + filePath: snapshot.filePath, + }); + process.send?.({ + type: "documentContent", + requestId, + content: document.content, + source: "file", + bindingToken: snapshot.bindingToken, + boundFilePath: snapshot.filePath, + boundRoot: snapshot.currentRoot, + boundRelativePath: snapshot.boundRelativePath, + revision: document.revision, + timestamp: Date.now(), + }); + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : "Unknown error"; + process.send?.({ + type: "documentContent", + requestId, + content: "", + source: "error", + error: errorMessage, + identityMismatch: /binding|workspace root/.test(errorMessage), + bindingToken: snapshot.bindingToken, + boundFilePath: snapshot.filePath, + boundRoot: snapshot.filePath ? snapshot.currentRoot : null, + boundRelativePath: snapshot.boundRelativePath, + revision: null, + timestamp: Date.now(), + }); + } } else if (message.type === "uiCommandResult") { // Handle UI command results from agent debug( diff --git a/ts/packages/agents/markdown/test/markdownActionHandler.spec.ts b/ts/packages/agents/markdown/test/markdownActionHandler.spec.ts index f1f31552af..6cbb885f0f 100644 --- a/ts/packages/agents/markdown/test/markdownActionHandler.spec.ts +++ b/ts/packages/agents/markdown/test/markdownActionHandler.spec.ts @@ -15,7 +15,7 @@ type TestAgentContext = { }; describe("markdown document creation", () => { - let workspace: string; + let workspace = ""; beforeEach(() => { workspace = fs.mkdtempSync( diff --git a/ts/packages/agents/markdown/test/markdownUpdatePersistence.spec.ts b/ts/packages/agents/markdown/test/markdownUpdatePersistence.spec.ts new file mode 100644 index 0000000000..fb23d6a6e1 --- /dev/null +++ b/ts/packages/agents/markdown/test/markdownUpdatePersistence.spec.ts @@ -0,0 +1,157 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import type { ChildProcess } from "node:child_process"; +import { EventEmitter } from "node:events"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { + computeContentRevision, + persistDocumentOperations, + readBoundDocument, + type DocumentBinding, +} from "../src/agent/documentUpdatePersistence.js"; +import { + getDocumentContentFromView, + sendOperationsToView, +} from "../src/agent/markdownActionHandler.js"; + +describe("markdown update persistence", () => { + let temporaryDirectory: string; + let workspace: string; + let filePath: string; + let binding: DocumentBinding; + + beforeEach(() => { + temporaryDirectory = fs.mkdtempSync( + path.join(os.tmpdir(), "typeagent-markdown-update-"), + ); + workspace = path.join(temporaryDirectory, "workspace"); + fs.mkdirSync(path.join(workspace, "notes"), { recursive: true }); + workspace = fs.realpathSync(workspace); + filePath = path.join(workspace, "notes", "plan.md"); + fs.writeFileSync(filePath, "original", "utf-8"); + binding = { + token: "binding-1", + root: workspace, + relativePath: "notes/plan.md", + filePath, + }; + }); + + afterEach(() => { + fs.rmSync(temporaryDirectory, { recursive: true, force: true }); + }); + + test("persists operations and treats a repeated apply as complete", () => { + const operation = { + type: "insert" as const, + position: 8, + content: [{ type: "text" as const, text: " updated" }], + }; + const expected = { + bindingToken: binding.token, + root: binding.root, + relativePath: binding.relativePath, + revision: computeContentRevision("original"), + updatedRevision: computeContentRevision("original updated"), + }; + + expect( + persistDocumentOperations(binding, [operation], expected) + .alreadyApplied, + ).toBe(false); + expect( + persistDocumentOperations(binding, [operation], expected) + .alreadyApplied, + ).toBe(true); + expect(fs.readFileSync(filePath, "utf-8")).toBe("original updated"); + }); + + test("rejects stale revisions and binding identities", () => { + fs.writeFileSync(filePath, "changed", "utf-8"); + expect(() => + persistDocumentOperations(binding, [], { + bindingToken: binding.token, + root: binding.root, + relativePath: binding.relativePath, + revision: computeContentRevision("original"), + updatedRevision: undefined, + }), + ).toThrow(/revision mismatch/); + expect(() => + persistDocumentOperations(binding, [], { + bindingToken: "binding-2", + root: binding.root, + relativePath: binding.relativePath, + revision: computeContentRevision("changed"), + updatedRevision: undefined, + }), + ).toThrow(/binding token changed/); + }); + + test("rejects a workspace replaced by a junction", () => { + const movedWorkspace = path.join(temporaryDirectory, "moved-workspace"); + const outside = path.join(temporaryDirectory, "outside"); + fs.mkdirSync(outside); + fs.writeFileSync(path.join(outside, "plan.md"), "outside", "utf-8"); + fs.renameSync(workspace, movedWorkspace); + fs.symlinkSync(outside, workspace, "junction"); + + expect(() => readBoundDocument(binding)).toThrow( + /workspace root changed/, + ); + expect(fs.readFileSync(path.join(outside, "plan.md"), "utf-8")).toBe( + "outside", + ); + fs.unlinkSync(workspace); + }); + + test("correlates concurrent view reads and applies", async () => { + const view = new EventEmitter() as EventEmitter & { + send: (message: Record) => void; + }; + view.send = (message) => { + const requestId = message.requestId as string; + queueMicrotask(() => { + view.emit("message", { + type: + message.type === "getDocumentContent" + ? "documentContent" + : "operationsApplied", + requestId: "unrelated", + success: false, + }); + view.emit("message", { + type: + message.type === "getDocumentContent" + ? "documentContent" + : "operationsApplied", + requestId, + content: "original", + bindingToken: binding.token, + revision: computeContentRevision("original"), + success: true, + }); + }); + }; + const child = view as unknown as ChildProcess; + const identity = { + expectedBindingToken: binding.token, + expectedRoot: binding.root, + expectedRelativePath: binding.relativePath, + }; + + await expect( + getDocumentContentFromView(child, identity), + ).resolves.toMatchObject({ content: "original" }); + await expect( + sendOperationsToView(child, [], { + ...identity, + expectedRevision: computeContentRevision("original"), + expectedUpdatedRevision: undefined, + }), + ).resolves.toMatchObject({ success: true }); + }); +});