1use std::{collections::BTreeMap, sync::Arc};
2
3use serde_json::Value;
4
5use crate::{
6 BudgetReservation, BudgetTracker, CancellationToken, CapabilityId, CapabilitySet, EventId,
7 Instant, Journal, JournalError, RunEvent, RunEventKind, RunId, RunRecorder,
8};
9use thiserror::Error;
10
11#[derive(Clone, Debug, Error, Eq, PartialEq)]
13#[error("child Run requested capability `{capability}` ({capability_id}) absent from its parent")]
14pub struct AuthorityAmplification {
15 pub capability_id: CapabilityId,
17 pub capability: String,
19}
20
21#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
23#[error("budget reservation does not belong to this run tree")]
24pub struct BudgetReservationMismatch;
25
26#[derive(Clone, Debug, Error, Eq, PartialEq)]
28#[non_exhaustive]
29pub enum ChildRunError {
30 #[error(transparent)]
32 Authority(#[from] AuthorityAmplification),
33 #[error(transparent)]
35 BudgetReservation(#[from] BudgetReservationMismatch),
36}
37
38pub type Metadata = BTreeMap<String, Value>;
40
41#[derive(Clone, Debug)]
43pub struct RunContext {
44 run_id: RunId,
45 parent_run_id: Option<RunId>,
46 root_run_id: RunId,
47 caused_by: Option<EventId>,
48 deadline: Option<Instant>,
49 cancellation: CancellationToken,
50 budget: BudgetTracker,
51 capabilities: CapabilitySet,
52 metadata: Metadata,
53 recorder: Option<RunRecorder>,
54}
55
56impl RunContext {
57 pub fn root(budget: BudgetTracker, capabilities: CapabilitySet) -> Self {
59 let run_id = RunId::new();
60 Self {
61 run_id,
62 parent_run_id: None,
63 root_run_id: run_id,
64 caused_by: None,
65 deadline: None,
66 cancellation: CancellationToken::new(),
67 budget,
68 capabilities,
69 metadata: Metadata::new(),
70 recorder: None,
71 }
72 }
73
74 pub const fn run_id(&self) -> RunId {
76 self.run_id
77 }
78
79 pub const fn parent_run_id(&self) -> Option<RunId> {
81 self.parent_run_id
82 }
83
84 pub const fn root_run_id(&self) -> RunId {
86 self.root_run_id
87 }
88
89 pub const fn caused_by(&self) -> Option<EventId> {
91 self.caused_by
92 }
93
94 pub const fn deadline(&self) -> Option<Instant> {
96 self.deadline
97 }
98
99 pub const fn cancellation(&self) -> &CancellationToken {
101 &self.cancellation
102 }
103
104 pub const fn budget(&self) -> &BudgetTracker {
106 &self.budget
107 }
108
109 pub const fn capabilities(&self) -> &CapabilitySet {
111 &self.capabilities
112 }
113
114 pub const fn metadata(&self) -> &Metadata {
116 &self.metadata
117 }
118
119 pub const fn metadata_mut(&mut self) -> &mut Metadata {
121 &mut self.metadata
122 }
123
124 pub const fn recorder(&self) -> Option<&RunRecorder> {
126 self.recorder.as_ref()
127 }
128
129 #[must_use]
131 pub fn with_journal(mut self, journal: Arc<dyn Journal>) -> Self {
132 self.recorder = Some(RunRecorder::new(journal, self.run_id, self.parent_run_id));
133 self
134 }
135
136 pub fn record(
142 &self,
143 kind: RunEventKind,
144 caused_by: Option<EventId>,
145 ) -> Result<Option<RunEvent>, JournalError> {
146 self.recorder
147 .as_ref()
148 .map(|recorder| recorder.record(kind, caused_by))
149 .transpose()
150 }
151
152 #[must_use]
154 pub fn with_deadline(mut self, deadline: Instant) -> Self {
155 self.deadline = Some(
156 self.deadline
157 .map_or(deadline, |current| current.min(deadline)),
158 );
159 self
160 }
161
162 #[must_use]
164 pub const fn with_cause(mut self, event_id: EventId) -> Self {
165 self.caused_by = Some(event_id);
166 self
167 }
168
169 pub fn child(&self, capabilities: CapabilitySet) -> Result<Self, AuthorityAmplification> {
179 self.validate_child_authority(&capabilities)?;
180 Ok(self.child_with_budget(capabilities, self.budget.clone()))
181 }
182
183 pub fn child_reserved(
190 &self,
191 capabilities: CapabilitySet,
192 reservation: &BudgetReservation,
193 ) -> Result<Self, ChildRunError> {
194 self.validate_child_authority(&capabilities)?;
195 if !reservation.belongs_to(&self.budget) {
196 return Err(BudgetReservationMismatch.into());
197 }
198 Ok(self.child_with_budget(capabilities, reservation.tracker()))
199 }
200
201 fn validate_child_authority(
202 &self,
203 capabilities: &CapabilitySet,
204 ) -> Result<(), AuthorityAmplification> {
205 if let Some(missing) = capabilities.first_missing_from(&self.capabilities) {
206 return Err(AuthorityAmplification {
207 capability_id: missing.id,
208 capability: missing.name.clone(),
209 });
210 }
211 Ok(())
212 }
213
214 fn child_with_budget(&self, capabilities: CapabilitySet, budget: BudgetTracker) -> Self {
215 let run_id = RunId::new();
216 Self {
217 run_id,
218 parent_run_id: Some(self.run_id),
219 root_run_id: self.root_run_id,
220 caused_by: None,
221 deadline: self.deadline,
222 cancellation: self.cancellation.child_token(),
223 budget,
224 capabilities,
225 metadata: Metadata::new(),
226 recorder: self
227 .recorder
228 .as_ref()
229 .map(|recorder| recorder.child(run_id, self.run_id)),
230 }
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use std::sync::Arc;
237
238 use super::{AuthorityAmplification, BudgetReservationMismatch, ChildRunError, RunContext};
239 use crate::{
240 Budget, BudgetTracker, CapabilityDescriptor, CapabilityId, CapabilityKind, CapabilitySet,
241 EffectClass, InMemoryJournal, LifecycleEvent, RiskLevel, RunEventKind, Usage,
242 };
243
244 #[test]
245 fn child_preserves_lineage_but_not_ambient_authority() {
246 let mut parent_capabilities = CapabilitySet::new();
247 parent_capabilities.grant(CapabilityDescriptor {
248 id: CapabilityId::new(),
249 name: "write-file".into(),
250 version: "1".into(),
251 kind: CapabilityKind::Tool,
252 input_schema: serde_json::json!({}),
253 output_schema: serde_json::json!({}),
254 effect: EffectClass::NonIdempotentWrite,
255 risk: RiskLevel::High,
256 metadata: std::collections::BTreeMap::default(),
257 });
258
259 let parent = RunContext::root(
260 BudgetTracker::new(Budget {
261 tokens: Some(10),
262 ..Budget::default()
263 }),
264 parent_capabilities,
265 );
266 let child = parent.child(CapabilitySet::new()).unwrap();
267
268 assert_eq!(child.parent_run_id(), Some(parent.run_id()));
269 assert_eq!(child.root_run_id(), parent.root_run_id());
270 assert!(child.capabilities().is_empty());
271
272 child
273 .budget()
274 .try_consume(Usage {
275 tokens: 4,
276 ..Usage::default()
277 })
278 .unwrap();
279 assert_eq!(parent.budget().usage().tokens, 4);
280
281 parent.cancellation().cancel();
282 assert!(child.cancellation().is_cancelled());
283 }
284
285 #[test]
286 fn child_recorders_share_a_journal_but_keep_per_run_sequences() {
287 let journal = InMemoryJournal::new();
288 let parent = RunContext::root(BudgetTracker::new(Budget::default()), CapabilitySet::new())
289 .with_journal(Arc::new(journal.clone()));
290 let child = parent.child(CapabilitySet::new()).unwrap();
291
292 parent
293 .record(RunEventKind::Lifecycle(LifecycleEvent::Started), None)
294 .unwrap();
295 child
296 .record(RunEventKind::Lifecycle(LifecycleEvent::Started), None)
297 .unwrap();
298 parent
299 .record(
300 RunEventKind::Lifecycle(LifecycleEvent::Completed {
301 output: serde_json::json!({}),
302 }),
303 None,
304 )
305 .unwrap();
306
307 let events = journal.events();
308 assert_eq!(events.len(), 3);
309 assert_eq!(events[0].meta.sequence, 0);
310 assert_eq!(events[1].meta.sequence, 0);
311 assert_eq!(events[2].meta.sequence, 1);
312 assert_eq!(events[1].meta.parent_run_id, Some(parent.run_id()));
313 }
314
315 #[test]
316 fn child_rejects_a_reservation_from_another_run_tree() {
317 let first = RunContext::root(
318 BudgetTracker::new(Budget {
319 turns: Some(1),
320 ..Budget::default()
321 }),
322 CapabilitySet::new(),
323 );
324 let second = RunContext::root(
325 BudgetTracker::new(Budget {
326 turns: Some(1),
327 ..Budget::default()
328 }),
329 CapabilitySet::new(),
330 );
331 let reservation = first
332 .budget()
333 .try_reserve(Usage {
334 turns: 1,
335 ..Usage::default()
336 })
337 .unwrap();
338
339 let error = second
340 .child_reserved(CapabilitySet::new(), &reservation)
341 .unwrap_err();
342
343 assert_eq!(
344 error,
345 ChildRunError::BudgetReservation(BudgetReservationMismatch)
346 );
347 }
348
349 #[test]
350 fn child_rejects_authority_absent_from_parent() {
351 let parent = RunContext::root(BudgetTracker::new(Budget::default()), CapabilitySet::new());
352 let requested = CapabilityDescriptor {
353 id: CapabilityId::new(),
354 name: "filesystem.write".into(),
355 version: "1".into(),
356 kind: CapabilityKind::Tool,
357 input_schema: serde_json::json!({}),
358 output_schema: serde_json::json!({}),
359 effect: EffectClass::NonIdempotentWrite,
360 risk: RiskLevel::High,
361 metadata: std::collections::BTreeMap::default(),
362 };
363 let mut capabilities = CapabilitySet::new();
364 capabilities.grant(requested.clone());
365
366 let error = parent.child(capabilities).unwrap_err();
367
368 assert_eq!(
369 error,
370 AuthorityAmplification {
371 capability_id: requested.id,
372 capability: requested.name,
373 }
374 );
375 }
376
377 #[test]
378 fn reserved_child_checks_authority_before_using_the_reservation() {
379 let parent = RunContext::root(
380 BudgetTracker::new(Budget {
381 turns: Some(1),
382 ..Budget::default()
383 }),
384 CapabilitySet::new(),
385 );
386 let reservation = parent
387 .budget()
388 .try_reserve(Usage {
389 turns: 1,
390 ..Usage::default()
391 })
392 .unwrap();
393 let requested = CapabilityDescriptor {
394 id: CapabilityId::new(),
395 name: "network.write".into(),
396 version: "1".into(),
397 kind: CapabilityKind::Tool,
398 input_schema: serde_json::json!({}),
399 output_schema: serde_json::json!({}),
400 effect: EffectClass::IdempotentWrite,
401 risk: RiskLevel::Medium,
402 metadata: std::collections::BTreeMap::default(),
403 };
404 let mut capabilities = CapabilitySet::new();
405 capabilities.grant(requested.clone());
406
407 let error = parent
408 .child_reserved(capabilities, &reservation)
409 .unwrap_err();
410
411 assert_eq!(
412 error,
413 ChildRunError::Authority(AuthorityAmplification {
414 capability_id: requested.id,
415 capability: requested.name,
416 })
417 );
418 assert_eq!(reservation.remaining().turns, 1);
419 }
420}