@mondaydotcomorg/atp-server 0.21.6 → 0.23.0

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@mondaydotcomorg/atp-server",
-	"version": "0.21.6",
+	"version": "0.23.0",
 	"description": "Server implementation for Agent Tool Protocol",
 	"type": "module",
 	"main": "./dist/index.cjs",
@@ -49,11 +49,11 @@
 		"@babel/parser": "^7.26.0",
 		"@babel/traverse": "^7.26.0",
 		"@babel/types": "^7.26.0",
-		"@mondaydotcomorg/atp-compiler": "0.19.26",
-		"@mondaydotcomorg/atp-protocol": "0.19.22",
-		"@mondaydotcomorg/atp-provenance": "0.19.21",
-		"@mondaydotcomorg/atp-providers": "0.19.22",
-		"@mondaydotcomorg/atp-runtime": "0.19.21",
+		"@mondaydotcomorg/atp-compiler": "0.21.0",
+		"@mondaydotcomorg/atp-protocol": "0.21.0",
+		"@mondaydotcomorg/atp-provenance": "0.21.0",
+		"@mondaydotcomorg/atp-providers": "0.21.0",
+		"@mondaydotcomorg/atp-runtime": "0.21.0",
 		"@opentelemetry/api": "^1.9.0",
 		"@opentelemetry/auto-instrumentations-node": "^0.66.0",
 		"@opentelemetry/core": "^2.2.0",

package/src/executor/compiler-config.ts

@@ -18,10 +18,15 @@ import {
 	resumablePromiseAll,
 	resumablePromiseAllSettled,
 	batchParallel,
+	initializeCheckpointRuntime,
+	initializeCheckpointRuntimeWithProvenance,
+	cleanupCheckpointRuntime,
+	getCheckpointRuntime,
+	getCheckpointDataForError,
 	type TransformResult,
 	type DetectionResult,
 	type ICompiler,
-	type CacheStats,
+	type CheckpointProvenanceMetadata,
 } from '@mondaydotcomorg/atp-compiler';
 import { ATP_COMPILER_ENABLED, ATP_BATCH_SIZE_THRESHOLD } from './constants.js';
 
