@nghyane/arcane 0.1.0

package/src/memories/storage.ts

@@ -0,0 +1,563 @@
+import { Database } from "bun:sqlite";
+
+export interface MemoryThread {
+	id: string;
+	updatedAt: number;
+	rolloutPath: string;
+	cwd: string;
+	sourceKind: string;
+}
+
+export interface Stage1OutputRow {
+	threadId: string;
+	sourceUpdatedAt: number;
+	rawMemory: string;
+	rolloutSummary: string;
+	rolloutSlug: string | null;
+	generatedAt: number;
+	cwd: string;
+}
+
+export interface Stage1Claim {
+	threadId: string;
+	ownershipToken: string;
+	inputWatermark: number;
+	sourceUpdatedAt: number;
+	rolloutPath: string;
+	cwd: string;
+}
+
+export interface GlobalClaim {
+	ownershipToken: string;
+	inputWatermark: number;
+}
+
+const STAGE1_KIND = "memory_stage1";
+const GLOBAL_KIND = "memory_consolidate_global";
+const GLOBAL_KEY = "global";
+const DEFAULT_RETRY_REMAINING = 3;
+
+export function openMemoryDb(dbPath: string): Database {
+	const db = new Database(dbPath);
+	db.exec(`
+PRAGMA journal_mode=WAL;
+PRAGMA synchronous=NORMAL;
+PRAGMA busy_timeout=5000;
+
+CREATE TABLE IF NOT EXISTS threads (
+	id TEXT PRIMARY KEY,
+	updated_at INTEGER NOT NULL,
+	rollout_path TEXT NOT NULL,
+	cwd TEXT NOT NULL,
+	source_kind TEXT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS stage1_outputs (
+	thread_id TEXT PRIMARY KEY,
+	source_updated_at INTEGER NOT NULL,
+	raw_memory TEXT NOT NULL,
+	rollout_summary TEXT NOT NULL,
+	rollout_slug TEXT,
+	generated_at INTEGER NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS jobs (
+	kind TEXT NOT NULL,
+	job_key TEXT NOT NULL,
+	status TEXT NOT NULL,
+	worker_id TEXT,
+	ownership_token TEXT,
+	started_at INTEGER,
+	finished_at INTEGER,
+	lease_until INTEGER,
+	retry_at INTEGER,
+	retry_remaining INTEGER NOT NULL,
+	last_error TEXT,
+	input_watermark INTEGER,
+	last_success_watermark INTEGER,
+	PRIMARY KEY (kind, job_key)
+);
+`);
+	return db;
+}
+
+export function closeMemoryDb(db: Database): void {
+	db.close();
+}
+
+export function clearMemoryData(db: Database): void {
+	db.exec(`
+DELETE FROM stage1_outputs;
+DELETE FROM threads;
+DELETE FROM jobs WHERE kind IN ('memory_stage1', 'memory_consolidate_global');
+`);
+}
+
+export function upsertThreads(db: Database, threads: MemoryThread[]): void {
+	if (threads.length === 0) return;
+	const stmt = db.prepare(`
+INSERT INTO threads (id, updated_at, rollout_path, cwd, source_kind)
+VALUES (?, ?, ?, ?, ?)
+ON CONFLICT(id) DO UPDATE SET
+	updated_at = excluded.updated_at,
+	rollout_path = excluded.rollout_path,
+	cwd = excluded.cwd,
+	source_kind = excluded.source_kind
+`);
+	const tx = db.transaction((rows: MemoryThread[]) => {
+		for (const row of rows) {
+			stmt.run(row.id, row.updatedAt, row.rolloutPath, row.cwd, row.sourceKind);
+		}
+	});
+	tx(threads);
+}
+
+function ensureStage1Job(db: Database, threadId: string): void {
+	db.prepare(`
+INSERT OR IGNORE INTO jobs (kind, job_key, status, retry_remaining, input_watermark, last_success_watermark)
+VALUES (?, ?, 'pending', ?, 0, 0)
+`).run(STAGE1_KIND, threadId, DEFAULT_RETRY_REMAINING);
+}
+
+function ensureGlobalJob(db: Database): void {
+	db.prepare(`
+INSERT OR IGNORE INTO jobs (kind, job_key, status, retry_remaining, input_watermark, last_success_watermark)
+VALUES (?, ?, 'pending', ?, 0, 0)
+`).run(GLOBAL_KIND, GLOBAL_KEY, DEFAULT_RETRY_REMAINING);
+}
+
+export function claimStage1Jobs(
+	db: Database,
+	params: {
+		nowSec: number;
+		threadScanLimit: number;
+		maxRolloutsPerStartup: number;
+		maxRolloutAgeDays: number;
+		minRolloutIdleHours: number;
+		leaseSeconds: number;
+		runningConcurrencyCap: number;
+		workerId: string;
+		excludeThreadIds?: string[];
+	},
+): Stage1Claim[] {
+	const {
+		nowSec,
+		threadScanLimit,
+		maxRolloutsPerStartup,
+		maxRolloutAgeDays,
+		minRolloutIdleHours,
+		leaseSeconds,
+		runningConcurrencyCap,
+		workerId,
+		excludeThreadIds = [],
+	} = params;
+	const maxAgeSec = maxRolloutAgeDays * 24 * 60 * 60;
+	const minIdleSec = minRolloutIdleHours * 60 * 60;
+	const runningCountRow = db
+		.prepare(
+			"SELECT COUNT(*) AS count FROM jobs WHERE kind = ? AND status = 'running' AND lease_until IS NOT NULL AND lease_until > ?",
+		)
+		.get(STAGE1_KIND, nowSec) as { count?: number } | undefined;
+	let runningCount = runningCountRow?.count ?? 0;
+	if (runningCount >= runningConcurrencyCap) return [];
+	const candidateRows = db
+		.prepare("SELECT id, updated_at, rollout_path, cwd, source_kind FROM threads ORDER BY updated_at DESC LIMIT ?")
+		.all(threadScanLimit) as Array<{
+		id: string;
+		updated_at: number;
+		rollout_path: string;
+		cwd: string;
+		source_kind: string;
+	}>;
+	const claims: Stage1Claim[] = [];
+	const excluded = new Set(excludeThreadIds);
+	for (const row of candidateRows) {
+		if (claims.length >= maxRolloutsPerStartup) break;
+		if (excluded.has(row.id)) continue;
+		if (row.source_kind !== "cli" && row.source_kind !== "app") continue;
+		if (nowSec - row.updated_at > maxAgeSec) continue;
+		if (nowSec - row.updated_at < minIdleSec) continue;
+		if (runningCount >= runningConcurrencyCap) break;
+		const stage1 = db.prepare("SELECT source_updated_at FROM stage1_outputs WHERE thread_id = ?").get(row.id) as
+			| { source_updated_at?: number }
+			| undefined;
+		if ((stage1?.source_updated_at ?? 0) >= row.updated_at) continue;
+		ensureStage1Job(db, row.id);
+		const ownershipToken = crypto.randomUUID();
+		const leaseUntil = nowSec + leaseSeconds;
+		const claimed = db
+			.prepare(`
+UPDATE jobs
+SET status = 'running', worker_id = ?, ownership_token = ?, started_at = ?, finished_at = NULL,
+	lease_until = ?, retry_at = NULL, last_error = NULL, input_watermark = ?,
+	retry_remaining = CASE
+		WHEN input_watermark IS NULL THEN ?
+		WHEN input_watermark < ? THEN ?
+		ELSE retry_remaining
+	END
+WHERE kind = ? AND job_key = ?
+	AND (last_success_watermark IS NULL OR last_success_watermark < ?)
+	AND NOT (status = 'running' AND lease_until IS NOT NULL AND lease_until > ?)
+	AND (
+		input_watermark IS NULL
+		OR input_watermark < ?
+		OR (
+			retry_remaining > 0
+			AND (retry_at IS NULL OR retry_at <= ?)
+		)
+	)
+`)
+			.run(
+				workerId,
+				ownershipToken,
+				nowSec,
+				leaseUntil,
+				row.updated_at,
+				DEFAULT_RETRY_REMAINING,
+				row.updated_at,
+				DEFAULT_RETRY_REMAINING,
+				STAGE1_KIND,
+				row.id,
+				row.updated_at,
+				nowSec,
+				row.updated_at,
+				nowSec,
+			);
+		if (Number(claimed.changes ?? 0) <= 0) continue;
+		claims.push({
+			threadId: row.id,
+			ownershipToken,
+			inputWatermark: row.updated_at,
+			sourceUpdatedAt: row.updated_at,
+			rolloutPath: row.rollout_path,
+			cwd: row.cwd,
+		});
+		runningCount += 1;
+	}
+	return claims;
+}
+
+export function enqueueGlobalWatermark(
+	db: Database,
+	sourceUpdatedAt: number,
+	params?: { forceDirtyWhenNotAdvanced?: boolean },
+): void {
+	const forceDirtyWhenNotAdvanced = params?.forceDirtyWhenNotAdvanced ?? false;
+	ensureGlobalJob(db);
+	db.prepare(`
+UPDATE jobs
+SET
+	input_watermark = CASE
+		WHEN input_watermark IS NULL THEN ?
+		WHEN input_watermark < ? THEN ?
+		WHEN ? = 1 AND (last_success_watermark IS NULL OR input_watermark <= last_success_watermark) THEN
+			CASE
+				WHEN last_success_watermark IS NULL THEN input_watermark + 1
+				ELSE last_success_watermark + 1
+			END
+		ELSE input_watermark
+	END,
+	retry_remaining = CASE
+		WHEN input_watermark IS NULL THEN ?
+		WHEN input_watermark < ? THEN ?
+		WHEN ? = 1 AND (last_success_watermark IS NULL OR input_watermark <= last_success_watermark) THEN ?
+		ELSE retry_remaining
+	END,
+	retry_at = CASE
+		WHEN input_watermark IS NULL THEN NULL
+		WHEN input_watermark < ? THEN NULL
+		WHEN ? = 1 AND (last_success_watermark IS NULL OR input_watermark <= last_success_watermark) THEN NULL
+		ELSE retry_at
+	END
+WHERE kind = ? AND job_key = ?
+`).run(
+		sourceUpdatedAt,
+		sourceUpdatedAt,
+		sourceUpdatedAt,
+		forceDirtyWhenNotAdvanced ? 1 : 0,
+		DEFAULT_RETRY_REMAINING,
+		sourceUpdatedAt,
+		DEFAULT_RETRY_REMAINING,
+		forceDirtyWhenNotAdvanced ? 1 : 0,
+		DEFAULT_RETRY_REMAINING,
+		sourceUpdatedAt,
+		forceDirtyWhenNotAdvanced ? 1 : 0,
+		GLOBAL_KIND,
+		GLOBAL_KEY,
+	);
+}
+
+export function markStage1SucceededWithOutput(
+	db: Database,
+	params: {
+		threadId: string;
+		ownershipToken: string;
+		sourceUpdatedAt: number;
+		rawMemory: string;
+		rolloutSummary: string;
+		rolloutSlug: string | null;
+		nowSec: number;
+	},
+): boolean {
+	const { threadId, ownershipToken, sourceUpdatedAt, rawMemory, rolloutSummary, rolloutSlug, nowSec } = params;
+	const tx = db.transaction(() => {
+		const matched = db
+			.prepare(
+				"SELECT 1 AS ok FROM jobs WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ?",
+			)
+			.get(STAGE1_KIND, threadId, ownershipToken) as { ok?: number } | undefined;
+		if (!matched?.ok) return false;
+
+		db.prepare(`
+UPDATE jobs
+SET status = 'done', finished_at = ?, lease_until = NULL, retry_at = NULL,
+	last_error = NULL, last_success_watermark = input_watermark
+WHERE kind = ? AND job_key = ? AND ownership_token = ?
+`).run(nowSec, STAGE1_KIND, threadId, ownershipToken);
+
+		db.prepare(`
+INSERT INTO stage1_outputs (thread_id, source_updated_at, raw_memory, rollout_summary, rollout_slug, generated_at)
+VALUES (?, ?, ?, ?, ?, ?)
+ON CONFLICT(thread_id) DO UPDATE SET
+	source_updated_at = excluded.source_updated_at,
+	raw_memory = excluded.raw_memory,
+	rollout_summary = excluded.rollout_summary,
+	rollout_slug = excluded.rollout_slug,
+	generated_at = excluded.generated_at
+WHERE excluded.source_updated_at >= stage1_outputs.source_updated_at
+`).run(threadId, sourceUpdatedAt, rawMemory, rolloutSummary, rolloutSlug, nowSec);
+
+		enqueueGlobalWatermark(db, sourceUpdatedAt, { forceDirtyWhenNotAdvanced: true });
+		return true;
+	});
+	return tx() as boolean;
+}
+
+export function markStage1SucceededNoOutput(
+	db: Database,
+	params: { threadId: string; ownershipToken: string; sourceUpdatedAt: number; nowSec: number },
+): boolean {
+	const { threadId, ownershipToken, sourceUpdatedAt, nowSec } = params;
+	const tx = db.transaction(() => {
+		const matched = db
+			.prepare(
+				"SELECT 1 AS ok FROM jobs WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ?",
+			)
+			.get(STAGE1_KIND, threadId, ownershipToken) as { ok?: number } | undefined;
+		if (!matched?.ok) return false;
+
+		db.prepare(`
+UPDATE jobs
+SET status = 'done', finished_at = ?, lease_until = NULL, retry_at = NULL,
+	last_error = NULL, last_success_watermark = input_watermark
+WHERE kind = ? AND job_key = ? AND ownership_token = ?
+`).run(nowSec, STAGE1_KIND, threadId, ownershipToken);
+
+		db.prepare("DELETE FROM stage1_outputs WHERE thread_id = ?").run(threadId);
+		enqueueGlobalWatermark(db, sourceUpdatedAt, { forceDirtyWhenNotAdvanced: true });
+		return true;
+	});
+	return tx() as boolean;
+}
+
+export function markStage1Failed(
+	db: Database,
+	params: { threadId: string; ownershipToken: string; retryDelaySeconds: number; reason: string; nowSec: number },
+): boolean {
+	const { threadId, ownershipToken, retryDelaySeconds, reason, nowSec } = params;
+	const result = db
+		.prepare(`
+UPDATE jobs
+SET status = 'error', finished_at = ?, lease_until = NULL, retry_at = ?,
+	retry_remaining = CASE WHEN retry_remaining > 0 THEN retry_remaining - 1 ELSE 0 END,
+	last_error = ?
+WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ?
+`)
+		.run(nowSec, nowSec + retryDelaySeconds, reason, STAGE1_KIND, threadId, ownershipToken);
+	return Number(result.changes ?? 0) > 0;
+}
+
+export function tryClaimGlobalPhase2Job(
+	db: Database,
+	params: { workerId: string; leaseSeconds: number; nowSec: number },
+): { kind: "claimed"; claim: GlobalClaim } | { kind: "skipped_not_dirty" } | { kind: "skipped_running" } {
+	const { workerId, leaseSeconds, nowSec } = params;
+	ensureGlobalJob(db);
+	const pre = db
+		.prepare(
+			"SELECT status, lease_until, input_watermark, last_success_watermark, retry_at, retry_remaining FROM jobs WHERE kind = ? AND job_key = ?",
+		)
+		.get(GLOBAL_KIND, GLOBAL_KEY) as
+		| {
+				status: string;
+				lease_until: number | null;
+				input_watermark: number | null;
+				last_success_watermark: number | null;
+				retry_at: number | null;
+				retry_remaining: number;
+		  }
+		| undefined;
+	if (!pre) return { kind: "skipped_not_dirty" };
+	const ownershipToken = crypto.randomUUID();
+	const claimed = db
+		.prepare(`
+UPDATE jobs
+SET status = 'running', worker_id = ?, ownership_token = ?, started_at = ?, finished_at = NULL,
+	lease_until = ?, retry_at = NULL, last_error = NULL
+WHERE kind = ? AND job_key = ?
+	AND NOT (status = 'running' AND lease_until IS NOT NULL AND lease_until > ?)
+	AND (input_watermark IS NOT NULL AND (last_success_watermark IS NULL OR input_watermark > last_success_watermark))
+	AND retry_remaining > 0
+	AND (retry_at IS NULL OR retry_at <= ?)
+`)
+		.run(workerId, ownershipToken, nowSec, nowSec + leaseSeconds, GLOBAL_KIND, GLOBAL_KEY, nowSec, nowSec);
+	if (Number(claimed.changes ?? 0) > 0) {
+		const row = db
+			.prepare("SELECT input_watermark FROM jobs WHERE kind = ? AND job_key = ? AND ownership_token = ?")
+			.get(GLOBAL_KIND, GLOBAL_KEY, ownershipToken) as { input_watermark: number | null } | undefined;
+		return {
+			kind: "claimed",
+			claim: {
+				ownershipToken,
+				inputWatermark: row?.input_watermark ?? 0,
+			},
+		};
+	}
+
+	if (pre.status === "running" && pre.lease_until !== null && pre.lease_until > nowSec) {
+		return { kind: "skipped_running" };
+	}
+	const preInput = pre.input_watermark ?? 0;
+	const preSuccess = pre.last_success_watermark ?? 0;
+	if (preInput <= preSuccess) {
+		return { kind: "skipped_not_dirty" };
+	}
+	if (pre.retry_remaining <= 0) {
+		return { kind: "skipped_not_dirty" };
+	}
+	if (pre.retry_at !== null && pre.retry_at > nowSec) {
+		return { kind: "skipped_not_dirty" };
+	}
+
+	const post = db
+		.prepare(
+			"SELECT status, lease_until, input_watermark, last_success_watermark, retry_at, retry_remaining FROM jobs WHERE kind = ? AND job_key = ?",
+		)
+		.get(GLOBAL_KIND, GLOBAL_KEY) as
+		| {
+				status: string;
+				lease_until: number | null;
+				input_watermark: number | null;
+				last_success_watermark: number | null;
+				retry_at: number | null;
+				retry_remaining: number;
+		  }
+		| undefined;
+	if (!post) return { kind: "skipped_not_dirty" };
+	if (post.status === "running" && post.lease_until !== null && post.lease_until > nowSec) {
+		return { kind: "skipped_running" };
+	}
+
+	return { kind: "skipped_not_dirty" };
+}
+
+export function heartbeatGlobalJob(
+	db: Database,
+	params: { ownershipToken: string; leaseSeconds: number; nowSec: number },
+): boolean {
+	const { ownershipToken, leaseSeconds, nowSec } = params;
+	const result = db
+		.prepare(`
+UPDATE jobs
+SET lease_until = ?
+WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ?
+`)
+		.run(nowSec + leaseSeconds, GLOBAL_KIND, GLOBAL_KEY, ownershipToken);
+	return Number(result.changes ?? 0) > 0;
+}
+
+export function listStage1OutputsForGlobal(db: Database, limit: number): Stage1OutputRow[] {
+	const rows = db
+		.prepare(`
+SELECT o.thread_id, o.source_updated_at, o.raw_memory, o.rollout_summary, o.rollout_slug, o.generated_at, t.cwd
+FROM stage1_outputs o
+LEFT JOIN threads t ON t.id = o.thread_id
+WHERE TRIM(COALESCE(o.raw_memory, '')) != '' OR TRIM(COALESCE(o.rollout_summary, '')) != ''
+ORDER BY o.source_updated_at DESC
+LIMIT ?
+`)
+		.all(limit) as Array<{
+		thread_id: string;
+		source_updated_at: number;
+		raw_memory: string;
+		rollout_summary: string;
+		rollout_slug: string | null;
+		generated_at: number;
+		cwd: string | null;
+	}>;
+	return rows.map(row => ({
+		threadId: row.thread_id,
+		sourceUpdatedAt: row.source_updated_at,
+		rawMemory: row.raw_memory,
+		rolloutSummary: row.rollout_summary,
+		rolloutSlug: row.rollout_slug,
+		generatedAt: row.generated_at,
+		cwd: row.cwd ?? "",
+	}));
+}
+
+export function markGlobalPhase2Succeeded(
+	db: Database,
+	params: { ownershipToken: string; newWatermark: number; nowSec: number },
+): boolean {
+	const { ownershipToken, newWatermark, nowSec } = params;
+	const result = db
+		.prepare(`
+UPDATE jobs
+SET status = 'done', finished_at = ?, lease_until = NULL, retry_at = NULL,
+	last_error = NULL,
+	last_success_watermark = CASE
+		WHEN last_success_watermark IS NULL THEN ?
+		WHEN last_success_watermark < ? THEN ?
+		ELSE last_success_watermark
+	END
+WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ?
+`)
+		.run(nowSec, newWatermark, newWatermark, newWatermark, GLOBAL_KIND, GLOBAL_KEY, ownershipToken);
+	return Number(result.changes ?? 0) > 0;
+}
+
+export function markGlobalPhase2Failed(
+	db: Database,
+	params: { ownershipToken: string; retryDelaySeconds: number; reason: string; nowSec: number },
+): boolean {
+	const { ownershipToken, retryDelaySeconds, reason, nowSec } = params;
+	const result = db
+		.prepare(`
+UPDATE jobs
+SET status = 'error', finished_at = ?, lease_until = NULL, retry_at = ?,
+	retry_remaining = CASE WHEN retry_remaining > 0 THEN retry_remaining - 1 ELSE 0 END,
+	last_error = ?
+WHERE kind = ? AND job_key = ? AND status = 'running' AND ownership_token = ?
+`)
+		.run(nowSec, nowSec + retryDelaySeconds, reason, GLOBAL_KIND, GLOBAL_KEY, ownershipToken);
+	return Number(result.changes ?? 0) > 0;
+}
+
+export function markGlobalPhase2FailedUnowned(
+	db: Database,
+	params: { retryDelaySeconds: number; reason: string; nowSec: number },
+): boolean {
+	const { retryDelaySeconds, reason, nowSec } = params;
+	const result = db
+		.prepare(`
+UPDATE jobs
+SET status = 'error', finished_at = ?, lease_until = NULL, retry_at = ?,
+	retry_remaining = CASE WHEN retry_remaining > 0 THEN retry_remaining - 1 ELSE 0 END,
+	last_error = ?
+WHERE kind = ? AND job_key = ? AND status = 'running'
+	AND (ownership_token IS NULL OR lease_until IS NULL OR lease_until <= ?)
+`)
+		.run(nowSec, nowSec + retryDelaySeconds, reason, GLOBAL_KIND, GLOBAL_KEY, nowSec);
+	return Number(result.changes ?? 0) > 0;
+}