takos-control 1.0.0

package/src/application/services/sync/snapshot-compressor.ts

@@ -0,0 +1,79 @@
+import { logWarn } from '../../../shared/utils/logger';
+
+/** Concatenate an array of Uint8Array chunks into a single Uint8Array. */
+function concatChunks(chunks: Uint8Array[], totalLength: number): Uint8Array {
+    const result = new Uint8Array(totalLength);
+    let offset = 0;
+    for (const chunk of chunks) {
+        result.set(chunk, offset);
+        offset += chunk.length;
+    }
+    return result;
+}
+
+const MAX_DECOMPRESSED_SIZE = 512 * 1024 * 1024; // 512MB
+const COMPRESSION_RATIO_LIMIT = 100; // 100:1 ratio limit
+
+/**
+ * Handles gzip compression/decompression with decompression bomb protection.
+ */
+export class SnapshotCompressor {
+    constructor(
+        private maxDecompressedSize: number = MAX_DECOMPRESSED_SIZE,
+        private compressionRatioLimit: number = COMPRESSION_RATIO_LIMIT,
+    ) {}
+
+    async compress(data: string): Promise<ArrayBuffer> {
+        const encoder = new TextEncoder();
+        const stream = new Blob([encoder.encode(data)]).stream();
+        const compressedStream = stream.pipeThrough(new CompressionStream('gzip'));
+        const reader = compressedStream.getReader();
+        const chunks: Uint8Array[] = [];
+
+        while (true) {
+            const { done, value } = await reader.read();
+            if (done) break;
+            chunks.push(value);
+        }
+
+        const totalLength = chunks.reduce((sum, chunk) => sum + chunk.length, 0);
+        return concatChunks(chunks, totalLength).buffer as ArrayBuffer;
+    }
+
+    async decompress(data: ArrayBuffer): Promise<string> {
+        const stream = new Blob([data]).stream();
+        const decompressedStream = stream.pipeThrough(new DecompressionStream('gzip'));
+        const reader = decompressedStream.getReader();
+        const chunks: Uint8Array[] = [];
+        let totalDecompressedSize = 0;
+        const compressedSize = data.byteLength;
+
+        try {
+            while (true) {
+                const { done, value } = await reader.read();
+                if (done) break;
+
+                totalDecompressedSize += value.length;
+
+                if (totalDecompressedSize > this.maxDecompressedSize) {
+                    reader.cancel();
+                    throw new Error(`Decompression bomb detected: decompressed size exceeds ${this.maxDecompressedSize / 1024 / 1024}MB limit`);
+                }
+
+                const currentRatio = totalDecompressedSize / compressedSize;
+                if (currentRatio > this.compressionRatioLimit && compressedSize > 1024) {
+                    reader.cancel();
+                    throw new Error(`Decompression bomb detected: compression ratio ${currentRatio.toFixed(0)}:1 exceeds limit of ${this.compressionRatioLimit}:1`);
+                }
+
+                chunks.push(value);
+            }
+        } catch (error) {
+            try { reader.releaseLock(); } catch (err) { logWarn('Failed to release reader lock (non-critical)', { module: 'snapshot-compressor', error: err instanceof Error ? err.message : String(err) }); }
+            throw error;
+        }
+
+        const decoder = new TextDecoder();
+        return decoder.decode(concatChunks(chunks, totalDecompressedSize));
+    }
+}

package/src/application/services/sync/snapshot-storage.ts

