a3s_code_core/capability/
runtime.rs

1use std::collections::BTreeMap;
2use std::fmt;
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use thiserror::Error;
7use tokio_util::sync::CancellationToken;
8
9use super::{
10    CapabilityCatalog, CapabilityCeiling, CapabilityCommitReceipt, CapabilityId, CapabilityKind,
11    CapabilityProjectionAdapter, CapabilityProjectionError, CapabilityProjectionLease,
12    CapabilityScope, CapabilityScopeError, CapabilitySet, CapabilityTxn, CapabilityValue, Prepared,
13    RetainedUseGeneration, Run, ScopeCloseReport, Session, Staged, UseCapabilityGeneration,
14    Validated,
15};
16
17const MAX_USE_LEASE_ERROR_BYTES: usize = 1_024;
18
19/// Bounded failure returned while retaining one exact A3S Use generation.
20#[derive(Clone, Debug, Eq, Error, PartialEq)]
21#[error("{message}")]
22pub struct UseGenerationLeaseError {
23    message: Box<str>,
24}
25
26impl UseGenerationLeaseError {
27    pub fn new(message: impl Into<String>) -> Self {
28        let message = message.into();
29        let message = if message.is_empty() {
30            "A3S Use generation lease acquisition failed".to_owned()
31        } else {
32            truncate_utf8(message, MAX_USE_LEASE_ERROR_BYTES)
33        };
34        Self {
35            message: message.into_boxed_str(),
36        }
37    }
38
39    pub fn message(&self) -> &str {
40        &self.message
41    }
42}
43
44/// Generation-bound host seam for the real non-clone A3S Use snapshot lease.
45///
46/// A provider is published in the same catalog CAS as its projection. Each Run
47/// calls [`Self::acquire`] again so A3S Use can reject a generation that has
48/// become hidden or stale since Code installed the projection. Implementations
49/// must retain the concrete `a3s_use::CapabilitySnapshotLease` inside the
50/// returned [`RetainedUseGeneration`] value.
51#[async_trait]
52pub trait UseGenerationLeaseProvider: Send + Sync + 'static {
53    fn use_generation(&self) -> &UseCapabilityGeneration;
54
55    async fn acquire(
56        &self,
57        cancellation: CancellationToken,
58    ) -> Result<Box<dyn RetainedUseGeneration>, UseGenerationLeaseError>;
59}
60
61/// Session-host failure that cannot expose a partial capability generation.
62#[derive(Clone, Debug, Eq, Error, PartialEq)]
63pub enum CapabilityRuntimeError {
64    #[error(transparent)]
65    Projection(#[from] CapabilityProjectionError),
66    #[error(transparent)]
67    Scope(#[from] CapabilityScopeError),
68    #[error("A Session capability batch with an A3S Use cursor requires a lease provider")]
69    MissingUseLeaseProvider,
70    #[error("A Session capability batch without an A3S Use cursor cannot carry a lease provider")]
71    UnexpectedUseLeaseProvider,
72    #[error(
73        "A3S Use lease provider does not match the batch cursor (generation {expected_generation} vs {actual_generation}, capability revision mismatch: {revision_mismatch}, Registry revision mismatch: {registry_revision_mismatch})"
74    )]
75    UseLeaseProviderMismatch {
76        expected_generation: u64,
77        actual_generation: u64,
78        revision_mismatch: bool,
79        registry_revision_mismatch: bool,
80    },
81    #[error("A3S Use generation lease acquisition failed: {message}")]
82    UseLeaseAcquisition { message: String },
83    #[error("Session capability Run admission was cancelled")]
84    Cancelled,
85    #[error("The owning Session is closed")]
86    SessionClosed,
87    #[error("Session capability kind '{kind}' is not migrated to the atomic host runtime")]
88    UnsupportedSessionKind { kind: CapabilityKind },
89    #[error("Session runtime {kind} name '{public_name}' conflicts with a compatibility value")]
90    RuntimeNameConflict {
91        kind: CapabilityKind,
92        public_name: String,
93    },
94    #[error("Session runtime {kind} value '{public_name}' is invalid: {message}")]
95    RuntimeValueInvalid {
96        kind: CapabilityKind,
97        public_name: String,
98        message: String,
99    },
100    #[error("Session recovery capability binding is unavailable: {message}")]
101    RecoveryBinding { message: String },
102    #[error(
103        "Capability Run close was incomplete (tasks failed: {tasks_failed}, tasks timed out: {tasks_timed_out}, child scopes failed: {child_scopes_failed}, child scopes timed out: {child_scopes_timed_out}, effects failed: {effects_failed}, effects timed out: {effects_timed_out})"
104    )]
105    RunCloseIncomplete {
106        tasks_failed: usize,
107        tasks_timed_out: usize,
108        child_scopes_failed: usize,
109        child_scopes_timed_out: usize,
110        effects_failed: usize,
111        effects_timed_out: usize,
112    },
113}
114
115/// One complete next-generation Tool/Skill/Agent/Command/Hook/MCP/Flow/
116/// Knowledge Surface/Knowledge/UI/Context projection for a Session.
117///
118/// The batch owns every adapter before preparation starts. A Use-backed batch
119/// also owns the generation-specific lease provider that will be published in
120/// the same catalog compare-and-swap as the prepared projection.
121#[must_use = "a Session capability batch must be applied or dropped"]
122pub struct SessionCapabilityBatch {
123    target: Arc<CapabilitySet>,
124    staged: BTreeMap<CapabilityId, Box<dyn CapabilityProjectionAdapter>>,
125    use_lease_provider: Option<Arc<dyn UseGenerationLeaseProvider>>,
126}
127
128impl SessionCapabilityBatch {
129    pub fn new(target: Arc<CapabilitySet>) -> Result<Self, CapabilityRuntimeError> {
130        if target.use_capability_generation().is_some() {
131            return Err(CapabilityRuntimeError::MissingUseLeaseProvider);
132        }
133        validate_session_kinds(&target)?;
134        Ok(Self {
135            target,
136            staged: BTreeMap::new(),
137            use_lease_provider: None,
138        })
139    }
140
141    pub fn from_use_projection(
142        target: Arc<CapabilitySet>,
143        provider: Arc<dyn UseGenerationLeaseProvider>,
144    ) -> Result<Self, CapabilityRuntimeError> {
145        let expected = target
146            .use_capability_generation()
147            .ok_or(CapabilityRuntimeError::UnexpectedUseLeaseProvider)?;
148        ensure_use_generation_matches(expected, provider.use_generation())?;
149        validate_session_kinds(&target)?;
150        Ok(Self {
151            target,
152            staged: BTreeMap::new(),
153            use_lease_provider: Some(provider),
154        })
155    }
156
157    pub fn target(&self) -> &CapabilitySet {
158        &self.target
159    }
160
161    pub fn stage<A>(
162        &mut self,
163        id: CapabilityId,
164        adapter: A,
165    ) -> Result<&mut Self, CapabilityRuntimeError>
166    where
167        A: CapabilityProjectionAdapter,
168    {
169        self.stage_boxed(id, Box::new(adapter))
170    }
171
172    pub fn stage_value(
173        &mut self,
174        id: CapabilityId,
175        value: CapabilityValue,
176    ) -> Result<&mut Self, CapabilityRuntimeError> {
177        struct ReadyValue(CapabilityValue);
178
179        #[async_trait]
180        impl CapabilityProjectionAdapter for ReadyValue {
181            async fn prepare(
182                self: Box<Self>,
183                _cancellation: CancellationToken,
184            ) -> Result<super::PreparedCapability, super::CapabilityAdapterError> {
185                Ok(super::PreparedCapability::new(self.0))
186            }
187        }
188
189        self.stage(id, ReadyValue(value))
190    }
191
192    pub fn len(&self) -> usize {
193        self.staged.len()
194    }
195
196    pub fn is_empty(&self) -> bool {
197        self.staged.is_empty()
198    }
199
200    fn stage_boxed(
201        &mut self,
202        id: CapabilityId,
203        adapter: Box<dyn CapabilityProjectionAdapter>,
204    ) -> Result<&mut Self, CapabilityRuntimeError> {
205        if !self.target.contains(&id) {
206            return Err(CapabilityProjectionError::UnknownStagedCapability {
207                capability: id.to_string(),
208            }
209            .into());
210        }
211        if self.staged.insert(id.clone(), adapter).is_some() {
212            return Err(CapabilityProjectionError::DuplicateStagedCapability {
213                capability: id.to_string(),
214            }
215            .into());
216        }
217        Ok(self)
218    }
219
220    pub(crate) async fn prepare(
221        self,
222        catalog: &CapabilityCatalog,
223        cancellation: CancellationToken,
224    ) -> Result<PreparedSessionCapabilityBatch, CapabilityRuntimeError> {
225        let Self {
226            target,
227            staged,
228            use_lease_provider,
229        } = self;
230        let mut transaction: CapabilityTxn<Staged> = catalog.begin(target)?;
231        for (id, adapter) in staged {
232            transaction.stage_boxed(id, adapter)?;
233        }
234        let transaction: CapabilityTxn<Prepared> = transaction.prepare(cancellation).await?;
235        let transaction: CapabilityTxn<Validated> = transaction.validate()?;
236        Ok(PreparedSessionCapabilityBatch {
237            transaction,
238            use_lease_provider,
239        })
240    }
241
242    pub(crate) async fn prepare_recovery_bootstrap(
243        self,
244        catalog: &CapabilityCatalog,
245        cancellation: CancellationToken,
246    ) -> Result<PreparedSessionCapabilityBatch, CapabilityRuntimeError> {
247        let Self {
248            target,
249            staged,
250            use_lease_provider,
251        } = self;
252        let mut transaction: CapabilityTxn<Staged> = catalog.begin_recovery_bootstrap(target)?;
253        for (id, adapter) in staged {
254            transaction.stage_boxed(id, adapter)?;
255        }
256        let transaction: CapabilityTxn<Prepared> = transaction.prepare(cancellation).await?;
257        let transaction: CapabilityTxn<Validated> = transaction.validate()?;
258        Ok(PreparedSessionCapabilityBatch {
259            transaction,
260            use_lease_provider,
261        })
262    }
263}
264
265impl fmt::Debug for SessionCapabilityBatch {
266    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
267        formatter
268            .debug_struct("SessionCapabilityBatch")
269            .field("target_generation", &self.target.generation())
270            .field("target_digest", &self.target.digest())
271            .field("staged", &self.staged.len())
272            .field("use_backed", &self.use_lease_provider.is_some())
273            .finish()
274    }
275}
276
277pub(crate) struct PreparedSessionCapabilityBatch {
278    transaction: CapabilityTxn<Validated>,
279    use_lease_provider: Option<Arc<dyn UseGenerationLeaseProvider>>,
280}
281
282impl PreparedSessionCapabilityBatch {
283    pub(crate) fn projection(
284        &self,
285    ) -> Result<&super::CapabilityProjection, CapabilityRuntimeError> {
286        self.transaction.projection().map_err(Into::into)
287    }
288
289    pub(crate) fn commit(self) -> Result<CapabilityCommitReceipt, CapabilityRuntimeError> {
290        self.transaction
291            .commit_with_use_lease_provider(self.use_lease_provider)
292            .map_err(Into::into)
293    }
294}
295
296/// Non-clone Run guard retaining one Code projection and one exact Use lease.
297///
298/// The Run and its generation-specific Session scope are closed together. The
299/// A3S Use lease lives in the Run supervisor and is released after children,
300/// tasks, and effects. The projection lease remains pinned until this guard is
301/// dropped, so model definitions and execution borrow the same runtime values.
302#[must_use = "a Session capability Run must remain alive for the complete execution"]
303pub struct SessionCapabilityRun {
304    run_scope: CapabilityScope<Run>,
305    session_scope: CapabilityScope<Session>,
306    projection: CapabilityProjectionLease,
307}
308
309impl SessionCapabilityRun {
310    pub(crate) async fn admit(
311        projection: CapabilityProjectionLease,
312        session_local_id: &str,
313        run_local_id: &str,
314        ceiling: CapabilityCeiling,
315        cancellation: CancellationToken,
316    ) -> Result<Self, CapabilityRuntimeError> {
317        if cancellation.is_cancelled() {
318            return Err(CapabilityRuntimeError::Cancelled);
319        }
320        let set = projection.projection().set();
321        // The host invocation token is the cancellation-tree root, not a
322        // resource owned by capability teardown. The Session owns a child so
323        // host cancellation still cascades downward, while normal Run close
324        // cannot make a successfully completed host lifecycle look cancelled.
325        let session_scope = CapabilityScope::new_session_with_cancellation(
326            session_local_id,
327            Arc::clone(projection.projection().set_arc()),
328            ceiling.clone(),
329            cancellation.child_token(),
330        )?;
331
332        let run_scope = match set.use_capability_generation() {
333            Some(expected) => {
334                let provider = projection
335                    .use_lease_provider()
336                    .ok_or(CapabilityRuntimeError::MissingUseLeaseProvider)?;
337                ensure_use_generation_matches(expected, provider.use_generation())?;
338                let acquire = provider.acquire(cancellation.clone());
339                tokio::pin!(acquire);
340                let lease = tokio::select! {
341                    biased;
342                    _ = cancellation.cancelled() => {
343                        return Err(CapabilityRuntimeError::Cancelled);
344                    }
345                    result = &mut acquire => result.map_err(|error| {
346                        CapabilityRuntimeError::UseLeaseAcquisition {
347                            message: error.message().to_owned(),
348                        }
349                    })?,
350                };
351                session_scope.admit_use_run(run_local_id, ceiling, lease)?
352            }
353            None => {
354                if projection.use_lease_provider().is_some() {
355                    return Err(CapabilityRuntimeError::UnexpectedUseLeaseProvider);
356                }
357                session_scope.admit_run(run_local_id, ceiling)?
358            }
359        };
360
361        Ok(Self {
362            run_scope,
363            session_scope,
364            projection,
365        })
366    }
367
368    pub fn projection(&self) -> &super::CapabilityProjection {
369        self.projection.projection()
370    }
371
372    pub fn run_scope(&self) -> &CapabilityScope<Run> {
373        &self.run_scope
374    }
375
376    pub(crate) fn task_spawner(&self) -> super::SupervisedTaskSpawner {
377        self.run_scope.task_spawner()
378    }
379
380    pub async fn close(&self) -> Result<ScopeCloseReport, CapabilityRuntimeError> {
381        let report = self.session_scope.close().await?;
382        if !report.is_clean() {
383            return Err(CapabilityRuntimeError::RunCloseIncomplete {
384                tasks_failed: report.tasks_failed,
385                tasks_timed_out: report.tasks_timed_out,
386                child_scopes_failed: report.child_scopes_failed,
387                child_scopes_timed_out: report.child_scopes_timed_out,
388                effects_failed: report.effects_failed,
389                effects_timed_out: report.effects_timed_out,
390            });
391        }
392        Ok(report)
393    }
394}
395
396impl fmt::Debug for SessionCapabilityRun {
397    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
398        formatter
399            .debug_struct("SessionCapabilityRun")
400            .field("stamp", self.projection.stamp())
401            .field("run_scope", &self.run_scope.id())
402            .finish_non_exhaustive()
403    }
404}
405
406fn validate_session_kinds(target: &CapabilitySet) -> Result<(), CapabilityRuntimeError> {
407    for (_, descriptor) in target.iter() {
408        if !matches!(
409            descriptor.id().kind(),
410            CapabilityKind::Tool
411                | CapabilityKind::Skill
412                | CapabilityKind::Agent
413                | CapabilityKind::Command
414                | CapabilityKind::Hook
415                | CapabilityKind::Mcp
416                | CapabilityKind::Flow
417                | CapabilityKind::KnowledgeSurface
418                | CapabilityKind::Knowledge
419                | CapabilityKind::Ui
420                | CapabilityKind::Context
421        ) {
422            return Err(CapabilityRuntimeError::UnsupportedSessionKind {
423                kind: descriptor.id().kind(),
424            });
425        }
426    }
427    Ok(())
428}
429
430fn ensure_use_generation_matches(
431    expected: &UseCapabilityGeneration,
432    actual: &UseCapabilityGeneration,
433) -> Result<(), CapabilityRuntimeError> {
434    if expected == actual {
435        return Ok(());
436    }
437    Err(CapabilityRuntimeError::UseLeaseProviderMismatch {
438        expected_generation: expected.generation(),
439        actual_generation: actual.generation(),
440        revision_mismatch: expected.revision() != actual.revision(),
441        registry_revision_mismatch: expected.registry_revision() != actual.registry_revision(),
442    })
443}
444
445fn truncate_utf8(mut value: String, max: usize) -> String {
446    if value.len() <= max {
447        return value;
448    }
449    let mut boundary = max;
450    while !value.is_char_boundary(boundary) {
451        boundary -= 1;
452    }
453    value.truncate(boundary);
454    value
455}
a3s_code_core/capability/runtime.rs

a3s_code_core/capability/
runtime.rs
