Skip to main content

lash_sqlite_store/
persistence.rs

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