@@ -0,0 +1,136 @@
+import type { Env } from '../../../shared/types';
+import type { SnapshotTree } from './models';
+import { SnapshotCompressor } from './snapshot-compressor';
+import { computeSHA256 } from '../../../shared/utils/hash';
+import { logError } from '../../../shared/utils/logger';
+
+/**
+ * Handles R2 (TENANT_SOURCE) storage operations for snapshot trees and blobs.
+ */
+export class SnapshotStorage {
+    private compressor: SnapshotCompressor;
+
+    constructor(
+        private env: Env,
+        private spaceId: string,
+        compressor?: SnapshotCompressor,
+    ) {
+        this.compressor = compressor ?? new SnapshotCompressor();
+    }
+
+    /** Ensure TENANT_SOURCE is available, throw otherwise. */
+    private requireStorage(): NonNullable<Env['TENANT_SOURCE']> {
+        if (!this.env.TENANT_SOURCE) {
+            throw new Error('Storage not configured (TENANT_SOURCE)');
+        }
+        return this.env.TENANT_SOURCE;
+    }
+
+    // ── Tree operations ──────────────────────────────────────────────
+
+    /** Build the R2 key for a snapshot tree. */
+    buildTreeKey(snapshotId: string): string {
+        return `trees/${this.spaceId}/${snapshotId}.json.gz`;
+    }
+
+    /** Put a compressed tree into R2 and return the key. */
+    async putTree(snapshotId: string, tree: SnapshotTree): Promise<{ treeKey: string; uncompressedSize: number }> {
+        const storage = this.requireStorage();
+        const treeKey = this.buildTreeKey(snapshotId);
+        const jsonData = JSON.stringify(tree);
+        const compressedData = await this.compressor.compress(jsonData);
+
+        await storage.put(treeKey, compressedData, {
+            httpMetadata: {
+                contentType: 'application/json',
+                contentEncoding: 'gzip',
+            },
+            customMetadata: {
+                'snapshot-id': snapshotId,
+                'workspace-id': this.spaceId,
+                'uncompressed-size': String(jsonData.length),
+            }
+        });
+
+        return { treeKey, uncompressedSize: jsonData.length };
+    }
+
+    /** Get a tree from R2, decompressing if needed. Returns null if not found. */
+    async getTree(treeKey: string): Promise<SnapshotTree | null> {
+        const storage = this.requireStorage();
+        const object = await storage.get(treeKey);
+        if (!object) return null;
+
+        if (treeKey.endsWith('.gz')) {
+            const compressedData = await object.arrayBuffer();
+            const jsonData = await this.compressor.decompress(compressedData);
+            return JSON.parse(jsonData) as SnapshotTree;
+        }
+
+        return await object.json<SnapshotTree>();
+    }
+
+    /** Delete a tree from R2. */
+    async deleteTree(treeKey: string): Promise<void> {
+        await this.env.TENANT_SOURCE?.delete(treeKey);
+    }
+
+    /** Validate that a tree key matches the expected prefix for a snapshot. */
+    validateTreeKeyIntegrity(treeKey: string, snapshotId: string): boolean {
+        const expectedKeyPrefix = `trees/${this.spaceId}/${snapshotId}`;
+        return treeKey.startsWith(expectedKeyPrefix);
+    }
+
+    // ── Blob operations ──────────────────────────────────────────────
+
+    /** Build the R2 key for a blob. */
+    private buildBlobKey(hash: string): string {
+        return `blobs/${this.spaceId}/${hash}`;
+    }
+
+    /** Check if a blob exists in R2. */
+    async blobExists(hash: string): Promise<boolean> {
+        const storage = this.requireStorage();
+        const existing = await storage.head(this.buildBlobKey(hash));
+        return existing !== null;
+    }
+
+    /** Put a blob into R2. */
+    async putBlob(hash: string, content: string, size: number): Promise<void> {
+        const storage = this.requireStorage();
+        await storage.put(this.buildBlobKey(hash), content, {
+            customMetadata: {
+                'workspace-id': this.spaceId,
+                'size': String(size),
+            }
+        });
+    }
+
+    /** Get blob content from R2, with integrity verification. Returns null if not found or integrity check fails. */
+    async getBlob(hash: string): Promise<string | null> {
+        if (!this.env.TENANT_SOURCE) {
+            return null;
+        }
+
+        const blob = await this.env.TENANT_SOURCE.get(this.buildBlobKey(hash));
+        if (!blob) {
+            return null;
+        }
+
+        const content = await blob.text();
+        const actualHash = await computeSHA256(content);
+        if (actualHash !== hash) {
+            logError(`Blob integrity check failed for ${this.spaceId}: expected ${hash}, got ${actualHash}`, undefined, { module: 'services/sync/snapshot' });
+            return null;
+        }
+
+        return content;
+    }
+
+    /** Delete blobs from R2 by their hashes. */
+    async deleteBlobs(hashes: string[]): Promise<void> {
+        if (hashes.length === 0 || !this.env.TENANT_SOURCE) return;
+        const blobKeys = hashes.map(hash => this.buildBlobKey(hash));
+        await this.env.TENANT_SOURCE.delete(blobKeys);
+    }
+}

