1use std::collections::BTreeMap;
2use std::path::PathBuf;
3use std::sync::Arc;
4
5use serde_json::Value;
6use sha2::{Digest, Sha256};
7
8use crate::sessions::Session;
9use crate::types::Metadata;
10
11#[derive(Debug, Clone, PartialEq)]
12pub struct ContextFragment {
13 pub id: String,
14 pub text: String,
15 pub stable: bool,
16 pub priority: i32,
17 pub source: Option<String>,
18 pub cache_hint: Option<String>,
19 pub metadata: Metadata,
20}
21
22impl ContextFragment {
23 pub fn new(id: impl Into<String>, text: impl Into<String>) -> Self {
24 Self {
25 id: id.into(),
26 text: text.into(),
27 stable: false,
28 priority: 100,
29 source: None,
30 cache_hint: None,
31 metadata: Metadata::new(),
32 }
33 }
34
35 pub fn stable(mut self, stable: bool) -> Self {
36 self.stable = stable;
37 self
38 }
39
40 pub fn priority(mut self, priority: i32) -> Self {
41 self.priority = priority;
42 self
43 }
44
45 pub fn source(mut self, source: impl Into<String>) -> Self {
46 self.source = Some(source.into());
47 self
48 }
49
50 pub fn cache_hint(mut self, cache_hint: impl Into<String>) -> Self {
51 self.cache_hint = Some(cache_hint.into());
52 self
53 }
54
55 pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
56 self.metadata.insert(key.into(), value);
57 self
58 }
59}
60
61pub struct ContextRequest<'a> {
62 pub agent_name: &'a str,
63 pub input: &'a str,
64 pub model: Option<String>,
65 pub trace_id: Option<String>,
66 pub session: Option<Arc<dyn Session>>,
67 pub workspace: Option<PathBuf>,
68 pub context: Option<Arc<dyn std::any::Any + Send + Sync>>,
69 pub max_prompt_chars: Option<usize>,
70 pub metadata: Metadata,
71}
72
73impl<'a> ContextRequest<'a> {
74 pub fn new(agent_name: &'a str, input: &'a str) -> Self {
75 Self {
76 agent_name,
77 input,
78 model: None,
79 trace_id: None,
80 session: None,
81 workspace: None,
82 context: None,
83 max_prompt_chars: None,
84 metadata: Metadata::new(),
85 }
86 }
87
88 pub fn for_test(agent_name: &'a str, input: &'a str) -> Self {
89 Self::new(agent_name, input)
90 }
91
92 pub fn max_prompt_chars(mut self, max_prompt_chars: usize) -> Self {
93 self.max_prompt_chars = Some(max_prompt_chars);
94 self
95 }
96
97 pub fn model(mut self, model: impl Into<String>) -> Self {
98 self.model = Some(model.into());
99 self
100 }
101
102 pub fn trace_id(mut self, trace_id: impl Into<String>) -> Self {
103 self.trace_id = Some(trace_id.into());
104 self
105 }
106
107 pub fn session(mut self, session: Arc<dyn Session>) -> Self {
108 self.session = Some(session);
109 self
110 }
111
112 pub fn workspace(mut self, workspace: impl Into<PathBuf>) -> Self {
113 self.workspace = Some(workspace.into());
114 self
115 }
116
117 pub fn context(mut self, context: Arc<dyn std::any::Any + Send + Sync>) -> Self {
118 self.context = Some(context);
119 self
120 }
121
122 pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
123 self.metadata.insert(key.into(), value);
124 self
125 }
126}
127
128#[derive(Debug, Clone, PartialEq)]
129pub struct ContextSection {
130 pub id: String,
131 pub text: String,
132 pub stable: bool,
133 pub priority: i32,
134 pub source: Option<String>,
135 pub cache_hint: Option<String>,
136 pub metadata: Metadata,
137}
138
139impl ContextSection {
140 pub fn to_metadata(&self) -> Value {
141 let mut payload = serde_json::Map::from_iter([
142 ("id".to_string(), Value::String(self.id.clone())),
143 ("text".to_string(), Value::String(self.text.clone())),
144 ("stable".to_string(), Value::Bool(self.stable)),
145 ]);
146 if let Some(source) = self.source.as_deref().filter(|value| !value.is_empty()) {
147 payload.insert("source".to_string(), Value::String(source.to_string()));
148 }
149 if let Some(cache_hint) = &self.cache_hint {
150 payload.insert("cache_hint".to_string(), Value::String(cache_hint.clone()));
151 }
152 if !self.metadata.is_empty() {
153 payload.insert(
154 "metadata".to_string(),
155 serde_json::to_value(&self.metadata).unwrap_or(Value::Null),
156 );
157 }
158 Value::Object(payload)
159 }
160}
161
162#[derive(Debug, Clone, PartialEq)]
163pub struct ContextBundle {
164 pub prompt: String,
165 pub sections: Vec<ContextSection>,
166 pub stable_hash: String,
167 pub sources: BTreeMap<String, String>,
168 pub total_chars: usize,
169 pub omitted_section_ids: Vec<String>,
170}
171
172impl ContextBundle {
173 pub fn metadata_sections(&self) -> Vec<Value> {
174 self.sections
175 .iter()
176 .map(ContextSection::to_metadata)
177 .collect()
178 }
179}
180
181pub trait ContextProvider: Send + Sync {
182 fn fragments(&self, request: &ContextRequest<'_>)
183 -> Result<Vec<ContextFragment>, ContextError>;
184}
185
186#[derive(Debug, Clone, PartialEq, Eq)]
187pub struct ContextError {
188 message: String,
189}
190
191impl ContextError {
192 pub fn new(message: impl Into<String>) -> Self {
193 Self {
194 message: message.into(),
195 }
196 }
197}
198
199impl std::fmt::Display for ContextError {
200 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201 formatter.write_str(&self.message)
202 }
203}
204
205impl std::error::Error for ContextError {}
206
207pub fn collect_context_fragments(
208 request: &ContextRequest<'_>,
209 providers: &[std::sync::Arc<dyn ContextProvider>],
210) -> Result<Vec<ContextFragment>, ContextError> {
211 let mut fragments = Vec::new();
212 for provider in providers {
213 fragments.extend(provider.fragments(request)?);
214 }
215 Ok(fragments)
216}
217
218pub fn assemble_context_fragments(
219 request: &ContextRequest<'_>,
220 mut fragments: Vec<ContextFragment>,
221) -> Result<ContextBundle, ContextError> {
222 fragments.sort_by(|left, right| {
223 left.priority
224 .cmp(&right.priority)
225 .then_with(|| right.stable.cmp(&left.stable))
226 .then_with(|| left.id.cmp(&right.id))
227 });
228
229 let mut prompt_parts = Vec::new();
230 let mut total_chars = 0usize;
231 let mut sections = Vec::new();
232 let mut stable_parts = Vec::new();
233 let mut sources = BTreeMap::new();
234 let mut omitted_section_ids = Vec::new();
235
236 for fragment in fragments {
237 let text = fragment.text.trim().to_string();
238 if text.is_empty() {
239 continue;
240 }
241 let separator_chars = usize::from(!prompt_parts.is_empty()) * 2;
242 let candidate_len = total_chars + separator_chars + text.chars().count();
243 if request
244 .max_prompt_chars
245 .is_some_and(|max_chars| candidate_len > max_chars)
246 {
247 omitted_section_ids.push(fragment.id);
248 continue;
249 }
250 if let Some(source) = fragment.source.as_ref() {
251 sources.insert(fragment.id.clone(), source.clone());
252 }
253 if fragment.stable {
254 stable_parts.push(text.clone());
255 }
256 prompt_parts.push(text.clone());
257 total_chars = candidate_len;
258 sections.push(ContextSection {
259 id: fragment.id,
260 text,
261 stable: fragment.stable,
262 priority: fragment.priority,
263 source: fragment.source,
264 cache_hint: fragment.cache_hint,
265 metadata: fragment.metadata,
266 });
267 }
268
269 Ok(ContextBundle {
270 prompt: prompt_parts.join("\n\n"),
271 sections,
272 stable_hash: sha256_hex(stable_parts.join("").as_bytes()),
273 sources,
274 total_chars,
275 omitted_section_ids,
276 })
277}
278
279fn sha256_hex(bytes: &[u8]) -> String {
280 let mut hasher = Sha256::new();
281 hasher.update(bytes);
282 hex_lower(&hasher.finalize())
283}
284
285fn hex_lower(bytes: &[u8]) -> String {
286 bytes.iter().map(|byte| format!("{byte:02x}")).collect()
287}