Function delegates_to_inner_guard

core/src/orchestration/workflow_budget.rs:201–225

()

    #[tokio::test]
    async fn delegates_to_inner_guard() {
        #[derive(Default)]
        struct Counting {
            checks: AtomicUsize,
            records: AtomicUsize,
        }
        #[async_trait]
        impl BudgetGuard for Counting {
            async fn check_before_llm(&self, _: &str, _: usize) -> BudgetDecision {
                self.checks.fetch_add(1, Ordering::SeqCst);
                BudgetDecision::Allow
            }
            async fn record_after_llm(&self, _: &str, _: &TokenUsage) {
                self.records.fetch_add(1, Ordering::SeqCst);
            }
        }

        let inner = Arc::new(Counting::default());
        let budget = WorkflowBudget::new(Some(1000)).with_inner(inner.clone());
        budget.check_before_llm("s", 0).await;
        budget.record_after_llm("s", &usage(10)).await;
        assert_eq!(inner.checks.load(Ordering::SeqCst), 1);
        assert_eq!(inner.records.load(Ordering::SeqCst), 1);
        assert_eq!(budget.consumed_tokens(), 10);
    }
}