@@ -31,8 +36,11 @@ import { ATP_COMPILER_ENABLED, ATP_BATCH_SIZE_THRESHOLD } from './constants.js';
 class ATPCompilerAdapter implements ICompiler {
 	private compiler: ATPCompiler;
 
-	constructor(config: { enableBatchParallel: boolean; batchSizeThreshold: number }) {
-		this.compiler = new ATPCompiler(config);
+	constructor(config: { enableBatchParallel: boolean; batchSizeThreshold: number; enableOperationCheckpoints?: boolean }) {
+		this.compiler = new ATPCompiler({
+			...config,
+			enableOperationCheckpoints: config.enableOperationCheckpoints ?? true,
+		});
 	}
 
 	detect(code: string): DetectionResult {
@@ -58,8 +66,11 @@ class ATPCompilerAdapter implements ICompiler {
 class PluggableCompilerAdapter implements ICompiler {
 	private compiler: ReturnType<typeof createDefaultCompiler>;
 
-	constructor(config: { enableBatchParallel: boolean; batchSizeThreshold: number }) {
-		this.compiler = createDefaultCompiler(config);
+	constructor(config: { enableBatchParallel: boolean; batchSizeThreshold: number; enableOperationCheckpoints?: boolean }) {
+		this.compiler = createDefaultCompiler({
+			...config,
+			enableOperationCheckpoints: config.enableOperationCheckpoints ?? true,
+		});
 	}
 
 	async detect(code: string): Promise<DetectionResult> {
@@ -84,7 +95,7 @@ class PluggableCompilerAdapter implements ICompiler {
  * This is where you can easily add new compiler types
  */
 class CompilerFactory {
-	static create(config: { enableBatchParallel: boolean; batchSizeThreshold: number }): ICompiler {
+	static create(config: { enableBatchParallel: boolean; batchSizeThreshold: number; enableOperationCheckpoints?: boolean }): ICompiler {
 		const compilerType = process.env.ATP_USE_PLUGGABLE_COMPILER === 'true' ? 'pluggable' : 'atp';
 
 		switch (compilerType) {
@@ -173,6 +184,16 @@ export function getCompilerRuntime() {
 	};
 }
 
+// Re-export checkpoint functions for executor use
+export {
+	initializeCheckpointRuntime,
+	initializeCheckpointRuntimeWithProvenance,
+	cleanupCheckpointRuntime,
+	getCheckpointRuntime,
+	getCheckpointDataForError,
+	type CheckpointProvenanceMetadata,
+};
+
 export async function transformCodeWithCompiler(
 	code: string,
 	executionId: string,
@@ -202,13 +223,19 @@ export async function transformCodeWithCompiler(
 		// Detect patterns (abstracted sync/async handling)
 		const detection = await compiler.detect(code);
 
+		// With enableOperationCheckpoints, we always need to transform to wrap API calls
+		const enableCheckpoints = process.env.ATP_DISABLE_CHECKPOINTS !== 'true';
+		const needsTransform = detection.needsTransform || enableCheckpoints;
+
 		executionLogger.info('ATP Compiler detection result', {
 			needsTransform: detection.needsTransform,
 			patterns: detection.patterns,
 			batchable: detection.batchableParallel,
+			checkpointsEnabled: enableCheckpoints,
+			willTransform: needsTransform,
 		});
 
-		if (detection.needsTransform) {
+		if (needsTransform) {
 			const codeHash = getCodeHash(code);
 			const cached = transformCache.get(codeHash);
 			if (cached) {
@@ -244,6 +271,7 @@ export async function transformCodeWithCompiler(
 				loopCount: transformed.metadata.loopCount,
 				arrayMethodCount: transformed.metadata.arrayMethodCount,
 				parallelCallCount: transformed.metadata.parallelCallCount,
+				checkpointCount: transformed.metadata.checkpointCount,
 				batchSizeThreshold: ATP_BATCH_SIZE_THRESHOLD,
 			};
 

package/src/executor/execution-error-handler.ts

@@ -1,6 +1,6 @@
 import ivm from 'isolated-vm';
 import { ExecutionStatus } from '@mondaydotcomorg/atp-protocol';
-import type { ExecutionResult } from '@mondaydotcomorg/atp-protocol';
+import type { ExecutionResult, ExecutionCheckpointData } from '@mondaydotcomorg/atp-protocol';
 import type { Logger } from '@mondaydotcomorg/atp-runtime';
 import {
 	isPauseError,
@@ -10,17 +10,20 @@ import {
 	getAPICallResults,
 	clearAPICallResults,
 	setCurrentExecutionId,
-	clearCurrentExecutionId,
 	type PauseExecutionError,
 } from '@mondaydotcomorg/atp-runtime';
-import { isBatchPauseError, type BatchPauseExecutionError } from '@mondaydotcomorg/atp-compiler';
+import {
+	isBatchPauseError,
+	getCheckpointDataForError,
+	type BatchPauseExecutionError,
+} from '@mondaydotcomorg/atp-compiler';
 import { randomUUID } from 'node:crypto';
 import type { CallbackRecord } from '../execution-state/index.js';
 import type { RuntimeContext } from './types.js';
 import { categorizeError } from './error-handler.js';
 import { PAUSE_EXECUTION_MARKER } from './constants.js';
 
-export function handleExecutionError(
+export async function handleExecutionError(
 	error: unknown,
 	pauseError: unknown,
 	context: RuntimeContext,
@@ -30,7 +33,7 @@ export function handleExecutionError(
 	executionLogger: Logger,
 	isolate: ivm.Isolate,
 	transformedCode?: string
-): ExecutionResult {
+): Promise<ExecutionResult> {
 	const errMsg = error instanceof Error ? error.message : String(error);
 
 	if (errMsg.includes(PAUSE_EXECUTION_MARKER) && pauseError) {
@@ -192,6 +195,35 @@ export function handleExecutionError(
 	const memoryAfter = process.memoryUsage().heapUsed;
 	const memoryUsed = Math.max(0, memoryAfter - memoryBefore);
 
+	// Collect checkpoint data if available
+	// This also flushes buffered checkpoints to cache for recovery
+	let checkpointData: ExecutionCheckpointData | undefined;
+	try {
+		const rawCheckpointData = await getCheckpointDataForError();
+		if (rawCheckpointData && rawCheckpointData.checkpoints.length > 0) {
+			checkpointData = {
+				checkpoints: rawCheckpointData.checkpoints.map((cp) => ({
+					id: cp.id,
+					type: cp.type,
+					operation: cp.operation,
+					description: cp.description,
+					timestamp: cp.timestamp,
+				})),
+				restoreInstructions: rawCheckpointData.restoreInstructions,
+				stats: rawCheckpointData.stats,
+			};
+			executionLogger.info('Checkpoint data included in error response', {
+				checkpointCount: checkpointData.checkpoints.length,
+				fullSnapshots: checkpointData.stats.fullSnapshots,
+				references: checkpointData.stats.references,
+			});
+		}
+	} catch (checkpointError) {
+		executionLogger.debug('No checkpoint data available', {
+			reason: checkpointError instanceof Error ? checkpointError.message : String(checkpointError),
+		});
+	}
+
 	try {
 		isolate.dispose();
 	} catch (e) {}
@@ -208,6 +240,7 @@ export function handleExecutionError(
 			stack: err.stack,
 			retryable: errorInfo.retryable,
 			suggestion: errorInfo.suggestion,
+			checkpointData,
 		},
 		stats: {
 			duration: Date.now() - context.startTime,

package/src/executor/executor.ts

@@ -28,24 +28,38 @@ import type { ExecutorConfig, RuntimeContext } from './types.js';
 import { SandboxBuilder } from './sandbox-builder.js';
 import { CodeInstrumentor, StateManager } from '../instrumentation/index.js';
 import { ATP_COMPILER_ENABLED } from './constants.js';
-import { getCompilerRuntime, transformCodeWithCompiler } from './compiler-config.js';
+import {
+	getCompilerRuntime,
+	transformCodeWithCompiler,
+	initializeCheckpointRuntime,
+	initializeCheckpointRuntimeWithProvenance,
+	cleanupCheckpointRuntime,
+	getCheckpointRuntime,
+	type CheckpointProvenanceMetadata,
+} from './compiler-config.js';
 import { setupResumeExecution } from './resume-handler.js';
 import {
 	injectSandbox,
 	injectTimerPolyfills,
 	setupAPINamespace,
 	setupRuntimeNamespace,
+	setupCheckpointNamespace,
 } from './sandbox-injector.js';
 import { handleExecutionError } from './execution-error-handler.js';
 import {
+	attachProvenanceMetaForCheckpoint,
 	captureProvenanceSnapshot,
 	cleanupProvenanceForExecution,
 	clearProvenanceExecutionId,
+	createProvenanceProxy,
 	createTrackingRuntime,
+	getProvenance,
+	getAllProvenance,
 	instrumentCode as astInstrumentCode,
 	registerProvenanceMetadata,
 	SecurityPolicyEngine,
 	setProvenanceExecutionId,
+	type ProvenanceMetadata,
 } from '@mondaydotcomorg/atp-provenance';
 import {
 	createASTProvenanceChecker,
@@ -297,11 +311,72 @@ export class SandboxExecutor {
 				};
 			}
 
-			if (ATP_COMPILER_ENABLED) {
-				sandbox.__runtime = getCompilerRuntime();
+		if (ATP_COMPILER_ENABLED) {
+			sandbox.__runtime = getCompilerRuntime();
+		}
+
+		// Initialize checkpoint runtime if cache provider is available
+		if (this.config.cacheProvider) {
+			if (provenanceMode !== ProvenanceMode.NONE) {
+				// Provenance-aware checkpoint initialization
+				initializeCheckpointRuntimeWithProvenance({
+					executionId,
+					cache: this.config.cacheProvider,
+					config: { enabled: true },
+					provenanceMetaAttacher: attachProvenanceMetaForCheckpoint,
+					provenanceExtractor: (value: unknown) => {
+						// Just return ProvenanceMetadata as-is - types are now compatible
+						return getProvenance(value);
+					},
+					provenanceAttacher: (
+						value: unknown,
+						metadata,
+						primitives?
+					): unknown => {
+						// Skip null values (primitive registration calls)
+						if (value === null) {
+							// Re-register primitive taints if present
+							if (primitives) {
+								for (const [key, primMeta] of primitives) {
+									registerProvenanceMetadata(key, primMeta, executionId);
+								}
+							}
+							return null;
+						}
+
+						// Re-attach provenance to restored value using the metadata as-is
+						const restored = createProvenanceProxy(
+							value,
+							metadata.source,
+							metadata.readers,
+							metadata.dependencies
+						);
+						
+						// Re-register primitive taints if present
+						if (primitives) {
+							for (const [key, primMeta] of primitives) {
+								registerProvenanceMetadata(key, primMeta, executionId);
+							}
+						}
+						
+						return restored;
+					},
+				});
+				executionLogger.debug('Checkpoint runtime initialized with provenance integration', {
+					provenanceMode,
+				});
+			} else {
+				// Standard checkpoint initialization (no provenance)
+				initializeCheckpointRuntime({
+					executionId,
+					cache: this.config.cacheProvider,
+					config: { enabled: true },
+				});
 			}
+			sandbox.__checkpoint = getCheckpointRuntime();
+		}
 
-			let hintMetadata: Map<string, any> | undefined;
+		let hintMetadata: Map<string, any> | undefined;
 			if (provenanceMode === ProvenanceMode.AST) {
 				hintMetadata = getHintMap(executionId);
 
@@ -336,11 +411,16 @@ export class SandboxExecutor {
 
 			await setupAPINamespace(ivmContext, sandbox, provenanceMode);
 
-			if (ATP_COMPILER_ENABLED) {
-				await setupRuntimeNamespace(ivmContext, sandbox);
-			}
+		if (ATP_COMPILER_ENABLED) {
+			await setupRuntimeNamespace(ivmContext, sandbox);
+		}
 
-			let useCompiler = false;
+		// Setup checkpoint namespace if available
+		if (this.config.cacheProvider) {
+			await setupCheckpointNamespace(ivmContext, sandbox);
+		}
+
+		let useCompiler = false;
 			let astInstrumented = false;
 
 			const isResume = resumeData !== undefined;
@@ -366,13 +446,44 @@ export class SandboxExecutor {
 				codePreview: code.substring(0, 100),
 			});
 
-			if (provenanceMode === ProvenanceMode.AST && !useCompiler && !alreadyTransformed) {
+			// STEP 1: Checkpoint transformation FIRST (before AST instrumentation)
+			// This ensures checkpoints are tracked even in AST provenance mode
+			if (
+				ATP_COMPILER_ENABLED &&
+				this.config.cacheProvider &&
+				!alreadyTransformed
+			) {
+				const compilerResult = await transformCodeWithCompiler(
+					code,
+					executionId,
+					this.config.cacheProvider,
+					executionLogger,
+					this.compiler
+				);
+				codeToExecute = compilerResult.code;
+				useCompiler = compilerResult.useCompiler;
+				executionLogger.debug('Checkpoint transformation applied', {
+					useCompiler,
+					originalLength: code.length,
+					transformedLength: codeToExecute.length,
+				});
+			} else if (alreadyTransformed) {
+				codeToExecute = code;
+				useCompiler = true;
+				executionLogger.debug('Using already-transformed code on resume');
+			}
+
+			// STEP 2: AST instrumentation AFTER checkpoint transformation
+			// This ensures provenance tracking works with checkpoint-wrapped code
+			if (provenanceMode === ProvenanceMode.AST && !alreadyTransformed) {
 				try {
-					const instrumentResult = astInstrumentCode(code);
+					// Instrument the (potentially checkpoint-transformed) code
+					const instrumentResult = astInstrumentCode(codeToExecute);
 					codeToExecute = instrumentResult.code;
 					astInstrumented = true;
 					executionLogger.info('Code instrumented for provenance tracking (AST mode)', {
 						trackingCalls: instrumentResult.metadata.trackingCalls,
+						checkpointsPreserved: useCompiler,
 						instrumentedCodeStart: codeToExecute.substring(0, 150),
 						instrumentedCodeEnd: codeToExecute.substring(codeToExecute.length - 150),
 					});
@@ -402,27 +513,6 @@ export class SandboxExecutor {
 				}
 			}
 
-			if (
-				ATP_COMPILER_ENABLED &&
-				this.config.cacheProvider &&
-				!astInstrumented &&
-				!alreadyTransformed
-			) {
-				const compilerResult = await transformCodeWithCompiler(
-					code,
-					executionId,
-					this.config.cacheProvider,
-					executionLogger,
-					this.compiler
-				);
-				codeToExecute = compilerResult.code;
-				useCompiler = compilerResult.useCompiler;
-			} else if (alreadyTransformed) {
-				codeToExecute = code;
-				useCompiler = true;
-				executionLogger.debug('Using already-transformed code on resume');
-			}
-
 			if (!useCompiler && !astInstrumented && stateManager) {
 				try {
 					const instrumentor = new CodeInstrumentor();
@@ -580,7 +670,7 @@ export class SandboxExecutor {
 				}
 			}
 
-			return handleExecutionError(
+			return await handleExecutionError(
 				error,
 				pauseError,
 				context,
@@ -643,6 +733,11 @@ export class SandboxExecutor {
 
 		clearVectorStoreExecutionId();
 
+		// Cleanup checkpoint runtime
+		try {
+			cleanupCheckpointRuntime();
+		} catch (e) {}
+
 		if (executionId) {
 			try {
 				cleanupExecutionState(executionId);

package/src/executor/sandbox-injector.ts

@@ -1,7 +1,7 @@
 import ivm from 'isolated-vm';
 import type { Logger } from '@mondaydotcomorg/atp-runtime';
 import { isPauseError, runInExecutionContext } from '@mondaydotcomorg/atp-runtime';
-import { isBatchPauseError } from '@mondaydotcomorg/atp-compiler';
+import { isBatchPauseError, CHECKPOINT_RUNTIME_NAMESPACE } from '@mondaydotcomorg/atp-compiler';
 import { PAUSE_EXECUTION_MARKER } from './constants.js';
 import { isInIsolateFunction, getInIsolateImplementation } from './in-isolate-runtime.js';
 
@@ -117,14 +117,16 @@ export async function injectSandbox(
 							// In AST mode, tag result with provenance ID before copying so tag survives
 							if (isASTMode && result && typeof result === 'object') {
 								try {
-									// Generate unique ID for this API result
-									const provId = `tracked_${Date.now()}_${Math.random().toString(36).substring(7)}`;
-									Object.defineProperty(result, '__prov_id__', {
-										value: provId,
-										writable: false,
-										enumerable: true,
-										configurable: true,
-									});
+									// Only add __prov_id__ if not already present (avoids overwriting UUID from createProvenanceProxy)
+									if (!Object.prototype.hasOwnProperty.call(result, '__prov_id__')) {
+										const provId = `tracked_${Date.now()}_${Math.random().toString(36).substring(7)}`;
+										Object.defineProperty(result, '__prov_id__', {
+											value: provId,
+											writable: false,
+											enumerable: true,
+											configurable: true,
+										});
+									}
 								} catch (e) {
 									// If can't define property, that's ok
 								}
@@ -176,6 +178,41 @@ export async function injectSandbox(
 			continue;
 		}
 
+		// Handle checkpoint namespace
+		if (namespace === '__checkpoint' && typeof value === 'object' && value !== null) {
+			for (const [key, fn] of Object.entries(value)) {
+				if (typeof fn === 'function') {
+					await jail.set(
+						`__checkpoint_${key}_impl`,
+						new ivm.Reference(async (...args: unknown[]) => {
+							try {
+								const execute = async () => {
+									const result = await fn(...args);
+									return new ivm.ExternalCopy(result).copyInto();
+								};
+
+								if (executionId) {
+									return await runInExecutionContext(executionId, execute);
+								} else {
+									return await execute();
+								}
+							} catch (error) {
+								const err = error as Error;
+								if (isPauseError(error) || err.message === PAUSE_EXECUTION_MARKER) {
+									if (isPauseError(error)) {
+										onPauseError(error);
+									}
+									throw new Error(PAUSE_EXECUTION_MARKER);
+								}
+								throw error;
+							}
+						})
+					);
+				}
+			}
+			continue;
+		}
+
 		if (namespace === '__runtime' && typeof value === 'object' && value !== null) {
 			for (const [key, fn] of Object.entries(value)) {
 				if (typeof fn === 'function') {
@@ -352,3 +389,32 @@ ${newAccessPath} = async function(...args) {
 	setupNestedAPI(apiObject, '', 'globalThis.api');
 	await ivmContext.eval(apiSetup);
 }
+
+export async function setupCheckpointNamespace(
+	ivmContext: ivm.Context,
+	sandbox: Record<string, unknown>
+): Promise<void> {
+	const checkpointObject = sandbox.__checkpoint as Record<string, unknown>;
+	if (!checkpointObject || typeof checkpointObject !== 'object') {
+		return;
+	}
+
+	const checkpointKeys = Object.keys(checkpointObject).filter(
+		(k) => typeof checkpointObject[k] === 'function'
+	);
+	if (checkpointKeys.length === 0) {
+		return;
+	}
+
+	// Setup __checkpoint namespace for internal use by transformed code
+	let checkpointSetup = `globalThis.${CHECKPOINT_RUNTIME_NAMESPACE} = {\n`;
+	checkpointSetup += checkpointKeys
+		.map(
+			(key) =>
+				`\t${key}: async (...args) => {\n\t\treturn await ${CHECKPOINT_RUNTIME_NAMESPACE}_${key}_impl.apply(undefined, args, { arguments: { copy: true }, result: { promise: true } });\n\t}`
+		)
+		.join(',\n');
+	checkpointSetup += '\n};';
+
+	await ivmContext.eval(checkpointSetup);
+}