Skip to main content

lash_sqlite_store/
persistence.rs

1//! The [`RuntimePersistence`] capability-segment implementations for
2//! [`Store`]: [`SessionCommitStore`], [`SessionExecutionLeaseStore`],
3//! [`QueuedWorkStore`], [`TurnInputStore`], and [`StoreMaintenance`].
4//!
5//! This is the tokio-rusqlite port of the prior store's `persistence.rs`. The
6//! public surface is byte-for-byte the prior store async trait: identical method
7//! names and signatures, so consumers swap backends with a path rename only.
8//!
9//! The translation rules (see `conn.rs`, `lifecycle.rs`, `blobs.rs`):
10//!
11//! * Pure reads run through `self.conn.call(move |conn| { ... })`.
12//! * Read-then-write paths run through `self.conn.write(move |tx| { ... })`
13//!   (`BEGIN IMMEDIATE`, commit on `Ok`, rollback on `Err`) — this is the
14//!   cross-process write-lock guard.
15//! * Paths that may abandon partially-applied writes (the queued-work claim)
16//!   run through `self.conn.write_flow`, deciding commit vs rollback via
17//!   [`TxOutcome`].
18//! * The shared `*_conn` helpers (`try_load_session_head_meta_from_conn`,
19//!   `Self::put_checkpoint_conn`, `Self::load_usage_deltas_conn`,
20//!   `Self::load_session_graph_from_conn`, the queued-work helpers, …) are
21//!   synchronous and take a `&rusqlite::Connection`, so they are reused from
22//!   inside these closures (a `&Transaction` derefs to `&Connection`).
23//! * Closures must be `'static` + `Send`: every borrow of `self`/caller data is
24//!   cloned into an owned value before being moved in.
25
26use super::*;
27
28#[async_trait::async_trait]
29impl SessionCommitStore for Store {
30    fn durability_tier(&self) -> DurabilityTier {
31        DurabilityTier::Durable
32    }
33
34    async fn load_session(
35        &self,
36        scope: SessionReadScope,
37    ) -> Result<Option<PersistedSessionRead>, StoreError> {
38        self.conn
39            .call(move |conn| {
40                let outcome: Result<Option<PersistedSessionRead>, StoreError> = (|| {
41                    let Some(meta) = try_load_session_head_meta_from_conn(conn)? else {
42                        return Ok(None);
43                    };
44                    let leaf_node_id = match &scope {
45                        SessionReadScope::FullGraph => meta.leaf_node_id.clone(),
46                        SessionReadScope::ActivePath { leaf_node_id } => {
47                            leaf_node_id.clone().or_else(|| meta.leaf_node_id.clone())
48                        }
49                    };
50                    let mut graph = match scope {
51                        SessionReadScope::FullGraph => {
52                            Self::load_session_graph_from_conn(conn, meta.leaf_node_id.clone())
53                        }
54                        SessionReadScope::ActivePath { .. } => {
55                            Self::load_active_path_session_graph_from_conn(
56                                conn,
57                                leaf_node_id.clone(),
58                            )
59                            .map_err(sqlite_error)?
60                        }
61                    };
62                    graph.set_leaf_node_id(leaf_node_id);
63                    let checkpoint = meta
64                        .checkpoint_ref
65                        .as_ref()
66                        .map(|blob_ref| Self::get_checkpoint_conn(conn, blob_ref))
67                        .transpose()?
68                        .flatten();
69                    Ok(Some(PersistedSessionRead {
70                        session_id: meta.session_id,
71                        head_revision: meta.head_revision,
72                        config: meta.config,
73                        agent_frames: meta.agent_frames,
74                        current_agent_frame_id: meta.current_agent_frame_id,
75                        graph,
76                        checkpoint_ref: meta.checkpoint_ref,
77                        checkpoint,
78                        token_ledger: merge_token_ledger_entries(Self::load_usage_deltas_conn(
79                            conn,
80                        )),
81                    }))
82                })(
83                );
84                Ok(outcome)
85            })
86            .await
87            .map_err(sqlite_error)?
88    }
89
90    async fn load_node(
91        &self,
92        node_id: &str,
93    ) -> Result<Option<lash_core::SessionNodeRecord>, StoreError> {
94        let node_id = node_id.to_string();
95        let row: Option<String> = self
96            .conn
97            .call(move |conn| {
98                conn.query_row(
99                    "SELECT node_json FROM graph_nodes WHERE node_id = ?1 AND tombstoned = 0",
100                    params![node_id],
101                    |row| row.get(0),
102                )
103                .optional()
104            })
105            .await
106            .map_err(sqlite_error)?;
107        Ok(row.and_then(|json| serde_json::from_str(&json).ok()))
108    }
109
110    async fn commit_runtime_state(
111        &self,
112        commit: RuntimeCommit,
113    ) -> Result<RuntimeCommitResult, StoreError> {
114        let blob_profile = self.options.blob_profile;
115        let now = self.clock.timestamp_ms();
116        let enqueue_nonce_start = self.commit_count.fetch_add(
117            commit.enqueued_queue_batches.len() as u64,
118            AtomicOrdering::Relaxed,
119        );
120        let result = self
121            .conn
122            .write_flow(move |tx| {
123                let outcome: Result<RuntimeCommitResult, StoreError> = (|| {
124                    let existing = try_load_session_head_meta_from_conn(tx)?;
125                    if let Some(bound_session_id) =
126                        existing.as_ref().map(|meta| meta.session_id.as_str())
127                        && bound_session_id != commit.session_id
128                    {
129                        return Err(StoreError::SessionBindingMismatch {
130                            bound_session_id: bound_session_id.to_string(),
131                            attempted_session_id: commit.session_id.clone(),
132                        });
133                    }
134                    if let Some(completed) = &commit.turn_commit {
135                        if completed.session_id != commit.session_id {
136                            return Err(StoreError::RuntimeTurnCommitConflict {
137                                session_id: completed.session_id.clone(),
138                                turn_id: completed.turn_id.clone(),
139                            });
140                        }
141                        let prior: Option<(String, String)> = tx
142                            .query_row(
143                                "SELECT turn_commit_hash, result_json FROM runtime_turn_commits
144                                 WHERE session_id = ?1 AND turn_id = ?2",
145                                params![completed.session_id, completed.turn_id],
146                                |row| Ok((row.get(0)?, row.get(1)?)),
147                            )
148                            .optional()
149                            .map_err(sqlite_error)?;
150                        if let Some((turn_commit_hash, result_json)) = prior {
151                            if turn_commit_hash == completed.turn_commit_hash {
152                                let result: RuntimeCommitResult =
153                                    serde_json::from_str(&result_json).map_err(|err| {
154                                        StoreError::Backend(format!(
155                                            "failed to decode runtime turn commit result: {err}"
156                                        ))
157                                    })?;
158                                if let Some(completion) =
159                                    commit.release_session_execution_lease.as_ref()
160                                {
161                                    release_session_execution_lease_conn(tx, completion)?;
162                                }
163                                return Ok(result);
164                            }
165                            return Err(StoreError::RuntimeTurnCommitConflict {
166                                session_id: completed.session_id.clone(),
167                                turn_id: completed.turn_id.clone(),
168                            });
169                        }
170                    }
171                    let Some(session_execution_lease) = commit.session_execution_lease.as_ref()
172                    else {
173                        return Err(StoreError::SessionExecutionLeaseExpired {
174                            session_id: commit.session_id.clone(),
175                        });
176                    };
177                    ensure_session_execution_lease_conn(
178                        tx,
179                        &commit.session_id,
180                        session_execution_lease,
181                        now,
182                    )?;
183                    let actual_revision = existing.as_ref().map_or(0, |meta| meta.head_revision);
184                    if commit.expected_head_revision.is_some()
185                        && commit.expected_head_revision != Some(actual_revision)
186                    {
187                        return Err(StoreError::HeadRevisionConflict {
188                            expected: commit.expected_head_revision,
189                            actual: actual_revision,
190                        });
191                    }
192                    for completed in &commit.completed_queue_claims {
193                        if completed.session_id != commit.session_id {
194                            return Err(StoreError::QueuedWorkClaimSuperseded {
195                                session_id: completed.session_id.clone(),
196                                claim_id: completed.claim_id.clone(),
197                            });
198                        }
199                        ensure_queued_work_completion_conn(tx, completed)?;
200                    }
201                    for completed in &commit.completed_turn_input_claims {
202                        if completed.session_id != commit.session_id {
203                            return Err(StoreError::TurnInputClaimSuperseded {
204                                session_id: completed.session_id.clone(),
205                                claim_id: completed.claim_id.clone(),
206                            });
207                        }
208                        let owned_rows: usize = tx
209                            .query_row(
210                                "SELECT COUNT(*)
211                                 FROM pending_turn_inputs
212                                 WHERE session_id = ?1
213                                   AND claim_id = ?2
214                                   AND claim_token = ?3",
215                                params![
216                                    completed.session_id,
217                                    completed.claim_id,
218                                    completed.lease_token
219                                ],
220                                |row| row.get::<_, i64>(0),
221                            )
222                            .map_err(sqlite_error)? as usize;
223                        ensure_turn_input_completion_owns_all_inputs(completed, owned_rows)?;
224                    }
225
226                    let stored_checkpoint =
227                        Self::put_checkpoint_conn(tx, &commit.checkpoint, blob_profile)
228                            .map_err(sqlite_error)?;
229
230                    if !commit.usage_deltas.is_empty() {
231                        let mut stmt = tx
232                            .prepare(
233                                "INSERT INTO usage_deltas (
234                                    source, model, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, reasoning_output_tokens
235                                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
236                            )
237                            .map_err(sqlite_error)?;
238                        for entry in &commit.usage_deltas {
239                            stmt.execute(params![
240                                entry.source,
241                                entry.model,
242                                entry.usage.input_tokens,
243                                entry.usage.output_tokens,
244                                entry.usage.cache_read_input_tokens,
245                                entry.usage.cache_write_input_tokens,
246                                entry.usage.reasoning_output_tokens,
247                            ])
248                            .map_err(sqlite_error)?;
249                        }
250                    }
251
252                    let leaf_node_id = match &commit.graph {
253                        GraphCommitDelta::Unchanged { leaf_node_id } => leaf_node_id.clone(),
254                        GraphCommitDelta::Append {
255                            nodes,
256                            leaf_node_id,
257                        } => {
258                            for node in nodes {
259                                let node_json = encode_json(node);
260                                tx.execute(
261                                    "INSERT INTO graph_nodes (node_id, node_json) VALUES (?1, ?2)",
262                                    params![node.node_id, node_json],
263                                )
264                                .map_err(sqlite_error)?;
265                            }
266                            leaf_node_id.clone()
267                        }
268                        GraphCommitDelta::ReplaceFull(graph) => {
269                            tx.execute("DELETE FROM graph_nodes", [])
270                                .map_err(sqlite_error)?;
271                            for node in &graph.nodes {
272                                let node_json = encode_json(node);
273                                tx.execute(
274                                    "INSERT INTO graph_nodes (node_id, node_json) VALUES (?1, ?2)",
275                                    params![node.node_id, node_json],
276                                )
277                                .map_err(sqlite_error)?;
278                            }
279                            graph.leaf_node_id.clone()
280                        }
281                    };
282                    let graph_node_count: usize = tx
283                        .query_row(
284                            "SELECT COUNT(*) FROM graph_nodes WHERE tombstoned = 0",
285                            [],
286                            |row| row.get::<_, i64>(0),
287                        )
288                        .map_err(sqlite_error)? as usize;
289                    let next_revision = actual_revision + 1;
290                    let meta = SessionHeadMeta {
291                        schema_version: lash_core::store::SESSION_HEAD_META_SCHEMA_VERSION,
292                        session_id: commit.session_id.clone(),
293                        head_revision: next_revision,
294                        config: commit.config.clone(),
295                        agent_frames: commit.agent_frames.clone(),
296                        current_agent_frame_id: commit.current_agent_frame_id.clone(),
297                        checkpoint_ref: Some(stored_checkpoint.checkpoint_ref.clone()),
298                        leaf_node_id,
299                        graph_node_count,
300                        token_ledger: Vec::new(),
301                    };
302                    tx.execute(
303                        "INSERT OR REPLACE INTO session_head (singleton, session_id, head_json, head_revision)
304                         VALUES (1, ?1, ?2, ?3)",
305                        params![
306                            meta.session_id,
307                            encode_json(&meta),
308                            meta.head_revision as i64
309                        ],
310                    )
311                    .map_err(sqlite_error)?;
312                    for completed in &commit.completed_queue_claims {
313                        for batch_id in &completed.batch_ids {
314                            tx.execute(
315                                "DELETE FROM queued_work_batches
316                                 WHERE session_id = ?1
317                                   AND batch_id = ?2
318                                   AND claim_id = ?3
319                                   AND claim_token = ?4",
320                                params![
321                                    completed.session_id,
322                                    batch_id,
323                                    completed.claim_id,
324                                    completed.lease_token
325                                ],
326                            )
327                            .map_err(sqlite_error)?;
328                        }
329                    }
330                    for completed in &commit.completed_turn_input_claims {
331                        for input_id in &completed.input_ids {
332                            tx.execute(
333                                "UPDATE pending_turn_inputs
334                                 SET state = ?5,
335                                     claim_id = NULL,
336                                     claim_owner_id = NULL,
337                                     claim_owner_incarnation_id = NULL,
338                                     claim_owner_liveness_json = NULL,
339                                     claim_token = NULL,
340                                     claim_session_lease_generation = 0
341                                 WHERE session_id = ?1
342                                   AND input_id = ?2
343                                   AND claim_id = ?3
344                                   AND claim_token = ?4",
345                                params![
346                                    completed.session_id,
347                                    input_id,
348                                    completed.claim_id,
349                                    completed.lease_token,
350                                    lash_core::TurnInputState::Completed.as_str(),
351                                ],
352                            )
353                            .map_err(sqlite_error)?;
354                        }
355                    }
356                    if let Some(turn_id) = commit.interrupted_turn_input_turn_id.as_deref() {
357                        let input_ids = {
358                            let mut stmt = tx
359                                .prepare(
360                                    "SELECT input_id, ingress_json
361                                     FROM pending_turn_inputs
362                                     WHERE session_id = ?1 AND state = ?2",
363                                )
364                                .map_err(sqlite_error)?;
365                            let rows = stmt
366                                .query_map(
367                                    params![
368                                        commit.session_id,
369                                        lash_core::TurnInputState::PendingActive.as_str()
370                                    ],
371                                    |row| {
372                                        Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
373                                    },
374                                )
375                                .map_err(sqlite_error)?;
376                            let mut input_ids = Vec::new();
377                            for row in rows {
378                                let (input_id, ingress_json) = row.map_err(sqlite_error)?;
379                                let ingress = decode_turn_input_ingress(ingress_json)?;
380                                if ingress
381                                    .active_turn_id()
382                                    .is_some_and(|active| active == turn_id)
383                                {
384                                    input_ids.push(input_id);
385                                }
386                            }
387                            input_ids
388                        };
389                        let next_turn_ingress = encode_json(&lash_core::TurnInputIngress::NextTurn);
390                        let mut stmt = tx
391                            .prepare(
392                                "UPDATE pending_turn_inputs
393                                 SET state = ?3,
394                                     ingress_json = ?4,
395                                     claim_id = NULL,
396                                     claim_owner_id = NULL,
397                                     claim_owner_incarnation_id = NULL,
398                                     claim_owner_liveness_json = NULL,
399                                     claim_token = NULL,
400                                     claim_session_lease_generation = 0
401                                 WHERE session_id = ?1 AND input_id = ?2",
402                            )
403                            .map_err(sqlite_error)?;
404                        for input_id in input_ids {
405                            stmt.execute(params![
406                                commit.session_id,
407                                input_id,
408                                lash_core::TurnInputState::DeferredNextTurn.as_str(),
409                                next_turn_ingress
410                            ])
411                            .map_err(sqlite_error)?;
412                        }
413                    }
414                    if !commit.committed_attachment_ids.is_empty() {
415                        let now = now as i64;
416                        let mut stmt = tx
417                            .prepare(
418                                "UPDATE attachment_manifest
419                                 SET committed_at_ms = COALESCE(committed_at_ms, ?1)
420                                 WHERE attachment_id = ?2 AND session_id = ?3",
421                            )
422                            .map_err(sqlite_error)?;
423                        for id in &commit.committed_attachment_ids {
424                            stmt.execute(params![now, id.as_str(), commit.session_id])
425                                .map_err(sqlite_error)?;
426                        }
427                    }
428                    let mut enqueued_queue_batches = Vec::new();
429                    for (index, batch) in commit.enqueued_queue_batches.iter().enumerate() {
430                        if batch.session_id != commit.session_id {
431                            return Err(StoreError::SessionBindingMismatch {
432                                bound_session_id: commit.session_id.clone(),
433                                attempted_session_id: batch.session_id.clone(),
434                            });
435                        }
436                        enqueued_queue_batches.push(enqueue_queued_work_conn(
437                            tx,
438                            batch,
439                            now,
440                            enqueue_nonce_start.saturating_add(index as u64),
441                        )?);
442                    }
443                    let result = RuntimeCommitResult {
444                        head_revision: next_revision,
445                        checkpoint_ref: stored_checkpoint.checkpoint_ref,
446                        manifest: stored_checkpoint.manifest,
447                        enqueued_queue_batches,
448                    };
449                    if let Some(completed) = &commit.turn_commit {
450                        tx.execute(
451                            "INSERT INTO runtime_turn_commits (
452                                session_id, turn_id, turn_commit_hash, result_json, committed_at_ms
453                             )
454                             VALUES (?1, ?2, ?3, ?4, ?5)",
455                            params![
456                                completed.session_id,
457                                completed.turn_id,
458                                completed.turn_commit_hash,
459                                encode_json(&result),
460                                now as i64
461                            ],
462                        )
463                        .map_err(sqlite_error)?;
464                    }
465                    if let Some(completion) = commit.release_session_execution_lease.as_ref() {
466                        release_session_execution_lease_conn(tx, completion)?;
467                    }
468                    Ok(result)
469                })();
470                // Roll back on a `StoreError` so a failure after the first
471                // write (e.g. a head-revision conflict surfaced mid-commit, or a
472                // backend write error) does not leave the partial transaction
473                // committed, while still carrying the typed error to the caller.
474                match outcome {
475                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
476                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
477                }
478            })
479            .await
480            .map_err(sqlite_error)??;
481        self.maybe_auto_gc().await;
482        Ok(result)
483    }
484
485    async fn save_session_meta(&self, meta: SessionMeta) -> Result<(), StoreError> {
486        Store::save_session_meta(self, meta).await;
487        Ok(())
488    }
489
490    async fn load_session_meta(&self) -> Result<Option<SessionMeta>, StoreError> {
491        Ok(Store::load_session_meta(self).await)
492    }
493}
494
495#[async_trait::async_trait]
496impl SessionExecutionLeaseStore for Store {
497    async fn try_claim_session_execution_lease(
498        &self,
499        session_id: &str,
500        owner: &LeaseOwnerIdentity,
501        lease_ttl_ms: u64,
502    ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError> {
503        let session_id = session_id.to_string();
504        let owner = owner.clone();
505        let now = self.clock.timestamp_ms();
506        self.conn
507            .write_flow(move |tx| {
508                let outcome: Result<SessionExecutionLeaseClaimOutcome, StoreError> = (|| {
509                    let current = load_session_execution_lease_row_conn(tx, &session_id)?;
510                    if current.as_ref().is_some_and(|lease| {
511                        lease.lease_token.is_some() && lease.expires_at_ms > now
512                    }) {
513                        let current = current.expect("checked current lease is present");
514                        if current
515                            .owner
516                            .as_ref()
517                            .is_some_and(|current_owner| current_owner.same_incarnation(&owner))
518                        {
519                            let expires_at = now.saturating_add(lease_ttl_ms);
520                            tx.execute(
521                                "UPDATE session_execution_leases
522                                 SET lease_expires_at_ms = ?2
523                                 WHERE session_id = ?1",
524                                params![session_id, expires_at as i64],
525                            )
526                            .map_err(sqlite_error)?;
527                            return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
528                                SessionExecutionLease {
529                                    session_id,
530                                    owner,
531                                    lease_token: current.lease_token.expect("live lease token set"),
532                                    fencing_token: current.fencing_token,
533                                    claimed_at_epoch_ms: current.claimed_at_ms,
534                                    expires_at_epoch_ms: expires_at,
535                                },
536                            ));
537                        }
538                        return Ok(SessionExecutionLeaseClaimOutcome::Busy {
539                            holder: row_to_session_execution_lease(&session_id, current)?,
540                        });
541                    }
542                    Ok(SessionExecutionLeaseClaimOutcome::Acquired(
543                        acquire_session_execution_lease_conn(
544                            tx,
545                            &session_id,
546                            &owner,
547                            current.as_ref().map_or(0, |lease| lease.fencing_token),
548                            now,
549                            lease_ttl_ms,
550                        )?,
551                    ))
552                })(
553                );
554                match outcome {
555                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
556                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
557                }
558            })
559            .await
560            .map_err(sqlite_error)?
561    }
562
563    async fn reclaim_session_execution_lease(
564        &self,
565        session_id: &str,
566        owner: &LeaseOwnerIdentity,
567        observed_holder: &SessionExecutionLeaseFence,
568        lease_ttl_ms: u64,
569    ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError> {
570        let session_id = session_id.to_string();
571        let owner = owner.clone();
572        let observed_holder = observed_holder.clone();
573        let now = self.clock.timestamp_ms();
574        self.conn
575            .write_flow(move |tx| {
576                let outcome: Result<SessionExecutionLeaseClaimOutcome, StoreError> = (|| {
577                    let current = load_session_execution_lease_row_conn(tx, &session_id)?;
578                    let Some(current) = current else {
579                        return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
580                            acquire_session_execution_lease_conn(
581                                tx,
582                                &session_id,
583                                &owner,
584                                0,
585                                now,
586                                lease_ttl_ms,
587                            )?,
588                        ));
589                    };
590                    if current.lease_token.is_none() || current.expires_at_ms <= now {
591                        return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
592                            acquire_session_execution_lease_conn(
593                                tx,
594                                &session_id,
595                                &owner,
596                                current.fencing_token,
597                                now,
598                                lease_ttl_ms,
599                            )?,
600                        ));
601                    }
602                    let holder = row_to_session_execution_lease(&session_id, current)?;
603                    if observed_holder.session_id == session_id
604                        && holder.owner.same_incarnation(&observed_holder.owner)
605                        && holder.lease_token == observed_holder.lease_token
606                        && holder.fencing_token == observed_holder.fencing_token
607                        && holder.owner.is_definitely_dead_for_claimant(&owner)
608                    {
609                        let fencing_token = holder.fencing_token.saturating_add(1);
610                        let lease_token = format!(
611                            "{}:{}:{}:{now}:{fencing_token}",
612                            session_id, owner.owner_id, owner.incarnation_id
613                        );
614                        let expires_at = now.saturating_add(lease_ttl_ms);
615                        let liveness_json = encode_liveness(&owner.liveness)?;
616                        let changed = tx
617                            .execute(
618                                "UPDATE session_execution_leases
619                                 SET lease_owner_id = ?1,
620                                     lease_owner_incarnation_id = ?2,
621                                     lease_owner_liveness_json = ?3,
622                                     lease_token = ?4,
623                                     lease_fencing_token = ?5,
624                                     lease_claimed_at_ms = ?6,
625                                     lease_expires_at_ms = ?7
626                                 WHERE session_id = ?8
627                                   AND lease_owner_id = ?9
628                                   AND lease_owner_incarnation_id = ?10
629                                   AND lease_token = ?11
630                                   AND lease_fencing_token = ?12",
631                                params![
632                                    owner.owner_id,
633                                    owner.incarnation_id,
634                                    liveness_json,
635                                    lease_token,
636                                    fencing_token as i64,
637                                    now as i64,
638                                    expires_at as i64,
639                                    session_id,
640                                    observed_holder.owner.owner_id,
641                                    observed_holder.owner.incarnation_id,
642                                    observed_holder.lease_token,
643                                    observed_holder.fencing_token as i64,
644                                ],
645                            )
646                            .map_err(sqlite_error)?;
647                        if changed == 1 {
648                            return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
649                                SessionExecutionLease {
650                                    session_id,
651                                    owner,
652                                    lease_token,
653                                    fencing_token,
654                                    claimed_at_epoch_ms: now,
655                                    expires_at_epoch_ms: expires_at,
656                                },
657                            ));
658                        }
659                        let current = load_session_execution_lease_row_conn(tx, &session_id)?;
660                        if current.as_ref().is_some_and(|lease| {
661                            lease.lease_token.is_some() && lease.expires_at_ms > now
662                        }) {
663                            let current = current.expect("checked current lease is present");
664                            return Ok(SessionExecutionLeaseClaimOutcome::Busy {
665                                holder: row_to_session_execution_lease(&session_id, current)?,
666                            });
667                        }
668                        let previous_fencing_token =
669                            current.as_ref().map_or(0, |lease| lease.fencing_token);
670                        return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
671                            acquire_session_execution_lease_conn(
672                                tx,
673                                &session_id,
674                                &owner,
675                                previous_fencing_token,
676                                now,
677                                lease_ttl_ms,
678                            )?,
679                        ));
680                    }
681                    Ok(SessionExecutionLeaseClaimOutcome::Busy { holder })
682                })(
683                );
684                match outcome {
685                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
686                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
687                }
688            })
689            .await
690            .map_err(sqlite_error)?
691    }
692
693    async fn renew_session_execution_lease(
694        &self,
695        fence: &SessionExecutionLeaseFence,
696        lease_ttl_ms: u64,
697    ) -> Result<SessionExecutionLease, StoreError> {
698        let fence = fence.clone();
699        let now = self.clock.timestamp_ms();
700        self.conn
701            .write_flow(move |tx| {
702                let outcome: Result<SessionExecutionLease, StoreError> = (|| {
703                    let current = load_session_execution_lease_row_conn(tx, &fence.session_id)?;
704                    let Some(current) = current else {
705                        return Err(StoreError::SessionExecutionLeaseExpired {
706                            session_id: fence.session_id.clone(),
707                        });
708                    };
709                    if !current
710                        .owner
711                        .as_ref()
712                        .is_some_and(|owner| owner.same_incarnation(&fence.owner))
713                        || current.lease_token.as_deref() != Some(fence.lease_token.as_str())
714                        || current.fencing_token != fence.fencing_token
715                        || current.expires_at_ms <= now
716                    {
717                        return Err(StoreError::SessionExecutionLeaseExpired {
718                            session_id: fence.session_id.clone(),
719                        });
720                    }
721                    let expires_at = now.saturating_add(lease_ttl_ms);
722                    tx.execute(
723                        "UPDATE session_execution_leases
724                         SET lease_expires_at_ms = ?5
725                         WHERE session_id = ?1
726                           AND lease_owner_id = ?2
727                           AND lease_owner_incarnation_id = ?3
728                           AND lease_token = ?4
729                           AND lease_fencing_token = ?6",
730                        params![
731                            fence.session_id,
732                            fence.owner.owner_id,
733                            fence.owner.incarnation_id,
734                            fence.lease_token,
735                            expires_at as i64,
736                            fence.fencing_token as i64
737                        ],
738                    )
739                    .map_err(sqlite_error)?;
740                    Ok(SessionExecutionLease {
741                        session_id: fence.session_id,
742                        owner: fence.owner,
743                        lease_token: fence.lease_token,
744                        fencing_token: fence.fencing_token,
745                        claimed_at_epoch_ms: current.claimed_at_ms,
746                        expires_at_epoch_ms: expires_at,
747                    })
748                })();
749                match outcome {
750                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
751                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
752                }
753            })
754            .await
755            .map_err(sqlite_error)?
756    }
757
758    async fn release_session_execution_lease(
759        &self,
760        completion: &SessionExecutionLeaseCompletion,
761    ) -> Result<(), StoreError> {
762        let completion = completion.clone();
763        self.conn
764            .write_flow(move |tx| {
765                let outcome = release_session_execution_lease_conn(tx, &completion);
766                match outcome {
767                    Ok(()) => Ok(TxOutcome::Commit(Ok(()))),
768                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
769                }
770            })
771            .await
772            .map_err(sqlite_error)?
773    }
774}
775
776#[async_trait::async_trait]
777impl QueuedWorkStore for Store {
778    async fn enqueue_queued_work(
779        &self,
780        batch: QueuedWorkBatchDraft,
781    ) -> Result<QueuedWorkBatch, StoreError> {
782        let nonce = self.commit_count.fetch_add(1, AtomicOrdering::Relaxed);
783        let now = self.clock.timestamp_ms();
784        self.conn
785            .write_flow(move |tx| {
786                let outcome = enqueue_queued_work_conn(tx, &batch, now, nonce);
787                // Roll back the partially-inserted batch/items on a
788                // `StoreError` while still returning the typed error.
789                match outcome {
790                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
791                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
792                }
793            })
794            .await
795            .map_err(sqlite_error)?
796    }
797
798    async fn claim_leading_ready_session_command(
799        &self,
800        session_id: &str,
801        session_execution_lease: &SessionExecutionLeaseFence,
802        owner: &LeaseOwnerIdentity,
803    ) -> Result<Option<QueuedWorkClaim>, StoreError> {
804        let session_id = session_id.to_string();
805        let session_execution_lease = session_execution_lease.clone();
806        let owner = owner.clone();
807        let now = self.clock.timestamp_ms();
808        self.conn
809            .write_flow(move |tx| {
810                let outcome: Result<TxOutcome<Option<QueuedWorkClaim>>, StoreError> = (|| {
811                    ensure_session_execution_lease_conn(
812                        tx,
813                        &session_id,
814                        &session_execution_lease,
815                        now,
816                    )?;
817                    // The fence is validated live, so its fencing token is the
818                    // currently-live session-lease generation; claims pin it and
819                    // are claimable only across a different generation (ADR 0029).
820                    let generation = session_execution_lease.fencing_token;
821                    let candidate_rows = {
822                        let mut stmt = tx
823                            .prepare(
824                                "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
825                                        slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
826                                        claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
827                                        claim_owner_liveness_json, claim_token, claim_session_lease_generation
828                                 FROM queued_work_batches
829                                 WHERE session_id = ?1
830                                   AND available_at_ms <= ?2
831                                   AND (
832                                        claim_token IS NULL
833                                        OR claim_session_lease_generation <> ?3
834                                   )
835                                 ORDER BY enqueue_seq ASC
836                                 LIMIT ?4",
837                            )
838                            .map_err(sqlite_error)?;
839                        let rows = stmt
840                            .query_map(
841                                params![
842                                    session_id,
843                                    now as i64,
844                                    generation as i64,
845                                    claim_scan_limit(1)
846                                ],
847                                queued_batch_row_from_sql,
848                            )
849                            .map_err(sqlite_error)?;
850                        rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
851                    };
852                    let candidate_rows = candidate_rows
853                        .into_iter()
854                        .filter(|row| {
855                            row.claim_token.is_none()
856                                || row.claim_session_lease_generation != generation
857                        })
858                        .collect::<Vec<_>>();
859                    let candidate_batches = candidate_rows
860                        .iter()
861                        .map(|row| queued_work_batch_from_conn(tx, row.clone()))
862                        .collect::<Result<Vec<_>, StoreError>>()?;
863                    let candidates = candidate_rows
864                        .iter()
865                        .zip(candidate_batches.iter())
866                        .map(|(row, batch)| {
867                            Ok(ClaimCandidate {
868                                enqueue_seq: row.enqueue_seq,
869                                claim_fencing_token: row.claim_fencing_token,
870                                work_class: batch.work_class().ok_or_else(|| {
871                                    StoreError::Backend(format!(
872                                        "queued-work batch `{}` has mixed or empty payload classes",
873                                        batch.batch_id
874                                    ))
875                                })?,
876                                delivery_policy: decode_delivery_policy(
877                                    row.delivery_policy.clone(),
878                                )?,
879                                slot_policy: decode_slot_policy(row.slot_policy.clone())?,
880                                merge_key: decode_merge_key(row.merge_key_json.clone())?,
881                            })
882                        })
883                        .collect::<Result<Vec<_>, StoreError>>()?;
884                    let selected_len = select_leading_session_command(&candidates);
885                    if selected_len == 0 {
886                        return Ok(TxOutcome::Commit(None));
887                    }
888                    let mut selected = candidate_rows;
889                    selected.truncate(selected_len);
890                    let mut selected_batches = candidate_batches;
891                    selected_batches.truncate(selected_len);
892                    let lease = QueuedWorkClaimLease::derive(
893                        &candidates[0],
894                        &session_id,
895                        &owner,
896                        now,
897                        generation,
898                    );
899                    let liveness_json = encode_liveness(&owner.liveness)?;
900                    for row in &selected {
901                        let claimed = tx
902                            .execute(
903                                "UPDATE queued_work_batches
904                                 SET claim_id = ?3,
905                                     claim_owner_id = ?4,
906                                     claim_owner_incarnation_id = ?5,
907                                     claim_owner_liveness_json = ?6,
908                                     claim_token = ?7,
909                                     claim_fencing_token = claim_fencing_token + 1,
910                                     claim_session_lease_generation = ?8
911                                 WHERE session_id = ?1
912                                   AND batch_id = ?2
913                                   AND (
914                                        claim_token IS NULL
915                                        OR claim_session_lease_generation <> ?8
916                                   )",
917                                params![
918                                    session_id,
919                                    row.batch_id,
920                                    lease.claim_id,
921                                    owner.owner_id.as_str(),
922                                    owner.incarnation_id.as_str(),
923                                    liveness_json.as_str(),
924                                    lease.lease_token,
925                                    lease.session_lease_generation as i64,
926                                ],
927                            )
928                            .map_err(sqlite_error)?;
929                        if claimed == 0 {
930                            return Ok(TxOutcome::Rollback(None));
931                        }
932                    }
933                    Ok(TxOutcome::Commit(Some(QueuedWorkClaim {
934                        session_id: session_id.clone(),
935                        claim_id: lease.claim_id,
936                        owner: owner.clone(),
937                        lease_token: lease.lease_token,
938                        fencing_token: lease.fencing_token,
939                        session_lease_generation: lease.session_lease_generation,
940                        batches: selected_batches,
941                    })))
942                })();
943                match outcome {
944                    Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
945                    Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
946                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
947                }
948            })
949            .await
950            .map_err(sqlite_error)?
951    }
952
953    async fn claim_ready_queued_work(
954        &self,
955        session_id: &str,
956        session_execution_lease: &SessionExecutionLeaseFence,
957        owner: &LeaseOwnerIdentity,
958        boundary: QueuedWorkClaimBoundary,
959        max_batches: usize,
960    ) -> Result<Option<QueuedWorkClaim>, StoreError> {
961        if max_batches == 0 {
962            return Ok(None);
963        }
964        let session_id = session_id.to_string();
965        let session_execution_lease = session_execution_lease.clone();
966        let owner = owner.clone();
967        let now = self.clock.timestamp_ms();
968        self.conn
969            .write_flow(move |tx| {
970                let outcome: Result<TxOutcome<Option<QueuedWorkClaim>>, StoreError> = (|| {
971                    ensure_session_execution_lease_conn(
972                        tx,
973                        &session_id,
974                        &session_execution_lease,
975                        now,
976                    )?;
977                    let generation = session_execution_lease.fencing_token;
978                    let candidate_rows = {
979                        let mut stmt = tx
980                            .prepare(
981                                "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
982                                        slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
983                                        claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
984                                        claim_owner_liveness_json, claim_token, claim_session_lease_generation
985                                 FROM queued_work_batches
986                                 WHERE session_id = ?1
987                                   AND available_at_ms <= ?2
988                                   AND (
989                                        claim_token IS NULL
990                                        OR claim_session_lease_generation <> ?3
991                                   )
992                                 ORDER BY enqueue_seq ASC
993                                 LIMIT ?4",
994                            )
995                            .map_err(sqlite_error)?;
996                        let rows = stmt
997                            .query_map(
998                                params![
999                                    session_id,
1000                                    now as i64,
1001                                    generation as i64,
1002                                    claim_scan_limit(max_batches)
1003                                ],
1004                                queued_batch_row_from_sql,
1005                            )
1006                            .map_err(sqlite_error)?;
1007                        rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1008                    };
1009                    let candidate_rows = candidate_rows
1010                        .into_iter()
1011                        .filter(|row| {
1012                            row.claim_token.is_none()
1013                                || row.claim_session_lease_generation != generation
1014                        })
1015                        .collect::<Vec<_>>();
1016                    let candidate_batches = candidate_rows
1017                        .iter()
1018                        .map(|row| queued_work_batch_from_conn(tx, row.clone()))
1019                        .collect::<Result<Vec<_>, StoreError>>()?;
1020                    let candidates = candidate_rows
1021                        .iter()
1022                        .zip(candidate_batches.iter())
1023                        .map(|(row, batch)| {
1024                            Ok(ClaimCandidate {
1025                                enqueue_seq: row.enqueue_seq,
1026                                claim_fencing_token: row.claim_fencing_token,
1027                                work_class: batch.work_class().ok_or_else(|| {
1028                                    StoreError::Backend(format!(
1029                                        "queued-work batch `{}` has mixed or empty payload classes",
1030                                        batch.batch_id
1031                                    ))
1032                                })?,
1033                                delivery_policy: decode_delivery_policy(
1034                                    row.delivery_policy.clone(),
1035                                )?,
1036                                slot_policy: decode_slot_policy(row.slot_policy.clone())?,
1037                                merge_key: decode_merge_key(row.merge_key_json.clone())?,
1038                            })
1039                        })
1040                        .collect::<Result<Vec<_>, StoreError>>()?;
1041                    let selected_len =
1042                        select_turn_work_claim_prefix(&candidates, boundary, max_batches);
1043                    if selected_len == 0 {
1044                        return Ok(TxOutcome::Commit(None));
1045                    }
1046                    let mut selected = candidate_rows;
1047                    selected.truncate(selected_len);
1048                    let mut selected_batches = candidate_batches;
1049                    selected_batches.truncate(selected_len);
1050                    let lease = QueuedWorkClaimLease::derive(
1051                        &candidates[0],
1052                        &session_id,
1053                        &owner,
1054                        now,
1055                        generation,
1056                    );
1057                    let liveness_json = encode_liveness(&owner.liveness)?;
1058                    for row in &selected {
1059                        // Under `BEGIN IMMEDIATE` this connection already holds
1060                        // the write lock, but the row could still have been
1061                        // claimed by an earlier committed writer (its
1062                        // `claim_token` set and not yet expired). The `WHERE`
1063                        // clause filters those out, so a 0-row update means we
1064                        // lost the race for this batch: treat the whole claim as
1065                        // not-won rather than returning a claim that doesn't
1066                        // actually own the row.
1067                        let claimed = tx
1068                            .execute(
1069                                "UPDATE queued_work_batches
1070                                 SET claim_id = ?3,
1071                                     claim_owner_id = ?4,
1072                                     claim_owner_incarnation_id = ?5,
1073                                     claim_owner_liveness_json = ?6,
1074                                     claim_token = ?7,
1075                                     claim_fencing_token = claim_fencing_token + 1,
1076                                     claim_session_lease_generation = ?8
1077                                 WHERE session_id = ?1
1078                                   AND batch_id = ?2
1079                                   AND (
1080                                        claim_token IS NULL
1081                                        OR claim_session_lease_generation <> ?8
1082                                   )",
1083                                params![
1084                                    session_id,
1085                                    row.batch_id,
1086                                    lease.claim_id,
1087                                    owner.owner_id.as_str(),
1088                                    owner.incarnation_id.as_str(),
1089                                    liveness_json.as_str(),
1090                                    lease.lease_token,
1091                                    lease.session_lease_generation as i64,
1092                                ],
1093                            )
1094                            .map_err(sqlite_error)?;
1095                        if claimed == 0 {
1096                            // Lost the race for this batch. Roll back any sibling
1097                            // rows we already claimed in this transaction so we
1098                            // never return a half-owned claim.
1099                            return Ok(TxOutcome::Rollback(None));
1100                        }
1101                    }
1102                    Ok(TxOutcome::Commit(Some(QueuedWorkClaim {
1103                        session_id: session_id.clone(),
1104                        claim_id: lease.claim_id,
1105                        owner: owner.clone(),
1106                        lease_token: lease.lease_token,
1107                        fencing_token: lease.fencing_token,
1108                        session_lease_generation: lease.session_lease_generation,
1109                        batches: selected_batches,
1110                    })))
1111                })();
1112                // Lower a `StoreError` into the rollback arm so the closure body
1113                // can keep using `?` while still propagating the error to the
1114                // caller. Encode it as a `Result` carried out of the flow.
1115                match outcome {
1116                    Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
1117                    Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
1118                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1119                }
1120            })
1121            .await
1122            .map_err(sqlite_error)?
1123    }
1124
1125    async fn claim_ready_queued_work_by_batch_ids(
1126        &self,
1127        session_id: &str,
1128        session_execution_lease: &SessionExecutionLeaseFence,
1129        owner: &LeaseOwnerIdentity,
1130        boundary: QueuedWorkClaimBoundary,
1131        batch_ids: &[String],
1132    ) -> Result<Option<QueuedWorkClaim>, StoreError> {
1133        if batch_ids.is_empty() {
1134            return Ok(None);
1135        }
1136        let session_id = session_id.to_string();
1137        let fence = session_execution_lease.clone();
1138        let owner = owner.clone();
1139        let batch_ids = batch_ids.to_vec();
1140        let now = self.clock.timestamp_ms();
1141        self.conn
1142            .write_flow(move |tx| {
1143                let outcome: Result<Option<QueuedWorkClaim>, StoreError> = (|| {
1144                    ensure_session_execution_lease_conn(tx, &session_id, &fence, now)?;
1145                    let generation = fence.fencing_token;
1146                    let mut rows = Vec::new();
1147                    let mut batches = Vec::new();
1148                    for batch_id in &batch_ids {
1149                        let row = tx
1150                            .query_row(
1151                                "SELECT enqueue_seq, batch_id, session_id, source_key,
1152                                        delivery_policy, slot_policy, merge_key_json,
1153                                        available_at_ms, enqueued_at_ms, claim_fencing_token,
1154                                        claim_owner_id, claim_owner_incarnation_id,
1155                                        claim_owner_liveness_json, claim_token,
1156                                        claim_session_lease_generation
1157                                 FROM queued_work_batches
1158                                 WHERE session_id = ?1 AND batch_id = ?2
1159                                   AND available_at_ms <= ?3
1160                                   AND (claim_token IS NULL
1161                                        OR claim_session_lease_generation <> ?4)",
1162                                params![session_id, batch_id, now as i64, generation as i64],
1163                                queued_batch_row_from_sql,
1164                            )
1165                            .optional()
1166                            .map_err(sqlite_error)?;
1167                        let Some(row) = row else {
1168                            return Ok(None);
1169                        };
1170                        let batch = queued_work_batch_from_conn(tx, row.clone())?;
1171                        if batch.work_class() != Some(lash_core::runtime::QueuedWorkClass::TurnWork)
1172                        {
1173                            return Ok(None);
1174                        }
1175                        rows.push(row);
1176                        batches.push(batch);
1177                    }
1178                    let candidates = rows
1179                        .iter()
1180                        .map(|row| {
1181                            Ok(ClaimCandidate {
1182                                enqueue_seq: row.enqueue_seq,
1183                                claim_fencing_token: row.claim_fencing_token,
1184                                work_class: lash_core::runtime::QueuedWorkClass::TurnWork,
1185                                delivery_policy: decode_delivery_policy(
1186                                    row.delivery_policy.clone(),
1187                                )?,
1188                                slot_policy: decode_slot_policy(row.slot_policy.clone())?,
1189                                merge_key: decode_merge_key(row.merge_key_json.clone())?,
1190                            })
1191                        })
1192                        .collect::<Result<Vec<_>, StoreError>>()?;
1193                    if select_turn_work_claim_prefix(&candidates, boundary, candidates.len())
1194                        != candidates.len()
1195                    {
1196                        return Ok(None);
1197                    }
1198                    let lease = QueuedWorkClaimLease::derive(
1199                        &candidates[0],
1200                        &session_id,
1201                        &owner,
1202                        now,
1203                        generation,
1204                    );
1205                    let owner_liveness_json = encode_liveness(&owner.liveness)?;
1206                    for row in &rows {
1207                        let changed = tx
1208                            .execute(
1209                                "UPDATE queued_work_batches
1210                                 SET claim_id = ?3, claim_owner_id = ?4,
1211                                     claim_owner_incarnation_id = ?5,
1212                                     claim_owner_liveness_json = ?6, claim_token = ?7,
1213                                     claim_fencing_token = claim_fencing_token + 1,
1214                                     claim_session_lease_generation = ?8
1215                                 WHERE session_id = ?1 AND batch_id = ?2
1216                                   AND (claim_token IS NULL
1217                                        OR claim_session_lease_generation <> ?8)",
1218                                params![
1219                                    session_id,
1220                                    row.batch_id,
1221                                    lease.claim_id,
1222                                    owner.owner_id,
1223                                    owner.incarnation_id,
1224                                    owner_liveness_json,
1225                                    lease.lease_token,
1226                                    generation as i64,
1227                                ],
1228                            )
1229                            .map_err(sqlite_error)?;
1230                        if changed != 1 {
1231                            return Ok(None);
1232                        }
1233                    }
1234                    Ok(Some(QueuedWorkClaim {
1235                        session_id,
1236                        claim_id: lease.claim_id,
1237                        owner,
1238                        lease_token: lease.lease_token,
1239                        fencing_token: lease.fencing_token,
1240                        session_lease_generation: lease.session_lease_generation,
1241                        batches,
1242                    }))
1243                })();
1244                match outcome {
1245                    Ok(Some(value)) => Ok(TxOutcome::Commit(Ok(Some(value)))),
1246                    Ok(None) => Ok(TxOutcome::Rollback(Ok(None))),
1247                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1248                }
1249            })
1250            .await
1251            .map_err(sqlite_error)?
1252    }
1253
1254    async fn abandon_queued_work_claim(&self, claim: &QueuedWorkClaim) -> Result<(), StoreError> {
1255        let session_id = claim.session_id.clone();
1256        let claim_id = claim.claim_id.clone();
1257        let lease_token = claim.lease_token.clone();
1258        self.conn
1259            .write(move |tx| {
1260                tx.execute(
1261                    "UPDATE queued_work_batches
1262                     SET claim_id = NULL,
1263                         claim_owner_id = NULL,
1264                         claim_owner_incarnation_id = NULL,
1265                         claim_owner_liveness_json = NULL,
1266                         claim_token = NULL,
1267                         claim_session_lease_generation = 0
1268                     WHERE session_id = ?1 AND claim_id = ?2 AND claim_token = ?3",
1269                    params![session_id, claim_id, lease_token],
1270                )
1271            })
1272            .await
1273            .map_err(sqlite_error)?;
1274        Ok(())
1275    }
1276
1277    async fn cancel_queued_work_batch(
1278        &self,
1279        session_id: &str,
1280        batch_id: &str,
1281    ) -> Result<Option<QueuedWorkBatch>, StoreError> {
1282        let session_id = session_id.to_string();
1283        let batch_id = batch_id.to_string();
1284        let now = self.clock.timestamp_ms() as i64;
1285        self.conn
1286            .write_flow(move |tx| {
1287                let outcome: Result<Option<QueuedWorkBatch>, StoreError> = (|| {
1288                    let row = tx
1289                        .query_row(
1290                            "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1291                                    slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1292                                    claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1293                                    claim_owner_liveness_json, claim_token, claim_session_lease_generation
1294                             FROM queued_work_batches
1295                             WHERE session_id = ?1
1296                               AND batch_id = ?2
1297                               AND (claim_token IS NULL OR NOT EXISTS (
1298                                        SELECT 1 FROM session_execution_leases sel
1299                                        WHERE sel.session_id = ?1
1300                                          AND sel.lease_token IS NOT NULL
1301                                          AND sel.lease_expires_at_ms > ?3
1302                                          AND sel.lease_fencing_token
1303                                              = queued_work_batches.claim_session_lease_generation
1304                                   ))",
1305                            params![session_id, batch_id, now],
1306                            queued_batch_row_from_sql,
1307                        )
1308                        .optional()
1309                        .map_err(sqlite_error)?;
1310                    let Some(row) = row else {
1311                        return Ok(None);
1312                    };
1313                    let batch = queued_work_batch_from_conn(tx, row)?;
1314                    tx.execute(
1315                        "DELETE FROM queued_work_batches
1316                         WHERE session_id = ?1
1317                           AND batch_id = ?2
1318                           AND (claim_token IS NULL OR NOT EXISTS (
1319                                SELECT 1 FROM session_execution_leases sel
1320                                WHERE sel.session_id = ?1
1321                                  AND sel.lease_token IS NOT NULL
1322                                  AND sel.lease_expires_at_ms > ?3
1323                                  AND sel.lease_fencing_token
1324                                      = queued_work_batches.claim_session_lease_generation
1325                           ))",
1326                        params![session_id, batch_id, now],
1327                    )
1328                    .map_err(sqlite_error)?;
1329                    Ok(Some(batch))
1330                })();
1331                match outcome {
1332                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1333                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1334                }
1335            })
1336            .await
1337            .map_err(sqlite_error)?
1338    }
1339
1340    async fn list_queued_work(&self, session_id: &str) -> Result<Vec<QueuedWorkBatch>, StoreError> {
1341        let session_id = session_id.to_string();
1342        self.conn
1343            .call(move |conn| {
1344                let outcome: Result<Vec<QueuedWorkBatch>, StoreError> = (|| {
1345                    let rows = {
1346                        let mut stmt = conn
1347                            .prepare(
1348                                "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1349                                        slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1350                                        claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1351                                        claim_owner_liveness_json, claim_token, claim_session_lease_generation
1352                                 FROM queued_work_batches
1353                                 WHERE session_id = ?1
1354                                 ORDER BY enqueue_seq ASC",
1355                            )
1356                            .map_err(sqlite_error)?;
1357                        let rows = stmt
1358                            .query_map(params![session_id], queued_batch_row_from_sql)
1359                            .map_err(sqlite_error)?;
1360                        rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1361                    };
1362                    rows.into_iter()
1363                        .map(|row| queued_work_batch_from_conn(conn, row))
1364                        .collect()
1365                })();
1366                Ok(outcome)
1367            })
1368            .await
1369            .map_err(sqlite_error)?
1370    }
1371
1372    async fn list_pending_queued_work(
1373        &self,
1374        session_id: &str,
1375    ) -> Result<Vec<QueuedWorkBatch>, StoreError> {
1376        let session_id = session_id.to_string();
1377        let now = self.clock.timestamp_ms();
1378        self.conn
1379            .call(move |conn| {
1380                let outcome: Result<Vec<QueuedWorkBatch>, StoreError> = (|| {
1381                    let rows = {
1382                        let mut stmt = conn
1383                            .prepare(
1384                                "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1385                                        slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1386                                        claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1387                                        claim_owner_liveness_json, claim_token, claim_session_lease_generation
1388                                 FROM queued_work_batches
1389                                 WHERE session_id = ?1
1390                                   AND (claim_token IS NULL OR NOT EXISTS (
1391                                        SELECT 1 FROM session_execution_leases sel
1392                                        WHERE sel.session_id = ?1
1393                                          AND sel.lease_token IS NOT NULL
1394                                          AND sel.lease_expires_at_ms > ?2
1395                                          AND sel.lease_fencing_token
1396                                              = queued_work_batches.claim_session_lease_generation
1397                                   ))
1398                                 ORDER BY enqueue_seq ASC",
1399                            )
1400                            .map_err(sqlite_error)?;
1401                        let rows = stmt
1402                            .query_map(
1403                                params![session_id, now as i64],
1404                                queued_batch_row_from_sql,
1405                            )
1406                            .map_err(sqlite_error)?;
1407                        rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1408                    };
1409                    rows.into_iter()
1410                        .map(|row| queued_work_batch_from_conn(conn, row))
1411                        .collect()
1412                })();
1413                Ok(outcome)
1414            })
1415            .await
1416            .map_err(sqlite_error)?
1417    }
1418}
1419
1420#[async_trait::async_trait]
1421impl TurnInputStore for Store {
1422    async fn enqueue_pending_turn_input(
1423        &self,
1424        draft: lash_core::PendingTurnInputDraft,
1425    ) -> Result<lash_core::PendingTurnInput, StoreError> {
1426        let nonce = self.commit_count.fetch_add(1, AtomicOrdering::Relaxed);
1427        let now = self.clock.timestamp_ms();
1428        self.conn
1429            .write_flow(move |tx| {
1430                let outcome: Result<lash_core::PendingTurnInput, StoreError> = (|| {
1431                    if let Some(source_key) = draft.source_key.as_deref() {
1432                        let existing_id: Option<String> = tx
1433                            .query_row(
1434                                "SELECT input_id
1435                                 FROM pending_turn_inputs
1436                                 WHERE session_id = ?1 AND source_key = ?2",
1437                                params![draft.session_id, source_key],
1438                                |row| row.get(0),
1439                            )
1440                            .optional()
1441                            .map_err(sqlite_error)?;
1442                        if let Some(input_id) = existing_id {
1443                            let existing = load_pending_turn_input_by_id_conn(
1444                                tx,
1445                                &draft.session_id,
1446                                &input_id,
1447                            )?
1448                            .ok_or_else(|| {
1449                                StoreError::Backend(
1450                                    "pending turn input source row disappeared".to_string(),
1451                                )
1452                            })?;
1453                            if !draft.submitted_content_matches(&existing).map_err(|err| {
1454                                StoreError::Backend(format!(
1455                                    "failed to compare pending turn input submission: {err}"
1456                                ))
1457                            })? {
1458                                return Err(StoreError::PendingTurnInputSourceKeyConflict {
1459                                    session_id: draft.session_id.clone(),
1460                                    source_key: source_key.to_string(),
1461                                    existing_input_id: existing.input_id.clone(),
1462                                });
1463                            }
1464                            return Ok(existing);
1465                        }
1466                    }
1467                    let input_id = draft.input_id.clone().unwrap_or_else(|| {
1468                        derive_pending_turn_input_id(
1469                            &draft.session_id,
1470                            draft.source_key.as_deref(),
1471                            now,
1472                            nonce,
1473                        )
1474                    });
1475                    let state = match draft.ingress {
1476                        lash_core::TurnInputIngress::ActiveTurn { .. } => {
1477                            lash_core::TurnInputState::PendingActive
1478                        }
1479                        lash_core::TurnInputIngress::NextTurn => {
1480                            lash_core::TurnInputState::DeferredNextTurn
1481                        }
1482                    };
1483                    tx.execute(
1484                        "INSERT INTO pending_turn_inputs (
1485                            input_id, session_id, source_key, ingress_json, state,
1486                            input_json, enqueued_at_ms
1487                         )
1488                         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
1489                        params![
1490                            input_id,
1491                            draft.session_id,
1492                            draft.source_key.as_deref(),
1493                            encode_json(&draft.ingress),
1494                            state.as_str(),
1495                            encode_json(&draft.input),
1496                            now as i64,
1497                        ],
1498                    )
1499                    .map_err(sqlite_error)?;
1500                    load_pending_turn_input_by_id_conn(tx, &draft.session_id, &input_id)?
1501                        .ok_or_else(|| {
1502                            StoreError::Backend("pending turn input insert disappeared".to_string())
1503                        })
1504                })();
1505                match outcome {
1506                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1507                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1508                }
1509            })
1510            .await
1511            .map_err(sqlite_error)?
1512    }
1513
1514    async fn list_pending_turn_inputs(
1515        &self,
1516        session_id: &str,
1517    ) -> Result<Vec<lash_core::PendingTurnInput>, StoreError> {
1518        let session_id = session_id.to_string();
1519        let now = self.clock.timestamp_ms();
1520        self.conn
1521            .call(move |conn| {
1522                let outcome: Result<Vec<lash_core::PendingTurnInput>, StoreError> = (|| {
1523                    let rows = {
1524                        let mut stmt = conn
1525                            .prepare(
1526                                "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1527                                        state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1528                                        claim_owner_id, claim_owner_incarnation_id,
1529                                        claim_owner_liveness_json, claim_token, claim_session_lease_generation
1530                                 FROM pending_turn_inputs
1531                                 WHERE session_id = ?1
1532                                   AND state IN (?2, ?3)
1533                                   AND (claim_token IS NULL OR NOT EXISTS (
1534                                        SELECT 1 FROM session_execution_leases sel
1535                                        WHERE sel.session_id = ?1
1536                                          AND sel.lease_token IS NOT NULL
1537                                          AND sel.lease_expires_at_ms > ?4
1538                                          AND sel.lease_fencing_token
1539                                              = pending_turn_inputs.claim_session_lease_generation
1540                                   ))
1541                                 ORDER BY enqueue_seq ASC",
1542                            )
1543                            .map_err(sqlite_error)?;
1544                        let rows = stmt
1545                            .query_map(
1546                                params![
1547                                    session_id,
1548                                    lash_core::TurnInputState::PendingActive.as_str(),
1549                                    lash_core::TurnInputState::DeferredNextTurn.as_str(),
1550                                    now as i64
1551                                ],
1552                                pending_turn_input_row_from_sql,
1553                            )
1554                            .map_err(sqlite_error)?;
1555                        rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1556                    };
1557                    rows.into_iter().map(pending_turn_input_from_row).collect()
1558                })(
1559                );
1560                Ok(outcome)
1561            })
1562            .await
1563            .map_err(sqlite_error)?
1564    }
1565
1566    async fn cancel_pending_turn_inputs(
1567        &self,
1568        session_id: &str,
1569        targets: &[lash_core::PendingTurnInputCancelTarget],
1570    ) -> Result<Vec<lash_core::PendingTurnInputCancelResult>, StoreError> {
1571        let session_id = session_id.to_string();
1572        let targets = targets.to_vec();
1573        let now = self.clock.timestamp_ms();
1574        self.conn
1575            .write_flow(move |tx| {
1576                let outcome: Result<Vec<lash_core::PendingTurnInputCancelResult>, StoreError> =
1577                    (|| {
1578                        let mut results = Vec::with_capacity(targets.len());
1579                        for target in targets {
1580                            let outcome = match load_pending_turn_input_row_by_target_conn(
1581                                tx,
1582                                &session_id,
1583                                &target,
1584                            )? {
1585                                Some(row) => cancel_pending_turn_input_row_conn(tx, row, now)?,
1586                                None => lash_core::PendingTurnInputCancelOutcome::NotFound,
1587                            };
1588                            results
1589                                .push(lash_core::PendingTurnInputCancelResult { target, outcome });
1590                        }
1591                        Ok(results)
1592                    })();
1593                match outcome {
1594                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1595                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1596                }
1597            })
1598            .await
1599            .map_err(sqlite_error)?
1600    }
1601
1602    async fn cancel_pending_turn_input_suffix(
1603        &self,
1604        session_id: &str,
1605        anchor: &lash_core::PendingTurnInputCancelTarget,
1606    ) -> Result<lash_core::PendingTurnInputSuffixCancelOutcome, StoreError> {
1607        let session_id = session_id.to_string();
1608        let anchor = anchor.clone();
1609        let now = self.clock.timestamp_ms();
1610        self.conn
1611            .write_flow(move |tx| {
1612                let outcome: Result<lash_core::PendingTurnInputSuffixCancelOutcome, StoreError> =
1613                    (|| {
1614                        let Some(anchor_row) =
1615                            load_pending_turn_input_row_by_target_conn(tx, &session_id, &anchor)?
1616                        else {
1617                            return Ok(
1618                                lash_core::PendingTurnInputSuffixCancelOutcome::AnchorNotFound {
1619                                    anchor,
1620                                },
1621                            );
1622                        };
1623                        let rows = {
1624                            let mut stmt = tx
1625                                .prepare(
1626                                    "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1627                                            state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1628                                            claim_owner_id, claim_owner_incarnation_id,
1629                                            claim_owner_liveness_json, claim_token, claim_session_lease_generation
1630                                     FROM pending_turn_inputs
1631                                     WHERE session_id = ?1 AND enqueue_seq >= ?2
1632                                     ORDER BY enqueue_seq ASC",
1633                                )
1634                                .map_err(sqlite_error)?;
1635                            let rows = stmt
1636                                .query_map(
1637                                    params![session_id, anchor_row.enqueue_seq as i64],
1638                                    pending_turn_input_row_from_sql,
1639                                )
1640                                .map_err(sqlite_error)?;
1641                            rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1642                        };
1643                        let mut outcomes = Vec::with_capacity(rows.len());
1644                        for row in rows {
1645                            outcomes.push(cancel_pending_turn_input_row_conn(tx, row, now)?);
1646                        }
1647                        Ok(lash_core::PendingTurnInputSuffixCancelOutcome::Outcomes {
1648                            anchor,
1649                            outcomes,
1650                        })
1651                    })();
1652                match outcome {
1653                    Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1654                    Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1655                }
1656            })
1657            .await
1658            .map_err(sqlite_error)?
1659    }
1660
1661    async fn claim_active_turn_inputs(
1662        &self,
1663        session_id: &str,
1664        session_execution_lease: &SessionExecutionLeaseFence,
1665        owner: &LeaseOwnerIdentity,
1666        turn_id: &str,
1667        checkpoint: lash_core::CheckpointKind,
1668        max_inputs: usize,
1669    ) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1670        claim_pending_turn_inputs_sqlite(
1671            &self.conn,
1672            self.clock.timestamp_ms(),
1673            session_id,
1674            session_execution_lease,
1675            owner,
1676            max_inputs,
1677            lash_core::TurnInputClaimMode::ActiveTurn {
1678                turn_id: turn_id.to_string(),
1679                checkpoint,
1680            },
1681        )
1682        .await
1683    }
1684
1685    async fn claim_next_turn_inputs(
1686        &self,
1687        session_id: &str,
1688        session_execution_lease: &SessionExecutionLeaseFence,
1689        owner: &LeaseOwnerIdentity,
1690        max_inputs: usize,
1691    ) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1692        claim_pending_turn_inputs_sqlite(
1693            &self.conn,
1694            self.clock.timestamp_ms(),
1695            session_id,
1696            session_execution_lease,
1697            owner,
1698            max_inputs,
1699            lash_core::TurnInputClaimMode::NextTurn,
1700        )
1701        .await
1702    }
1703
1704    async fn abandon_turn_input_claim(
1705        &self,
1706        claim: &lash_core::TurnInputClaim,
1707    ) -> Result<(), StoreError> {
1708        let session_id = claim.session_id.clone();
1709        let claim_id = claim.claim_id.clone();
1710        let lease_token = claim.lease_token.clone();
1711        let restored_state = match claim.mode {
1712            lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1713                lash_core::TurnInputState::PendingActive
1714            }
1715            lash_core::TurnInputClaimMode::NextTurn => lash_core::TurnInputState::DeferredNextTurn,
1716        };
1717        self.conn
1718            .write(move |tx| {
1719                tx.execute(
1720                    "UPDATE pending_turn_inputs
1721                     SET state = CASE
1722                             WHEN state = ?4 THEN ?5
1723                             ELSE state
1724                         END,
1725                         claim_id = NULL,
1726                         claim_owner_id = NULL,
1727                         claim_owner_incarnation_id = NULL,
1728                         claim_owner_liveness_json = NULL,
1729                         claim_token = NULL,
1730                         claim_session_lease_generation = 0
1731                     WHERE session_id = ?1 AND claim_id = ?2 AND claim_token = ?3",
1732                    params![
1733                        session_id,
1734                        claim_id,
1735                        lease_token,
1736                        lash_core::TurnInputState::Accepted.as_str(),
1737                        restored_state.as_str(),
1738                    ],
1739                )
1740            })
1741            .await
1742            .map_err(sqlite_error)?;
1743        Ok(())
1744    }
1745}
1746
1747#[async_trait::async_trait]
1748impl StoreMaintenance for Store {
1749    async fn tombstone_nodes(&self, ids: &[String]) -> Result<(), StoreError> {
1750        if ids.is_empty() {
1751            return Ok(());
1752        }
1753        let ids = ids.to_vec();
1754        self.conn
1755            .write(move |tx| {
1756                let mut stmt =
1757                    tx.prepare("UPDATE graph_nodes SET tombstoned = 1 WHERE node_id = ?1")?;
1758                for id in &ids {
1759                    stmt.execute(params![id])?;
1760                }
1761                Ok(())
1762            })
1763            .await
1764            .map_err(sqlite_error)
1765    }
1766
1767    async fn vacuum(&self) -> Result<VacuumReport, StoreError> {
1768        let (removed_node_count, removed_pending_turn_input_tombstone_count) = self
1769            .conn
1770            .write(move |tx| {
1771                let removed_node_count =
1772                    tx.execute("DELETE FROM graph_nodes WHERE tombstoned = 1", [])?;
1773                let removed_pending_turn_input_tombstone_count = tx.execute(
1774                    "DELETE FROM pending_turn_inputs
1775                     WHERE state IN (?1, ?2)",
1776                    params![
1777                        lash_core::TurnInputState::Cancelled.as_str(),
1778                        lash_core::TurnInputState::Completed.as_str()
1779                    ],
1780                )?;
1781                Ok((
1782                    removed_node_count,
1783                    removed_pending_turn_input_tombstone_count,
1784                ))
1785            })
1786            .await
1787            .map_err(sqlite_error)?;
1788        Ok(VacuumReport {
1789            removed_node_count,
1790            removed_pending_turn_input_tombstone_count,
1791        })
1792    }
1793
1794    async fn gc_unreachable(&self) -> Result<GcReport, StoreError> {
1795        Ok(Store::gc_unreachable(self).await)
1796    }
1797}
1798
1799fn derive_pending_turn_input_id(
1800    session_id: &str,
1801    source_key: Option<&str>,
1802    now_epoch_ms: u64,
1803    nonce: u64,
1804) -> String {
1805    format!(
1806        "ti:{:x}",
1807        Sha256::digest(format!("{session_id}:{source_key:?}:{now_epoch_ms}:{nonce}").as_bytes())
1808    )
1809}
1810
1811fn cancel_pending_turn_input_row_conn(
1812    conn: &Connection,
1813    row: PendingTurnInputRow,
1814    now_epoch_ms: u64,
1815) -> Result<lash_core::PendingTurnInputCancelOutcome, StoreError> {
1816    let mut input = pending_turn_input_from_row(row.clone())?;
1817    match input.state {
1818        lash_core::TurnInputState::Cancelled => Ok(
1819            lash_core::PendingTurnInputCancelOutcome::AlreadyCancelled(input),
1820        ),
1821        lash_core::TurnInputState::Completed => Ok(
1822            lash_core::PendingTurnInputCancelOutcome::AlreadyCompleted(input),
1823        ),
1824        lash_core::TurnInputState::Accepted => {
1825            Ok(lash_core::PendingTurnInputCancelOutcome::AlreadyClaimed {
1826                claim: pending_turn_input_claim_diagnostics_from_row(&row, input.state),
1827                input,
1828            })
1829        }
1830        lash_core::TurnInputState::PendingActive | lash_core::TurnInputState::DeferredNextTurn => {
1831            // A claim is live only while the session-execution-lease generation it
1832            // pins still holds the session lease (ADR 0029).
1833            let live_claim = row.claim_token.is_some()
1834                && load_session_execution_lease_row_conn(conn, &row.session_id)?.is_some_and(
1835                    |lease| {
1836                        lease.lease_token.is_some()
1837                            && lease.expires_at_ms > now_epoch_ms
1838                            && lease.fencing_token == row.claim_session_lease_generation
1839                    },
1840                );
1841            if live_claim {
1842                return Ok(lash_core::PendingTurnInputCancelOutcome::AlreadyClaimed {
1843                    claim: pending_turn_input_claim_diagnostics_from_row(&row, input.state),
1844                    input,
1845                });
1846            }
1847            conn.execute(
1848                "UPDATE pending_turn_inputs
1849                 SET state = ?3,
1850                     claim_id = NULL,
1851                     claim_owner_id = NULL,
1852                     claim_owner_incarnation_id = NULL,
1853                     claim_owner_liveness_json = NULL,
1854                     claim_token = NULL,
1855                     claim_session_lease_generation = 0
1856                 WHERE session_id = ?1 AND input_id = ?2",
1857                params![
1858                    row.session_id,
1859                    row.input_id,
1860                    lash_core::TurnInputState::Cancelled.as_str(),
1861                ],
1862            )
1863            .map_err(sqlite_error)?;
1864            input.state = lash_core::TurnInputState::Cancelled;
1865            Ok(lash_core::PendingTurnInputCancelOutcome::Cancelled(input))
1866        }
1867    }
1868}
1869
1870async fn claim_pending_turn_inputs_sqlite(
1871    conn: &SqliteConnection,
1872    now: u64,
1873    session_id: &str,
1874    session_execution_lease: &SessionExecutionLeaseFence,
1875    owner: &LeaseOwnerIdentity,
1876    max_inputs: usize,
1877    mode: lash_core::TurnInputClaimMode,
1878) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1879    if max_inputs == 0 {
1880        return Ok(None);
1881    }
1882    let session_id = session_id.to_string();
1883    let session_execution_lease = session_execution_lease.clone();
1884    let owner = owner.clone();
1885    conn.write_flow(move |tx| {
1886        let outcome: Result<TxOutcome<Option<lash_core::TurnInputClaim>>, StoreError> = (|| {
1887            ensure_session_execution_lease_conn(
1888                tx,
1889                &session_id,
1890                &session_execution_lease,
1891                now,
1892            )?;
1893            let generation = session_execution_lease.fencing_token;
1894            let wanted_state = match &mode {
1895                lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1896                    lash_core::TurnInputState::PendingActive
1897                }
1898                lash_core::TurnInputClaimMode::NextTurn => {
1899                    lash_core::TurnInputState::DeferredNextTurn
1900                }
1901            };
1902            let candidate_rows = {
1903                let mut stmt = tx
1904                    .prepare(
1905                        "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1906                                state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1907                                claim_owner_id, claim_owner_incarnation_id,
1908                                claim_owner_liveness_json, claim_token, claim_session_lease_generation
1909                         FROM pending_turn_inputs
1910                         WHERE session_id = ?1 AND state = ?2
1911                           AND (
1912                                claim_token IS NULL
1913                                OR claim_session_lease_generation <> ?3
1914                           )
1915                         ORDER BY enqueue_seq ASC
1916                         LIMIT ?4",
1917                    )
1918                    .map_err(sqlite_error)?;
1919                let rows = stmt
1920                    .query_map(
1921                        params![
1922                            session_id,
1923                            wanted_state.as_str(),
1924                            generation as i64,
1925                            (max_inputs as i64).saturating_add(32)
1926                        ],
1927                        pending_turn_input_row_from_sql,
1928                    )
1929                    .map_err(sqlite_error)?;
1930                rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1931            };
1932            let mut selected = Vec::new();
1933            for row in candidate_rows {
1934                let claim_available = row.claim_token.is_none()
1935                    || row.claim_session_lease_generation != generation;
1936                if !claim_available {
1937                    continue;
1938                }
1939                let input = pending_turn_input_from_row(row.clone())?;
1940                let matches_mode = match &mode {
1941                    lash_core::TurnInputClaimMode::ActiveTurn {
1942                        turn_id,
1943                        checkpoint,
1944                    } => {
1945                        input
1946                            .ingress
1947                            .active_turn_id()
1948                            .is_some_and(|active| active == turn_id)
1949                            && input.ingress.admits_checkpoint(*checkpoint)
1950                    }
1951                    lash_core::TurnInputClaimMode::NextTurn => input.state.is_next_turn_pending(),
1952                };
1953                if matches_mode {
1954                    selected.push((row, input));
1955                    if selected.len() >= max_inputs {
1956                        break;
1957                    }
1958                }
1959            }
1960            let Some((head, _)) = selected.first() else {
1961                return Ok(TxOutcome::Commit(None));
1962            };
1963            let lease = TurnInputClaimLease::derive(head, &session_id, &owner, now, generation);
1964            let liveness_json = encode_liveness(&owner.liveness)?;
1965            let state_after_claim = match &mode {
1966                lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1967                    lash_core::TurnInputState::Accepted
1968                }
1969                lash_core::TurnInputClaimMode::NextTurn => {
1970                    lash_core::TurnInputState::DeferredNextTurn
1971                }
1972            };
1973            let mut inputs = Vec::new();
1974            for (row, mut input) in selected {
1975                let claimed = tx
1976                    .execute(
1977                        "UPDATE pending_turn_inputs
1978                         SET state = ?3,
1979                             claim_id = ?4,
1980                             claim_owner_id = ?5,
1981                             claim_owner_incarnation_id = ?6,
1982                             claim_owner_liveness_json = ?7,
1983                             claim_token = ?8,
1984                             claim_fencing_token = claim_fencing_token + 1,
1985                             claim_session_lease_generation = ?9
1986                         WHERE session_id = ?1
1987                           AND input_id = ?2
1988                           AND (
1989                                claim_token IS NULL
1990                                OR claim_session_lease_generation <> ?9
1991                           )",
1992                        params![
1993                            session_id,
1994                            row.input_id,
1995                            state_after_claim.as_str(),
1996                            lease.claim_id,
1997                            owner.owner_id.as_str(),
1998                            owner.incarnation_id.as_str(),
1999                            liveness_json.as_str(),
2000                            lease.lease_token,
2001                            lease.session_lease_generation as i64,
2002                        ],
2003                    )
2004                    .map_err(sqlite_error)?;
2005                if claimed == 0 {
2006                    return Ok(TxOutcome::Rollback(None));
2007                }
2008                input.state = state_after_claim;
2009                inputs.push(input);
2010            }
2011            Ok(TxOutcome::Commit(Some(lash_core::TurnInputClaim {
2012                session_id: session_id.clone(),
2013                claim_id: lease.claim_id,
2014                owner: owner.clone(),
2015                lease_token: lease.lease_token,
2016                fencing_token: lease.fencing_token,
2017                session_lease_generation: lease.session_lease_generation,
2018                mode,
2019                inputs,
2020            })))
2021        })(
2022        );
2023        match outcome {
2024            Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
2025            Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
2026            Err(err) => Ok(TxOutcome::Rollback(Err(err))),
2027        }
2028    })
2029    .await
2030    .map_err(sqlite_error)?
2031}
2032
2033struct SessionExecutionLeaseRow {
2034    owner: Option<LeaseOwnerIdentity>,
2035    lease_token: Option<String>,
2036    fencing_token: u64,
2037    claimed_at_ms: u64,
2038    expires_at_ms: u64,
2039}
2040
2041fn load_session_execution_lease_row_conn(
2042    conn: &Connection,
2043    session_id: &str,
2044) -> Result<Option<SessionExecutionLeaseRow>, StoreError> {
2045    let row = conn
2046        .query_row(
2047            "SELECT lease_owner_id, lease_token, lease_fencing_token,
2048                    lease_claimed_at_ms, lease_expires_at_ms,
2049                    lease_owner_incarnation_id, lease_owner_liveness_json
2050             FROM session_execution_leases
2051             WHERE session_id = ?1",
2052            params![session_id],
2053            |row| {
2054                let owner_id: Option<String> = row.get(0)?;
2055                let incarnation_id: Option<String> = row.get(5)?;
2056                let liveness_json: Option<String> = row.get(6)?;
2057                Ok(SessionExecutionLeaseRow {
2058                    owner: lease_owner_from_columns(owner_id, incarnation_id, liveness_json),
2059                    lease_token: row.get(1)?,
2060                    fencing_token: row.get::<_, i64>(2)? as u64,
2061                    claimed_at_ms: row.get::<_, i64>(3)? as u64,
2062                    expires_at_ms: row.get::<_, i64>(4)? as u64,
2063                })
2064            },
2065        )
2066        .optional()
2067        .map_err(sqlite_error)?;
2068    Ok(row)
2069}
2070
2071fn lease_owner_from_columns(
2072    owner_id: Option<String>,
2073    incarnation_id: Option<String>,
2074    liveness_json: Option<String>,
2075) -> Option<LeaseOwnerIdentity> {
2076    owner_id.map(|owner_id| LeaseOwnerIdentity {
2077        incarnation_id: incarnation_id.unwrap_or_else(|| owner_id.clone()),
2078        owner_id,
2079        liveness: liveness_json
2080            .as_deref()
2081            .and_then(|json| serde_json::from_str(json).ok())
2082            .unwrap_or(LeaseOwnerLiveness::Opaque),
2083    })
2084}
2085
2086fn encode_liveness(liveness: &LeaseOwnerLiveness) -> Result<String, StoreError> {
2087    serde_json::to_string(liveness)
2088        .map_err(|err| StoreError::Backend(format!("failed to encode lease liveness: {err}")))
2089}
2090
2091fn row_to_session_execution_lease(
2092    session_id: &str,
2093    row: SessionExecutionLeaseRow,
2094) -> Result<SessionExecutionLease, StoreError> {
2095    Ok(SessionExecutionLease {
2096        session_id: session_id.to_string(),
2097        owner: row
2098            .owner
2099            .ok_or_else(|| StoreError::Backend("live session lease missing owner".to_string()))?,
2100        lease_token: row.lease_token.ok_or_else(|| {
2101            StoreError::Backend("live session lease missing lease token".to_string())
2102        })?,
2103        fencing_token: row.fencing_token,
2104        claimed_at_epoch_ms: row.claimed_at_ms,
2105        expires_at_epoch_ms: row.expires_at_ms,
2106    })
2107}
2108
2109fn acquire_session_execution_lease_conn(
2110    conn: &Connection,
2111    session_id: &str,
2112    owner: &LeaseOwnerIdentity,
2113    previous_fencing_token: u64,
2114    now: u64,
2115    lease_ttl_ms: u64,
2116) -> Result<SessionExecutionLease, StoreError> {
2117    let fencing_token = previous_fencing_token.saturating_add(1);
2118    let lease_token = format!(
2119        "{}:{}:{}:{now}:{fencing_token}",
2120        session_id, owner.owner_id, owner.incarnation_id
2121    );
2122    let expires_at = now.saturating_add(lease_ttl_ms);
2123    let liveness_json = encode_liveness(&owner.liveness)?;
2124    conn.execute(
2125        "INSERT INTO session_execution_leases (
2126            session_id, lease_owner_id, lease_owner_incarnation_id, lease_owner_liveness_json,
2127            lease_token, lease_fencing_token, lease_claimed_at_ms, lease_expires_at_ms
2128         )
2129         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
2130         ON CONFLICT(session_id) DO UPDATE SET
2131            lease_owner_id = excluded.lease_owner_id,
2132            lease_owner_incarnation_id = excluded.lease_owner_incarnation_id,
2133            lease_owner_liveness_json = excluded.lease_owner_liveness_json,
2134            lease_token = excluded.lease_token,
2135            lease_fencing_token = excluded.lease_fencing_token,
2136            lease_claimed_at_ms = excluded.lease_claimed_at_ms,
2137            lease_expires_at_ms = excluded.lease_expires_at_ms",
2138        params![
2139            session_id,
2140            owner.owner_id,
2141            owner.incarnation_id,
2142            liveness_json,
2143            lease_token,
2144            fencing_token as i64,
2145            now as i64,
2146            expires_at as i64
2147        ],
2148    )
2149    .map_err(sqlite_error)?;
2150    Ok(SessionExecutionLease {
2151        session_id: session_id.to_string(),
2152        owner: owner.clone(),
2153        lease_token,
2154        fencing_token,
2155        claimed_at_epoch_ms: now,
2156        expires_at_epoch_ms: expires_at,
2157    })
2158}
2159
2160fn ensure_session_execution_lease_conn(
2161    conn: &Connection,
2162    session_id: &str,
2163    fence: &SessionExecutionLeaseFence,
2164    now: u64,
2165) -> Result<(), StoreError> {
2166    if fence.session_id != session_id {
2167        return Err(StoreError::SessionExecutionLeaseExpired {
2168            session_id: session_id.to_string(),
2169        });
2170    }
2171    let current = load_session_execution_lease_row_conn(conn, session_id)?;
2172    let Some(current) = current else {
2173        return Err(StoreError::SessionExecutionLeaseExpired {
2174            session_id: session_id.to_string(),
2175        });
2176    };
2177    if current
2178        .owner
2179        .as_ref()
2180        .is_some_and(|owner| owner.same_incarnation(&fence.owner))
2181        && current.lease_token.as_deref() == Some(fence.lease_token.as_str())
2182        && current.fencing_token == fence.fencing_token
2183        && current.expires_at_ms > now
2184    {
2185        Ok(())
2186    } else {
2187        Err(StoreError::SessionExecutionLeaseExpired {
2188            session_id: session_id.to_string(),
2189        })
2190    }
2191}
2192
2193fn release_session_execution_lease_conn(
2194    conn: &Connection,
2195    completion: &SessionExecutionLeaseCompletion,
2196) -> Result<(), StoreError> {
2197    conn.execute(
2198        "UPDATE session_execution_leases
2199         SET lease_owner_id = NULL,
2200             lease_owner_incarnation_id = NULL,
2201             lease_owner_liveness_json = NULL,
2202             lease_token = NULL,
2203             lease_claimed_at_ms = 0,
2204             lease_expires_at_ms = 0
2205         WHERE session_id = ?1
2206           AND lease_owner_id = ?2
2207           AND lease_owner_incarnation_id = ?3
2208           AND lease_token = ?4
2209           AND lease_fencing_token = ?5",
2210        params![
2211            completion.session_id,
2212            completion.owner.owner_id,
2213            completion.owner.incarnation_id,
2214            completion.lease_token,
2215            completion.fencing_token as i64
2216        ],
2217    )
2218    .map_err(sqlite_error)?;
2219    Ok(())
2220}