1use crate::Error;
4use serde::{Deserialize, Serialize};
5
6#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
8#[serde(rename_all = "snake_case")]
9pub enum ContextLane {
10 SystemContract,
12 CurrentTask,
14 ToolSchemas,
16 RecentTurns,
18 RetrievedContext,
20 ArtifactSummary,
22 Policy,
24}
25
26#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
28#[serde(deny_unknown_fields)]
29pub struct ContextFragment {
30 pub lane: ContextLane,
32 pub source: String,
34 pub content: String,
36 pub estimated_tokens: u32,
38}
39
40#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
42#[serde(deny_unknown_fields)]
43pub struct ContextPack {
44 pub token_budget: u32,
46 pub fragments: Vec<ContextFragment>,
48}
49
50impl ContextPack {
51 pub fn estimated_tokens(&self) -> u32 {
53 self.estimated_tokens_u64().min(u64::from(u32::MAX)) as u32
54 }
55
56 fn estimated_tokens_u64(&self) -> u64 {
57 self.fragments
58 .iter()
59 .map(|fragment| u64::from(fragment.estimated_tokens))
60 .sum()
61 }
62
63 pub fn validate_budget(&self) -> Result<(), Error> {
65 let used = self.estimated_tokens_u64();
66 if used > u64::from(self.token_budget) {
67 return Err(Error::ContextBudgetExceeded {
68 used: used.min(u64::from(u32::MAX)) as u32,
69 budget: self.token_budget,
70 });
71 }
72 Ok(())
73 }
74}
75
76#[cfg(test)]
77mod tests {
78 use super::*;
79
80 #[test]
81 fn context_budget_is_enforced() {
82 let pack = ContextPack {
83 token_budget: 10,
84 fragments: vec![ContextFragment {
85 lane: ContextLane::CurrentTask,
86 source: "user".into(),
87 content: "test".into(),
88 estimated_tokens: 11,
89 }],
90 };
91
92 assert!(matches!(
93 pack.validate_budget(),
94 Err(Error::ContextBudgetExceeded {
95 used: 11,
96 budget: 10
97 })
98 ));
99 }
100
101 #[test]
102 fn context_budget_rejects_overflowing_fragment_sum() {
103 let large_fragment = u32::MAX / 2 + 1;
104 let pack = ContextPack {
105 token_budget: u32::MAX,
106 fragments: vec![
107 ContextFragment {
108 lane: ContextLane::CurrentTask,
109 source: "first".into(),
110 content: "test".into(),
111 estimated_tokens: large_fragment,
112 },
113 ContextFragment {
114 lane: ContextLane::RetrievedContext,
115 source: "second".into(),
116 content: "test".into(),
117 estimated_tokens: large_fragment,
118 },
119 ],
120 };
121
122 assert_eq!(pack.estimated_tokens(), u32::MAX);
123 assert!(matches!(
124 pack.validate_budget(),
125 Err(Error::ContextBudgetExceeded {
126 used: u32::MAX,
127 budget: u32::MAX
128 })
129 ));
130 }
131}