bamboo_engine/runtime/managers/
lifecycle.rs

1use async_trait::async_trait;
2use bamboo_agent_core::tools::ToolExecutor;
3use bamboo_agent_core::{AgentError, AgentEvent, Session};
4use bamboo_domain::AgentRuntimeState;
5use bamboo_llm::LLMProvider;
6use tokio::sync::mpsc;
7use tokio_util::sync::CancellationToken;
8
9use crate::runtime::config::AgentLoopConfig;
10use crate::runtime::task_context::TaskLoopContext;
11
12/// Manages agent run state machine and round transitions.
13#[async_trait]
14pub trait LifecycleManager: Send + Sync {
15    /// Initialize runtime state for a new agent run.
16    fn initialize_run(&self, session: &Session, config: &AgentLoopConfig) -> AgentRuntimeState;
17
18    /// Prepare for a new round. Returns the round ID.
19    ///
20    /// `max_rounds` is the run's round cap; `None` = unlimited.
21    #[allow(clippy::too_many_arguments)]
22    async fn prepare_round(
23        &self,
24        session: &mut Session,
25        task_context: &mut Option<TaskLoopContext>,
26        runtime_state: &mut AgentRuntimeState,
27        round: usize,
28        max_rounds: Option<usize>,
29        config: &AgentLoopConfig,
30        cancel_token: &CancellationToken,
31        metrics_collector: Option<&bamboo_metrics::MetricsCollector>,
32        session_id: &str,
33        model_name: &str,
34        tools: &dyn ToolExecutor,
35        llm: &dyn LLMProvider,
36    ) -> Result<String, AgentError>;
37
38    /// Optional data-only companion. Older implementations execute exactly once.
39    async fn prepare_round_with_observation(
40        &self,
41        session: &mut Session,
42        task_context: &mut Option<TaskLoopContext>,
43        runtime_state: &mut AgentRuntimeState,
44        round: usize,
45        max_rounds: Option<usize>,
46        config: &AgentLoopConfig,
47        cancel_token: &CancellationToken,
48        metrics_collector: Option<&bamboo_metrics::MetricsCollector>,
49        session_id: &str,
50        model_name: &str,
51        tools: &dyn ToolExecutor,
52        llm: &dyn LLMProvider,
53    ) -> Result<ObservedRoundPreparation, AgentError> {
54        let round_id = self
55            .prepare_round(
56                session,
57                task_context,
58                runtime_state,
59                round,
60                max_rounds,
61                config,
62                cancel_token,
63                metrics_collector,
64                session_id,
65                model_name,
66                tools,
67                llm,
68            )
69            .await?;
70        Ok(ObservedRoundPreparation {
71            round_id,
72            observation: InputObservation::default(),
73        })
74    }
75
76    /// Handle post-round processing and determine next action.
77    /// Returns `true` if the agent run should break out of the round loop.
78    async fn handle_round_outcome(
79        &self,
80        session: &mut Session,
81        runtime_state: &mut AgentRuntimeState,
82        task_context: &mut Option<TaskLoopContext>,
83        round: usize,
84        should_break: bool,
85    ) -> Result<bool, AgentError>;
86
87    /// Finalize the agent run.
88    #[allow(clippy::too_many_arguments)]
89    async fn finalize_run(
90        &self,
91        session: &mut Session,
92        runtime_state: &mut AgentRuntimeState,
93        event_tx: &mpsc::Sender<AgentEvent>,
94        session_id: &str,
95        config: &AgentLoopConfig,
96        metrics_collector: Option<&bamboo_metrics::MetricsCollector>,
97        task_context: Option<TaskLoopContext>,
98    );
99}
100
101pub use crate::runtime::runner::state_bridge::{
102    project_input_request_batch, BorrowedInputRequestRecord,
103};
104
105/// Original round result plus opaque, untrusted execution-local input data.
106#[derive(Debug)]
107pub struct ObservedRoundPreparation {
108    pub round_id: String,
109    pub observation: InputObservation,
110}
111
112/// Only checked admission creates positive data; this is never a Skill grant.
113/// The caller owns finite current data and supplies its existing run identity.
114#[derive(Debug, Default)]
115pub struct InputObservation(ObservationState);
116
117#[derive(Debug, Default)]
118enum ObservationState {
119    #[default]
120    Unavailable,
121    SuccessfulNoNewInput(BoundedInputRequestBatch),
122    New(BoundedInputRequestBatch),
123}
124
125impl InputObservation {
126    pub(crate) fn projected(projection: InputRequestProjection) -> Self {
127        Self(match projection {
128            InputRequestProjection::Available(batch) if batch.records.is_empty() => {
129                ObservationState::SuccessfulNoNewInput(batch)
130            }
131            InputRequestProjection::Available(batch) => ObservationState::New(batch),
132            InputRequestProjection::Unavailable(_) => ObservationState::Unavailable,
133        })
134    }
135
136    pub fn new_inputs(&self) -> Option<&BoundedInputRequestBatch> {
137        match &self.0 {
138            ObservationState::New(batch) => Some(batch),
139            _ => None,
140        }
141    }
142
143    pub fn is_unavailable(&self) -> bool {
144        matches!(self.0, ObservationState::Unavailable)
145    }
146
147    pub fn is_successful_no_new_input(&self) -> bool {
148        matches!(self.0, ObservationState::SuccessfulNoNewInput(_))
149    }
150
151    pub(crate) fn matches_session(&self, session_id: &str) -> bool {
152        match &self.0 {
153            ObservationState::New(batch) | ObservationState::SuccessfulNoNewInput(batch) => {
154                batch.session_id == session_id
155            }
156            ObservationState::Unavailable => false,
157        }
158    }
159
160    /// Move a whole new batch, retain only same-run NoNew, or clear on unknown.
161    /// A new input with no request replaces the old IDs and selections too.
162    pub fn update_current(
163        self,
164        current: &mut Option<BoundedInputRequestBatch>,
165        session_id: &str,
166        execution_id: &str,
167    ) {
168        let matches = |batch: &BoundedInputRequestBatch| {
169            batch.session_id == session_id && batch.execution_id == execution_id
170        };
171        match self.0 {
172            ObservationState::New(batch) if matches(&batch) => *current = Some(batch),
173            ObservationState::SuccessfulNoNewInput(batch)
174                if matches(&batch) && current.as_ref().is_none_or(matches) => {}
175            _ => *current = None,
176        }
177    }
178}
179
180/// Untrusted request data only; neither current-input evidence nor permission.
181#[derive(Debug, PartialEq)]
182pub struct ProjectedInputRequest {
183    pub input_id: String,
184    pub source: bamboo_domain::SessionMessageSource,
185    pub kind: bamboo_domain::SessionMessageKind,
186    pub wrapper: Option<String>,
187    pub created_at: chrono::DateTime<chrono::Utc>,
188    pub request: Option<bamboo_domain::SessionSkillRequest>,
189}
190
191/// One execution-private data value, borrowed by the SDK provider/Tool round.
192/// Compact bytes bound the measurement, not total AST allocation or RSS.
193///
194/// ```
195/// use bamboo_engine::runtime::managers::lifecycle::{project_input_request_batch, InputRequestProjection};
196/// let InputRequestProjection::Available(batch) = project_input_request_batch("s", "e", &[]) else { panic!() };
197/// assert!(batch.records().is_empty());
198/// ```
199///
200/// ```compile_fail
201/// use bamboo_engine::runtime::managers::lifecycle::BoundedInputRequestBatch;
202/// fn requires_clone<T: Clone>() {}
203/// requires_clone::<BoundedInputRequestBatch>();
204/// ```
205///
206/// ```compile_fail
207/// use bamboo_engine::runtime::managers::lifecycle::BoundedInputRequestBatch;
208/// fn wire<T: serde::Serialize>() {}
209/// wire::<BoundedInputRequestBatch>();
210/// ```
211#[derive(Debug, PartialEq)]
212pub struct BoundedInputRequestBatch {
213    pub(crate) session_id: String,
214    pub(crate) execution_id: String,
215    pub(crate) records: Box<[ProjectedInputRequest]>,
216    pub(crate) compact_bytes: usize,
217}
218
219impl BoundedInputRequestBatch {
220    pub fn session_id(&self) -> &str {
221        &self.session_id
222    }
223    pub fn execution_id(&self) -> &str {
224        &self.execution_id
225    }
226    pub fn records(&self) -> &[ProjectedInputRequest] {
227        &self.records
228    }
229    pub fn compact_bytes(&self) -> usize {
230        self.compact_bytes
231    }
232}
233
234tokio::task_local! {
235    static INPUT_REQUEST_DATA: Option<BoundedInputRequestBatch>;
236}
237
238/// Lend one existing, uniquely owned input-data batch during a future's polls.
239///
240/// The optional batch is moved, never cloned or persisted. A synchronous reader
241/// can borrow it only in this scope. Normal completion, including an error
242/// result, returns the original owner. Dropping or unwinding the future drops
243/// that owner and restores the surrounding scope; it cannot resume the cancelled
244/// work. `None` shadows an outer value. Detached tasks do not inherit the scope;
245/// explicitly moving the whole scoped future carries only its own value.
246///
247/// This is unwired data transport, not an admission seal, caller identity, Source
248/// lease or Skill permission. The host must apply `InputObservation::update_current`
249/// between rounds and independently validate current authority before each use.
250/// No `Send` bound is imposed on inline futures.
251///
252/// ```
253/// use bamboo_engine::runtime::managers::lifecycle::{
254///     project_input_request_batch, scope_input_request_data,
255///     with_scoped_input_request_data, InputRequestProjection,
256/// };
257/// # async fn example() {
258/// let InputRequestProjection::Available(batch) = project_input_request_batch("s", "e", &[]) else { panic!() };
259/// let (count, returned) = scope_input_request_data(Some(batch), async {
260///     with_scoped_input_request_data("s", "e", |batch| batch.records().len())
261/// }).await;
262/// assert_eq!(count, Some(0));
263/// assert!(returned.is_some());
264/// # }
265/// ```
266pub async fn scope_input_request_data<F: std::future::Future>(
267    batch: Option<BoundedInputRequestBatch>,
268    future: F,
269) -> (F::Output, Option<BoundedInputRequestBatch>) {
270    let mut scoped = Box::pin(INPUT_REQUEST_DATA.scope(batch, future));
271    let output = scoped.as_mut().await;
272    let batch = scoped
273        .as_mut()
274        .take_value()
275        .expect("completed input-data scope retains its unique owner");
276    (output, batch)
277}
278
279/// Synchronously borrow data for exactly the supplied Session and execution.
280///
281/// Absent data, empty IDs or either mismatch returns `None` without calling the
282/// reader. Matching IDs identify data only; they authenticate no caller and do
283/// not turn projected/unsealed data into New input. Return an independently
284/// owned result; do not poll another scoped future while holding this borrow.
285///
286/// ```compile_fail
287/// use bamboo_engine::runtime::managers::lifecycle::{
288///     with_scoped_input_request_data, BoundedInputRequestBatch,
289/// };
290/// fn escape() -> &'static BoundedInputRequestBatch {
291///     with_scoped_input_request_data("s", "e", |batch| batch).unwrap()
292/// }
293/// ```
294///
295/// ```compile_fail
296/// use bamboo_engine::runtime::managers::lifecycle::with_scoped_input_request_data;
297/// async fn escape_into_future() {
298///     let future = with_scoped_input_request_data("s", "e", |batch| async move {
299///         tokio::task::yield_now().await;
300///         batch.records().len()
301///     }).unwrap();
302///     future.await;
303/// }
304/// ```
305pub fn with_scoped_input_request_data<R>(
306    session_id: &str,
307    execution_id: &str,
308    reader: impl FnOnce(&BoundedInputRequestBatch) -> R,
309) -> Option<R> {
310    if session_id.is_empty() || execution_id.is_empty() {
311        return None;
312    }
313    INPUT_REQUEST_DATA
314        .try_with(|batch| {
315            batch
316                .as_ref()
317                .filter(|batch| {
318                    batch.session_id == session_id && batch.execution_id == execution_id
319                })
320                .map(reader)
321        })
322        .ok()
323        .flatten()
324}
325
326/// Embedded record/selection headers counted once, separate from dynamic data.
327pub const INPUT_REQUEST_FIXED_SLOT_BYTES: usize = 128
328    * std::mem::size_of::<ProjectedInputRequest>()
329    + 4096 * std::mem::size_of::<bamboo_domain::SessionSkillSelection>()
330    + std::mem::size_of::<BoundedInputRequestBatch>();
331
332#[derive(Debug, PartialEq, Eq)]
333pub enum InputRequestUnavailable {
334    RecordLimit { count: usize, limit: usize },
335    InvalidData,
336    CompactLimit { limit: usize },
337}
338
339#[derive(Debug, PartialEq)]
340pub enum InputRequestProjection {
341    Available(BoundedInputRequestBatch),
342    Unavailable(InputRequestUnavailable),
343}
344
345#[cfg(test)]
346mod input_request_batch_tests {
347    fn ql_observation(session: &str, execution: &str, id: Option<&str>) -> InputObservation {
348        let source = bamboo_domain::SessionMessageSource::User;
349        let records = id.map(|id| BorrowedInputRequestRecord {
350            input_id: id,
351            source: &source,
352            kind: bamboo_domain::SessionMessageKind::UserInput,
353            wrapper: None,
354            created_at: chrono::DateTime::from_timestamp(1, 0).unwrap(),
355            request: None,
356        });
357        InputObservation::projected(project_input_request_batch(
358            session,
359            execution,
360            records.as_slice(),
361        ))
362    }
363
364    #[test]
365    fn ql_finite_owner_retains_nonew_only_for_exact_session_and_execution() {
366        let mut current = None;
367        ql_observation("s", "e", Some("N")).update_current(&mut current, "s", "e");
368        let pointer = current.as_ref().unwrap().records().as_ptr();
369        ql_observation("s", "e", None).update_current(&mut current, "s", "e");
370        assert_eq!(current.as_ref().unwrap().records().as_ptr(), pointer);
371        ql_observation("s", "old-execution", None).update_current(&mut current, "s", "e");
372        assert!(
373            current.is_none(),
374            "foreign NoNew cannot preserve this owner's data"
375        );
376        ql_observation("foreign-session", "e", Some("wrong")).update_current(
377            &mut current,
378            "s",
379            "e",
380        );
381        assert!(current.is_none());
382        ql_observation("s", "e", Some("N")).update_current(&mut current, "s", "e");
383        ql_observation("s", "e", None).update_current(&mut current, "successor", "e");
384        assert!(current.is_none());
385        ql_observation("s", "e", Some("N")).update_current(&mut current, "s", "e");
386        InputObservation::default().update_current(&mut current, "s", "e");
387        ql_observation("s", "e", None).update_current(&mut current, "s", "e");
388        assert!(
389            current.is_none(),
390            "NoNew after Unavailable never resurrects N"
391        );
392    }
393
394    use super::*;
395    use crate::runtime::runner::state_bridge::{
396        project_input_request_batch, BorrowedInputRequestRecord,
397    };
398
399    #[test]
400    fn input_request_batch_compiled_fixed_slots_and_moves_preserve_one_owned_value() {
401        let request = bamboo_domain::SessionSkillRequest {
402            mode: None,
403            selections: (0..32)
404                .map(|i| bamboo_domain::SessionSkillSelection {
405                    id: format!("{i:x}"),
406                    source: "user".into(),
407                    revision: 1,
408                    args: serde_json::Value::Null,
409                })
410                .collect(),
411        };
412        let records = [BorrowedInputRequestRecord {
413            input_id: "u",
414            source: &bamboo_domain::SessionMessageSource::User,
415            kind: bamboo_domain::SessionMessageKind::UserInput,
416            wrapper: None,
417            created_at: chrono::DateTime::from_timestamp(1, 0).unwrap(),
418            request: Some(&request),
419        }; 128];
420        let make = |execution| match project_input_request_batch("s", execution, &records) {
421            InputRequestProjection::Available(batch) => batch,
422            _ => panic!("full slot fixture must actually fit"),
423        };
424        let mut owner = Some(make("old"));
425        let old_pointer = owner.as_ref().unwrap().records().as_ptr();
426        let prospective = make("new");
427        let new_pointer = prospective.records().as_ptr();
428        assert_ne!(old_pointer, new_pointer);
429        let old = owner.replace(prospective).unwrap();
430        assert_eq!(old.records().as_ptr(), old_pointer);
431        assert_eq!(owner.as_ref().unwrap().records().as_ptr(), new_pointer);
432        drop(old);
433        let batch = owner.take().unwrap();
434        assert!(owner.is_none());
435        assert_eq!(batch.execution_id(), "new");
436        assert_eq!(batch.records().len(), 128);
437        assert_eq!(
438            batch
439                .records
440                .iter()
441                .map(|r| r.request.as_ref().unwrap().selections.capacity())
442                .sum::<usize>(),
443            4096
444        );
445        assert_eq!(
446            INPUT_REQUEST_FIXED_SLOT_BYTES,
447            128 * std::mem::size_of::<ProjectedInputRequest>()
448                + 4096 * std::mem::size_of::<bamboo_domain::SessionSkillSelection>()
449                + std::mem::size_of::<BoundedInputRequestBatch>()
450        );
451        println!(
452            "compiled fixed slots={} Record={} Selection={} Batch={}",
453            INPUT_REQUEST_FIXED_SLOT_BYTES,
454            std::mem::size_of::<ProjectedInputRequest>(),
455            std::mem::size_of::<bamboo_domain::SessionSkillSelection>(),
456            std::mem::size_of::<BoundedInputRequestBatch>()
457        );
458    }
459}
460
461#[cfg(test)]
462mod input_request_scope_tests {
463    use super::*;
464    use bamboo_agent_core::tools::{
465        observed_tool_output_cap, scope_tool_output_cap, FunctionCall, ToolCall, ToolCtx,
466        ToolError, ToolResult, ToolSchema,
467    };
468    use futures::{future::join_all, FutureExt};
469    use std::sync::{
470        atomic::{AtomicUsize, Ordering},
471        Arc,
472    };
473
474    fn batch(session: &str, execution: &str, id: &str) -> BoundedInputRequestBatch {
475        let request = bamboo_domain::SessionSkillRequest {
476            mode: Some("plan".into()),
477            selections: vec![bamboo_domain::SessionSkillSelection {
478                id: "untrusted-selection".into(),
479                source: "user".into(),
480                revision: 7,
481                args: serde_json::json!({"text": "bounded arguments"}),
482            }],
483        };
484        let records = [BorrowedInputRequestRecord {
485            input_id: id,
486            source: &bamboo_domain::SessionMessageSource::User,
487            kind: bamboo_domain::SessionMessageKind::UserInput,
488            wrapper: None,
489            created_at: chrono::DateTime::from_timestamp(1, 0).unwrap(),
490            request: Some(&request),
491        }];
492        let InputRequestProjection::Available(batch) =
493            project_input_request_batch(session, execution, &records)
494        else {
495            panic!("bounded fixture must fit")
496        };
497        batch
498    }
499
500    fn pointer(session: &str, execution: &str) -> Option<usize> {
501        with_scoped_input_request_data(session, execution, |batch| {
502            batch.records().as_ptr() as usize
503        })
504    }
505
506    #[tokio::test]
507    async fn owned_scope_pending_and_ready_return_the_same_batch_and_args_allocations() {
508        let original = batch("s", "e", "N");
509        let records = original.records().as_ptr() as usize;
510        let args = original.records()[0].request.as_ref().unwrap().selections[0].args["text"]
511            .as_str()
512            .unwrap()
513            .as_ptr() as usize;
514        let bytes = original.compact_bytes();
515        let mut future = Box::pin(scope_input_request_data(Some(original), async {
516            assert_eq!(pointer("s", "e"), Some(records));
517            tokio::task::yield_now().await;
518            with_scoped_input_request_data("s", "e", |batch| {
519                assert_eq!(batch.compact_bytes(), bytes);
520                assert_eq!(batch.records()[0].input_id, "N");
521                let request = batch.records()[0].request.as_ref().unwrap();
522                assert_eq!(request.mode.as_deref(), Some("plan"));
523                assert_eq!(
524                    request.selections[0].args["text"]
525                        .as_str()
526                        .unwrap()
527                        .as_ptr() as usize,
528                    args
529                );
530                43
531            })
532        }));
533        assert!(futures::poll!(future.as_mut()).is_pending());
534        assert_eq!(pointer("s", "e"), None, "restored between polls");
535        let (result, returned) = future.await;
536        assert_eq!(result, Some(43));
537        assert_eq!(returned.unwrap().records().as_ptr() as usize, records);
538        assert_eq!(pointer("s", "e"), None);
539    }
540
541    #[tokio::test]
542    async fn owned_scope_absent_and_mismatched_ids_never_call_the_reader() {
543        let calls = AtomicUsize::new(0);
544        assert_eq!(
545            with_scoped_input_request_data("s", "e", |_| calls.fetch_add(1, Ordering::SeqCst)),
546            None
547        );
548        scope_input_request_data(Some(batch("s", "e", "N")), async {
549            for (session, execution) in [
550                ("", "e"),
551                ("s", ""),
552                ("other", "e"),
553                ("s", "other"),
554                (" s", "e"),
555            ] {
556                assert_eq!(
557                    with_scoped_input_request_data(session, execution, |_| calls
558                        .fetch_add(1, Ordering::SeqCst)),
559                    None
560                );
561            }
562            assert!(pointer("s", "e").is_some());
563        })
564        .await;
565        scope_input_request_data(Some(batch("", "", "N")), async {
566            assert_eq!(pointer("", ""), None);
567        })
568        .await;
569        scope_input_request_data(None, async {
570            assert_eq!(pointer("s", "e"), None);
571        })
572        .await;
573        assert_eq!(calls.load(Ordering::SeqCst), 0);
574    }
575
576    #[tokio::test]
577    async fn owned_scope_nested_values_and_explicit_absence_restore_the_outer_value() {
578        let original = batch("s", "e", "outer");
579        let outer = original.records().as_ptr() as usize;
580        let (_, returned) = scope_input_request_data(Some(original), async {
581            let inner = batch("s", "inner", "inner");
582            let inner_pointer = inner.records().as_ptr() as usize;
583            let (_, inner) = scope_input_request_data(Some(inner), async {
584                assert_eq!(pointer("s", "e"), None);
585                assert_eq!(pointer("s", "inner"), Some(inner_pointer));
586                tokio::task::yield_now().await;
587                assert_eq!(pointer("s", "inner"), Some(inner_pointer));
588            })
589            .await;
590            assert_eq!(inner.unwrap().records().as_ptr() as usize, inner_pointer);
591            assert_eq!(pointer("s", "e"), Some(outer));
592            let (_, absent) = scope_input_request_data(None, async {
593                assert_eq!(pointer("s", "e"), None);
594                tokio::task::yield_now().await;
595                assert_eq!(pointer("s", "e"), None);
596            })
597            .await;
598            assert!(absent.is_none());
599            assert_eq!(pointer("s", "e"), Some(outer));
600        })
601        .await;
602        assert_eq!(returned.unwrap().records().as_ptr() as usize, outer);
603    }
604
605    struct DataExecutor {
606        entered: tokio::sync::Barrier,
607    }
608
609    #[async_trait]
610    impl ToolExecutor for DataExecutor {
611        async fn execute(&self, call: &ToolCall) -> Result<ToolResult, ToolError> {
612            let mut ctx = ToolCtx::none(call.id.clone());
613            ctx.session_id = Some("s".into());
614            let before = pointer("s", "e").expect("real executor future sees the owned batch");
615            let cap = observed_tool_output_cap(&ctx).unwrap();
616            self.entered.wait().await;
617            tokio::task::yield_now().await;
618            assert_eq!(pointer("s", "e"), Some(before));
619            assert_eq!(observed_tool_output_cap(&ctx), Some(cap));
620            Ok(ToolResult {
621                success: true,
622                result: serde_json::json!({"pointer": before, "cap": cap}).to_string(),
623                ..Default::default()
624            })
625        }
626
627        fn list_tools(&self) -> Vec<ToolSchema> {
628            Vec::new()
629        }
630    }
631
632    #[tokio::test]
633    async fn owned_scope_real_inline_executor_futures_share_data_with_independent_dispatch_caps() {
634        let executor = DataExecutor {
635            entered: tokio::sync::Barrier::new(2),
636        };
637        let original = batch("s", "e", "N");
638        let records = original.records().as_ptr() as usize;
639        let calls = [7u32, 19].map(|cap| ToolCall {
640            id: cap.to_string(),
641            tool_type: "function".into(),
642            function: FunctionCall {
643                name: "data-probe".into(),
644                arguments: "{}".into(),
645            },
646        });
647        let (results, returned) = scope_input_request_data(Some(original), async {
648            join_all(calls.iter().zip([7, 19]).map(|(call, cap)| {
649                scope_tool_output_cap("s", &call.id, Some(cap), executor.execute(call))
650            }))
651            .await
652        })
653        .await;
654        for (result, cap) in results.into_iter().zip([7, 19]) {
655            let value: serde_json::Value = serde_json::from_str(&result.unwrap().result).unwrap();
656            assert_eq!(value["pointer"], records);
657            assert_eq!(value["cap"], cap);
658        }
659        assert_eq!(returned.unwrap().records().as_ptr() as usize, records);
660        assert_eq!(pointer("s", "e"), None);
661    }
662
663    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
664    async fn owned_scope_whole_send_futures_with_identical_ids_stay_task_local() {
665        let barrier = Arc::new(tokio::sync::Barrier::new(2));
666        let mut tasks = Vec::new();
667        for id in ["A", "B"] {
668            let barrier = barrier.clone();
669            let original = batch("s", "e", id);
670            let records = original.records().as_ptr() as usize;
671            tasks.push(tokio::spawn(scope_input_request_data(
672                Some(original),
673                async move {
674                    barrier.wait().await;
675                    for _ in 0..10 {
676                        tokio::task::yield_now().await;
677                        assert_eq!(pointer("s", "e"), Some(records));
678                        assert_eq!(
679                            with_scoped_input_request_data("s", "e", |batch| batch.records()[0]
680                                .input_id
681                                .clone())
682                            .as_deref(),
683                            Some(id)
684                        );
685                    }
686                    records
687                },
688            )));
689        }
690        for task in tasks {
691            let (records, returned) = task.await.unwrap();
692            assert_eq!(returned.unwrap().records().as_ptr() as usize, records);
693        }
694        assert_eq!(pointer("s", "e"), None);
695    }
696
697    #[tokio::test]
698    async fn owned_scope_unwrapped_spawn_and_spawn_blocking_do_not_inherit_data() {
699        scope_input_request_data(Some(batch("s", "e", "N")), async {
700            let records = pointer("s", "e");
701            assert!(records.is_some());
702            assert_eq!(
703                tokio::spawn(async { pointer("s", "e") }).await.unwrap(),
704                None
705            );
706            assert_eq!(
707                tokio::task::spawn_blocking(|| pointer("s", "e"))
708                    .await
709                    .unwrap(),
710                None
711            );
712            assert_eq!(pointer("s", "e"), records);
713        })
714        .await;
715    }
716
717    #[tokio::test]
718    async fn owned_scope_inline_non_send_future_is_supported() {
719        let local = std::rc::Rc::new(41);
720        let (result, returned) = scope_input_request_data(Some(batch("s", "e", "N")), async {
721            let local = local.clone();
722            tokio::task::yield_now().await;
723            assert!(pointer("s", "e").is_some());
724            *local
725        })
726        .await;
727        assert_eq!(result, 41);
728        assert!(returned.is_some());
729        assert_eq!(pointer("s", "e"), None);
730    }
731
732    struct DropProbe(Arc<AtomicUsize>);
733    impl Drop for DropProbe {
734        fn drop(&mut self) {
735            self.0.fetch_add(1, Ordering::SeqCst);
736        }
737    }
738
739    #[tokio::test]
740    async fn owned_scope_cancellation_drops_work_once_and_restores_outer_data() {
741        let drops = Arc::new(AtomicUsize::new(0));
742        let probe = DropProbe(drops.clone());
743        let (_, returned) = scope_input_request_data(Some(batch("s", "outer", "outer")), async {
744            let outer = pointer("s", "outer");
745            let mut cancelled = Box::pin(scope_input_request_data(
746                Some(batch("s", "e", "cancelled")),
747                async move {
748                    let _probe = probe;
749                    assert!(pointer("s", "e").is_some());
750                    assert_eq!(pointer("s", "outer"), None);
751                    std::future::pending::<()>().await;
752                },
753            ));
754            assert!(futures::poll!(cancelled.as_mut()).is_pending());
755            assert_eq!(pointer("s", "e"), None);
756            assert_eq!(pointer("s", "outer"), outer);
757            drop(cancelled);
758            assert_eq!(drops.load(Ordering::SeqCst), 1);
759            assert_eq!(pointer("s", "e"), None);
760            assert_eq!(pointer("s", "outer"), outer);
761        })
762        .await;
763        assert!(returned.is_some());
764        assert_eq!(drops.load(Ordering::SeqCst), 1);
765        assert_eq!(pointer("s", "outer"), None);
766    }
767
768    #[tokio::test]
769    async fn owned_scope_unwind_drops_work_once_and_restores_outer_data() {
770        let drops = Arc::new(AtomicUsize::new(0));
771        let probe = DropProbe(drops.clone());
772        scope_input_request_data(Some(batch("s", "outer", "outer")), async {
773            let outer = pointer("s", "outer");
774            let result = std::panic::AssertUnwindSafe(scope_input_request_data(
775                Some(batch("s", "e", "panicked")),
776                async move {
777                    let _probe = probe;
778                    tokio::task::yield_now().await;
779                    assert!(pointer("s", "e").is_some());
780                    panic!("intentional scoped-work panic");
781                },
782            ))
783            .catch_unwind()
784            .await;
785            assert!(result.is_err());
786            assert_eq!(drops.load(Ordering::SeqCst), 1);
787            assert_eq!(pointer("s", "e"), None);
788            assert_eq!(pointer("s", "outer"), outer);
789        })
790        .await;
791        assert_eq!(pointer("s", "outer"), None);
792    }
793
794    #[tokio::test]
795    async fn owned_scope_error_result_returns_ownership_and_never_retains_a_scope() {
796        let original = batch("s", "e", "N");
797        let records = original.records().as_ptr() as usize;
798        let (result, returned) = scope_input_request_data(Some(original), async {
799            tokio::task::yield_now().await;
800            Err::<(), _>("real error result")
801        })
802        .await;
803        assert_eq!(result, Err("real error result"));
804        assert_eq!(returned.unwrap().records().as_ptr() as usize, records);
805        assert_eq!(pointer("s", "e"), None);
806    }
807
808    #[tokio::test]
809    async fn owned_scope_round_updates_retain_nonew_replace_new_and_clear_unavailable() {
810        let mut current = None;
811        InputObservation::projected(InputRequestProjection::Available(batch("s", "e", "N")))
812            .update_current(&mut current, "s", "e");
813        let original = current.as_ref().unwrap().records().as_ptr() as usize;
814        let (_, returned) = scope_input_request_data(current.take(), async {
815            assert_eq!(pointer("s", "e"), Some(original));
816        })
817        .await;
818        current = returned;
819        InputObservation::projected(project_input_request_batch("s", "e", &[])).update_current(
820            &mut current,
821            "s",
822            "e",
823        );
824        assert_eq!(
825            current.as_ref().unwrap().records().as_ptr() as usize,
826            original
827        );
828        let mut newer = batch("s", "e", "N2");
829        newer.records[0].request = None;
830        InputObservation::projected(InputRequestProjection::Available(newer)).update_current(
831            &mut current,
832            "s",
833            "e",
834        );
835        let (_, returned) = scope_input_request_data(current.take(), async {
836            with_scoped_input_request_data("s", "e", |batch| {
837                assert_eq!(batch.records()[0].input_id, "N2");
838                assert!(batch.records()[0].request.is_none());
839            })
840            .unwrap();
841        })
842        .await;
843        current = returned;
844        InputObservation::default().update_current(&mut current, "s", "e");
845        InputObservation::projected(project_input_request_batch("s", "e", &[])).update_current(
846            &mut current,
847            "s",
848            "e",
849        );
850        let (_, returned) = scope_input_request_data(current, async {
851            assert_eq!(
852                pointer("s", "e"),
853                None,
854                "NoNew cannot resurrect cleared data"
855            );
856        })
857        .await;
858        assert!(returned.is_none());
859    }
860
861    #[tokio::test]
862    async fn owned_scope_aborting_a_real_task_drops_work_without_inheriting_data() {
863        let drops = Arc::new(AtomicUsize::new(0));
864        let probe = DropProbe(drops.clone());
865        let entered = Arc::new(tokio::sync::Notify::new());
866        let task_entered = entered.clone();
867        let task = tokio::spawn(scope_input_request_data(
868            Some(batch("s", "e", "aborted")),
869            async move {
870                let _probe = probe;
871                assert!(pointer("s", "e").is_some());
872                task_entered.notify_one();
873                std::future::pending::<()>().await;
874            },
875        ));
876        entered.notified().await;
877        assert_eq!(pointer("s", "e"), None);
878        task.abort();
879        assert!(task.await.unwrap_err().is_cancelled());
880        assert_eq!(drops.load(Ordering::SeqCst), 1);
881        assert_eq!(pointer("s", "e"), None);
882    }
883
884    #[tokio::test]
885    async fn owned_scope_maximum_record_batch_remains_one_allocation() {
886        let ids = (0..128).map(|i| format!("input-{i}")).collect::<Vec<_>>();
887        let source = bamboo_domain::SessionMessageSource::User;
888        let records = ids
889            .iter()
890            .map(|id| BorrowedInputRequestRecord {
891                input_id: id,
892                source: &source,
893                kind: bamboo_domain::SessionMessageKind::UserInput,
894                wrapper: None,
895                created_at: chrono::DateTime::from_timestamp(1, 0).unwrap(),
896                request: None,
897            })
898            .collect::<Vec<_>>();
899        let InputRequestProjection::Available(original) =
900            project_input_request_batch("s", "e", &records)
901        else {
902            panic!("maximum slots must fit")
903        };
904        let address = original.records().as_ptr() as usize;
905        let (_, returned) = scope_input_request_data(Some(original), async {
906            tokio::task::yield_now().await;
907            with_scoped_input_request_data("s", "e", |batch| {
908                assert_eq!(batch.records().len(), 128);
909                assert_eq!(batch.records().as_ptr() as usize, address);
910                assert_eq!(batch.records()[127].input_id, "input-127");
911            })
912            .unwrap();
913        })
914        .await;
915        assert_eq!(returned.unwrap().records().as_ptr() as usize, address);
916    }
917}
bamboo_engine/runtime/managers/lifecycle.rs

bamboo_engine/runtime/managers/
lifecycle.rs
