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