@axiom-lattice/local-stores 3.1.24 → 3.2.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/local-stores",
-  "version": "3.1.24",
+  "version": "3.2.1",
   "description": "Local SQLite-based stores for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -25,8 +25,8 @@
     "@types/sql.js": "^1.4.11",
     "sql.js": "^1.14.1",
     "uuid": "^14.0.1",
-    "@axiom-lattice/core": "6.2.0",
-    "@axiom-lattice/protocols": "4.6.1"
+    "@axiom-lattice/core": "6.3.1",
+    "@axiom-lattice/protocols": "4.7.0"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/LocalScheduleStorage.test.ts

@@ -0,0 +1,143 @@
+import { afterEach, describe, expect, it } from "@jest/globals";
+import { tmpdir } from "os";
+import path from "path";
+
+import {
+  ScheduleExecutionType,
+  ScheduledTaskDefinition,
+  ScheduledTaskStatus,
+} from "@axiom-lattice/protocols";
+
+import { closeDatabase, initDatabase } from "../database";
+import { LocalScheduleStorage } from "../stores/LocalScheduleStorage";
+
+const nextDbPath = (): string =>
+  path.join(
+    tmpdir(),
+    `lt-schedule-${Date.now()}-${Math.random().toString(36).slice(2)}.db`
+  );
+
+const buildTask = (
+  taskId: string,
+  overrides: Partial<ScheduledTaskDefinition> = {}
+): ScheduledTaskDefinition => {
+  const now = Date.now();
+  return {
+    taskId,
+    taskType: "test",
+    payload: {},
+    tenantId: "tenant-1",
+    executionType: ScheduleExecutionType.ONCE,
+    executeAt: now - 1000,
+    status: ScheduledTaskStatus.PENDING,
+    runCount: 0,
+    retryCount: 0,
+    maxRetries: 0,
+    createdAt: now,
+    updatedAt: now,
+    ...overrides,
+  };
+};
+
+const createStore = async (): Promise<{
+  db: Awaited<ReturnType<typeof initDatabase>>;
+  store: LocalScheduleStorage;
+}> => {
+  const db = await initDatabase({ dbPath: nextDbPath() });
+  return { db, store: new LocalScheduleStorage(db) };
+};
+
+describe("LocalScheduleStorage.claimScheduledRun", () => {
+  afterEach(() => closeDatabase());
+
+  it("claims a pending task exactly once", async () => {
+    const { store } = await createStore();
+    await store.save(buildTask("task-1"));
+
+    expect(await store.claimScheduledRun("task-1")).toBe(true);
+    expect(await store.claimScheduledRun("task-1")).toBe(false);
+
+    const task = await store.get("task-1");
+    expect(task?.status).toBe(ScheduledTaskStatus.RUNNING);
+  });
+
+  it("returns false for an unknown task", async () => {
+    const { store } = await createStore();
+    expect(await store.claimScheduledRun("missing")).toBe(false);
+  });
+
+  it("returns false when the task is not pending", async () => {
+    const { store } = await createStore();
+    await store.save(buildTask("task-1"));
+    await store.update("task-1", { status: ScheduledTaskStatus.COMPLETED });
+
+    expect(await store.claimScheduledRun("task-1")).toBe(false);
+  });
+
+  it("can claim again after a retry re-arms the task to pending", async () => {
+    const { store } = await createStore();
+    await store.save(buildTask("task-1"));
+
+    expect(await store.claimScheduledRun("task-1")).toBe(true);
+    await store.update("task-1", { status: ScheduledTaskStatus.PENDING });
+    expect(await store.claimScheduledRun("task-1")).toBe(true);
+  });
+});
+
+describe("LocalScheduleStorage.requeueStaleRunning", () => {
+  afterEach(() => closeDatabase());
+
+  it("requeues a running task older than the threshold", async () => {
+    const { store } = await createStore();
+    await store.save(
+      buildTask("task-stale", {
+        status: ScheduledTaskStatus.RUNNING,
+        updatedAt: Date.now() - 60_000,
+      })
+    );
+
+    expect(await store.requeueStaleRunning(30_000)).toBe(1);
+
+    const task = await store.get("task-stale");
+    expect(task?.status).toBe(ScheduledTaskStatus.PENDING);
+  });
+
+  it("leaves a recently updated running task alone", async () => {
+    const { store } = await createStore();
+    await store.save(
+      buildTask("task-fresh", {
+        status: ScheduledTaskStatus.RUNNING,
+        updatedAt: Date.now() - 1_000,
+      })
+    );
+
+    expect(await store.requeueStaleRunning(30_000)).toBe(0);
+
+    const task = await store.get("task-fresh");
+    expect(task?.status).toBe(ScheduledTaskStatus.RUNNING);
+  });
+
+  it("never touches pending or terminal tasks", async () => {
+    const { store } = await createStore();
+    await store.save(
+      buildTask("pending", {
+        status: ScheduledTaskStatus.PENDING,
+        updatedAt: Date.now() - 60_000,
+      })
+    );
+    await store.save(
+      buildTask("completed", {
+        status: ScheduledTaskStatus.COMPLETED,
+        updatedAt: Date.now() - 60_000,
+      })
+    );
+
+    expect(await store.requeueStaleRunning(30_000)).toBe(0);
+    expect((await store.get("pending"))?.status).toBe(
+      ScheduledTaskStatus.PENDING
+    );
+    expect((await store.get("completed"))?.status).toBe(
+      ScheduledTaskStatus.COMPLETED
+    );
+  });
+});

