cogsbox-sync 0.0.1

package/templates/worker/src/index.ts

@@ -0,0 +1,1860 @@
+import { DurableObject } from 'cloudflare:workers';
+import { sign, verify } from '@tsndr/cloudflare-worker-jwt';
+import { UserPresenceDO } from './UserObject';
+import { Container } from '@cloudflare/containers';
+import { buildShadowNode } from './utility';
+import { AuthService, AuthError } from './auth';
+// These types should match the ones in cogsbox-shape for clarity
+type SerializableFieldMetadata = {
+	type: 'field';
+	sql: any; // The SQLType object
+};
+type SerializableRelationMetadata = {
+	type: 'relation';
+	relationType: 'hasMany' | 'hasOne' | 'manyToMany';
+	fromKey: string;
+	toKey: string;
+	schema: SerializableSchemaMetadata;
+};
+type SerializableSchemaMetadata = {
+	_tableName: string;
+	primaryKey: string | null;
+	fields: Record<string, SerializableFieldMetadata>;
+	relations: Record<string, SerializableRelationMetadata>;
+};
+
+type DataRoutes = {
+	fetch: string;
+	update?: string;
+	log?: string;
+};
+
+type SyncTokenPayload = {
+	clientId: string;
+	serviceId: number;
+	boundaryId: string;
+	scopes: string[];
+	exp: number;
+	iat: number;
+	dataRoutes?: DataRoutes;
+	appApiKey: string;
+	tenantId: number;
+	context: Record<string, any>;
+	metadataTemplate?: Record<string, any>;
+};
+
+type AuthResult = {
+	success: boolean;
+	valid: boolean;
+	tenantId: number;
+	serviceId: number;
+	scopes: string[];
+};
+
+export class SchemaProcessorContainer extends Container<Env> {
+	defaultPort = 8080;
+	sleepAfter = '30m';
+
+	async fetch(request: Request): Promise<Response> {
+		const url = new URL(request.url);
+		let requestToForward = request;
+		if (url.pathname === '/process-notification') {
+			const tenantId = url.searchParams.get('tenantId');
+			if (!tenantId) return new Response('Missing tenantId', { status: 400 });
+
+			const bundlePath = `schemas/tenant-${tenantId}/bundle.js`;
+			const r2Object = await this.env.SCHEMA_STORAGE.get(bundlePath);
+			if (!r2Object) return new Response('Bundle not found', { status: 404 });
+			const bundledCode = await r2Object.text();
+
+			const requestBody = (await request.json()) as Record<string, any>;
+
+			// Create a new request object with the bundled code injected into the body.
+			requestToForward = new Request(request.url, {
+				method: 'POST',
+				headers: { 'Content-Type': 'application/json' },
+				body: JSON.stringify({ bundledCode, ...requestBody }),
+			});
+		}
+
+		// STEP 2: UNIFIED ROUTER for all allowed paths.
+		// It will forward either the original request or the modified one.
+		if (
+			url.pathname.startsWith('/health') ||
+			url.pathname.startsWith('/load-schema') ||
+			url.pathname.startsWith('/validate') ||
+			url.pathname.startsWith('/process-notification') ||
+			url.pathname.startsWith('/get-api-routes')
+		) {
+			// Always forward the final `requestToForward` object.
+			return this.containerFetch(requestToForward);
+		}
+
+		// Return 404 for unknown paths instead of infinite forwarding
+		return new Response('Not Found', { status: 404 });
+	}
+}
+// Sync Engine Shadow Types
+export type ValidationStatus = 'NOT_VALIDATED' | 'VALIDATING' | 'VALID' | 'INVALID';
+export type ValidationError = {
+	path: (string | number)[];
+	message: string;
+	code?: string;
+};
+
+export type ValidationState = {
+	status: ValidationStatus;
+	errors: ValidationError[];
+	lastValidated?: number;
+	validatedValue?: any;
+};
+
+export type TypeInfo = {
+	type: 'string' | 'number' | 'boolean' | 'array' | 'object' | 'date' | 'unknown';
+	schema: any; // The Zod schema object
+	source: 'sync' | 'zod' | 'runtime' | 'default';
+	default: any;
+	nullable?: boolean;
+	optional?: boolean;
+};
+
+export type SyncEngineShadowMetadata = {
+	arrayKeys?: string[];
+	value?: any;
+	isDirty?: boolean;
+	lastSyncTimestamp?: number;
+	version?: number;
+
+	typeInfo?: TypeInfo;
+	validation?: ValidationState;
+};
+export type SyncEngineShadowNode = {
+	_meta?: SyncEngineShadowMetadata;
+	[key: string]: SyncEngineShadowNode | SyncEngineShadowMetadata | undefined;
+};
+
+type VersionRecord = {
+	value: number;
+};
+
+type MetadataTemplate = Record<string, any>;
+
+class WebSocketSyncEngine extends DurableObject<Env, unknown> {
+	private readonly DEBOUNCE_DELAY = 2000; // 2 seconds
+	private readonly MAX_SAVE_INTERVAL = 10000; // 10 seconds
+	private syncKey: string = '';
+	private expectedBoundaryId: string = '';
+	private schemaKey: string = '';
+	private stateId: string = '';
+	private params: Record<string, any> | undefined;
+	private instanceId: string = '';
+	private operationCounter: number = 0;
+	private clientActivities: Map<string, any> = new Map();
+
+	constructor(
+		readonly ctx: DurableObjectState,
+		readonly env: Env,
+	) {
+		super(ctx, env);
+
+		this.ctx.blockConcurrencyWhile(async () => {
+			this.syncKey = (await this.ctx.storage.get<string>('syncKey')) || '';
+			this.expectedBoundaryId = (await this.ctx.storage.get<string>('boundaryId')) || '';
+
+			const storedData = await this.ctx.storage.get(['instanceId', 'operationCounter']);
+			this.instanceId = (storedData.get('instanceId') as string) || `inst_${Date.now()}_${Math.random().toString(36).substring(2, 9)}`;
+			this.operationCounter = (storedData.get('operationCounter') as number) || 0;
+
+			if (!storedData.get('instanceId')) {
+				await this.ctx.storage.put({ instanceId: this.instanceId, operationCounter: 0 });
+				console.log(`[DO] New instance created: ${this.instanceId}`);
+			}
+
+			if (this.syncKey) {
+				this.parseSyncKey();
+			}
+		});
+	}
+	async alarm() {
+		await this.processUpdateQueue();
+		this.rebuildCachedState();
+	}
+	private async ensureContainerReady(): Promise<boolean> {
+		try {
+			const tenantId = await this.ctx.storage.get<number>('tenantId');
+			if (!tenantId) {
+				console.log('[DO] No tenantId found, skipping container readiness check.');
+				return false;
+			}
+
+			const containerStub = this.env.SCHEMA_PROCESSOR.get(this.env.SCHEMA_PROCESSOR.idFromName(`tenant:${tenantId}`));
+
+			// 1. CHECK if the container has the schema cached already.
+			const checkRequest = new Request(`https://container/load-schema?tenantId=${tenantId}`);
+			const checkResponse = await this.fetchWithReadyCheck(containerStub, checkRequest);
+
+			if (checkResponse.ok) {
+				const { hasSchema } = await checkResponse.json<{ hasSchema: boolean }>();
+				if (hasSchema) {
+					console.log(`[DO] Container for tenant ${tenantId} confirmed schema is ready.`);
+					return true; // It's ready, we're done.
+				}
+			}
+
+			// 2. If not, LOAD it from R2 and send it to the container.
+			console.log(`[DO] Schema not cached in container for tenant ${tenantId}. Loading from R2...`);
+			const bundlePath = `schemas/tenant-${tenantId}/bundle.js`;
+			const r2Object = await this.env.SCHEMA_STORAGE.get(bundlePath);
+			if (!r2Object) {
+				console.error(`[DO] CRITICAL: Schema bundle not found in R2 at ${bundlePath}`);
+				return false;
+			}
+			const bundledCode = await r2Object.text();
+
+			// 3. Use the POST /load-schema endpoint you already built.
+			const loadRequest = new Request(`https://container/load-schema?tenantId=${tenantId}`, {
+				method: 'POST',
+				headers: { 'Content-Type': 'application/json' },
+				body: JSON.stringify({ bundledCode }),
+			});
+
+			const loadResponse = await this.fetchWithReadyCheck(containerStub, loadRequest);
+			if (loadResponse.ok) {
+				console.log(`[DO] Successfully sent schema to container for tenant ${tenantId}.`);
+				return true;
+			} else {
+				console.error(`[DO] Failed to load schema into container for tenant ${tenantId}.`);
+				return false;
+			}
+		} catch (error) {
+			console.error('[DO] The ensureContainerReady check failed:', error);
+			return false;
+		}
+	}
+
+	private async rebuildCachedState(): Promise<void> {
+		try {
+			const { state: freshState, shadow: freshShadow } = await this.reconstructObjectFromStorage(this.ctx.storage, 'shadow');
+			if (freshState) {
+				const masterVersion = await this.ctx.storage.get<VersionRecord>('version');
+				await this.ctx.storage.put({
+					state_cache: freshState,
+					cache_version: masterVersion,
+				});
+			}
+		} catch (error) {
+			console.error('[DO Alarm] Failed to rebuild state_cache:', error);
+		}
+	}
+	private parseSyncKey() {
+		const bracketStart = this.syncKey.indexOf('[');
+		const bracketEnd = this.syncKey.indexOf(']');
+
+		if (bracketStart > -1 && bracketEnd > bracketStart) {
+			this.schemaKey = this.syncKey.substring(0, bracketStart);
+			const paramsString = this.syncKey.substring(bracketStart + 1, bracketEnd);
+
+			// Parse "userId:1" into {userId: 1}
+			const [key, value] = paramsString.split(':');
+			this.params = { [key]: isNaN(Number(value)) ? value : Number(value) };
+		} else {
+			this.schemaKey = this.syncKey.split('::')[0];
+			this.params = undefined;
+		}
+
+		const stateIdStart = this.syncKey.indexOf('::');
+		this.stateId = stateIdStart > -1 ? this.syncKey.substring(stateIdStart + 2) : '';
+	}
+	private async sendCurrentState(ws: WebSocket) {
+		// 1. Get the most up-to-date state by rebuilding it from storage.
+		const { shadow: currentShadow } = await this.reconstructObjectFromStorage(this.ctx.storage, 'shadow');
+
+		if (currentShadow) {
+			// 2. Build the message with the correct, full version string.
+			const message = JSON.stringify({
+				type: 'updateInitialState',
+				syncKey: this.syncKey,
+				data: currentShadow,
+				version: `${this.instanceId}:${this.operationCounter}`,
+			});
+
+			// 3. Send it.
+			ws.send(message);
+		} else {
+			// If there's no state at all, tell the client to provide it (for a "cold start").
+			// This handles the case where the DO is totally empty.
+			ws.send(JSON.stringify({ type: 'requestInitialState', syncKey: this.syncKey }));
+		}
+	}
+	async fetch(request: Request) {
+		const url = new URL(request.url);
+		const sessionId = url.searchParams.get('sessionId');
+		if (!sessionId) {
+			return new Response('Missing sessionId', { status: 400 });
+		}
+		const pathParts = url.pathname.split('/');
+		let pathSyncKey = '';
+		if (pathParts.length >= 3 && pathParts[1] === 'sync') {
+			pathSyncKey = pathParts[2];
+		}
+
+		if (!this.syncKey && pathSyncKey) {
+			this.syncKey = pathSyncKey;
+			await this.ctx.storage.put('syncKey', pathSyncKey);
+			this.parseSyncKey();
+		}
+
+		if (request.headers.get('Upgrade') !== 'websocket') {
+			return new Response('Expected WebSocket', { status: 400 });
+		}
+
+		const token = url.searchParams.get('token');
+		if (!token) {
+			return new Response('Missing token', { status: 401 });
+		}
+
+		let payload: SyncTokenPayload;
+		let isExpired = false;
+		try {
+			const isValid = await verify(token, this.env.JWT_SECRET);
+			if (!isValid) throw new Error('Invalid token');
+
+			payload = JSON.parse(atob(token.split('.')[1]));
+			isExpired = payload.exp < Math.floor(Date.now() / 1000);
+
+			const storedtenantId = await this.ctx.storage.get<number>('tenantId');
+			if (!storedtenantId) {
+				if (!payload.tenantId) {
+					throw new Error("JWT is missing required 'tenantId' claim.");
+				}
+				await this.ctx.storage.put('tenantId', payload.tenantId);
+			}
+
+			if (payload.context) {
+				await this.ctx.storage.put('clientContext', payload.context);
+			}
+			const appApiKey = payload.appApiKey;
+			if (!appApiKey) {
+				throw new Error('appApiKey not found in token payload. The /sync-token endpoint might be misconfigured.');
+			}
+
+			const existingKey = await this.ctx.storage.get<string>('app_api_key');
+			if (!existingKey) {
+				await this.ctx.storage.put('app_api_key', appApiKey);
+			}
+		} catch (error) {
+			return new Response('Invalid token', { status: 401 });
+		}
+
+		if (!this.expectedBoundaryId) {
+			this.expectedBoundaryId = payload.boundaryId;
+			await this.ctx.storage.put('boundaryId', payload.boundaryId);
+		} else if (this.expectedBoundaryId !== payload.boundaryId) {
+			console.error(`SECURITY VIOLATION: Boundary mismatch. Expected: ${this.expectedBoundaryId}, Got: ${payload.boundaryId}`);
+			return new Response('Access denied: boundary violation', { status: 403 });
+		}
+		this.ensureContainerReady();
+		const pair = new WebSocketPair();
+		const [client, server] = Object.values(pair);
+
+		this.ctx.acceptWebSocket(server);
+		if (payload.metadataTemplate) {
+			const existingTemplate = await this.ctx.storage.get<MetadataTemplate>('metadataTemplate');
+			if (!existingTemplate) {
+				await this.ctx.storage.put('metadataTemplate', payload.metadataTemplate);
+			}
+		}
+		server.serializeAttachment({
+			serviceId: payload.serviceId,
+			clientId: payload.clientId,
+			sessionId: sessionId,
+			syncKey: this.syncKey,
+			boundaryId: payload.boundaryId,
+			context: payload.context,
+		});
+
+		if (isExpired) {
+			server.send(
+				JSON.stringify({
+					type: 'refreshToken',
+					syncKey: this.syncKey,
+				}),
+			);
+		}
+
+		this.broadcastSubscribers();
+
+		return new Response(null, { status: 101, webSocket: client });
+	}
+
+	async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer) {
+		const attachment = ws.deserializeAttachment();
+		const data = typeof message === 'string' ? JSON.parse(message) : JSON.parse(new TextDecoder().decode(message));
+		const messageId = Math.random().toString(36).substring(7);
+		console.log(`[MSG-${messageId}] Received message:`, message);
+		switch (data.type) {
+			case 'initialSend':
+				await this.handleInitialSetup(ws, data);
+				break;
+
+			case 'providingInitialState':
+				const currentState = await this.ctx.storage.get('state');
+				console.log('currentState', currentState);
+				if (!currentState) {
+					await this.ctx.storage.put('state', data.state);
+					// Broadcast the newly adopted state to everyone
+					const updateMessage = JSON.stringify({
+						type: 'updateInitialState',
+						syncKey: this.syncKey,
+						data: data.state,
+					});
+					this.ctx.getWebSockets().forEach((socket) => socket.send(updateMessage));
+				}
+				break;
+
+			case 'updateContext':
+				await this.ctx.storage.put('clientContext', data.ctx);
+				// You might want to broadcast this change to other clients if needed
+				break;
+			case 'requestState': {
+				await this.sendCurrentState(ws);
+				break;
+			}
+			case 'initialSyncState':
+				if (data.data && !data.data.error) {
+					await this.ctx.storage.put('state', data.data);
+					ws.send(
+						JSON.stringify({
+							type: 'syncReady',
+							syncKey: this.syncKey, // Normal sync key
+						}),
+					);
+				}
+				break;
+
+			case 'clientStatus':
+				this.handleClientActivity(ws, data.status);
+				break;
+
+			case 'queueUpdate':
+				console.log(`[MSG-${messageId}] Processing queueUpdate`);
+				await this.handleStateUpdate(ws, data.data, attachment);
+				console.log(`[MSG-${messageId}] Completed queueUpdate`);
+				break;
+
+			case 'getSubscribers':
+				this.sendSubscribers(ws);
+				break;
+
+			case 'clearStorage':
+				await this.ctx.storage.delete('state');
+				await this.ctx.storage.delete('schema');
+				ws.send(
+					JSON.stringify({
+						type: 'storageCleared',
+						syncKey: this.syncKey, // Normal sync key
+					}),
+				);
+				break;
+		}
+	}
+	private async evaluateMetadataTemplate(
+		template: MetadataTemplate,
+		receivingClientContext: any,
+		updatingClientContext: any,
+	): Promise<Record<string, any>> {
+		const result: Record<string, any> = {};
+
+		for (const [key, value] of Object.entries(template)) {
+			if (typeof value === 'string' && value.startsWith('__ctx_') && value.endsWith('__')) {
+				// Extract the field name from the marker
+				const fieldName = value.slice(6, -2); // Remove __ctx_ prefix and __ suffix
+
+				// Get the value from the updating client's context
+				result[key] = updatingClientContext[fieldName];
+			} else if (typeof value === 'string' && value.startsWith('__currentCtx_') && value.endsWith('__')) {
+				// Handle references to the receiving/current client's context
+				const fieldName = value.slice(13, -2); // Remove __currentCtx_ prefix and __ suffix
+
+				// Get the value from the receiving client's context
+				result[key] = receivingClientContext[fieldName];
+			} else if (typeof value === 'object' && value !== null && !Array.isArray(value)) {
+				// Recursively handle nested objects
+				result[key] = await this.evaluateMetadataTemplate(value, receivingClientContext, updatingClientContext);
+			} else {
+				// IMPORTANT: It's a literal value - preserve it exactly as is!
+				// This includes strings, numbers, booleans, arrays, null, etc.
+				result[key] = value;
+			}
+		}
+
+		return result;
+	}
+	private async getFilteredMetadata(receivingClientAttachment: any, updatingClientContext: any): Promise<Record<string, any> | null> {
+		try {
+			// Get the metadata template stored during token creation
+			const metadataTemplate = await this.ctx.storage.get<MetadataTemplate>('metadataTemplate');
+
+			if (!metadataTemplate) {
+				// No metadata template defined for this sync schema
+				return null;
+			}
+
+			// Get the receiving client's context from their attachment or stored data
+			const receivingClientContext =
+				receivingClientAttachment?.context || (await this.ctx.storage.get<any>(`clientContext:${receivingClientAttachment?.clientId}`));
+
+			if (!receivingClientContext) {
+				console.warn(`[DO] No context found for receiving client ${receivingClientAttachment?.clientId}`);
+				return null;
+			}
+
+			// Evaluate the template with both contexts
+			const filteredMetadata = await this.evaluateMetadataTemplate(metadataTemplate, receivingClientContext, updatingClientContext);
+
+			return filteredMetadata;
+		} catch (error) {
+			console.error('[DO] Error filtering metadata:', error);
+			return null;
+		}
+	}
+	async webSocketClose(ws: WebSocket) {
+		this.broadcastSubscribers(ws);
+		const attachment = ws.deserializeAttachment();
+		if (attachment && attachment.sessionId) {
+			this.clientActivities.delete(attachment.sessionId);
+
+			// Optional: Broadcast that the user has left
+			const clientContext = await this.ctx.storage.get<any>('clientContext');
+
+			// Broadcast that the user has left WITH METADATA
+			const leftMessage = JSON.stringify({
+				type: 'clientActivity',
+				syncKey: this.syncKey,
+				activity: {
+					clientId: attachment.clientId,
+					event: { type: 'disconnect' },
+					metaData: clientContext || {}, // ADD METADATA HERE
+				},
+			});
+			this.ctx.getWebSockets().forEach((socket) => socket.send(leftMessage));
+		}
+	}
+	private async getApiRoutes(params?: Record<string, any>): Promise<{
+		queryData?: { url: string; method: 'GET' | 'POST'; data?: any };
+		mutateData?: { url: string; method: 'GET' | 'POST'; data?: any };
+	} | null> {
+		const tenantId = await this.ctx.storage.get<number>('tenantId');
+		const schemaKey = this.schemaKey;
+
+		if (!params) {
+			try {
+				const routesPath = `schemas/tenant-${tenantId}/routes/${schemaKey}.json`;
+				const r2Routes = await this.env.SCHEMA_STORAGE.get(routesPath);
+
+				if (r2Routes) {
+					return await r2Routes.json();
+				}
+			} catch (error) {
+				console.error('[DO] Failed to get routes from R2:', error);
+			}
+		}
+
+		try {
+			const bundlePath = `schemas/tenant-${tenantId}/bundle.js`;
+			const r2Object = await this.env.SCHEMA_STORAGE.get(bundlePath);
+			if (!r2Object) {
+				console.error(`[DO] Schema bundle not found in R2 at ${bundlePath} while trying to get API routes.`);
+				return null;
+			}
+			const bundledCode = await r2Object.text();
+
+			const containerStub = this.env.SCHEMA_PROCESSOR.get(this.env.SCHEMA_PROCESSOR.idFromName(`tenant:${tenantId}`));
+
+			const request = new Request(`https://container/get-api-routes`, {
+				method: 'POST',
+				headers: { 'Content-Type': 'application/json' },
+				body: JSON.stringify({
+					tenantId,
+					schemaKey,
+					bundledCode,
+					params,
+				}),
+			});
+
+			const response = await this.fetchWithReadyCheck(containerStub, request);
+			console.log('response', response);
+			if (!response.ok) {
+				const errorText = await response.text();
+				console.error(`[DO] Failed to get API routes from container for ${schemaKey}: ${response.status} - ${errorText}`);
+				return null;
+			}
+
+			const apiRoutes = await response.json<{
+				queryData?: { url: string; method: 'GET' | 'POST'; data?: any };
+				mutateData?: { url: string; method: 'GET' | 'POST'; data?: any };
+			}>();
+
+			if (!params) {
+				try {
+					await this.env.SCHEMA_STORAGE.put(`schemas/tenant-${tenantId}/routes/${schemaKey}.json`, JSON.stringify(apiRoutes), {
+						httpMetadata: { contentType: 'application/json' },
+					});
+				} catch (error) {
+					console.error('[DO] Failed to cache routes in R2:', error);
+				}
+			}
+
+			return apiRoutes;
+		} catch (error) {
+			console.error('[DO] Error fetching API routes from container:', error);
+			return null;
+		}
+	}
+	private async handleInitialSetup(ws: WebSocket, data: any) {
+		try {
+			const { isArray, inMemoryState } = data;
+			await this.ctx.storage.put('isArray', isArray);
+			if (inMemoryState) {
+				await this.ctx.storage.put('inMemoryState', true);
+			}
+			if (!inMemoryState) {
+				this.ensureContainerReady(); // Fire and forget
+			}
+
+			const isInitialized = await this.ctx.storage.get('shadow');
+
+			if (isInitialized) {
+				await this.sendCurrentState(ws);
+				return;
+			}
+
+			if (inMemoryState) {
+				ws.send(JSON.stringify({ type: 'requestInitialState', syncKey: this.syncKey }));
+			} else {
+				await this.fetchInitialState(ws);
+			}
+		} catch (error: any) {
+			console.error(`[DO] Initial setup failed for syncKey ${this.syncKey}:`, error);
+			ws.close(1011, 'Initial setup failed');
+		}
+	}
+	private async fetchWithReadyCheck(
+		containerStub: DurableObjectStub,
+		request: Request, // The actual request we WANT to send
+		options: { retries?: number; delay?: number } = {},
+	): Promise<Response> {
+		const { retries = 5, delay = 400 } = options;
+
+		for (let i = 0; i < retries; i++) {
+			try {
+				if (i > 0) {
+					// Only log retries
+				}
+				return await containerStub.fetch(request.clone());
+			} catch (error: any) {
+				if (error.message.includes('Connection refused') || error.message.includes('container port not found')) {
+					await new Promise((resolve) => setTimeout(resolve, delay));
+				} else {
+					console.error(`[DO] Unrecoverable fetch error:`, error);
+					throw error;
+				}
+			}
+		}
+
+		throw new Error(`Container did not respond after ${retries} attempts.`);
+	}
+	private async handleClientActivity(ws: WebSocket, status: any) {
+		const attachment = ws.deserializeAttachment();
+		if (!attachment || !attachment.clientId || !attachment.sessionId) return;
+
+		// 1. Store the latest activity for this client
+		this.clientActivities.set(attachment.sessionId, {
+			clientId: attachment.clientId,
+			event: status,
+		});
+		const clientContext = await this.ctx.storage.get<any>('clientContext');
+
+		// 3. Prepare the broadcast message WITH METADATA
+		const message = JSON.stringify({
+			type: 'clientActivity',
+			syncKey: this.syncKey,
+			activity: {
+				clientId: attachment.clientId,
+				event: status.event,
+				metaData: clientContext || {}, // ADD METADATA HERE
+			},
+		});
+
+		// 3. Broadcast to all OTHER clients
+		this.ctx.getWebSockets().forEach((socket) => {
+			const socketAttachment = socket.deserializeAttachment();
+			if (socketAttachment?.sessionId !== attachment.sessionId) {
+				try {
+					socket.send(message);
+				} catch (e) {
+					console.error('[DO] Failed to send activity to socket:', e);
+				}
+			}
+		});
+	}
+	private async processUpdateQueue() {
+		const isInMemory = await this.ctx.storage.get<boolean>('inMemoryState');
+		if (isInMemory) {
+			await this.ctx.storage.delete('firstDirtyTimestamp');
+			await this.ctx.storage.deleteAlarm();
+			return;
+		}
+
+		await this.ctx.storage.transaction(async (txn) => {
+			const firstDirtyTimestamp = await txn.get<number>('firstDirtyTimestamp');
+			if (!firstDirtyTimestamp) {
+				return;
+			}
+
+			await txn.delete('firstDirtyTimestamp');
+
+			const stateToPersist = await txn.get<any>('state');
+			if (stateToPersist === undefined || stateToPersist === null) {
+				return;
+			}
+
+			try {
+				// Get the API routes configuration
+
+				const apiRoutes = await this.getApiRoutes(this.params);
+				const mutateConfig = apiRoutes?.mutateData;
+
+				if (!mutateConfig || !mutateConfig.url) {
+					console.warn(`[DO] Persistence skipped: 'mutateData' configuration not found for syncKey ${this.syncKey}.`);
+					await txn.put('firstDirtyTimestamp', firstDirtyTimestamp);
+					return;
+				}
+
+				const authToken = await this.getAuthToken();
+				let response: Response;
+
+				if (mutateConfig.method === 'GET') {
+					// GET request for mutations (unusual but supported)
+					const finalUrl = new URL(mutateConfig.url);
+					finalUrl.searchParams.append('action', `${this.schemaKey}.update`); // Use this.schemaKey
+					finalUrl.searchParams.append('stateId', this.stateId); // Use this.stateId
+
+					response = await fetch(finalUrl.href, {
+						method: 'GET',
+						headers: {
+							'Content-Type': 'application/json',
+							Authorization: `Bearer ${authToken}`,
+						},
+					});
+				} else {
+					// POST request for mutations (standard)
+					const bodyData = {
+						action: `${this.schemaKey}.update`, // Use this.schemaKey
+						stateId: this.stateId, // Use this.stateId
+						payload: stateToPersist,
+						...(mutateConfig.data || {}),
+					};
+
+					response = await fetch(mutateConfig.url, {
+						method: 'POST',
+						headers: {
+							'Content-Type': 'application/json',
+							Authorization: `Bearer ${authToken}`,
+						},
+						body: JSON.stringify(bodyData),
+					});
+				}
+
+				if (!response.ok) {
+					const errorText = await response.text();
+					throw new Error(`DB sync failed: ${response.status} ${errorText}`);
+				}
+			} catch (error) {
+				console.error('Processing update queue failed, will retry:', error);
+				await txn.put('firstDirtyTimestamp', firstDirtyTimestamp);
+				const currentAlarm = await this.ctx.storage.getAlarm();
+				if (!currentAlarm || currentAlarm < Date.now()) {
+					await this.ctx.storage.setAlarm(Date.now() + this.DEBOUNCE_DELAY);
+				}
+			}
+		});
+	}
+
+	private async checkNotifications(operation: any, fullState: any) {
+		const stateKey = this.schemaKey;
+		const tenantId = await this.ctx.storage.get<number>('tenantId');
+		const boundaryId = await this.ctx.storage.get<string>('boundaryId');
+
+		if (!tenantId || !boundaryId) return;
+
+		const allUserContexts = this.ctx
+			.getWebSockets()
+			.map((ws) => {
+				const attachment = ws.deserializeAttachment();
+				return {
+					userId: attachment?.clientId,
+					tenantId,
+					boundaryId,
+				};
+			})
+			.filter((ctx) => ctx.userId);
+
+		if (allUserContexts.length === 0) {
+			return;
+		}
+
+		try {
+			const bundlePath = `schemas/tenant-${tenantId}/bundle.js`;
+			const r2Object = await this.env.SCHEMA_STORAGE.get(bundlePath);
+			if (!r2Object) {
+				console.error(`[DO] Schema bundle not found for notifications at ${bundlePath}`);
+				return;
+			}
+			const bundledCode = await r2Object.text();
+
+			const containerStub = this.env.SCHEMA_PROCESSOR.get(this.env.SCHEMA_PROCESSOR.idFromName(`tenant:${tenantId}`));
+
+			const requestUrl = `https://container/process-notifications-batch?tenantId=${tenantId}`;
+
+			const response = await this.fetchWithReadyCheck(
+				containerStub,
+				new Request(requestUrl, {
+					method: 'POST',
+					headers: { 'Content-Type': 'application/json' },
+					body: JSON.stringify({
+						bundledCode,
+						fullState: { [stateKey]: fullState },
+						operation,
+						contexts: allUserContexts,
+					}),
+				}),
+			);
+
+			if (!response.ok) {
+				const errorBody = await response.text();
+				console.error(`[DO] Batch notification processing failed with status: ${response.status}. Body: ${errorBody}`);
+				return;
+			}
+
+			const notificationsByUserId = await response.json<Record<string, any[]>>();
+
+			for (const userId in notificationsByUserId) {
+				const notifications = notificationsByUserId[userId];
+				if (notifications && notifications.length > 0) {
+					await this.sendNotificationToUser(userId, notifications, boundaryId);
+				}
+			}
+		} catch (error) {
+			console.error(`[DO] Failed to process notification batch:`, error);
+		}
+	}
+
+	private async sendNotificationToUser(userId: string, notifications: any[], boundaryId: string) {
+		try {
+			if (!notifications || notifications.length === 0) {
+				return;
+			}
+
+			const doName = `${boundaryId}:user:${userId}`;
+			const id = this.env.USER_PRESENCE.idFromName(doName);
+			const stub = this.env.USER_PRESENCE.get(id);
+
+			const batchMessage = {
+				type: 'notifications_batch',
+				payload: notifications,
+				timestamp: Date.now(),
+			};
+
+			console.log(`[WebSocketSyncEngine] Forwarding EFFICIENT batch notification:`, JSON.stringify(batchMessage));
+
+			await stub.forwardNotification(batchMessage);
+		} catch (error) {
+			console.error(`[DO] Failed to send notification batch to user ${userId}:`, error);
+		}
+	}
+
+	private async resolvePathToKey(txn: DurableObjectTransaction, path: (string | number)[]): Promise<string> {
+		let currentKey = 'shadow';
+		for (const segment of path) {
+			if (typeof segment === 'number') {
+				const parentMeta = await txn.get<any>(currentKey);
+				if (!parentMeta?.arrayKeys || parentMeta.arrayKeys.length <= segment) {
+					throw new Error(`Could not resolve path. Invalid index ${segment} at key ${currentKey}`);
+				}
+				const canonicalId = parentMeta.arrayKeys[segment];
+				currentKey = `${currentKey}:${canonicalId}`;
+			} else {
+				currentKey = `${currentKey}:${segment}`;
+			}
+		}
+		return currentKey;
+	}
+
+	private async reconstructObjectFromStorage(
+		storage: DurableObjectStorage | DurableObjectTransaction,
+		startKey: string,
+	): Promise<{ state: any; shadow: any }> {
+		const entries = await storage.list<any>({ prefix: `${startKey}:` });
+		const rootMeta = await storage.get<any>(startKey);
+
+		if (!rootMeta) return { state: undefined, shadow: undefined };
+
+		if (rootMeta.hasOwnProperty('value')) {
+			const shadow = { _meta: rootMeta };
+			return {
+				state: rootMeta.value,
+				shadow: shadow,
+			};
+		}
+
+		if (rootMeta.arrayKeys) {
+			const array: any[] = [];
+			const shadow: any = { _meta: rootMeta };
+
+			for (const itemKey of rootMeta.arrayKeys) {
+				const result = await this.reconstructObjectFromStorage(storage, `${startKey}:${itemKey}`);
+				array.push(result.state);
+				shadow[itemKey] = result.shadow;
+			}
+
+			return { state: array, shadow };
+		}
+
+		const obj: { [key: string]: any } = {};
+		const shadow: any = { _meta: rootMeta };
+		const keyPrefix = `${startKey}:`;
+
+		const childKeys = new Set<string>();
+		for (const key of entries.keys()) {
+			const propName = key.substring(keyPrefix.length).split(':')[0];
+			childKeys.add(propName);
+		}
+
+		for (const propName of childKeys) {
+			const result = await this.reconstructObjectFromStorage(storage, `${keyPrefix}${propName}`);
+			obj[propName] = result.state;
+			shadow[propName] = result.shadow;
+		}
+
+		return { state: obj, shadow };
+	}
+
+	private async handleStateUpdate(ws: WebSocket, updateData: any, attachment: any) {
+		const stateKey = this.schemaKey;
+		let validationErrors: { errors: any[]; sessionId: string | null } | null = null;
+		console.log('new updateData', updateData);
+		if (!updateData.operation) return;
+		await this.ctx.blockConcurrencyWhile(async () => {
+			await this.ctx.storage.transaction(async (txn) => {
+				try {
+					const { operation } = updateData;
+
+					const isInMemory = await txn.get<boolean>('inMemoryState');
+
+					console.log('new operation 22222');
+					const newVersion = Date.now();
+					await txn.put('version', { value: newVersion });
+
+					if (Array.isArray(operation)) {
+						for (const patch of operation) {
+							if (patch.op === 'replace') {
+								const pathSegments = patch.path
+									.split('/')
+									.slice(1)
+									.map((p: any) => (isNaN(Number(p)) ? p : Number(p)));
+								const targetKey = await this.resolvePathToKey(txn, pathSegments);
+								const oldMeta = (await txn.get<any>(targetKey)) || {};
+								await txn.put(targetKey, { ...oldMeta, value: patch.value, isDirty: true });
+							}
+						}
+					} else {
+						console.log('new operation 33333');
+
+						switch (operation.updateType) {
+							case 'update': {
+								const targetKey = await this.resolvePathToKey(txn, operation.path);
+								const oldMeta = (await txn.get<any>(targetKey)) || {};
+								await txn.put(targetKey, { ...oldMeta, value: operation.newValue, isDirty: true });
+								break;
+							}
+							case 'insert': {
+								const { itemId, insertAfterId, newValue } = operation;
+								console.log(`[DO] Insert operation: adding item ${itemId} to array at path [${operation.path.join(', ')}]`);
+								if (!itemId || !itemId.startsWith('id:')) {
+									throw new Error(`Invalid insert operation: missing or invalid item ID`);
+								}
+
+								let arrayKey: string;
+								if (operation.path.length === 0) {
+									arrayKey = 'shadow';
+								} else {
+									arrayKey = await this.resolvePathToKey(txn, operation.path);
+								}
+
+								let arrayMeta = await txn.get<any>(arrayKey);
+								if (!arrayMeta) {
+									arrayMeta = { arrayKeys: [] };
+								}
+								if (!arrayMeta.arrayKeys) {
+									arrayMeta.arrayKeys = [];
+								}
+
+								let insertIndex = 0;
+								if (insertAfterId && arrayMeta.arrayKeys.includes(insertAfterId)) {
+									insertIndex = arrayMeta.arrayKeys.indexOf(insertAfterId) + 1;
+								} else if (insertAfterId === null) {
+									insertIndex = 0;
+								} else {
+									insertIndex = arrayMeta.arrayKeys.length;
+								}
+
+								arrayMeta.arrayKeys.splice(insertIndex, 0, itemId);
+
+								arrayMeta.isDirty = true;
+								await txn.put(arrayKey, arrayMeta);
+
+								const newItemKey = `${arrayKey}:${itemId}`;
+
+								const storeNode = async (value: any, storageKey: string) => {
+									if (value === null || value === undefined || typeof value !== 'object') {
+										await txn.put(storageKey, { value, isDirty: true });
+										return;
+									}
+
+									if (Array.isArray(value)) {
+										const subArrayKeys: string[] = [];
+										for (let i = 0; i < value.length; i++) {
+											const subItemId = `id:${this.instanceId}_${Math.random().toString(36).substring(2, 9)}`;
+											subArrayKeys.push(subItemId);
+											await storeNode(value[i], `${storageKey}:${subItemId}`);
+										}
+										await txn.put(storageKey, { arrayKeys: subArrayKeys, isDirty: true });
+									} else {
+										await txn.put(storageKey, { isDirty: true });
+										for (const [key, val] of Object.entries(value)) {
+											await storeNode(val, `${storageKey}:${key}`);
+										}
+									}
+								};
+
+								await storeNode(operation.newValue, newItemKey);
+
+								console.log(
+									`[DO] Insert operation: added item ${itemId} ${
+										insertAfterId ? `after ${insertAfterId}` : 'at beginning'
+									} to array at path [${operation.path.join(', ')}]`,
+								);
+								break;
+							}
+							case 'cut': {
+								const itemIdToCut = operation.path[operation.path.length - 1];
+								const arrayPath = operation.path.slice(0, -1);
+
+								if (!itemIdToCut || !itemIdToCut.startsWith('id:')) {
+									throw new Error(`Invalid cut operation: missing or invalid item ID in path: ${operation.path.join('.')}`);
+								}
+
+								let arrayKey: string;
+								if (arrayPath.length === 0) {
+									arrayKey = 'shadow';
+								} else {
+									arrayKey = await this.resolvePathToKey(txn, arrayPath);
+								}
+
+								const arrayMeta = await txn.get<any>(arrayKey);
+								console.log('arrayMeta', arrayMeta);
+								if (!arrayMeta?.arrayKeys) {
+									console.error(`[DO] Cut operation failed: array not found at path [${arrayPath.join(', ')}]`);
+									return;
+								}
+
+								const indexToRemove = arrayMeta.arrayKeys.indexOf(itemIdToCut);
+								if (indexToRemove === -1) {
+									console.error(`[DO] Cut operation failed: item ${itemIdToCut} not found in array at path [${arrayPath.join(', ')}]`);
+									return;
+								}
+
+								arrayMeta.arrayKeys.splice(indexToRemove, 1);
+								arrayMeta.isDirty = true;
+								await txn.put(arrayKey, arrayMeta);
+
+								const deleteRecursive = async (keyToDelete: string) => {
+									const entries = await txn.list({ prefix: `${keyToDelete}:` });
+									for (const key of entries.keys()) {
+										await txn.delete(key);
+									}
+									await txn.delete(keyToDelete);
+								};
+
+								const itemKey = `${arrayKey}:${itemIdToCut}`;
+								await deleteRecursive(itemKey);
+
+								console.log(`[DO] Cut operation: removed item ${itemIdToCut} from array at path [${arrayPath.join(', ')}]`);
+								break;
+							}
+						}
+					}
+
+					if (!Array.isArray(operation) && !isInMemory) {
+						const { state: fullStateForNotifications, shadow: fullShadowForNotifications } = await this.reconstructObjectFromStorage(
+							txn,
+							'shadow',
+						);
+						if (fullStateForNotifications) {
+							await this.checkNotifications(operation, fullStateForNotifications);
+						} else {
+							console.warn(`[DO] Could not reconstruct state for notification check on syncKey: ${this.syncKey}`);
+						}
+					}
+					if (!isInMemory) {
+						const firstDirtyTimestamp = await txn.get<number>('firstDirtyTimestamp');
+						if (!firstDirtyTimestamp) {
+							await txn.put('firstDirtyTimestamp', Date.now());
+							await this.ctx.storage.setAlarm(Date.now() + this.MAX_SAVE_INTERVAL);
+						}
+						const currentAlarm = await this.ctx.storage.getAlarm();
+						const timeForDebounceAlarm = Date.now() + this.DEBOUNCE_DELAY;
+						if (!currentAlarm || timeForDebounceAlarm < currentAlarm) {
+							await this.ctx.storage.setAlarm(timeForDebounceAlarm);
+						}
+					}
+
+					const clientOperation = updateData.operation;
+
+					this.operationCounter++;
+
+					const allSockets = this.ctx.getWebSockets();
+					const updatingClientAttachment = ws.deserializeAttachment();
+					const updatingClientContext = updatingClientAttachment?.context;
+
+					for (const socket of allSockets) {
+						try {
+							const socketAttachment = socket.deserializeAttachment();
+							const receivingClientContext = socketAttachment?.context;
+
+							const messagePayload: any = {
+								type: 'applyPatch',
+								syncKey: this.syncKey,
+								data: { ...clientOperation },
+								sessionId: updatingClientAttachment?.sessionId,
+								version: `${this.instanceId}:${this.operationCounter}`,
+							};
+
+							if (updatingClientContext && receivingClientContext) {
+								const metadataTemplate = await this.ctx.storage.get<MetadataTemplate>('metadataTemplate');
+
+								if (metadataTemplate) {
+									const filteredMetadata = await this.evaluateMetadataTemplate(
+										metadataTemplate,
+										receivingClientContext,
+										updatingClientContext,
+									);
+
+									if (filteredMetadata) {
+										messagePayload.data.metaData = filteredMetadata;
+									}
+								}
+							}
+
+							socket.send(JSON.stringify(messagePayload));
+						} catch (e) {
+							console.error('[DO] Failed to send update to socket:', e);
+						}
+					}
+
+					await txn.put('operationCounter', this.operationCounter);
+				} catch (error) {
+					console.error('[DO] Fatal error in handleStateUpdate:', error);
+					ws.send(JSON.stringify({ type: 'error', message: 'An unexpected error occurred while processing the update.' }));
+				}
+			});
+		});
+	}
+
+	private async getAuthToken(): Promise<string> {
+		let token = await this.ctx.storage.get<string>('s2s_auth_token');
+		if (token) {
+			const payload = JSON.parse(atob(token.split('.')[1]));
+			if (payload.exp > Math.floor(Date.now() / 1000)) {
+				return token;
+			}
+		}
+
+		const boundaryId = await this.ctx.storage.get<string>('boundaryId');
+		if (!boundaryId) {
+			throw new Error('BUG: boundaryId was not found in storage. Cannot create auth token.');
+		}
+
+		const payload = {
+			iss: 'cogs-sync-engine',
+			sub: `do:${this.syncKey}`,
+			boundaryId: boundaryId,
+			iat: Math.floor(Date.now() / 1000),
+			exp: Math.floor(Date.now() / 1000) + 24 * 60 * 60,
+		};
+		const appApiKey = await this.ctx.storage.get<string>('app_api_key');
+		if (!appApiKey) {
+			throw new Error('BUG: app_api_key was not found in storage. The fetch() handler should have stored it.');
+		}
+		const newToken = await sign(payload, appApiKey);
+		await this.ctx.storage.put('s2s_auth_token', newToken);
+		return newToken;
+	}
+
+	private async initializeShadowState(state: any, schema?: any): Promise<any> {
+		const writes: { [key: string]: object } = {};
+		const rootKey = 'shadow';
+
+		function traverse(node: SyncEngineShadowNode, path: string) {
+			if (node._meta) {
+				writes[path] = node._meta;
+			}
+			for (const key in node) {
+				if (key === '_meta') continue;
+				const childNode = node[key] as SyncEngineShadowNode;
+				traverse(childNode, `${path}:${key}`);
+			}
+		}
+
+		const shadowTree = buildShadowNode(state, schema);
+		traverse(shadowTree, rootKey);
+
+		const currentVersion = Date.now();
+		const versionRecord: VersionRecord = { value: currentVersion };
+
+		writes['state_cache'] = state;
+		writes['version'] = versionRecord;
+		writes['cache_version'] = versionRecord;
+
+		await this.ctx.storage.put(writes);
+		return writes;
+	}
+	private async fetchInitialState(ws: WebSocket) {
+		const existingState = await this.ctx.storage.get('shadow');
+		if (existingState) {
+			console.warn(`[DO] ABORTING fetchInitialState: State already exists. This prevents overwriting live data.`);
+
+			await this.sendCurrentState(ws);
+
+			return;
+		}
+
+		try {
+			const apiRoutes = await this.getApiRoutes(this.params);
+			const queryConfig = apiRoutes?.queryData;
+
+			if (!queryConfig || !queryConfig.url) {
+				console.error(
+					`[fetchInitialState] 2. ERROR: 'queryData' configuration could not be resolved for syncKey "${this.syncKey}". Aborting.`,
+				);
+				return;
+			}
+
+			let url = queryConfig.url;
+
+			if (url.includes('localhost') || url.includes('127.0.0.1')) {
+				url = url.replace(/localhost|127\.0\.0\.1/, '172.17.32.1');
+			} else {
+			}
+
+			const authToken = await this.getAuthToken();
+			let response: Response;
+
+			if (queryConfig.method === 'GET') {
+				const finalUrl = new URL(url);
+				finalUrl.searchParams.append('action', `${this.schemaKey}.fetch`);
+				finalUrl.searchParams.append('stateId', this.stateId);
+
+				response = await fetch(finalUrl.href, {
+					method: 'GET',
+					headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${authToken}` },
+				});
+			} else {
+				const bodyData = {
+					action: `${this.schemaKey}.fetch`,
+					stateId: this.stateId,
+					...(queryConfig.data || {}),
+				};
+
+				response = await fetch(url, {
+					method: 'POST',
+					headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${authToken}` },
+					body: JSON.stringify(bodyData),
+				});
+			}
+
+			if (!response.ok) {
+				const errorText = await response.text();
+				throw new Error(`Fetch failed with status ${response.status}: ${errorText}`);
+			}
+			const responseData = (await response.json()) as any;
+			console.log('responseData', responseData);
+			const initialState = responseData?.data || responseData;
+
+			if (initialState) {
+				await this.initializeShadowState(initialState);
+				console.log(`[DO] Successfully fetched and initialized state for syncKey: ${this.syncKey}`);
+			} else {
+				throw new Error('Fetched initial state is null or undefined.');
+			}
+
+			await this.sendCurrentState(ws);
+			ws.send(JSON.stringify({ type: 'syncReady', syncKey: this.syncKey }));
+		} catch (error) {
+			console.error('[fetchInitialState] FATAL ERROR during fetch:', error);
+			ws.send(JSON.stringify({ type: 'error', message: 'Failed to fetch initial state.' }));
+		}
+	}
+
+	private broadcastSubscribers(closingSocket?: WebSocket) {
+		let allSockets = this.ctx.getWebSockets();
+
+		if (closingSocket) {
+			const closingSessionId = closingSocket.deserializeAttachment()?.sessionId;
+
+			allSockets = allSockets.filter((socket) => {
+				return socket.deserializeAttachment()?.sessionId !== closingSessionId;
+			});
+		}
+
+		const subscribers = allSockets.map((ws) => {
+			const attachment = ws.deserializeAttachment();
+			return {
+				clientId: attachment?.clientId || 'unknown',
+				sessionId: attachment?.sessionId || 'unknown',
+				syncKey: attachment?.syncKey || 'unknown',
+			};
+		});
+
+		const message = JSON.stringify({
+			type: 'subscribers',
+			syncKey: this.syncKey,
+			subscribers: subscribers,
+		});
+
+		allSockets.forEach((ws) => {
+			try {
+				ws.send(message);
+			} catch (e) {}
+		});
+	}
+
+	private sendSubscribers(ws: WebSocket) {
+		const sockets = this.ctx.getWebSockets();
+		const clientIds = new Set(
+			sockets.map((socket) => {
+				const attachment = socket.deserializeAttachment();
+
+				return attachment?.clientId || 'unknown';
+			}),
+		);
+
+		ws.send(
+			JSON.stringify({
+				type: 'subscribers',
+				syncKey: this.syncKey,
+				subscribers: clientIds,
+			}),
+		);
+	}
+}
+
+export { WebSocketSyncEngine, UserPresenceDO };
+
+// CHANGED: now takes auth param
+async function handleSchemaUpload(request: Request, env: Env, auth: AuthService): Promise<Response> {
+	// CHANGED: replaced entire auth block with AuthService
+	let tenantId: number;
+	try {
+		const apiKey = auth.extractBearerToken(request);
+		const authResult = await auth.validateApiKey(apiKey);
+
+		if (!authResult.success || !authResult.valid) {
+			return new Response(JSON.stringify({ success: false, message: 'Authentication failed: Invalid API key' }), {
+				status: 401,
+				headers: { 'Content-Type': 'application/json' },
+			});
+		}
+		tenantId = authResult.tenantId;
+	} catch (error) {
+		if (error instanceof AuthError) {
+			return new Response(JSON.stringify({ success: false, message: error.message }), {
+				status: error.status,
+				headers: { 'Content-Type': 'application/json' },
+			});
+		}
+		console.error('Auth failed:', error);
+		return new Response(JSON.stringify({ success: false, message: 'Internal server error during authentication' }), {
+			status: 500,
+			headers: { 'Content-Type': 'application/json' },
+		});
+	}
+	// CHANGED: end of auth replacement
+
+	let payload: { bundledSchema: string; timestamp: number; schemaPath: string };
+
+	try {
+		payload = await request.json();
+	} catch (e: any) {
+		console.error('Failed to parse request body', e);
+		return new Response(
+			JSON.stringify({
+				success: false,
+				message: 'Failed to parse request body as JSON.',
+				error: e.message,
+			}),
+			{
+				status: 400,
+				headers: { 'Content-Type': 'application/json' },
+			},
+		);
+	}
+
+	if (!payload.bundledSchema) {
+		return new Response(JSON.stringify({ success: false, message: 'bundledSchema is required.' }), {
+			status: 400,
+			headers: { 'Content-Type': 'application/json' },
+		});
+	}
+
+	const bundleStoragePath = `schemas/tenant-${tenantId}/bundle.js`;
+	try {
+		await env.SCHEMA_STORAGE.put(bundleStoragePath, payload.bundledSchema, {
+			httpMetadata: { contentType: 'application/javascript; charset=utf-8' },
+		});
+	} catch (error: any) {
+		console.error(`R2 upload failed for tenant ${tenantId}:`, error);
+		return new Response(JSON.stringify({ success: false, message: 'Failed to store schema bundle' }), {
+			status: 500,
+			headers: { 'Content-Type': 'application/json' },
+		});
+	}
+	const routesPrefix = `schemas/tenant-${tenantId}/routes/`;
+	const listed = await env.SCHEMA_STORAGE.list({ prefix: routesPrefix });
+
+	for (const key of listed.objects) {
+		await env.SCHEMA_STORAGE.delete(key.key);
+	}
+	return new Response(
+		JSON.stringify({
+			success: true,
+			message: `Schema bundle for tenant ${tenantId} uploaded successfully.`,
+		}),
+		{
+			status: 200,
+			headers: { 'Content-Type': 'application/json' },
+		},
+	);
+}
+async function handleUpdateContext(request: Request, env: Env) {
+	try {
+		const authHeader = request.headers.get('Authorization');
+
+		if (!authHeader || !authHeader.startsWith('Bearer ')) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'Missing or invalid authorization header',
+				}),
+				{
+					status: 401,
+					headers: { 'Content-Type': 'application/json' },
+				},
+			);
+		}
+
+		const apiKey = authHeader.substring(7);
+
+		const body = (await request.json()) as {
+			syncKey: string;
+			ctx: Record<string, any>;
+		};
+
+		const { syncKey, ctx } = body;
+
+		if (!syncKey) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'syncKey is required',
+				}),
+				{
+					status: 400,
+					headers: { 'Content-Type': 'application/json' },
+				},
+			);
+		}
+
+		const boundaryId = syncKey.split(':')[0];
+		const boundaryDOKey = `${boundaryId}:${syncKey}`;
+		const id = env.WEBSOCKET_SYNC_ENGINE.idFromName(boundaryDOKey);
+		const stub = env.WEBSOCKET_SYNC_ENGINE.get(id);
+
+		return stub.fetch(
+			new Request(request.url, {
+				method: 'POST',
+				headers: request.headers,
+				body: JSON.stringify({ type: 'updateContext', ctx }),
+			}),
+		);
+	} catch (error: any) {
+		console.error('Error updating context:', error);
+		return new Response(
+			JSON.stringify({
+				success: false,
+				message: error.message || 'Internal server error',
+			}),
+			{
+				status: 500,
+				headers: { 'Content-Type': 'application/json' },
+			},
+		);
+	}
+}
+export default {
+	async fetch(request: Request, env: Env, ctx: ExecutionContext) {
+		const url = new URL(request.url);
+
+		// Handle CORS
+		if (request.method === 'OPTIONS') {
+			return new Response(null, {
+				headers: {
+					'Access-Control-Allow-Origin': '*',
+					'Access-Control-Allow-Methods': 'POST, GET, OPTIONS',
+					'Access-Control-Allow-Headers': 'Content-Type, Authorization',
+					'Access-Control-Max-Age': '86400',
+				},
+			});
+		}
+
+		// CHANGED: create auth service once
+		const auth = new AuthService(env.JWT_SECRET, env.AUTH_SECRET);
+
+		if (url.pathname === '/sync-token' && request.method === 'POST') {
+			return handleSyncToken(request, env, auth);
+		}
+		if (url.pathname === '/refresh-token' && request.method === 'POST') {
+			return handleRefreshToken(request, env, auth);
+		}
+		if (url.pathname === '/upload-schema' && request.method === 'POST') {
+			return handleSchemaUpload(request, env, auth); // CHANGED: pass auth
+		}
+		if (url.pathname === '/update-context' && request.method === 'POST') {
+			return handleUpdateContext(request, env);
+		}
+		if (url.pathname.startsWith('/sync/') && request.headers.get('Upgrade') === 'websocket') {
+			const syncKey = url.pathname.split('/')[2];
+			if (!syncKey) {
+				return new Response('Missing sync key', { status: 400 });
+			}
+
+			const token = url.searchParams.get('token');
+			if (!token) {
+				return new Response('Missing token', {
+					status: 401,
+					headers: { 'X-Error-Reason': 'token_missing' },
+				});
+			}
+
+			try {
+				const isValid = await verify(token, env.JWT_SECRET);
+				if (!isValid) throw new Error('Invalid token signature');
+
+				const payload = JSON.parse(atob(token.split('.')[1]));
+
+				if (payload.exp && payload.exp < Math.floor(Date.now() / 1000)) {
+					throw new Error('TokenExpired');
+				}
+
+				const boundaryDOKey = `${payload.boundaryId}:${syncKey}`;
+				const id = env.WEBSOCKET_SYNC_ENGINE.idFromName(boundaryDOKey);
+				const stub = env.WEBSOCKET_SYNC_ENGINE.get(id);
+				return stub.fetch(request);
+			} catch (e: any) {
+				const errorReason = e.message || 'unknown_auth_error';
+
+				const pair = new WebSocketPair();
+				const [client, server] = Object.values(pair);
+
+				server.accept();
+
+				server.send(
+					JSON.stringify({
+						type: 'auth_failed',
+						reason: errorReason,
+						message: `Authentication failed: ${errorReason.replace(/_/g, ' ')}.`,
+					}),
+				);
+
+				server.close(4001, errorReason);
+
+				return new Response(null, { status: 101, webSocket: client });
+			}
+		}
+		if (url.pathname.startsWith('/user-presence/') && request.headers.get('Upgrade') === 'websocket') {
+			const token = url.searchParams.get('token');
+			if (!token) return new Response('Missing token', { status: 401 });
+
+			try {
+				const isValid = await verify(token, env.JWT_SECRET);
+				if (!isValid) throw new Error('Invalid token');
+				const payload = JSON.parse(atob(token.split('.')[1]));
+
+				const boundaryId = payload.boundaryId;
+				const userId = payload.clientId;
+
+				if (!boundaryId || !userId) {
+					return new Response('Token missing required claims', { status: 400 });
+				}
+
+				const doName = `${boundaryId}:user:${userId}`;
+				const id = env.USER_PRESENCE.idFromName(doName);
+				const stub = env.USER_PRESENCE.get(id);
+
+				return stub.fetch(request);
+			} catch (e) {
+				return new Response('Invalid token', { status: 401 });
+			}
+		}
+		return new Response(`<html><body><h1>Sync Engine</h1><p>Status: Online</p></body></html>`, {
+			headers: { 'Content-Type': 'text/html' },
+		});
+	},
+};
+
+// CHANGED: now takes auth param, replaced external fetch with auth.validateApiKey
+async function handleSyncToken(request: Request, env: Env, auth: AuthService) {
+	try {
+		// CHANGED: use auth service to extract and validate
+		const apiKey = auth.extractBearerToken(request);
+
+		const body = (await request.json()) as {
+			clientId: string;
+			boundaryId: string;
+			ctx: Record<string, any>;
+			metadataTemplate?: Record<string, any>;
+		};
+
+		const { clientId: realClientId, boundaryId, ctx, metadataTemplate } = body;
+
+		if (!realClientId) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'clientKey is required',
+				}),
+				{
+					status: 400,
+					headers: { 'Content-Type': 'application/json' },
+				},
+			);
+		}
+
+		// CHANGED: replaced external fetch('https://goot.co.uk:60002/...') with auth service
+		const authResult = await auth.validateApiKey(apiKey);
+
+		if (!authResult.success || !authResult.valid) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'Authentication failed: Invalid API key',
+				}),
+				{
+					status: 401,
+					headers: { 'Content-Type': 'application/json' },
+				},
+			);
+		}
+		// CHANGED: end of auth replacement
+
+		const sessionClientId = `session_${realClientId}_${Math.random().toString(36).substring(2, 9)}`;
+		const payload: SyncTokenPayload = {
+			clientId: sessionClientId,
+			serviceId: authResult.serviceId,
+			boundaryId,
+			scopes: authResult.scopes,
+			tenantId: authResult.tenantId,
+			exp: Math.floor(Date.now() / 1000) + 3600, // 1 hour
+			iat: Math.floor(Date.now() / 1000),
+			appApiKey: apiKey,
+			context: ctx,
+			metadataTemplate: metadataTemplate,
+		};
+		const refreshPayload = {
+			clientId: sessionClientId,
+			boundaryId,
+			tenantId: authResult.tenantId,
+			type: 'refresh',
+			exp: Math.floor(Date.now() / 1000) + 30 * 24 * 60 * 60, // 30 days
+			iat: Math.floor(Date.now() / 1000),
+		};
+		const token = await sign(payload, env.JWT_SECRET);
+
+		const refreshToken = await sign(refreshPayload, env.JWT_SECRET);
+		return new Response(
+			JSON.stringify({
+				success: true,
+				sessionToken: token,
+				refreshToken,
+				expiresIn: 3600,
+				clientId: sessionClientId,
+			}),
+			{
+				status: 200,
+				headers: {
+					'Content-Type': 'application/json',
+					'Cache-Control': 'no-store',
+					'Access-Control-Allow-Origin': '*',
+					'Access-Control-Allow-Methods': 'POST, OPTIONS',
+					'Access-Control-Allow-Headers': 'Content-Type, Authorization',
+				},
+			},
+		);
+	} catch (error: any) {
+		// CHANGED: handle AuthError specifically
+		if (error instanceof AuthError) {
+			return new Response(JSON.stringify({ success: false, message: error.message }), {
+				status: error.status,
+				headers: { 'Content-Type': 'application/json' },
+			});
+		}
+		console.error('Error generating sync token:', error);
+		return new Response(
+			JSON.stringify({
+				success: false,
+				message: error.message || 'Internal server error',
+			}),
+			{
+				status: 500,
+				headers: { 'Content-Type': 'application/json' },
+			},
+		);
+	}
+}
+
+// CHANGED: now takes auth param, replaced external fetch with auth.validateApiKey
+async function handleRefreshToken(request: Request, env: Env, auth: AuthService) {
+	try {
+		const body = (await request.json()) as {
+			refreshToken: string;
+			clientId: string;
+			ctx: Record<string, any>;
+		};
+
+		const { refreshToken, clientId } = body;
+
+		if (!refreshToken || !clientId) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'refreshToken and clientId are required',
+				}),
+				{ status: 400, headers: { 'Content-Type': 'application/json' } },
+			);
+		}
+
+		// Verify the refresh token
+		const isValid = await verify(refreshToken, env.JWT_SECRET);
+		if (!isValid) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'Invalid refresh token',
+				}),
+				{ status: 401, headers: { 'Content-Type': 'application/json' } },
+			);
+		}
+
+		const payload = JSON.parse(atob(refreshToken.split('.')[1]));
+
+		// Check if it's a refresh token and not expired
+		if (payload.type !== 'refresh' || payload.exp < Math.floor(Date.now() / 1000)) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'Invalid or expired refresh token',
+				}),
+				{ status: 401, headers: { 'Content-Type': 'application/json' } },
+			);
+		}
+
+		// Verify clientId matches
+		if (payload.clientId !== clientId) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'Client ID mismatch',
+				}),
+				{ status: 401, headers: { 'Content-Type': 'application/json' } },
+			);
+		}
+
+		// CHANGED: use auth service instead of manual header extraction
+		const apiKey = auth.extractBearerToken(request);
+
+		// CHANGED: replaced external fetch('https://goot.co.uk:60002/...') with auth service
+		const authResult = await auth.validateApiKey(apiKey);
+
+		if (!authResult.success || !authResult.valid) {
+			return new Response(
+				JSON.stringify({
+					success: false,
+					message: 'Failed to validate API key',
+				}),
+				{ status: 401, headers: { 'Content-Type': 'application/json' } },
+			);
+		}
+		// CHANGED: end of auth replacement
+
+		// Generate new access token
+		const newAccessPayload: SyncTokenPayload = {
+			clientId: payload.clientId,
+			serviceId: authResult.serviceId,
+			boundaryId: payload.boundaryId,
+			scopes: authResult.scopes,
+			tenantId: payload.tenantId,
+			exp: Math.floor(Date.now() / 1000) + 3600, // CHANGED: was 5 seconds, fixed to 1 hour
+			iat: Math.floor(Date.now() / 1000),
+			appApiKey: apiKey,
+			context: payload.context,
+		};
+
+		const newAccessToken = await sign(newAccessPayload, env.JWT_SECRET);
+
+		return new Response(
+			JSON.stringify({
+				success: true,
+				sessionToken: newAccessToken,
+				expiresIn: 3600,
+			}),
+			{
+				status: 200,
+				headers: {
+					'Content-Type': 'application/json',
+					'Cache-Control': 'no-store',
+					'Access-Control-Allow-Origin': '*',
+				},
+			},
+		);
+	} catch (error: any) {
+		// CHANGED: handle AuthError specifically
+		if (error instanceof AuthError) {
+			return new Response(JSON.stringify({ success: false, message: error.message }), {
+				status: error.status,
+				headers: { 'Content-Type': 'application/json' },
+			});
+		}
+		console.error('Error refreshing token:', error);
+		return new Response(
+			JSON.stringify({
+				success: false,
+				message: error.message || 'Internal server error',
+			}),
+			{
+				status: 500,
+				headers: { 'Content-Type': 'application/json' },
+			},
+		);
+	}
+}