Skip to main content

platonic_core/
context.rs

1//! Context assembly primitives with lane labels and budget validation.
2
3use crate::Error;
4use serde::{Deserialize, Serialize};
5
6/// Lane accounting for context assembly.
7#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
8#[serde(rename_all = "snake_case")]
9pub enum ContextLane {
10    /// Stable system contract.
11    SystemContract,
12    /// Current user task.
13    CurrentTask,
14    /// Selected tool schemas only.
15    ToolSchemas,
16    /// Recent turns preserved verbatim.
17    RecentTurns,
18    /// Retrieved memories or project facts.
19    RetrievedContext,
20    /// Artifact summaries instead of raw large blobs.
21    ArtifactSummary,
22    /// Runtime policy and approval constraints.
23    Policy,
24}
25
26/// One accountable context fragment.
27#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
28#[serde(deny_unknown_fields)]
29pub struct ContextFragment {
30    /// Context lane this fragment belongs to.
31    pub lane: ContextLane,
32    /// Human-readable source path, URL, event id, or synthetic label.
33    pub source: String,
34    /// Text injected into the model prompt.
35    pub content: String,
36    /// Estimated token count used for budget checks.
37    pub estimated_tokens: u32,
38}
39
40/// A bounded prompt/context bundle.
41#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
42#[serde(deny_unknown_fields)]
43pub struct ContextPack {
44    /// Maximum allowed prompt tokens for this pack.
45    pub token_budget: u32,
46    /// Context fragments selected for the next model call.
47    pub fragments: Vec<ContextFragment>,
48}
49
50impl ContextPack {
51    /// Sums fragment estimates, saturating at [`u32::MAX`].
52    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    /// Rejects a fragment sum that exceeds the declared token budget.
64    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}