package/src/__tests__/LocalThreadMessageQueueStore.test.ts

@@ -198,4 +198,44 @@ describe("LocalThreadMessageQueueStore scope isolation", () => {
 
     await expect(store.getPendingMessages("thread")).rejects.toThrow("Invalid trusted agent run context");
   });
+
+  describe("claimPendingMessages / reclaimExpiredProcessing", () => {
+    it("claims pending messages atomically and marks them processing", async () => {
+      const message = await store.addMessage({
+        threadId: "claim", tenantId: "tenant", assistantId: "assistant", content: "hi",
+      });
+
+      const claimed = await store.claimPendingMessages("claim", undefined, { owner: "r1", ttlMs: 1000 });
+
+      expect(claimed.map((m) => m.id)).toEqual([message.id]);
+      expect(await store.getPendingMessages("claim")).toHaveLength(0);
+      expect(await store.getProcessingMessages("claim")).toHaveLength(1);
+    });
+
+    it("gives sequential claimers disjoint sets bounded by the limit", async () => {
+      for (let i = 0; i < 4; i += 1) {
+        await store.addMessage({ threadId: "claim-disjoint", tenantId: "tenant", assistantId: "assistant", content: `hi${i}`, id: `m${i}` });
+      }
+
+      const a = await store.claimPendingMessages("claim-disjoint", undefined, { owner: "r1", ttlMs: 1000, limit: 2 });
+      const b = await store.claimPendingMessages("claim-disjoint", undefined, { owner: "r2", ttlMs: 1000, limit: 2 });
+
+      expect(a).toHaveLength(2);
+      expect(b).toHaveLength(2);
+      expect(new Set([...a, ...b].map((m) => m.id)).size).toBe(4);
+      expect(await store.getPendingMessages("claim-disjoint")).toHaveLength(0);
+    });
+
+    it("reclaims only expired processing rows", async () => {
+      await store.addMessage({ threadId: "lease-live", tenantId: "tenant", assistantId: "assistant", content: "live", id: "live" });
+      await store.addMessage({ threadId: "lease-expired", tenantId: "tenant", assistantId: "assistant", content: "expired", id: "expired" });
+      await store.claimPendingMessages("lease-live", undefined, { owner: "r1", ttlMs: 60000 });
+      await store.claimPendingMessages("lease-expired", undefined, { owner: "r1", ttlMs: 1 });
+      await new Promise((resolve) => setTimeout(resolve, 30));
+
+      expect(await store.reclaimExpiredProcessing("lease-live")).toBe(0);
+      expect(await store.reclaimExpiredProcessing("lease-expired")).toBe(1);
+      expect(await store.getPendingMessages("lease-expired")).toHaveLength(1);
+    });
+  });
 });

package/src/stores/LocalScheduleStorage.ts

@@ -154,6 +154,29 @@ export class LocalScheduleStorage implements ScheduleStorage {
     ).run(...values);
   }
 
+  async claimScheduledRun(taskId: string): Promise<boolean> {
+    const result = this.db.prepare(
+      `UPDATE lt_scheduled_tasks SET status = 'running', updated_at = ? WHERE task_id = ? AND status = 'pending'`,
+    ).run(new Date().toISOString(), taskId);
+    return result.changes > 0;
+  }
+
+  /**
+   * Requeues `running` tasks whose `updated_at` is older than `olderThanMs` back
+   * to `pending` (recovery after a crash that left a claim unsettled).
+   *
+   * @param olderThanMs Age threshold in milliseconds.
+   * @returns The number of tasks requeued.
+   */
+  async requeueStaleRunning(olderThanMs: number): Promise<number> {
+    const now = new Date();
+    const cutoff = new Date(now.getTime() - olderThanMs).toISOString();
+    const result = this.db.prepare(
+      `UPDATE lt_scheduled_tasks SET status = 'pending', updated_at = ? WHERE status = 'running' AND updated_at < ?`,
+    ).run(now.toISOString(), cutoff);
+    return result.changes;
+  }
+
   async delete(taskId: string): Promise<void> {
     this.db.prepare(`DELETE FROM lt_scheduled_tasks WHERE task_id = ?`).run(taskId);
   }

package/src/stores/LocalThreadMessageQueueStore.ts

