stakpak_server/context/
builder.rs

1use crate::context::{
2    ContextBudget,
3    budget::{apply_budget, truncate_with_marker},
4    environment::EnvironmentContext,
5    project::{ContextFile, ProjectContext},
6};
7
8#[derive(Debug, Clone, Default)]
9pub struct SessionContext {
10    pub system_prompt: String,
11    pub user_context_block: Option<String>,
12}
13
14#[derive(Debug, Clone, Default)]
15pub struct SessionContextBuilder {
16    environment: Option<EnvironmentContext>,
17    project: Option<ProjectContext>,
18    base_system_prompt: Option<String>,
19    tool_summaries: Vec<String>,
20    budget: ContextBudget,
21}
22
23impl SessionContextBuilder {
24    pub fn new() -> Self {
25        Self::default()
26    }
27
28    pub fn environment(mut self, environment: EnvironmentContext) -> Self {
29        self.environment = Some(environment);
30        self
31    }
32
33    pub fn project(mut self, project: ProjectContext) -> Self {
34        self.project = Some(project);
35        self
36    }
37
38    pub fn base_system_prompt(mut self, prompt: impl Into<String>) -> Self {
39        self.base_system_prompt = Some(prompt.into());
40        self
41    }
42
43    pub fn tools(mut self, tools: &[stakai::Tool]) -> Self {
44        self.tool_summaries = tools
45            .iter()
46            .map(|tool| {
47                let description = tool.function.description.trim();
48                if description.is_empty() {
49                    format!("- {}", tool.function.name)
50                } else {
51                    format!("- {}: {}", tool.function.name, description)
52                }
53            })
54            .collect();
55        self
56    }
57
58    pub fn budget(mut self, budget: ContextBudget) -> Self {
59        self.budget = budget;
60        self
61    }
62
63    pub fn build(self) -> SessionContext {
64        let system_prompt = self.build_system_prompt();
65        let user_context_block = self.build_user_context_block();
66
67        SessionContext {
68            system_prompt,
69            user_context_block,
70        }
71    }
72
73    fn build_system_prompt(&self) -> String {
74        let mut sections = Vec::new();
75
76        if let Some(base) = self.base_system_prompt.as_ref().map(|prompt| prompt.trim())
77            && !base.is_empty()
78        {
79            sections.push(base.to_string());
80        }
81
82        if !self.tool_summaries.is_empty() {
83            sections.push(format!(
84                "## Available Tools\n{}",
85                self.tool_summaries.join("\n")
86            ));
87        }
88
89        let combined = sections.join("\n\n");
90        let (truncated, _) = truncate_with_marker(
91            &combined,
92            self.budget.system_prompt_max_chars,
93            "system prompt",
94        );
95        truncated
96    }
97
98    fn build_user_context_block(&self) -> Option<String> {
99        let mut sections = Vec::new();
100
101        if let Some(environment) = &self.environment {
102            sections.push(format!(
103                "<local_context>\n{}\n</local_context>",
104                environment.to_local_context_block()
105            ));
106        }
107
108        let mut files = self
109            .project
110            .as_ref()
111            .map(|project| project.files.clone())
112            .unwrap_or_default();
113
114        if !files.is_empty() {
115            apply_budget(&mut files, &self.budget);
116            for file in files {
117                sections.push(format_context_file(&file));
118            }
119        }
120
121        if sections.is_empty() {
122            return None;
123        }
124
125        Some(sections.join("\n\n"))
126    }
127}
128
129fn format_context_file(file: &ContextFile) -> String {
130    if file.name.eq_ignore_ascii_case("AGENTS.md") {
131        return format!(
132            "<agents_md>\n# AGENTS.md (from {})\n\n{}\n</agents_md>",
133            file.path, file.content
134        );
135    }
136
137    if file.name.eq_ignore_ascii_case("APPS.md") {
138        return format!(
139            "<apps_md>\n# APPS.md (from {})\n\n{}\n</apps_md>",
140            file.path, file.content
141        );
142    }
143
144    format!(
145        "<context_file name=\"{}\" path=\"{}\">\n{}\n</context_file>",
146        escape_xml_attribute(&file.name),
147        escape_xml_attribute(&file.path),
148        file.content
149    )
150}
151
152fn escape_xml_attribute(value: &str) -> String {
153    value
154        .replace('&', "&amp;")
155        .replace('"', "&quot;")
156        .replace('\'', "&apos;")
157        .replace('<', "&lt;")
158        .replace('>', "&gt;")
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use crate::context::{ContextPriority, project::ContextFile};
165
166    fn test_environment() -> EnvironmentContext {
167        EnvironmentContext {
168            machine_name: "test-machine".to_string(),
169            operating_system: "Linux".to_string(),
170            shell_type: "bash".to_string(),
171            is_container: false,
172            working_directory: "/tmp".to_string(),
173            current_datetime_utc: chrono::Utc::now(),
174            directory_tree: "├── src".to_string(),
175            git: None,
176        }
177    }
178
179    #[test]
180    fn builds_context_with_local_context_and_agents_file() {
181        let project = ProjectContext {
182            files: vec![ContextFile::new(
183                "AGENTS.md",
184                "/tmp/AGENTS.md",
185                "Follow project conventions",
186                ContextPriority::Critical,
187            )],
188        };
189
190        let context = SessionContextBuilder::new()
191            .environment(test_environment())
192            .project(project)
193            .build();
194
195        assert!(context.user_context_block.is_some());
196        if let Some(block) = context.user_context_block {
197            assert!(block.contains("<local_context>"));
198            assert!(block.contains("<agents_md>"));
199        }
200    }
201
202    #[test]
203    fn system_prompt_includes_base_prompt() {
204        let context = SessionContextBuilder::new()
205            .base_system_prompt("You are an expert DevOps agent.")
206            .build();
207
208        assert!(context.system_prompt.contains("expert DevOps agent"));
209    }
210
211    #[test]
212    fn system_prompt_includes_tool_summaries() {
213        let tools = vec![stakai::Tool::function(
214            "run_command",
215            "Execute a shell command",
216        )];
217
218        let context = SessionContextBuilder::new().tools(&tools).build();
219
220        assert!(context.system_prompt.contains("run_command"));
221        assert!(context.system_prompt.contains("Execute a shell command"));
222    }
223
224    #[test]
225    fn empty_builder_produces_empty_system_prompt_and_no_user_context() {
226        let context = SessionContextBuilder::new().build();
227
228        assert!(context.system_prompt.is_empty());
229        assert!(context.user_context_block.is_none());
230    }
231
232    #[test]
233    fn apps_md_formatted_with_apps_md_tag() {
234        let project = ProjectContext {
235            files: vec![ContextFile::new(
236                "APPS.md",
237                "/workspace/APPS.md",
238                "App configuration guide",
239                ContextPriority::High,
240            )],
241        };
242
243        let context = SessionContextBuilder::new().project(project).build();
244
245        if let Some(block) = context.user_context_block {
246            assert!(block.contains("<apps_md>"));
247            assert!(block.contains("App configuration guide"));
248        } else {
249            panic!("expected user context block with APPS.md");
250        }
251    }
252
253    #[test]
254    fn generic_context_file_uses_context_file_tag() {
255        let project = ProjectContext {
256            files: vec![ContextFile::new(
257                "notes.txt",
258                "caller://notes.txt",
259                "custom caller notes",
260                ContextPriority::CallerSupplied,
261            )],
262        };
263
264        let context = SessionContextBuilder::new().project(project).build();
265
266        if let Some(block) = context.user_context_block {
267            assert!(block.contains("<context_file"));
268            assert!(block.contains("custom caller notes"));
269        } else {
270            panic!("expected user context block with context file");
271        }
272    }
273
274    #[test]
275    fn generic_context_file_escapes_xml_attributes() {
276        let project = ProjectContext {
277            files: vec![ContextFile::new(
278                "bad\"name<>",
279                "caller://path?x=1&y='2'",
280                "content",
281                ContextPriority::CallerSupplied,
282            )],
283        };
284
285        let context = SessionContextBuilder::new().project(project).build();
286
287        if let Some(block) = context.user_context_block {
288            assert!(block.contains("name=\"bad&quot;name&lt;&gt;\""));
289            assert!(block.contains("path=\"caller://path?x=1&amp;y=&apos;2&apos;\""));
290        } else {
291            panic!("expected user context block with escaped attributes");
292        }
293    }
294
295    #[test]
296    fn system_prompt_truncated_by_budget() {
297        let budget = ContextBudget {
298            system_prompt_max_chars: 50,
299            ..Default::default()
300        };
301
302        let context = SessionContextBuilder::new()
303            .base_system_prompt("A".repeat(1_000))
304            .budget(budget)
305            .build();
306
307        assert!(
308            context.system_prompt.chars().count() <= 50,
309            "system prompt should be truncated to budget"
310        );
311    }
312}
stakpak_server/context/builder.rs

stakpak_server/context/
builder.rs