package/src/application/services/sync/snapshot.ts

@@ -0,0 +1,426 @@
+import type { Env } from '../../../shared/types';
+import type { Snapshot, SnapshotTree, BlobFetcher } from './models';
+import { generateId, now, toIsoString } from '../../../shared/utils';
+import { computeSHA256 } from '../../../shared/utils/hash';
+import { getDb, snapshots, blobs, files } from '../../../infra/db';
+import { eq, and, inArray, lte, ne, sql } from 'drizzle-orm';
+import { logError, logInfo, logWarn } from '../../../shared/utils/logger';
+import { SnapshotStorage } from './snapshot-storage';
+import {
+    cleanupPendingSnapshots as cleanupPendingSnapshotsImpl,
+    getReachableSnapshots as getReachableSnapshotsImpl,
+    runGC as runGCImpl,
+} from './snapshot-cleanup';
+
+/** Extract non-empty blob hashes from a snapshot tree. */
+function extractTreeHashes(tree: SnapshotTree): string[] {
+    return Object.values(tree).map(entry => entry.hash).filter(Boolean);
+}
+
+/** Parse a JSON-encoded parent ID array, returning [] on failure. */
+function parseParentIds(raw: string | null, snapshotId?: string): string[] {
+    if (!raw) return [];
+    try {
+        return JSON.parse(raw) as string[];
+    } catch (err) {
+        if (snapshotId) {
+            logError(`Failed to parse parent_ids for snapshot ${snapshotId}`, err, { module: 'services/sync/snapshot' });
+        }
+        return [];
+    }
+}
+
+export class SnapshotManager {
+    private storage: SnapshotStorage;
+
+    constructor(
+        private env: Env,
+        private spaceId: string,
+        storage?: SnapshotStorage,
+    ) {
+        this.storage = storage ?? new SnapshotStorage(env, spaceId);
+    }
+
+    /**
+     * Create a new snapshot from a tree.
+     * Creates snapshot in 'pending' status; call completeSnapshot() after
+     * successful DB updates to mark as 'complete'.
+     */
+    async createSnapshot(
+        tree: SnapshotTree,
+        parentIds: string[],
+        message?: string,
+        author: 'user' | 'ai' = 'ai'
+    ): Promise<Snapshot> {
+        const db = getDb(this.env.DB);
+        const snapshotId = generateId();
+
+        const { treeKey } = await this.storage.putTree(snapshotId, tree);
+
+        const timestamp = now();
+        await db.insert(snapshots)
+            .values({
+                id: snapshotId,
+                accountId: this.spaceId,
+                parentIds: JSON.stringify(parentIds),
+                treeKey: treeKey,
+                message: message || null,
+                author,
+                status: 'pending',
+                createdAt: timestamp,
+            })
+            .run();
+
+        const blobHashes = extractTreeHashes(tree);
+        if (blobHashes.length > 0) {
+            await this.increaseBlobRefcount(blobHashes);
+        }
+
+        return {
+            id: snapshotId,
+            space_id: this.spaceId,
+            parent_ids: parentIds,
+            tree_key: treeKey,
+            message,
+            author,
+            status: 'pending',
+            created_at: timestamp
+        };
+    }
+
+    /**
+     * Mark a pending snapshot as complete after related DB updates have succeeded.
+     */
+    async completeSnapshot(snapshotId: string): Promise<void> {
+        const db = getDb(this.env.DB);
+        const result = await db.update(snapshots)
+            .set({ status: 'complete' })
+            .where(
+                and(
+                    eq(snapshots.id, snapshotId),
+                    eq(snapshots.accountId, this.spaceId),
+                    eq(snapshots.status, 'pending'),
+                )
+            )
+            .run();
+
+        if ((result.meta.changes ?? 0) === 0) {
+            logWarn(`Snapshot ${snapshotId} was not in pending status or not found`, { module: 'services/sync/snapshot' });
+        }
+    }
+
+    /**
+     * Mark a pending snapshot as failed and clean up.
+     * Call this if DB updates fail after snapshot creation.
+     */
+    async failSnapshot(snapshotId: string): Promise<void> {
+        const db = getDb(this.env.DB);
+        const snapshot = await this.getSnapshot(snapshotId);
+        if (!snapshot || snapshot.status !== 'pending') {
+            return;
+        }
+
+        try {
+            const tree = await this.getTree(snapshotId);
+            const hashes = extractTreeHashes(tree);
+
+            if (hashes.length > 0) {
+                await this.decreaseBlobRefcount(hashes);
+            }
+
+            await this.storage.deleteTree(snapshot.tree_key);
+
+            await db.delete(snapshots)
+                .where(eq(snapshots.id, snapshotId))
+                .run();
+
+            logInfo(`Cleaned up failed snapshot ${snapshotId}`, { module: 'services/sync/snapshot' });
+        } catch (err) {
+            logError(`Failed to clean up snapshot ${snapshotId}`, err, { module: 'services/sync/snapshot' });
+            await db.update(snapshots)
+                .set({ status: 'failed' })
+                .where(eq(snapshots.id, snapshotId))
+                .run();
+        }
+    }
+
+    /**
+     * Clean up old pending/failed snapshots (call periodically or on startup).
+     * Delegates to standalone cleanupPendingSnapshots function.
+     */
+    async cleanupPendingSnapshots(maxAgeMinutes: number = 30): Promise<number> {
+        return cleanupPendingSnapshotsImpl(this, this.env, this.spaceId, maxAgeMinutes);
+    }
+
+    /**  Get a snapshot by ID. */
+    async getSnapshot(snapshotId: string): Promise<Snapshot | null> {
+        const db = getDb(this.env.DB);
+        const snapshot = await db.select()
+            .from(snapshots)
+            .where(
+                and(
+                    eq(snapshots.id, snapshotId),
+                    eq(snapshots.accountId, this.spaceId),
+                )
+            )
+            .get();
+
+        if (!snapshot) return null;
+
+        return {
+            id: snapshot.id,
+            space_id: snapshot.accountId,
+            parent_ids: parseParentIds(snapshot.parentIds, snapshotId),
+            tree_key: snapshot.treeKey,
+            message: snapshot.message ?? undefined,
+            author: (snapshot.author as 'user' | 'ai' | null) ?? undefined,
+            status: snapshot.status as 'pending' | 'complete' | 'failed',
+            created_at: toIsoString(snapshot.createdAt),
+        };
+    }
+
+    /** Validate tree structure (checks a sample of entries). */
+    private validateTree(tree: unknown): tree is SnapshotTree {
+        if (typeof tree !== 'object' || tree === null) {
+            return false;
+        }
+        const entries = Object.entries(tree as Record<string, unknown>);
+        for (const [path, entry] of entries.slice(0, 10)) {
+            if (typeof path !== 'string' || path.length === 0) return false;
+            if (typeof entry !== 'object' || entry === null) return false;
+            const e = entry as Record<string, unknown>;
+            if (typeof e.hash !== 'string') return false;
+            if (typeof e.type !== 'string') return false;
+        }
+        return true;
+    }
+
+    /** Get the tree for a snapshot (with integrity verification). */
+    async getTree(snapshotId: string): Promise<SnapshotTree> {
+        const db = getDb(this.env.DB);
+        const snapshot = await this.getSnapshot(snapshotId);
+        if (!snapshot) {
+            throw new Error(`Snapshot not found: ${snapshotId}`);
+        }
+
+        let tree: SnapshotTree | null;
+        try {
+            tree = await this.storage.getTree(snapshot.tree_key);
+        } catch (parseError) {
+            // Re-throw configuration errors (e.g. missing TENANT_SOURCE)
+            if (parseError instanceof Error && parseError.message.includes('not configured')) {
+                throw parseError;
+            }
+            logError(`Failed to parse tree JSON for snapshot ${snapshotId}`, parseError, { module: 'services/sync/snapshot' });
+            return {};
+        }
+
+        if (!tree) {
+            logError(`Tree object not found for snapshot ${snapshotId}: ${snapshot.tree_key}`, undefined, { module: 'services/sync/snapshot' });
+            try {
+                await db.delete(snapshots)
+                    .where(eq(snapshots.id, snapshotId))
+                    .run();
+                logInfo(`Cleaned up orphaned snapshot record: ${snapshotId}`, { module: 'services/sync/snapshot' });
+            } catch (cleanupError) {
+                logError(`Failed to cleanup orphaned snapshot`, cleanupError, { module: 'services/sync/snapshot' });
+            }
+            return {};
+        }
+
+        if (!this.validateTree(tree)) {
+            logError(`Invalid tree structure for snapshot ${snapshotId}`, undefined, { module: 'services/sync/snapshot' });
+            return {};
+        }
+
+        if (!this.storage.validateTreeKeyIntegrity(snapshot.tree_key, snapshotId)) {
+            logError(`Snapshot integrity check failed: tree_key ${snapshot.tree_key} does not match expected prefix`, undefined, { module: 'services/sync/snapshot' });
+            return {};
+        }
+
+        return tree;
+    }
+
+    /** Create a tree from current workspace files table. */
+    async createTreeFromWorkspace(): Promise<SnapshotTree> {
+        const db = getDb(this.env.DB);
+        const fileRows = await db.select({
+            path: files.path,
+            sha256: files.sha256,
+            size: files.size,
+        })
+            .from(files)
+            .where(
+                and(
+                    eq(files.accountId, this.spaceId),
+                    ne(files.origin, 'system'),
+                )
+            )
+            .all();
+
+        const tree: SnapshotTree = {};
+
+        for (const file of fileRows) {
+            tree[file.path] = {
+                hash: file.sha256 || '', // Should not be null for source files
+                mode: 0o644, // Default to non-executable
+                type: 'file',
+                size: file.size
+            };
+        }
+
+        return tree;
+    }
+
+    /** Create a blob fetcher for this workspace (with integrity check). */
+    createBlobFetcher(): BlobFetcher {
+        return async (hash: string): Promise<string | null> => {
+            return this.storage.getBlob(hash);
+        };
+    }
+
+    /** Write a blob and return its hash and size. Refcount starts at 1. */
+    async writeBlob(content: string): Promise<{ hash: string; size: number }> {
+        const db = getDb(this.env.DB);
+        const hash = await computeSHA256(content);
+        const size = new TextEncoder().encode(content).length;
+
+        const exists = await this.storage.blobExists(hash);
+        if (!exists) {
+            await this.storage.putBlob(hash, content, size);
+        }
+
+        const existingBlob = await db.select()
+            .from(blobs)
+            .where(
+                and(
+                    eq(blobs.accountId, this.spaceId),
+                    eq(blobs.hash, hash),
+                )
+            )
+            .get();
+        if (existingBlob) {
+            await db.update(blobs)
+                .set({ refcount: sql`${blobs.refcount} + 1` })
+                .where(
+                    and(
+                        eq(blobs.accountId, this.spaceId),
+                        eq(blobs.hash, hash),
+                    )
+                )
+                .run();
+        } else {
+            try {
+                await db.insert(blobs)
+                    .values({
+                        accountId: this.spaceId,
+                        hash: hash,
+                        size: size,
+                        isBinary: false,
+                        refcount: 1,
+                        createdAt: now(),
+                    })
+                    .run();
+            } catch {
+                await db.update(blobs)
+                    .set({ refcount: sql`${blobs.refcount} + 1` })
+                    .where(
+                        and(
+                            eq(blobs.accountId, this.spaceId),
+                            eq(blobs.hash, hash),
+                        )
+                    )
+                    .run();
+            }
+        }
+
+        return { hash, size };
+    }
+
+    /** Increase refcount for blobs included in a snapshot. */
+    async increaseBlobRefcount(hashes: string[]): Promise<void> {
+        if (hashes.length === 0) return;
+
+        const db = getDb(this.env.DB);
+        for (const hash of hashes) {
+            await db.update(blobs)
+                .set({ refcount: sql`${blobs.refcount} + 1` })
+                .where(
+                    and(
+                        eq(blobs.accountId, this.spaceId),
+                        eq(blobs.hash, hash),
+                    )
+                )
+                .run();
+        }
+    }
+
+    /** Decrease refcount for blobs and delete those that reach zero. */
+    async decreaseBlobRefcount(hashes: string[]): Promise<void> {
+        if (hashes.length === 0) return;
+
+        const db = getDb(this.env.DB);
+
+        for (const hash of hashes) {
+            await db.update(blobs)
+                .set({ refcount: sql`${blobs.refcount} - 1` })
+                .where(
+                    and(
+                        eq(blobs.accountId, this.spaceId),
+                        eq(blobs.hash, hash),
+                        sql`${blobs.refcount} > 0`,
+                    )
+                )
+                .run();
+        }
+
+        const blobsToDelete = await db.select({ hash: blobs.hash })
+            .from(blobs)
+            .where(
+                and(
+                    eq(blobs.accountId, this.spaceId),
+                    inArray(blobs.hash, hashes),
+                    lte(blobs.refcount, 0),
+                )
+            )
+            .all();
+
+        const deletedHashes = blobsToDelete.map(b => b.hash);
+
+        if (deletedHashes.length > 0) {
+            await db.delete(blobs)
+                .where(
+                    and(
+                        eq(blobs.accountId, this.spaceId),
+                        inArray(blobs.hash, deletedHashes),
+                        lte(blobs.refcount, 0),
+                    )
+                )
+                .run();
+
+            await this.storage.deleteBlobs(deletedHashes);
+        }
+    }
+
+    /** Get all snapshots reachable from a head snapshot (DAG traversal). Delegates to standalone function. */
+    async getReachableSnapshots(headSnapshotId: string): Promise<Set<string>> {
+        return getReachableSnapshotsImpl(this.env, this.spaceId, headSnapshotId);
+    }
+
+    /** Run garbage collection on orphaned blobs and session data. Delegates to standalone function. */
+    async runGC(): Promise<{ deletedBlobs: number; deletedSnapshots: number; deletedSessionFiles: number }> {
+        return runGCImpl(this, this.env, this.spaceId);
+    }
+
+    // ── Methods exposed for use by snapshot-cleanup ─────────────────
+
+    /** Delete blobs from storage (exposed for cleanup functions). */
+    async deleteBlobs(hashes: string[]): Promise<void> {
+        await this.storage.deleteBlobs(hashes);
+    }
+
+    /** Delete a tree from storage (exposed for cleanup functions). */
+    async deleteTree(treeKey: string): Promise<void> {
+        await this.storage.deleteTree(treeKey);
+    }
+}