Skip to main content

a3s_code_core/capability/
projection.rs

1use std::collections::{BTreeMap, VecDeque};
2use std::fmt;
3use std::sync::{Arc, Mutex};
4
5use tokio::time::Instant;
6
7use crate::cognitive_context::CognitiveContextSession;
8use crate::commands::SlashCommand;
9use crate::context::ContextProvider;
10use crate::hooks::HookBinding;
11use crate::skills::Skill;
12use crate::subagent::AgentDefinition;
13use crate::tools::Tool;
14
15use super::{
16    CapabilityEffect, CapabilityId, CapabilityKind, CapabilityProjectionError,
17    CapabilityReadinessPlan, CapabilitySet, CapabilityValue, CodeCatalogGeneration, FlowBinding,
18    KnowledgeSurfaceBinding, McpBinding, ScopeClosePolicy, Sha256Digest, UiBinding,
19    UseGenerationLeaseProvider,
20};
21
22/// Immutable pairing of one identity set with exactly one typed runtime value
23/// for every descriptor.
24#[derive(Debug)]
25pub struct CapabilityProjection {
26    set: Arc<CapabilitySet>,
27    readiness: Arc<CapabilityReadinessPlan>,
28    values: BTreeMap<CapabilityId, CapabilityValue>,
29}
30
31impl CapabilityProjection {
32    pub fn new(
33        set: Arc<CapabilitySet>,
34        values: impl IntoIterator<Item = (CapabilityId, CapabilityValue)>,
35    ) -> Result<Arc<Self>, CapabilityProjectionError> {
36        let readiness = Arc::new(CapabilityReadinessPlan::from_set(&set)?);
37        Self::with_readiness(set, readiness, values)
38    }
39
40    pub(super) fn with_readiness(
41        set: Arc<CapabilitySet>,
42        readiness: Arc<CapabilityReadinessPlan>,
43        values: impl IntoIterator<Item = (CapabilityId, CapabilityValue)>,
44    ) -> Result<Arc<Self>, CapabilityProjectionError> {
45        if !readiness.matches(&set) {
46            return Err(CapabilityProjectionError::ReadinessPlanMismatch {
47                expected_generation: set.generation().get(),
48                actual_generation: readiness.generation().get(),
49                digest_mismatch: readiness.digest() != set.digest(),
50            });
51        }
52        let mut canonical = BTreeMap::new();
53        for (id, value) in values {
54            if canonical.insert(id.clone(), value).is_some() {
55                return Err(CapabilityProjectionError::DuplicateValue {
56                    capability: id.to_string(),
57                });
58            }
59        }
60
61        for (id, descriptor) in set.iter() {
62            let Some(value) = canonical.get(id) else {
63                return Err(CapabilityProjectionError::MissingValue {
64                    capability: id.to_string(),
65                });
66            };
67            if value.kind() != descriptor.id().kind() {
68                return Err(CapabilityProjectionError::KindMismatch {
69                    capability: id.to_string(),
70                    descriptor_kind: descriptor.id().kind(),
71                    value_kind: value.kind(),
72                });
73            }
74            if let Some(actual) = value.public_name() {
75                if actual != descriptor.public_name() {
76                    return Err(CapabilityProjectionError::PublicNameMismatch {
77                        capability: id.to_string(),
78                        expected: descriptor.public_name().to_owned(),
79                        actual: actual.to_owned(),
80                    });
81                }
82            }
83            if let CapabilityValue::Ui(binding) = value {
84                if descriptor.surface_digest() != binding.surface_digest() {
85                    return Err(CapabilityProjectionError::SurfaceDigestMismatch {
86                        capability: id.to_string(),
87                        expected: descriptor.surface_digest().to_string(),
88                        actual: binding.surface_digest().to_string(),
89                    });
90                }
91                for dependency in descriptor.dependencies() {
92                    if !matches!(
93                        dependency.kind(),
94                        CapabilityKind::Tool
95                            | CapabilityKind::Skill
96                            | CapabilityKind::Mcp
97                            | CapabilityKind::Flow
98                    ) {
99                        return Err(CapabilityProjectionError::UnsupportedUiDependencyKind {
100                            capability: id.to_string(),
101                            dependency: dependency.to_string(),
102                            dependency_kind: dependency.kind(),
103                        });
104                    }
105                }
106            }
107            if let CapabilityValue::KnowledgeSurface(binding) = value {
108                if descriptor.surface_digest() != binding.surface_digest() {
109                    return Err(CapabilityProjectionError::SurfaceDigestMismatch {
110                        capability: id.to_string(),
111                        expected: descriptor.surface_digest().to_string(),
112                        actual: binding.surface_digest().to_string(),
113                    });
114                }
115            }
116        }
117        for id in canonical.keys() {
118            if !set.contains(id) {
119                return Err(CapabilityProjectionError::UnexpectedValue {
120                    capability: id.to_string(),
121                });
122            }
123        }
124
125        Ok(Arc::new(Self {
126            set,
127            readiness,
128            values: canonical,
129        }))
130    }
131
132    pub fn set(&self) -> &CapabilitySet {
133        &self.set
134    }
135
136    pub(crate) fn set_arc(&self) -> &Arc<CapabilitySet> {
137        &self.set
138    }
139
140    pub fn readiness_plan(&self) -> &CapabilityReadinessPlan {
141        &self.readiness
142    }
143
144    pub fn len(&self) -> usize {
145        self.values.len()
146    }
147
148    pub fn is_empty(&self) -> bool {
149        self.values.is_empty()
150    }
151
152    pub fn contains(&self, id: &CapabilityId) -> bool {
153        self.values.contains_key(id)
154    }
155
156    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&CapabilityId, &CapabilityValue)> {
157        self.values.iter()
158    }
159
160    pub fn tool(&self, id: &CapabilityId) -> Option<&dyn Tool> {
161        match self.values.get(id) {
162            Some(CapabilityValue::Tool(value)) => Some(value.as_ref()),
163            _ => None,
164        }
165    }
166
167    pub fn skill(&self, id: &CapabilityId) -> Option<&Skill> {
168        match self.values.get(id) {
169            Some(CapabilityValue::Skill(value)) => Some(value.as_ref()),
170            _ => None,
171        }
172    }
173
174    pub fn agent(&self, id: &CapabilityId) -> Option<&AgentDefinition> {
175        match self.values.get(id) {
176            Some(CapabilityValue::Agent(value)) => Some(value.as_ref()),
177            _ => None,
178        }
179    }
180
181    pub fn command(&self, id: &CapabilityId) -> Option<&dyn SlashCommand> {
182        match self.values.get(id) {
183            Some(CapabilityValue::Command(value)) => Some(value.as_ref()),
184            _ => None,
185        }
186    }
187
188    pub fn hook(&self, id: &CapabilityId) -> Option<&HookBinding> {
189        match self.values.get(id) {
190            Some(CapabilityValue::Hook(value)) => Some(value.as_ref()),
191            _ => None,
192        }
193    }
194
195    pub fn mcp(&self, id: &CapabilityId) -> Option<&McpBinding> {
196        match self.values.get(id) {
197            Some(CapabilityValue::Mcp(value)) => Some(value.as_ref()),
198            _ => None,
199        }
200    }
201
202    pub fn flow(&self, id: &CapabilityId) -> Option<&FlowBinding> {
203        match self.values.get(id) {
204            Some(CapabilityValue::Flow(value)) => Some(value.as_ref()),
205            _ => None,
206        }
207    }
208
209    pub fn knowledge(&self, id: &CapabilityId) -> Option<&CognitiveContextSession> {
210        match self.values.get(id) {
211            Some(CapabilityValue::Knowledge(value)) => Some(value.as_ref()),
212            _ => None,
213        }
214    }
215
216    pub fn knowledge_surface(&self, id: &CapabilityId) -> Option<&KnowledgeSurfaceBinding> {
217        match self.values.get(id) {
218            Some(CapabilityValue::KnowledgeSurface(value)) => Some(value.as_ref()),
219            _ => None,
220        }
221    }
222
223    pub fn ui(&self, id: &CapabilityId) -> Option<&UiBinding> {
224        match self.values.get(id) {
225            Some(CapabilityValue::Ui(value)) => Some(value.as_ref()),
226            _ => None,
227        }
228    }
229
230    pub fn context(&self, id: &CapabilityId) -> Option<&dyn ContextProvider> {
231        match self.values.get(id) {
232            Some(CapabilityValue::Context(value)) => Some(value.as_ref()),
233            _ => None,
234        }
235    }
236}
237
238/// Exact local generation and identity digest used by catalog CAS publication.
239#[derive(Clone, Debug, Eq, PartialEq)]
240pub struct CapabilityCatalogStamp {
241    generation: CodeCatalogGeneration,
242    digest: Sha256Digest,
243}
244
245impl CapabilityCatalogStamp {
246    fn from_projection(projection: &CapabilityProjection) -> Self {
247        Self {
248            generation: projection.set().generation(),
249            digest: projection.set().digest().clone(),
250        }
251    }
252
253    pub const fn generation(&self) -> CodeCatalogGeneration {
254        self.generation
255    }
256
257    pub fn digest(&self) -> &Sha256Digest {
258        &self.digest
259    }
260}
261
262/// Successful all-or-nothing projection publication evidence.
263#[derive(Clone, Debug, Eq, PartialEq)]
264pub struct CapabilityCommitReceipt {
265    previous: CapabilityCatalogStamp,
266    committed: CapabilityCatalogStamp,
267}
268
269impl CapabilityCommitReceipt {
270    pub fn previous(&self) -> &CapabilityCatalogStamp {
271        &self.previous
272    }
273
274    pub fn committed(&self) -> &CapabilityCatalogStamp {
275        &self.committed
276    }
277}
278
279#[derive(Clone, Copy)]
280enum CleanupReason {
281    Rollback,
282    Retired,
283}
284
285struct CleanupBatch {
286    reason: CleanupReason,
287    effects: Vec<Box<dyn CapabilityEffect>>,
288}
289
290#[derive(Default)]
291struct CleanupQueue {
292    batches: Mutex<VecDeque<CleanupBatch>>,
293}
294
295impl CleanupQueue {
296    fn enqueue(&self, reason: CleanupReason, effects: Vec<Box<dyn CapabilityEffect>>) {
297        if effects.is_empty() {
298            return;
299        }
300        self.batches
301            .lock()
302            .unwrap_or_else(std::sync::PoisonError::into_inner)
303            .push_back(CleanupBatch { reason, effects });
304    }
305
306    fn take_all(&self) -> VecDeque<CleanupBatch> {
307        std::mem::take(
308            &mut *self
309                .batches
310                .lock()
311                .unwrap_or_else(std::sync::PoisonError::into_inner),
312        )
313    }
314
315    fn len(&self) -> usize {
316        self.batches
317            .lock()
318            .unwrap_or_else(std::sync::PoisonError::into_inner)
319            .len()
320    }
321}
322
323struct PublishedGeneration {
324    projection: Arc<CapabilityProjection>,
325    stamp: CapabilityCatalogStamp,
326    use_lease_provider: Option<Arc<dyn UseGenerationLeaseProvider>>,
327    effects: Mutex<Vec<Box<dyn CapabilityEffect>>>,
328    cleanup: Arc<CleanupQueue>,
329}
330
331impl PublishedGeneration {
332    fn new(
333        projection: Arc<CapabilityProjection>,
334        use_lease_provider: Option<Arc<dyn UseGenerationLeaseProvider>>,
335        effects: Vec<Box<dyn CapabilityEffect>>,
336        cleanup: Arc<CleanupQueue>,
337    ) -> Self {
338        let stamp = CapabilityCatalogStamp::from_projection(&projection);
339        Self {
340            projection,
341            stamp,
342            use_lease_provider,
343            effects: Mutex::new(effects),
344            cleanup,
345        }
346    }
347}
348
349impl Drop for PublishedGeneration {
350    fn drop(&mut self) {
351        let effects = std::mem::take(
352            self.effects
353                .get_mut()
354                .unwrap_or_else(std::sync::PoisonError::into_inner),
355        );
356        self.cleanup.enqueue(CleanupReason::Retired, effects);
357    }
358}
359
360struct CatalogState {
361    current: Arc<PublishedGeneration>,
362}
363
364pub(super) struct CatalogInner {
365    state: Mutex<CatalogState>,
366    cleanup: Arc<CleanupQueue>,
367}
368
369impl CatalogInner {
370    pub(super) fn current_stamp(&self) -> CapabilityCatalogStamp {
371        self.state
372            .lock()
373            .unwrap_or_else(std::sync::PoisonError::into_inner)
374            .current
375            .stamp
376            .clone()
377    }
378
379    pub(super) fn enqueue_rollback(&self, effects: Vec<Box<dyn CapabilityEffect>>) {
380        self.cleanup.enqueue(CleanupReason::Rollback, effects);
381    }
382
383    pub(super) fn publish(
384        &self,
385        base: &CapabilityCatalogStamp,
386        projection: Arc<CapabilityProjection>,
387        use_lease_provider: Option<Arc<dyn UseGenerationLeaseProvider>>,
388        effects: Vec<Box<dyn CapabilityEffect>>,
389    ) -> Result<CapabilityCommitReceipt, CapabilityProjectionError> {
390        let committed = CapabilityCatalogStamp::from_projection(&projection);
391        let old = {
392            let mut state = self
393                .state
394                .lock()
395                .unwrap_or_else(std::sync::PoisonError::into_inner);
396            let actual = &state.current.stamp;
397            if actual != base {
398                let error = CapabilityProjectionError::CommitConflict {
399                    expected_generation: base.generation().get(),
400                    expected_digest: base.digest().to_string(),
401                    actual_generation: actual.generation().get(),
402                    actual_digest: actual.digest().to_string(),
403                };
404                drop(state);
405                self.enqueue_rollback(effects);
406                return Err(error);
407            }
408            let published = Arc::new(PublishedGeneration::new(
409                projection,
410                use_lease_provider,
411                effects,
412                Arc::clone(&self.cleanup),
413            ));
414            std::mem::replace(&mut state.current, published)
415        };
416        let previous = old.stamp.clone();
417        drop(old);
418        Ok(CapabilityCommitReceipt {
419            previous,
420            committed,
421        })
422    }
423}
424
425/// Session-local immutable capability publication catalog.
426///
427/// Readers only clone one `Arc` under a short mutex and then resolve through a
428/// pinned [`CapabilityProjectionLease`]. Writers use an exact generation and
429/// digest compare-and-swap; a losing writer cannot mutate the current value.
430pub struct CapabilityCatalog {
431    pub(super) inner: Arc<CatalogInner>,
432}
433
434impl CapabilityCatalog {
435    pub fn new(initial: Arc<CapabilityProjection>) -> Self {
436        let cleanup = Arc::new(CleanupQueue::default());
437        let current = Arc::new(PublishedGeneration::new(
438            initial,
439            None,
440            Vec::new(),
441            Arc::clone(&cleanup),
442        ));
443        Self {
444            inner: Arc::new(CatalogInner {
445                state: Mutex::new(CatalogState { current }),
446                cleanup,
447            }),
448        }
449    }
450
451    pub fn current_stamp(&self) -> CapabilityCatalogStamp {
452        self.inner.current_stamp()
453    }
454
455    pub fn pin(&self) -> CapabilityProjectionLease {
456        let generation = Arc::clone(
457            &self
458                .inner
459                .state
460                .lock()
461                .unwrap_or_else(std::sync::PoisonError::into_inner)
462                .current,
463        );
464        CapabilityProjectionLease { generation }
465    }
466
467    pub fn pending_cleanup_batches(&self) -> usize {
468        self.inner.cleanup.len()
469    }
470
471    pub(crate) fn retire_current_effects(&self) {
472        let generation = Arc::clone(
473            &self
474                .inner
475                .state
476                .lock()
477                .unwrap_or_else(std::sync::PoisonError::into_inner)
478                .current,
479        );
480        let effects = std::mem::take(
481            &mut *generation
482                .effects
483                .lock()
484                .unwrap_or_else(std::sync::PoisonError::into_inner),
485        );
486        self.inner.cleanup.enqueue(CleanupReason::Retired, effects);
487    }
488
489    pub async fn drain_cleanup(&self) -> CapabilityCleanupReport {
490        self.drain_cleanup_with_policy(ScopeClosePolicy::default())
491            .await
492    }
493
494    pub async fn drain_cleanup_with_policy(
495        &self,
496        policy: ScopeClosePolicy,
497    ) -> CapabilityCleanupReport {
498        let mut batches = self.inner.cleanup.take_all();
499        let deadline = Instant::now() + policy.timeout();
500        let mut report = CapabilityCleanupReport::default();
501
502        while let Some(mut batch) = batches.pop_front() {
503            match batch.reason {
504                CleanupReason::Rollback => report.rollback_batches += 1,
505                CleanupReason::Retired => report.retired_batches += 1,
506            }
507            while let Some(effect) = batch.effects.pop() {
508                if Instant::now() >= deadline {
509                    report.effects_timed_out += 1 + batch.effects.len();
510                    break;
511                }
512                match tokio::time::timeout_at(deadline, effect.close()).await {
513                    Ok(Ok(())) => report.effects_closed += 1,
514                    Ok(Err(_)) => report.effects_failed += 1,
515                    Err(_) => report.effects_timed_out += 1,
516                }
517            }
518        }
519        report
520    }
521}
522
523impl fmt::Debug for CapabilityCatalog {
524    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
525        formatter
526            .debug_struct("CapabilityCatalog")
527            .field("current", &self.current_stamp())
528            .field("pending_cleanup_batches", &self.pending_cleanup_batches())
529            .finish()
530    }
531}
532
533/// Non-clone reader lease retaining one exact projected generation.
534///
535/// The contained values are exposed by borrow, so ordinary execution cannot
536/// accidentally switch to the catalog's latest generation. When the last
537/// lease and catalog pointer to a retired generation disappear, Rust `Arc`
538/// ownership moves its effects to the asynchronous cleanup queue.
539#[must_use = "a projection lease pins one exact catalog generation"]
540pub struct CapabilityProjectionLease {
541    generation: Arc<PublishedGeneration>,
542}
543
544impl CapabilityProjectionLease {
545    pub fn stamp(&self) -> &CapabilityCatalogStamp {
546        &self.generation.stamp
547    }
548
549    pub fn projection(&self) -> &CapabilityProjection {
550        &self.generation.projection
551    }
552
553    pub(super) fn use_lease_provider(&self) -> Option<&Arc<dyn UseGenerationLeaseProvider>> {
554        self.generation.use_lease_provider.as_ref()
555    }
556}
557
558impl fmt::Debug for CapabilityProjectionLease {
559    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
560        formatter
561            .debug_struct("CapabilityProjectionLease")
562            .field("stamp", self.stamp())
563            .finish_non_exhaustive()
564    }
565}
566
567/// Bounded reverse-teardown result for retired or rolled-back effects.
568#[derive(Clone, Debug, Default, Eq, PartialEq)]
569pub struct CapabilityCleanupReport {
570    pub rollback_batches: usize,
571    pub retired_batches: usize,
572    pub effects_closed: usize,
573    pub effects_failed: usize,
574    pub effects_timed_out: usize,
575}
576
577impl CapabilityCleanupReport {
578    pub const fn is_clean(&self) -> bool {
579        self.effects_failed == 0 && self.effects_timed_out == 0
580    }
581}