Skip to main content

lash_sqlite_store/
persistence.rs

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