a3s_code_core/capability/
transaction.rs

1use std::collections::BTreeMap;
2use std::fmt;
3use std::marker::PhantomData;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use tokio_util::sync::CancellationToken;
8
9use super::projection::CatalogInner;
10use super::{
11    CapabilityAdapterError, CapabilityCatalog, CapabilityCatalogStamp, CapabilityCommitReceipt,
12    CapabilityEffect, CapabilityId, CapabilityProjection, CapabilityProjectionError,
13    CapabilityReadinessPlan, CapabilitySet, CapabilityValue, UseGenerationLeaseProvider,
14};
15
16pub const MAX_CAPABILITY_TRANSACTION_EFFECTS: usize = 4_096;
17
18/// Transaction state before fallible runtime preparation.
19#[derive(Debug)]
20pub struct Staged;
21
22/// Transaction state after every adapter prepared successfully.
23#[derive(Debug)]
24pub struct Prepared;
25
26/// Transaction state after the complete value projection passed validation.
27#[derive(Debug)]
28pub struct Validated;
29
30/// One atomically returned runtime value and its reversible resources.
31///
32/// Adapters must return acquired resources in the same successful value. If
33/// their future is cancelled before returning, they remain responsible for
34/// cancellation-safe local RAII cleanup. Once returned, the transaction owns
35/// every effect and transfers it to the catalog rollback queue on any failure.
36#[must_use = "prepared capability effects must be transferred into a transaction"]
37pub struct PreparedCapability {
38    value: CapabilityValue,
39    effects: Vec<Box<dyn CapabilityEffect>>,
40}
41
42impl PreparedCapability {
43    pub fn new(value: CapabilityValue) -> Self {
44        Self {
45            value,
46            effects: Vec::new(),
47        }
48    }
49
50    pub fn push_effect<E>(&mut self, effect: E) -> Result<(), CapabilityAdapterError>
51    where
52        E: CapabilityEffect,
53    {
54        self.push_boxed_effect(Box::new(effect))
55    }
56
57    pub fn push_boxed_effect(
58        &mut self,
59        effect: Box<dyn CapabilityEffect>,
60    ) -> Result<(), CapabilityAdapterError> {
61        // The transaction checks the aggregate immediately after an adapter
62        // returns. Keeping this append infallible ensures an oversized batch
63        // is first transferred into catalog-owned rollback instead of dropping
64        // already acquired effects inside the adapter.
65        self.effects.push(effect);
66        Ok(())
67    }
68
69    fn into_parts(self) -> (CapabilityValue, Vec<Box<dyn CapabilityEffect>>) {
70        (self.value, self.effects)
71    }
72}
73
74/// Surface-owned fallible preparation boundary.
75///
76/// Tool, Skill, MCP, and other concerns implement this trait beside their
77/// native runtime types. The adapter does not resolve packages or dependencies;
78/// it projects one descriptor from the already selected A3S Use snapshot. A
79/// successful return is the surface readiness barrier: the adapter must not
80/// report success while its value still depends on unfinished initialization.
81#[async_trait]
82pub trait CapabilityProjectionAdapter: Send + 'static {
83    async fn prepare(
84        self: Box<Self>,
85        cancellation: CancellationToken,
86    ) -> Result<PreparedCapability, CapabilityAdapterError>;
87}
88
89struct ReadyValueAdapter(CapabilityValue);
90
91#[async_trait]
92impl CapabilityProjectionAdapter for ReadyValueAdapter {
93    async fn prepare(
94        self: Box<Self>,
95        _cancellation: CancellationToken,
96    ) -> Result<PreparedCapability, CapabilityAdapterError> {
97        Ok(PreparedCapability::new(self.0))
98    }
99}
100
101struct TransactionBody {
102    catalog: Arc<CatalogInner>,
103    base: CapabilityCatalogStamp,
104    target: Arc<CapabilitySet>,
105    readiness: Arc<CapabilityReadinessPlan>,
106    effects: Vec<Box<dyn CapabilityEffect>>,
107    rollback_armed: bool,
108}
109
110impl Drop for TransactionBody {
111    fn drop(&mut self) {
112        if self.rollback_armed {
113            let effects = std::mem::take(&mut self.effects);
114            self.catalog.enqueue_rollback(effects);
115        }
116    }
117}
118
119/// Atomic capability contribution transaction guarded by Rust typestate.
120///
121/// Only [`CapabilityTxn<Validated>`] exposes `commit`. A prepared transaction
122/// cannot publish, and dropping any uncommitted state transfers all completed
123/// effects to the catalog-owned asynchronous rollback queue.
124///
125/// ```compile_fail
126/// use a3s_code_core::capability::{CapabilityTxn, Prepared};
127///
128/// fn publish_without_validation(txn: CapabilityTxn<Prepared>) {
129///     let _ = txn.commit();
130/// }
131/// ```
132#[must_use = "capability transactions must be committed or drained as rollback"]
133pub struct CapabilityTxn<S> {
134    body: Option<TransactionBody>,
135    staged: BTreeMap<CapabilityId, Box<dyn CapabilityProjectionAdapter>>,
136    prepared: BTreeMap<CapabilityId, CapabilityValue>,
137    projection: Option<Arc<CapabilityProjection>>,
138    _state: PhantomData<S>,
139}
140
141impl CapabilityCatalog {
142    pub fn begin(
143        &self,
144        target: Arc<CapabilitySet>,
145    ) -> Result<CapabilityTxn<Staged>, CapabilityProjectionError> {
146        let base = self.current_stamp();
147        let expected = base
148            .generation()
149            .checked_next()
150            .ok_or(CapabilityProjectionError::GenerationExhausted)?;
151        if target.generation() != expected {
152            return Err(CapabilityProjectionError::TargetGenerationMismatch {
153                expected: expected.get(),
154                actual: target.generation().get(),
155            });
156        }
157        let readiness = Arc::new(CapabilityReadinessPlan::from_set(&target)?);
158        Ok(CapabilityTxn {
159            body: Some(TransactionBody {
160                catalog: Arc::clone(&self.inner),
161                base,
162                target,
163                readiness,
164                effects: Vec::new(),
165                rollback_armed: true,
166            }),
167            staged: BTreeMap::new(),
168            prepared: BTreeMap::new(),
169            projection: None,
170            _state: PhantomData,
171        })
172    }
173
174    /// Start a one-time recovery transaction from the untouched empty catalog
175    /// to an exact historical generation. This is intentionally separate from
176    /// ordinary N -> N+1 publication and is only usable before the Session has
177    /// published any scoped capability generation.
178    pub(crate) fn begin_recovery_bootstrap(
179        &self,
180        target: Arc<CapabilitySet>,
181    ) -> Result<CapabilityTxn<Staged>, CapabilityProjectionError> {
182        let current = self.pin();
183        let current_set = current.projection().set();
184        if current_set.generation().get() != 0 || !current_set.is_empty() {
185            return Err(CapabilityProjectionError::BootstrapUnavailable {
186                actual_generation: current_set.generation().get(),
187                actual_capabilities: current_set.len(),
188            });
189        }
190        if target.generation().get() == 0 {
191            return Err(CapabilityProjectionError::BootstrapTargetGeneration {
192                actual: target.generation().get(),
193            });
194        }
195        let base = current.stamp().clone();
196        drop(current);
197        let readiness = Arc::new(CapabilityReadinessPlan::from_set(&target)?);
198        Ok(CapabilityTxn {
199            body: Some(TransactionBody {
200                catalog: Arc::clone(&self.inner),
201                base,
202                target,
203                readiness,
204                effects: Vec::new(),
205                rollback_armed: true,
206            }),
207            staged: BTreeMap::new(),
208            prepared: BTreeMap::new(),
209            projection: None,
210            _state: PhantomData,
211        })
212    }
213}
214
215impl CapabilityTxn<Staged> {
216    pub fn stage<A>(
217        &mut self,
218        id: CapabilityId,
219        adapter: A,
220    ) -> Result<&mut Self, CapabilityProjectionError>
221    where
222        A: CapabilityProjectionAdapter,
223    {
224        let body = self
225            .body
226            .as_ref()
227            .ok_or(CapabilityProjectionError::InvalidTransactionState)?;
228        if !body.target.contains(&id) {
229            return Err(CapabilityProjectionError::UnknownStagedCapability {
230                capability: id.to_string(),
231            });
232        }
233        if self.staged.contains_key(&id) {
234            return Err(CapabilityProjectionError::DuplicateStagedCapability {
235                capability: id.to_string(),
236            });
237        }
238        self.staged.insert(id, Box::new(adapter));
239        Ok(self)
240    }
241
242    pub fn stage_value(
243        &mut self,
244        id: CapabilityId,
245        value: CapabilityValue,
246    ) -> Result<&mut Self, CapabilityProjectionError> {
247        self.stage(id, ReadyValueAdapter(value))
248    }
249
250    pub(crate) fn stage_boxed(
251        &mut self,
252        id: CapabilityId,
253        adapter: Box<dyn CapabilityProjectionAdapter>,
254    ) -> Result<&mut Self, CapabilityProjectionError> {
255        let body = self
256            .body
257            .as_ref()
258            .ok_or(CapabilityProjectionError::InvalidTransactionState)?;
259        if !body.target.contains(&id) {
260            return Err(CapabilityProjectionError::UnknownStagedCapability {
261                capability: id.to_string(),
262            });
263        }
264        if self.staged.contains_key(&id) {
265            return Err(CapabilityProjectionError::DuplicateStagedCapability {
266                capability: id.to_string(),
267            });
268        }
269        self.staged.insert(id, adapter);
270        Ok(self)
271    }
272
273    pub async fn prepare(
274        mut self,
275        cancellation: CancellationToken,
276    ) -> Result<CapabilityTxn<Prepared>, CapabilityProjectionError> {
277        if cancellation.is_cancelled() {
278            return Err(CapabilityProjectionError::Cancelled);
279        }
280        let body = self
281            .body
282            .as_ref()
283            .ok_or(CapabilityProjectionError::InvalidTransactionState)?;
284        if let Some((id, _)) = body
285            .target
286            .iter()
287            .find(|(id, _)| !self.staged.contains_key(*id))
288        {
289            return Err(CapabilityProjectionError::MissingStagedCapability {
290                capability: id.to_string(),
291            });
292        }
293        let activation_order = body.readiness.activation_order().to_vec();
294        for id in activation_order {
295            let adapter = self.staged.remove(&id).ok_or_else(|| {
296                CapabilityProjectionError::MissingStagedCapability {
297                    capability: id.to_string(),
298                }
299            })?;
300            let result = tokio::select! {
301                biased;
302                _ = cancellation.cancelled() => {
303                    return Err(CapabilityProjectionError::Cancelled);
304                }
305                result = adapter.prepare(cancellation.clone()) => result,
306            };
307            let prepared = result.map_err(|error| CapabilityProjectionError::PrepareFailed {
308                capability: id.to_string(),
309                message: error.message().to_owned(),
310            })?;
311            let (value, mut effects) = prepared.into_parts();
312            let body = self
313                .body
314                .as_mut()
315                .ok_or(CapabilityProjectionError::InvalidTransactionState)?;
316            body.effects.append(&mut effects);
317            if body.effects.len() > MAX_CAPABILITY_TRANSACTION_EFFECTS {
318                return Err(CapabilityProjectionError::EffectBoundExceeded {
319                    max: MAX_CAPABILITY_TRANSACTION_EFFECTS,
320                });
321            }
322            self.prepared.insert(id, value);
323        }
324        if cancellation.is_cancelled() {
325            return Err(CapabilityProjectionError::Cancelled);
326        }
327        self.transition()
328    }
329}
330
331impl CapabilityTxn<Prepared> {
332    pub fn validate(mut self) -> Result<CapabilityTxn<Validated>, CapabilityProjectionError> {
333        let target = Arc::clone(
334            &self
335                .body
336                .as_ref()
337                .ok_or(CapabilityProjectionError::InvalidTransactionState)?
338                .target,
339        );
340        let readiness = Arc::clone(
341            &self
342                .body
343                .as_ref()
344                .ok_or(CapabilityProjectionError::InvalidTransactionState)?
345                .readiness,
346        );
347        let values = std::mem::take(&mut self.prepared);
348        self.projection = Some(CapabilityProjection::with_readiness(
349            target, readiness, values,
350        )?);
351        self.transition()
352    }
353}
354
355impl CapabilityTxn<Validated> {
356    pub(crate) fn projection(&self) -> Result<&CapabilityProjection, CapabilityProjectionError> {
357        self.projection
358            .as_deref()
359            .ok_or(CapabilityProjectionError::InvalidTransactionState)
360    }
361
362    pub fn commit(mut self) -> Result<CapabilityCommitReceipt, CapabilityProjectionError> {
363        self.commit_inner(None)
364    }
365
366    pub(crate) fn commit_with_use_lease_provider(
367        mut self,
368        provider: Option<Arc<dyn UseGenerationLeaseProvider>>,
369    ) -> Result<CapabilityCommitReceipt, CapabilityProjectionError> {
370        self.commit_inner(provider)
371    }
372
373    fn commit_inner(
374        &mut self,
375        provider: Option<Arc<dyn UseGenerationLeaseProvider>>,
376    ) -> Result<CapabilityCommitReceipt, CapabilityProjectionError> {
377        let projection = self
378            .projection
379            .take()
380            .ok_or(CapabilityProjectionError::InvalidTransactionState)?;
381        let mut body = self
382            .body
383            .take()
384            .ok_or(CapabilityProjectionError::InvalidTransactionState)?;
385        let effects = std::mem::take(&mut body.effects);
386        let result = body
387            .catalog
388            .publish(&body.base, projection, provider, effects);
389        // `publish` owns the effect batch on both success and CAS conflict.
390        body.rollback_armed = false;
391        result
392    }
393}
394
395impl<S> CapabilityTxn<S> {
396    fn transition<T>(mut self) -> Result<CapabilityTxn<T>, CapabilityProjectionError> {
397        let body = self
398            .body
399            .take()
400            .ok_or(CapabilityProjectionError::InvalidTransactionState)?;
401        Ok(CapabilityTxn {
402            body: Some(body),
403            staged: std::mem::take(&mut self.staged),
404            prepared: std::mem::take(&mut self.prepared),
405            projection: self.projection.take(),
406            _state: PhantomData,
407        })
408    }
409
410    pub fn base(&self) -> Result<&CapabilityCatalogStamp, CapabilityProjectionError> {
411        self.body
412            .as_ref()
413            .map(|body| &body.base)
414            .ok_or(CapabilityProjectionError::InvalidTransactionState)
415    }
416
417    pub fn target(&self) -> Result<&CapabilitySet, CapabilityProjectionError> {
418        self.body
419            .as_ref()
420            .map(|body| body.target.as_ref())
421            .ok_or(CapabilityProjectionError::InvalidTransactionState)
422    }
423}
424
425impl<S> fmt::Debug for CapabilityTxn<S> {
426    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
427        formatter
428            .debug_struct("CapabilityTxn")
429            .field("base", &self.body.as_ref().map(|body| &body.base))
430            .field(
431                "target_generation",
432                &self.body.as_ref().map(|body| body.target.generation()),
433            )
434            .field("staged", &self.staged.len())
435            .field("prepared", &self.prepared.len())
436            .field("validated", &self.projection.is_some())
437            .finish()
438    }
439}
a3s_code_core/capability/transaction.rs

a3s_code_core/capability/
transaction.rs