@@ -9,6 +9,7 @@ import type {
   AddMessageParams,
   ThreadInfo,
   QueueScope,
+  ClaimMessagesOptions,
 } from "@axiom-lattice/core";
 import {
   parseQueuedExecutionMode,
@@ -18,6 +19,7 @@ import {
 } from "@axiom-lattice/protocols";
 import { ensureTable, nowISO, parseISO } from "../database";
 import { randomUUID } from "crypto";
+import { transaction, write } from "./project-room-schema";
 
 const DDL = `
 CREATE TABLE IF NOT EXISTS lt_thread_message_queue (
@@ -36,6 +38,9 @@ CREATE TABLE IF NOT EXISTS lt_thread_message_queue (
   custom_run_config TEXT,
   trusted_run_context TEXT,
   execution_mode TEXT CHECK (execution_mode IS NULL OR execution_mode = 'followup'),
+  locked_by TEXT,
+  locked_at TEXT,
+  lease_expires_at TEXT,
   created_at TEXT NOT NULL
 );
 CREATE INDEX IF NOT EXISTS idx_lt_tmq_thread ON lt_thread_message_queue(thread_id, status);
@@ -58,6 +63,9 @@ interface QueueRow {
   custom_run_config: string | null;
   trusted_run_context: string | null;
   execution_mode: string | null;
+  locked_by: string | null;
+  locked_at: string | null;
+  lease_expires_at: string | null;
   created_at: string;
 }
 
@@ -76,6 +84,16 @@ export class LocalThreadMessageQueueStore implements IMessageQueueStore {
     if (!names.has("execution_mode")) {
       db.run("ALTER TABLE lt_thread_message_queue ADD COLUMN execution_mode TEXT CHECK (execution_mode IS NULL OR execution_mode = 'followup')");
     }
+    if (!names.has("locked_by")) {
+      db.run("ALTER TABLE lt_thread_message_queue ADD COLUMN locked_by TEXT");
+    }
+    if (!names.has("locked_at")) {
+      db.run("ALTER TABLE lt_thread_message_queue ADD COLUMN locked_at TEXT");
+    }
+    if (!names.has("lease_expires_at")) {
+      db.run("ALTER TABLE lt_thread_message_queue ADD COLUMN lease_expires_at TEXT");
+    }
+    db.run("CREATE INDEX IF NOT EXISTS idx_lt_tmq_lease ON lt_thread_message_queue(status, lease_expires_at)");
   }
 
   async addMessage(params: AddMessageParams): Promise<PendingMessage> {
@@ -189,6 +207,43 @@ export class LocalThreadMessageQueueStore implements IMessageQueueStore {
     return rows.map(rowToMessage);
   }
 
+  async claimPendingMessages(
+    threadId: string,
+    scope: QueueScope | undefined,
+    options: ClaimMessagesOptions,
+  ): Promise<PendingMessage[]> {
+    const filter = scopeClause(scope);
+    const limitClause = options.limit === undefined ? "" : ` LIMIT ${Math.max(0, Math.floor(options.limit))}`;
+    const claimedRows = transaction(this.db, () => {
+      const candidates = this.db.prepare(
+        `SELECT * FROM lt_thread_message_queue WHERE thread_id = ? AND status = 'pending'${filter.sql} ORDER BY priority DESC, sequence_order ASC${limitClause}`,
+      ).all(threadId, ...filter.params) as unknown as QueueRow[];
+      const startedAt = Date.now();
+      const lockedAt = new Date(startedAt).toISOString();
+      const leaseExpiresAt = new Date(startedAt + options.ttlMs).toISOString();
+      const claimed: QueueRow[] = [];
+      for (const row of candidates) {
+        const changes = write(
+          this.db,
+          `UPDATE lt_thread_message_queue SET status = 'processing', locked_by = ?, locked_at = ?, lease_expires_at = ? WHERE id = ? AND status = 'pending'`,
+          [options.owner, lockedAt, leaseExpiresAt, row.id],
+        );
+        if (changes > 0) claimed.push(row);
+      }
+      return claimed;
+    });
+    return claimedRows.map(rowToMessage);
+  }
+
+  async reclaimExpiredProcessing(threadId: string, scope?: QueueScope): Promise<number> {
+    const filter = scopeClause(scope);
+    return transaction(this.db, () => write(
+      this.db,
+      `UPDATE lt_thread_message_queue SET status = 'pending', locked_by = NULL, locked_at = NULL, lease_expires_at = NULL WHERE thread_id = ? AND status = 'processing' AND lease_expires_at IS NOT NULL AND lease_expires_at <= ?${filter.sql}`,
+      [threadId, nowISO(), ...(filter.params as Array<string | number | null>)],
+    ));
+  }
+
   async getQueueSize(threadId: string, scope?: QueueScope): Promise<number> {
     const filter = scopeClause(scope);
     const row = this.db.prepare(
@@ -235,7 +290,7 @@ export class LocalThreadMessageQueueStore implements IMessageQueueStore {
   async resetProcessingToPending(threadId: string, scope?: QueueScope): Promise<number> {
     const filter = scopeClause(scope);
     const result = this.db.prepare(
-      `UPDATE lt_thread_message_queue SET status = 'pending' WHERE thread_id = ? AND status = 'processing'${filter.sql}`,
+      `UPDATE lt_thread_message_queue SET status = 'pending', locked_by = NULL, locked_at = NULL, lease_expires_at = NULL WHERE thread_id = ? AND status = 'processing'${filter.sql}`,
     ).run(threadId, ...filter.params);
     return result.changes;
   }