(
&self,
session_id: &str,
estimated_prompt_tokens: usize,
)
| 98 | #[async_trait] |
| 99 | impl BudgetGuard for WorkflowBudget { |
| 100 | async fn check_before_llm( |
| 101 | &self, |
| 102 | session_id: &str, |
| 103 | estimated_prompt_tokens: usize, |
| 104 | ) -> BudgetDecision { |
| 105 | if self.is_exhausted() { |
| 106 | return BudgetDecision::Deny { |
| 107 | resource: "workflow_tokens".to_string(), |
| 108 | reason: format!( |
| 109 | "workflow token budget exhausted ({} / {} tokens)", |
| 110 | self.consumed_tokens(), |
| 111 | self.limit_tokens.unwrap_or(0) |
| 112 | ), |
| 113 | }; |
| 114 | } |
| 115 | match &self.inner { |
| 116 | Some(inner) => { |
| 117 | inner |
| 118 | .check_before_llm(session_id, estimated_prompt_tokens) |
| 119 | .await |
| 120 | } |
| 121 | None => BudgetDecision::Allow, |
| 122 | } |
| 123 | } |
| 124 | |
| 125 | async fn record_after_llm(&self, session_id: &str, usage: &TokenUsage) { |
| 126 | self.consumed_tokens |
no test coverage detected