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#[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#[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#[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
425pub 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#[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#[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}