@jini-ai/chat 0.6.5 → 0.6.7

package/dist/store/runtime/durable-run-store.js

@@ -1,25 +1,33 @@
 import { recoveredRunEvents } from "../../core/index.js";
 import { chatKernelTables } from "../sql/boundary.js";
+import { sql } from "kysely";
 import { hasUnknownToolCall, CONTINUATION_DIVIDER } from "../../core/durable-runs/continuation.js";
 import { RECOVERY_WINDOW_MS } from "../../core/durable-runs/policies.js";
 import { createChatHistoryStore } from "./chat-history-store.js";
 import { runLockKey } from "../../core/durable-runs/run-lock.js";
 import { runContentFromEvents } from "../../core/run-events/entry.js";
+import { toMessage } from "../sql/store.js";
 export const RUN_SLOT_BUSY = "Another answer in this chat is still running. This turn is queued in your browser.";
 export class RunSlotBusyError extends Error {
     constructor() { super(RUN_SLOT_BUSY); }
 }
+function isAcceptedUserShape({ message }, _optional = {}) {
+    return Boolean(message) && typeof message.id === "string" && Boolean(message.id) && message.role === "user" && typeof message.content === "string"
+        && (message.createdAt === undefined || Number.isFinite(message.createdAt));
+}
+function areAcceptedAttachments({ attachments }, _optional = {}) {
+    return Array.isArray(attachments) && attachments.every((attachment) => attachment && typeof attachment.path === "string" && typeof attachment.name === "string"
+        && (attachment.kind === "file" || attachment.kind === "image"));
+}
 function acceptedUserMessage({ request, now }, _optional = {}) {
     const context = JSON.parse(request.contextRef);
     // Older clients supplied only a prompt; never invent a user identity from that transcript.
     const message = context.userMessage;
     if (message === undefined)
         return null;
-    if (!message || typeof message.id !== "string" || !message.id || message.role !== "user" || typeof message.content !== "string"
-        || (message.createdAt !== undefined && !Number.isFinite(message.createdAt)))
+    if (!isAcceptedUserShape({ message }, {}))
         throw new Error("Invalid accepted user message");
-    if (message.attachments !== undefined && (!Array.isArray(message.attachments) || !message.attachments.every((attachment) => attachment && typeof attachment.path === "string" && typeof attachment.name === "string"
-        && (attachment.kind === "file" || attachment.kind === "image"))))
+    if (message.attachments !== undefined && !areAcceptedAttachments({ attachments: message.attachments }, {}))
         throw new Error("Invalid accepted user attachments");
     // User input cannot claim an assistant run's status/events. Keep the same identity and payload
     // as its browser PUT, including the original clock and attachment order.
@@ -27,9 +35,10 @@ function acceptedUserMessage({ request, now }, _optional = {}) {
         ...(message.attachments ? { attachments: message.attachments } : {}) };
 }
 function recoveryFields({ row, metadata }, _optional = {}) {
-    return { recoveryCount: metadata?.recovery_count ?? 0, recoveryDeadline: metadata?.recovery_deadline ?? null,
-        attemptStartedAt: metadata?.attempt_started_at ?? 0, lastProgressAt: metadata?.last_progress_at ?? row.started_at ?? row.created_at,
-        cancelReason: metadata?.cancel_reason ?? null };
+    const saved = metadata ?? {};
+    return { recoveryCount: saved.recovery_count ?? 0, recoveryDeadline: saved.recovery_deadline ?? null,
+        attemptStartedAt: saved.attempt_started_at ?? 0, lastProgressAt: saved.last_progress_at ?? row.started_at ?? row.created_at,
+        cancelReason: saved.cancel_reason ?? null };
 }
 function sessionFields({ metadata }, _optional = {}) {
     return { sessionId: metadata?.session_id ?? null, sessionConfirmed: metadata?.session_confirmed === 1,
@@ -42,6 +51,16 @@ function acceptedRequest({ row, metadata, transcript }, _optional = {}) {
     return { agentId: row.agent_id ?? "claude", contextRef: JSON.stringify({ conversationId: row.conversation_id,
             prompt: transcript.map((item) => `${item.role}: ${item.content}`).join("\n\n") }) };
 }
+function runHead({ row, runId, metadata, message, transcript }, { isDaemonRunId }) {
+    return {
+        messageId: row.id, conversationId: row.conversation_id, runId,
+        workspaceId: row.scope_id,
+        principalId: row.owner_kind === "user" ? row.owner_id : null,
+        engine: metadata?.engine ?? (isDaemonRunId({ runId }) ? "daemon" : "request"), message,
+        request: acceptedRequest({ row, metadata, transcript }, {}),
+        ...recoveryFields({ row, metadata }, {}), ...sessionFields({ metadata }, {}),
+    };
+}
 function metadataForRun({ run }, _optional = {}) {
     return { message_id: run.messageId, engine: run.engine, accepted_json: JSON.stringify(run.request),
         recovery_count: run.recoveryCount, recovery_deadline: run.recoveryDeadline, recovery_elapsed_ms: 0,
@@ -61,33 +80,31 @@ export function savedRunEvents({ json }, _optional = {}) {
         return [];
     }
 }
-function continuationEvents({ current, incoming }, _optional = {}) {
-    let separator = -1;
-    // Scan backward to keep the latest separator without requiring ES2023 array methods.
-    for (let index = incoming.length - 1; index >= 0; index--) {
-        const event = incoming[index];
-        if (event.kind === "text" && event.text === CONTINUATION_DIVIDER) {
-            separator = index;
-            break;
-        }
-    }
+/** Index of the latest matching event, or -1. */
+function lastIndexOf({ events, matches }, _optional = {}) {
+    // Scan backward to keep the latest match without requiring ES2023 array methods.
+    for (let index = events.length - 1; index >= 0; index--)
+        if (matches(events[index]))
+            return index;
+    return -1;
+}
+/** Saved projection, with legacy content restored as text when no text event carries it. */
+function savedProjection({ current }, _optional = {}) {
     const saved = [...(current.message.events ?? [])];
     if (!saved.some((event) => event.kind === "text") && current.message.content)
         saved.unshift({ kind: "text", text: current.message.content });
+    return saved;
+}
+function continuationEvents({ current, incoming }, _optional = {}) {
+    const separator = lastIndexOf({ events: incoming, matches: (event) => event.kind === "text" && event.text === CONTINUATION_DIVIDER });
+    const saved = savedProjection({ current });
     // A trailing checkpoint/tool barrier may land after the recovery probe but before its CAS.
     // Build the continuation segment from the projection inside the transaction, never that probe.
     if (separator >= 0)
         return [...saved, ...incoming.slice(separator)];
     // Empty startup retries have a recovery marker but no divider. Rebase their saved work under
     // the same lock too: a checkpoint may have arrived since the recovery probe.
-    let marker = -1;
-    for (let index = incoming.length - 1; index >= 0; index--) {
-        const event = incoming[index];
-        if (event.kind === "status" && event.code === "run_recovering") {
-            marker = index;
-            break;
-        }
-    }
+    const marker = lastIndexOf({ events: incoming, matches: (event) => event.kind === "status" && event.code === "run_recovering" });
     if (marker >= 0) {
         const event = incoming[marker];
         const notice = event.kind === 'status' && event.label === '' ? 'silent' : 'visible';
@@ -95,87 +112,132 @@ function continuationEvents({ current, incoming }, _optional = {}) {
     }
     return runContentFromEvents({ events: incoming }).length >= current.message.content.length ? incoming : saved;
 }
+/** `events_json` with one more event, built in SQL: no dialect parses or re-serializes the saved
+ * array in JS. Writers store `JSON.stringify` output (no whitespace), so a non-empty array is
+ * `[...]` and gains `,event`; `[]`, NULL and anything not array-shaped (a truncated legacy value,
+ * which the old JS read also treated as no events) become `[event]`. Plain `substr`/`length`/`||`
+ * run unchanged on SQLite and PostgreSQL. */
+function appendedEventsJson({ event }, _optional = {}) {
+    const item = sql `cast(${JSON.stringify(event)} as text)`;
+    return sql `case when events_json like '[_%]' then substr(events_json, 1, length(events_json) - 1) || ',' || ${item} || ']' else '[' || ${item} || ']' end`;
+}
+function sameOwner({ run, required }, _optional = {}) {
+    return run.principalId === required.principalId && run.conversationId === required.conversationId;
+}
 /** Same kernel and message lock as checkpoint/settle. Claim changes run_id before any spawn;
  * old attempts cannot write projections or dispatch tools after that transaction commits. */
 export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _optional = {}) {
     const kernel = chatKernelTables({ kernel: hostKernel }, {});
-    async function loadSnapshot({ messageId }, _options = {}) {
-        const row = await kernel.run((db) => db.selectFrom("ai_chat_messages")
+    /** The run's own message row, with its events only when asked: never another message's. */
+    async function runRow({ match }, { events = false } = {}) {
+        const eventsColumn = events ? sql `${sql.ref("ai_chat_messages.events_json")}` : sql `null`;
+        return kernel.run((db) => db.selectFrom("ai_chat_messages")
             .innerJoin("ai_chats", "ai_chats.id", "ai_chat_messages.conversation_id")
-            .selectAll("ai_chat_messages").select(["owner_kind", "owner_id", "scope_id"])
-            .where("ai_chat_messages.id", "=", messageId).where("role", "=", "assistant").executeTakeFirst());
+            .select(["ai_chat_messages.id", "conversation_id", "role", "content", "agent_id", "agent_name", "attachments_json",
+            "run_id", "run_status", "position", "ai_chat_messages.created_at", "started_at", "ended_at", "owner_kind", "owner_id", "scope_id"])
+            .select(eventsColumn.as("events_json"))
+            .where(match.column, "=", match.value).where("role", "=", "assistant").executeTakeFirst());
+    }
+    /** Every message of the run's conversation; only `load` and legacy rows without attempt metadata pay this. */
+    async function conversationTranscript({ row }, _options = {}) {
+        const history = createChatHistoryStore({ kernel, isDaemonRunId, scope: { scopeId: row.scope_id, ownerKind: row.owner_kind, ownerId: row.owner_id } }, {});
+        return history.messages({ conversationId: row.conversation_id });
+    }
+    /** One run's head, plus its transcript when that was read (`[]` otherwise). */
+    async function readRun({ match }, options = {}) {
+        const row = await runRow({ match }, options);
         if (!row || !row.run_id)
             return null;
+        const messageId = row.id;
         const metadata = await kernel.run((db) => db.selectFrom("assistant_run_attempts").selectAll().where("message_id", "=", messageId).executeTakeFirst());
-        const history = createChatHistoryStore({ kernel, isDaemonRunId, scope: { scopeId: row.scope_id, ownerKind: row.owner_kind, ownerId: row.owner_id } }, {});
-        const transcript = await history.messages({ conversationId: row.conversation_id });
-        const message = transcript.find((item) => item.id === messageId);
+        // A legacy row's accepted request is derived from its transcript, so only it reads one here.
+        const transcript = options.transcript || !metadata ? await conversationTranscript({ row }) : [];
+        const message = options.transcript ? transcript.find((item) => item.id === messageId) : toMessage(row);
         if (!message)
             return null; // Deletion may have cascaded between the two reads.
-        const request = acceptedRequest({ row, metadata, transcript }, {});
-        return {
-            messageId, conversationId: row.conversation_id, runId: row.run_id,
-            workspaceId: row.scope_id,
-            principalId: row.owner_kind === "user" ? row.owner_id : null,
-            engine: metadata?.engine ?? (isDaemonRunId({ runId: row.run_id }) ? "daemon" : "request"), message, transcript, request,
-            ...recoveryFields({ row, metadata }, {}), ...sessionFields({ metadata }, {}),
-        };
+        return { head: runHead({ row, runId: row.run_id, metadata, message, transcript }, { isDaemonRunId }), transcript };
+    }
+    /** Locked read of one message; `transcript` decides whether it is a whole run or only its head. */
+    async function lockedRead({ messageId, conversationId }, options = {}) {
+        return kernel.transaction(async () => {
+            // Message, attempt metadata and transcript must describe the same generation. A CAS in
+            // the gap between independent reads must never pair a new projection with an old session.
+            await kernel.lockKey(runLockKey({ messageId, conversationId, runId: "" }));
+            return readRun({ match: { column: "ai_chat_messages.id", value: messageId } }, options);
+        });
     }
     async function load({ messageId }, _options = {}) {
         const binding = await kernel.run((db) => db.selectFrom("ai_chat_messages").select("conversation_id")
             .where("id", "=", messageId).executeTakeFirst());
         if (!binding)
             return null;
-        return kernel.transaction(async () => {
-            // Message, attempt metadata and transcript must describe the same generation. A CAS in
-            // the gap between independent reads must never pair a new projection with an old session.
-            await kernel.lockKey(runLockKey({ messageId, conversationId: binding.conversation_id, runId: "" }));
-            return loadSnapshot({ messageId }, {});
-        });
+        // The message comes from the transcript, so its row's events are not selected twice.
+        const read = await lockedRead({ messageId, conversationId: binding.conversation_id }, { transcript: true });
+        return read && { ...read.head, transcript: read.transcript };
+    }
+    /** Unlocked probe: the message and conversation an attempt id currently names. */
+    async function probeRun({ runId }, _options = {}) {
+        return kernel.run((db) => db.selectFrom("ai_chat_messages").select(["id", "conversation_id"])
+            .where("run_id", "=", runId).where("role", "=", "assistant").executeTakeFirst());
     }
     async function find({ runId, principalId }, _options = {}) {
-        const row = await kernel.run((db) => db.selectFrom("ai_chat_messages").select("id").where("run_id", "=", runId).where("role", "=", "assistant").executeTakeFirst());
+        const row = await probeRun({ runId });
         if (!row)
             return null;
-        const run = await load({ messageId: row.id }, {});
+        const run = (await lockedRead({ messageId: row.id, conversationId: row.conversation_id }, { events: true }))?.head ?? null;
         return principalId !== undefined && run?.principalId !== principalId ? null : run;
     }
+    /** Slot and run locks, then the owner-scoped conversation; null when the caller does not own it. */
+    async function lockOwnedConversation({ required }, _options = {}) {
+        await kernel.lockKey(`chat-slot:${required.conversationId}`);
+        await kernel.lockKey(runLockKey(required));
+        const history = createChatHistoryStore({ kernel, isDaemonRunId, scope: { scopeId: required.workspaceId, ownerKind: "user", ownerId: required.principalId } }, {});
+        return await history.get({ id: required.conversationId }) ? history : null;
+    }
+    async function assertSlotFree({ conversationId }, _options = {}) {
+        const occupied = await kernel.run((db) => db.selectFrom("ai_chat_messages").select("id")
+            .where("conversation_id", "=", conversationId).where("role", "=", "assistant")
+            .where("run_status", "in", ["queued", "running"]).executeTakeFirst());
+        if (occupied)
+            throw new RunSlotBusyError();
+    }
+    async function appendAcceptedQuestion({ history, required, userMessage }, _options = {}) {
+        if (userMessage.id === required.messageId)
+            throw new Error("User and assistant messages must have different IDs");
+        // A pending/failed browser PUT must never let the assistant take the first position.
+        // Nested appends join this transaction and retain the conversation lock through both
+        // writes, so either arrival order of the idempotent PUT yields user then assistant.
+        const saved = await history.appendMessage({ conversationId: required.conversationId, message: userMessage });
+        if (!saved || saved.role !== "user")
+            throw new Error("Accepted user message could not be persisted");
+    }
+    /** @returns false when the answer id collides and no question was written. */
+    async function appendQueuedAnswer({ history, required, hasQuestion }, _options = {}) {
+        const message = { id: required.messageId, role: "assistant", content: "", events: [],
+            runId: required.runId, runStatus: "queued", createdAt: required.now, startedAt: required.now,
+            ...(required.request.agentId ? { agentId: required.request.agentId } : {}) };
+        const savedMessage = await history.appendMessage({ conversationId: required.conversationId, message });
+        if (savedMessage?.role === "assistant")
+            return true;
+        // Returning null here would commit the question alone after an ID collision.
+        if (hasQuestion)
+            throw new Error("Accepted assistant message could not be persisted");
+        return false;
+    }
     async function accept(required, _options = {}) {
         return kernel.transaction(async () => {
-            await kernel.lockKey(`chat-slot:${required.conversationId}`);
-            await kernel.lockKey(runLockKey(required));
-            const history = createChatHistoryStore({ kernel, isDaemonRunId, scope: { scopeId: required.workspaceId, ownerKind: "user", ownerId: required.principalId } }, {});
-            if (!await history.get({ id: required.conversationId }))
+            const history = await lockOwnedConversation({ required });
+            if (!history)
                 return null;
             const existing = await load({ messageId: required.messageId }, {});
             if (existing)
-                return existing.principalId === required.principalId && existing.conversationId === required.conversationId ? existing : null;
-            const occupied = await kernel.run((db) => db.selectFrom("ai_chat_messages").select("id")
-                .where("conversation_id", "=", required.conversationId).where("role", "=", "assistant")
-                .where("run_status", "in", ["queued", "running"]).executeTakeFirst());
-            if (occupied)
-                throw new RunSlotBusyError();
+                return sameOwner({ run: existing, required }) ? existing : null;
+            await assertSlotFree({ conversationId: required.conversationId });
             const userMessage = acceptedUserMessage({ request: required.request, now: required.now }, {});
-            if (userMessage) {
-                if (userMessage.id === required.messageId)
-                    throw new Error("User and assistant messages must have different IDs");
-                // A pending/failed browser PUT must never let the assistant take the first position.
-                // Nested appends join this transaction and retain the conversation lock through both
-                // writes, so either arrival order of the idempotent PUT yields user then assistant.
-                const saved = await history.appendMessage({ conversationId: required.conversationId, message: userMessage });
-                if (!saved || saved.role !== "user")
-                    throw new Error("Accepted user message could not be persisted");
-            }
-            const message = { id: required.messageId, role: "assistant", content: "", events: [],
-                runId: required.runId, runStatus: "queued", createdAt: required.now, startedAt: required.now,
-                ...(required.request.agentId ? { agentId: required.request.agentId } : {}) };
-            const savedMessage = await history.appendMessage({ conversationId: required.conversationId, message });
-            if (!savedMessage || savedMessage.role !== "assistant") {
-                // Returning null here would commit the question alone after an ID collision.
-                if (userMessage)
-                    throw new Error("Accepted assistant message could not be persisted");
+            if (userMessage)
+                await appendAcceptedQuestion({ history, required, userMessage });
+            if (!await appendQueuedAnswer({ history, required, hasQuestion: userMessage !== null }))
                 return null;
-            }
             await kernel.run((db) => db.insertInto("assistant_run_attempts").values({
                 message_id: required.messageId, engine: "daemon", accepted_json: JSON.stringify(required.request),
                 recovery_count: 0, recovery_deadline: null, last_progress_at: required.now, cancel_reason: null,
@@ -188,7 +250,8 @@ export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _op
     async function advance({ run, nextRunId, now, events }, _options = {}) {
         return kernel.transaction(async () => {
             await kernel.lockKey(runLockKey(run));
-            const current = await load({ messageId: run.messageId }, {});
+            // The projection rebase needs this run's events, never the rest of its conversation.
+            const current = (await lockedRead(run, { events: true }))?.head;
             if (!current || current.cancelReason || current.runId !== run.runId)
                 return false;
             const projection = continuationEvents({ current, incoming: events }, {});
@@ -212,28 +275,33 @@ export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _op
             return true;
         });
     }
+    /** Probe without the run lock, take it, then read: a run deleted or advanced in the gap reads as gone.
+     * One read under the lock serves the whole operation; the unlocked probe only names the lock.
+     * @returns The run as of the lock, or null when the attempt id is unknown or no longer current. */
+    async function lockedRun({ runId }, options = {}) {
+        const probe = await probeRun({ runId });
+        if (!probe)
+            return null;
+        await kernel.lockKey(runLockKey({ messageId: probe.id, conversationId: probe.conversation_id, runId }));
+        return (await readRun({ match: { column: "ai_chat_messages.run_id", value: runId } }, options))?.head ?? null;
+    }
     async function cancel({ runId, reason }, _options = {}) {
         await kernel.transaction(async () => {
-            const run = await find({ runId }, {});
+            const run = await lockedRun({ runId });
             if (!run)
                 return;
-            await kernel.lockKey(runLockKey(run));
-            if (!await find({ runId }, {}))
-                return;
             await kernel.run((db) => db.insertInto("assistant_run_attempts").values({ ...metadataForRun({ run }, {}), cancel_reason: reason })
                 .onConflict((oc) => oc.column("message_id").doUpdateSet({ cancel_reason: reason })).execute());
         });
     }
     async function cancelPending(required, _options = {}) {
         return kernel.transaction(async () => {
-            await kernel.lockKey(`chat-slot:${required.conversationId}`);
-            await kernel.lockKey(runLockKey(required));
-            const history = createChatHistoryStore({ kernel, isDaemonRunId, scope: { scopeId: required.workspaceId, ownerKind: "user", ownerId: required.principalId } }, {});
-            if (!await history.get({ id: required.conversationId }))
+            const history = await lockOwnedConversation({ required });
+            if (!history)
                 return null;
             const existing = await load({ messageId: required.messageId }, {});
             if (existing) {
-                if (existing.principalId !== required.principalId || existing.conversationId !== required.conversationId)
+                if (!sameOwner({ run: existing, required }))
                     return null;
                 await cancel({ runId: existing.runId, reason: "user-stop" }, {});
                 return load({ messageId: existing.messageId }, {});
@@ -247,15 +315,11 @@ export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _op
     }
     async function captureSession({ runId, sessionId, child, confirmed = true }, _options = {}) {
         return kernel.transaction(async () => {
-            const run = await find({ runId }, {});
-            if (!run)
-                return false;
-            await kernel.lockKey(runLockKey(run));
-            const current = await find({ runId }, {});
-            if (!attemptCanExecute({ run: current }, {}) || !current)
+            const run = await lockedRun({ runId });
+            if (!run || !attemptCanExecute({ run }, {}))
                 return false;
             const fields = { session_id: sessionId, session_confirmed: confirmed ? 1 : 0, ...(child ? { child_pid: child.pid, child_started_at: child.startedAt } : {}) };
-            await kernel.run((db) => db.insertInto("assistant_run_attempts").values({ ...metadataForRun({ run: current }, {}), ...fields })
+            await kernel.run((db) => db.insertInto("assistant_run_attempts").values({ ...metadataForRun({ run }, {}), ...fields })
                 .onConflict((oc) => oc.column("message_id").doUpdateSet(fields)).execute());
             // Session and generation share the same lock/transaction; an old callback cannot write
             // the conversation's locator in the gap between a successful guard and a later CAS.
@@ -267,11 +331,7 @@ export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _op
     }
     async function clearSession({ runId, sessionId }, _options = {}) {
         await kernel.transaction(async () => {
-            const before = await find({ runId }, {});
-            if (!before)
-                return;
-            await kernel.lockKey(runLockKey(before));
-            const run = await find({ runId }, {});
+            const run = await lockedRun({ runId });
             if (!run)
                 return;
             await kernel.run((db) => db.deleteFrom("assistant_agent_sessions").where("conversation_id", "=", run.conversationId)
@@ -282,14 +342,9 @@ export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _op
     }
     async function guardTool(required, _options = {}) {
         return kernel.transaction(async () => {
-            const before = await find({ runId: required.runId }, {});
-            if (!before)
-                return "stale";
-            await kernel.lockKey(runLockKey(before));
-            const run = await find({ runId: required.runId }, {});
-            if (!attemptCanExecute({ run }, {}))
-                return "stale";
-            if (!run)
+            // The unknown-call check needs this run's saved events (one row), never other messages'.
+            const run = await lockedRun({ runId: required.runId }, { events: true });
+            if (!run || !attemptCanExecute({ run }, {}))
                 return "stale";
             const events = run.message.events ?? [];
             const canonical = events.filter((event) => event.kind !== "tool_use" || event.name === required.toolId);
@@ -301,7 +356,7 @@ export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _op
         });
     }
     async function writeToolEvent(run, event) {
-        await kernel.run((db) => db.updateTable("ai_chat_messages").set({ events_json: JSON.stringify([...(run.message.events ?? []), event]) })
+        await kernel.run((db) => db.updateTable("ai_chat_messages").set({ events_json: appendedEventsJson({ event }) })
             .where("id", "=", run.messageId).where("run_id", "=", run.runId).where("run_status", "in", ["queued", "running"]).execute());
         await kernel.run((db) => db.updateTable("assistant_run_attempts")
             .set({ recovery_count: 0, recovery_deadline: null, recovery_elapsed_ms: 0, last_progress_at: Date.now() })
@@ -309,23 +364,16 @@ export function createDurableRunStore({ kernel: hostKernel, isDaemonRunId }, _op
     }
     async function completeTool({ runId, toolUseId, content }, _options = {}) {
         await kernel.transaction(async () => {
-            const before = await find({ runId }, {});
-            if (!before)
-                return;
-            await kernel.lockKey(runLockKey(before));
-            const run = await find({ runId }, {});
+            const run = await lockedRun({ runId });
             if (run)
                 await writeToolEvent(run, { kind: "tool_result", toolUseId, content, isError: false });
         });
     }
     async function recoveryClock({ runId, now, active }, _options = {}) {
         await kernel.transaction(async () => {
-            const run = await find({ runId }, {});
+            const run = await lockedRun({ runId });
             if (!run)
                 return;
-            await kernel.lockKey(runLockKey(run));
-            if (!await find({ runId }, {}))
-                return;
             const row = await kernel.run((db) => db.selectFrom("assistant_run_attempts").selectAll().where("message_id", "=", run.messageId).executeTakeFirst())
                 ?? metadataForRun({ run }, {});
             let deadline = row.recovery_deadline;