aether_core/core/
prompt.rs

1use crate::core::{AgentError, Result};
2use crate::mcp::PROGRESSIVE_DISCOVERY_INSTRUCTIONS;
3use glob::glob;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use std::collections::BTreeMap;
7use std::path::{Path, PathBuf};
8use thiserror::Error;
9use tokio::fs;
10use tracing::warn;
11use utils::shell_expander::expand_shell;
12use utils::variables::VarError;
13use utils::{PathOrObject, ResourcePath};
14
15#[derive(Debug, Clone, PartialEq)]
16pub enum Prompt {
17    Text(String),
18    File { path: PathBuf, cwd: PathBuf },
19    McpInstructions(McpInstructions),
20}
21
22#[derive(Debug, Clone, Default, PartialEq)]
23pub struct McpInstructions {
24    pub servers: BTreeMap<String, String>,
25    pub progressive_discovery: bool,
26}
27
28/// Authored description of a prompt source — text, a file path, or a glob pattern.
29///
30/// Used by configuration layers to declare prompts before resolution. Convert into
31/// runtime [`Prompt`] values with [`Prompt::from_sources`].
32#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
33#[serde(from = "PathOrObject<PromptSourceObject>", into = "PathOrObject<PromptSourceObject>")]
34#[schemars(with = "PathOrObject<PromptSourceObject>", description = include_str!("../docs/prompt_source.md"))]
35pub enum PromptSource {
36    Text { text: String },
37    File { path: ResourcePath, optional: bool },
38    Glob { pattern: ResourcePath, optional: bool },
39}
40
41impl PromptSource {
42    pub fn file(path: impl Into<ResourcePath>) -> Self {
43        Self::File { path: path.into(), optional: false }
44    }
45
46    pub fn glob(pattern: impl Into<ResourcePath>) -> Self {
47        Self::Glob { pattern: pattern.into(), optional: false }
48    }
49
50    #[must_use]
51    pub fn optional(self) -> Self {
52        match self {
53            Self::File { path, .. } => Self::File { path, optional: true },
54            Self::Glob { pattern, .. } => Self::Glob { pattern, optional: true },
55            Self::Text { .. } => self,
56        }
57    }
58
59    /// The authored path/pattern string, if this source has one.
60    pub fn path(&self) -> Option<&str> {
61        match self {
62            Self::File { path, .. } => Some(path.as_authored()),
63            Self::Glob { pattern, .. } => Some(pattern.as_authored()),
64            Self::Text { .. } => None,
65        }
66    }
67
68    pub fn is_optional(&self) -> bool {
69        match self {
70            Self::File { optional, .. } | Self::Glob { optional, .. } => *optional,
71            Self::Text { .. } => false,
72        }
73    }
74}
75
76impl From<&str> for PromptSource {
77    fn from(value: &str) -> Self {
78        Self::file(value)
79    }
80}
81
82impl From<String> for PromptSource {
83    fn from(value: String) -> Self {
84        Self::file(value)
85    }
86}
87
88impl From<PromptSourceObject> for PromptSource {
89    fn from(object: PromptSourceObject) -> Self {
90        match object {
91            PromptSourceObject::Text { text } => Self::Text { text },
92            PromptSourceObject::File { path, optional } => Self::File { path, optional },
93            PromptSourceObject::Glob { pattern, optional } => Self::Glob { pattern, optional },
94        }
95    }
96}
97
98impl From<PathOrObject<PromptSourceObject>> for PromptSource {
99    fn from(source: PathOrObject<PromptSourceObject>) -> Self {
100        match source {
101            PathOrObject::Path(path) => Self::file(path),
102            PathOrObject::Object(object) => object.into(),
103        }
104    }
105}
106
107impl From<PromptSource> for PathOrObject<PromptSourceObject> {
108    fn from(source: PromptSource) -> Self {
109        match source {
110            PromptSource::File { path, optional: false } => Self::Path(path),
111            PromptSource::File { path, optional } => Self::Object(PromptSourceObject::File { path, optional }),
112            PromptSource::Text { text } => Self::Object(PromptSourceObject::Text { text }),
113            PromptSource::Glob { pattern, optional } => Self::Object(PromptSourceObject::Glob { pattern, optional }),
114        }
115    }
116}
117
118#[derive(schemars::JsonSchema, serde::Deserialize, serde::Serialize)]
119#[serde(tag = "type", rename_all = "camelCase", deny_unknown_fields)]
120enum PromptSourceObject {
121    Text {
122        /// Literal prompt text included verbatim.
123        text: String,
124    },
125    File {
126        /// Path to a prompt file, resolved as a resource path.
127        path: ResourcePath,
128        /// When true, a missing file is skipped instead of raising an error.
129        #[serde(default, skip_serializing_if = "std::ops::Not::not")]
130        optional: bool,
131    },
132    Glob {
133        /// Glob pattern matching prompt files, resolved as a resource path.
134        pattern: ResourcePath,
135        /// When true, a zero-match glob is skipped instead of raising an error.
136        #[serde(default, skip_serializing_if = "std::ops::Not::not")]
137        optional: bool,
138    },
139}
140
141/// Errors raised while resolving [`PromptSource`] values into [`Prompt`]s.
142#[derive(Debug, Clone, PartialEq, Eq, Error)]
143pub enum PromptSourceError {
144    /// A glob pattern is syntactically invalid.
145    #[error("Invalid glob pattern '{pattern}': {error}")]
146    InvalidGlobPattern { pattern: String, error: String },
147
148    /// A prompt file does not exist on disk.
149    #[error("Prompt file '{path}' does not exist")]
150    Missing { path: String },
151
152    /// A prompt glob matched no files.
153    #[error("Prompt glob '{pattern}' matched no files")]
154    ZeroMatch { pattern: String },
155
156    /// A `${VAR}` reference in a prompt path could not be resolved.
157    #[error("Prompt entry '{pattern}' references undefined variable '{variable}'")]
158    UnresolvedVariable { pattern: String, variable: String },
159}
160
161impl Prompt {
162    pub fn text(str: &str) -> Self {
163        Self::Text(str.to_string())
164    }
165
166    pub fn file(path: impl Into<PathBuf>, cwd: impl Into<PathBuf>) -> Self {
167        Self::File { path: path.into(), cwd: cwd.into() }
168    }
169
170    /// Resolve a slice of [`PromptSource`] declarations into runtime [`Prompt`] values.
171    pub fn from_sources(
172        workspace_root: &Path,
173        sources: &[PromptSource],
174    ) -> std::result::Result<Vec<Prompt>, PromptSourceError> {
175        let mut prompts = Vec::new();
176        for source in sources {
177            if let PromptSource::Text { text } = source {
178                prompts.push(Prompt::text(text));
179                continue;
180            }
181            match resolve_source_files(workspace_root, source) {
182                Ok(paths) => {
183                    for path in paths {
184                        prompts.push(Prompt::file(path, workspace_root.to_path_buf()));
185                    }
186                }
187                Err(PromptSourceError::Missing { .. }) if source.is_optional() => {}
188                Err(PromptSourceError::UnresolvedVariable { variable, .. }) if source.is_optional() => {
189                    warn!(
190                        "Skipping optional prompt entry '{}': variable '{variable}' is not defined",
191                        source.path().unwrap_or_default()
192                    );
193                }
194                Err(error) => return Err(error),
195            }
196        }
197        Ok(prompts)
198    }
199
200    /// Resolve this `SystemPrompt` to a String
201    pub async fn build(&self) -> Result<String> {
202        match self {
203            Prompt::Text(text) => Ok(text.clone()),
204            Prompt::File { path, cwd } => Ok(expand_shell(&Self::resolve_file(path).await?, cwd).await),
205            Prompt::McpInstructions(instructions) => Ok(format_mcp_instructions(instructions)),
206        }
207    }
208
209    async fn resolve_file(path: &Path) -> Result<String> {
210        fs::read_to_string(path)
211            .await
212            .map_err(|e| AgentError::IoError(format!("Failed to read file '{}': {e}", path.display())))
213    }
214}
215
216fn resolve_source_files(
217    workspace_root: &Path,
218    source: &PromptSource,
219) -> std::result::Result<Vec<PathBuf>, PromptSourceError> {
220    match source {
221        PromptSource::Text { .. } => Ok(Vec::new()),
222        PromptSource::File { path, .. } => {
223            let full_path = resolve_path(path, workspace_root)?;
224            if full_path.is_file() {
225                Ok(vec![full_path])
226            } else {
227                Err(PromptSourceError::Missing { path: path.as_authored().to_string() })
228            }
229        }
230        PromptSource::Glob { pattern, optional } => {
231            let full_pattern = resolve_path(pattern, workspace_root)?;
232            let mut paths: Vec<PathBuf> = glob(&full_pattern.to_string_lossy())
233                .map_err(|e| PromptSourceError::InvalidGlobPattern {
234                    pattern: pattern.as_authored().to_string(),
235                    error: e.to_string(),
236                })?
237                .filter_map(std::result::Result::ok)
238                .filter(|path| path.is_file())
239                .collect();
240            paths.sort();
241            if paths.is_empty() && !*optional {
242                Err(PromptSourceError::ZeroMatch { pattern: pattern.as_authored().to_string() })
243            } else {
244                Ok(paths)
245            }
246        }
247    }
248}
249
250fn resolve_path(path: &ResourcePath, workspace_root: &Path) -> std::result::Result<PathBuf, PromptSourceError> {
251    path.resolve(workspace_root).map_err(|VarError::NotFound(variable)| PromptSourceError::UnresolvedVariable {
252        pattern: path.as_authored().to_string(),
253        variable,
254    })
255}
256
257pub struct PromptCache {
258    prompts: Vec<Prompt>,
259    entries: Vec<(Prompt, String)>,
260}
261
262impl PromptCache {
263    pub fn new(mut prompts: Vec<Prompt>) -> Self {
264        if !prompts.iter().any(|p| matches!(p, Prompt::McpInstructions(_))) {
265            prompts.push(Prompt::McpInstructions(McpInstructions::default()));
266        }
267        Self { prompts, entries: Vec::new() }
268    }
269
270    pub fn set_mcp_instructions(&mut self, instructions: McpInstructions) -> bool {
271        let Some(Prompt::McpInstructions(current)) =
272            self.prompts.iter_mut().find(|prompt| matches!(prompt, Prompt::McpInstructions(_)))
273        else {
274            return false;
275        };
276        if *current == instructions {
277            return false;
278        }
279        *current = instructions;
280        true
281    }
282
283    pub async fn render(&mut self) -> Result<String> {
284        self.entries.truncate(self.prompts.len());
285        let mut rendered_prompt = String::new();
286        for i in 0..self.prompts.len() {
287            let prompt = &self.prompts[i];
288            match self.entries.get_mut(i) {
289                Some((cached, _)) if *cached == *prompt => {}
290                Some(entry) => *entry = (prompt.clone(), prompt.build().await?),
291                None => self.entries.push((prompt.clone(), prompt.build().await?)),
292            }
293
294            let (_, body) = &self.entries[i];
295            if !body.is_empty() {
296                if !rendered_prompt.is_empty() {
297                    rendered_prompt.push_str("\n\n");
298                }
299                rendered_prompt.push_str(body);
300            }
301        }
302        Ok(rendered_prompt)
303    }
304}
305
306fn format_mcp_instructions(instructions: &McpInstructions) -> String {
307    let mut sections = Vec::new();
308    if !instructions.servers.is_empty() {
309        let mut parts = vec!["# MCP Server Instructions\n".to_string()];
310        parts.push("You are connected to the following MCP servers:\n".to_string());
311        for (server_name, body) in &instructions.servers {
312            parts.push(format!("<mcp-server name=\"{server_name}\">\n{body}\n</mcp-server>\n"));
313        }
314        sections.push(parts.join("\n"));
315    }
316    if instructions.progressive_discovery {
317        sections.push(format!("# Deferred MCP Tools\n\n{PROGRESSIVE_DISCOVERY_INSTRUCTIONS}"));
318    }
319    sections.join("\n\n")
320}
321
322#[cfg(test)]
323mod tests {
324    use std::fs::{create_dir_all, write};
325
326    use super::*;
327    use crate::testing::mcp_instructions as instructions;
328
329    #[tokio::test]
330    async fn build_text_prompt() {
331        let prompt = Prompt::text("Hello, world!");
332        let result = prompt.build().await.unwrap();
333        assert_eq!(result, "Hello, world!");
334    }
335
336    #[tokio::test]
337    async fn prompt_cache_concatenates_files() {
338        let dir = tempfile::tempdir().unwrap();
339        write(dir.path().join("AGENTS.md"), "Agent instructions").unwrap();
340        write(dir.path().join("SYSTEM.md"), "System prompt").unwrap();
341
342        let mut cache = PromptCache::new(vec![
343            Prompt::file(dir.path().join("AGENTS.md"), dir.path()),
344            Prompt::file(dir.path().join("SYSTEM.md"), dir.path()),
345        ]);
346
347        assert_eq!(cache.render().await.unwrap(), "Agent instructions\n\nSystem prompt");
348    }
349
350    #[tokio::test]
351    async fn prompt_cache_skips_empty_prompts() {
352        let mut cache = PromptCache::new(vec![Prompt::text("Part one"), Prompt::text(""), Prompt::text("Part two")]);
353        assert_eq!(cache.render().await.unwrap(), "Part one\n\nPart two");
354    }
355
356    #[tokio::test]
357    async fn prompt_cache_renders_mcp_instructions_in_place() {
358        let mcp = Prompt::McpInstructions(instructions(&[("srv", "body")]));
359        let expected = format!("first\n\n{}\n\nlast", mcp.build().await.unwrap());
360
361        let mut cache = PromptCache::new(vec![Prompt::text("first"), mcp, Prompt::text("last")]);
362
363        assert_eq!(cache.render().await.unwrap(), expected);
364    }
365
366    #[tokio::test]
367    async fn prompt_cache_reuses_unchanged_slots() {
368        use std::fs::{remove_file, write};
369
370        let dir = tempfile::tempdir().unwrap();
371        write(dir.path().join("AGENTS.md"), "cached body").unwrap();
372        let mut cache = PromptCache::new(vec![
373            Prompt::file(dir.path().join("AGENTS.md"), dir.path()),
374            Prompt::McpInstructions(McpInstructions::default()),
375        ]);
376
377        cache.render().await.unwrap();
378
379        // Remove the source file to prove we cached things
380        remove_file(dir.path().join("AGENTS.md")).unwrap();
381        cache.set_mcp_instructions(instructions(&[("srv", "instr")]));
382
383        let rendered = cache.render().await.unwrap();
384        assert!(rendered.contains("cached body"));
385        assert!(rendered.contains("instr"));
386    }
387
388    #[tokio::test]
389    async fn prompt_cache_empty_renders_empty() {
390        assert_eq!(PromptCache::new(vec![]).render().await.unwrap(), "");
391    }
392
393    #[tokio::test]
394    async fn prompt_cache_drops_empty_slots() {
395        let mut cache = PromptCache::new(vec![Prompt::text("a"), Prompt::text("b")]);
396        assert_eq!(cache.render().await.unwrap(), "a\n\nb");
397    }
398
399    #[tokio::test]
400    async fn build_file_expands_shell_commands() {
401        let dir = tempfile::tempdir().unwrap();
402        write(dir.path().join("AGENTS.md"), "Instructions\n\nbranch: !`echo main`\n\nRules").unwrap();
403
404        let prompt = Prompt::file(dir.path().join("AGENTS.md"), dir.path().to_path_buf());
405        let result = prompt.build().await.unwrap();
406        assert!(result.contains("Instructions"));
407        assert!(result.contains("branch: main"));
408        assert!(result.contains("Rules"));
409        assert!(!result.contains("!`"));
410    }
411
412    #[tokio::test]
413    async fn build_file_runs_shell_in_cwd() {
414        let dir = tempfile::tempdir().unwrap();
415        write(dir.path().join("sentinel.txt"), "").unwrap();
416        let prompt_path = dir.path().join("AGENTS.md");
417        write(&prompt_path, "files: !`ls`").unwrap();
418
419        let prompt = Prompt::file(prompt_path, dir.path().to_path_buf());
420        let result = prompt.build().await.unwrap();
421        assert!(result.contains("sentinel.txt"), "expected sentinel.txt in output: {result}");
422    }
423
424    #[tokio::test]
425    async fn build_file_handles_multiple_commands() {
426        let dir = tempfile::tempdir().unwrap();
427        let prompt_path = dir.path().join("AGENTS.md");
428        write(&prompt_path, "a=!`echo one`, b=!`echo two`").unwrap();
429
430        let prompt = Prompt::file(prompt_path, dir.path().to_path_buf());
431        let result = prompt.build().await.unwrap();
432        assert_eq!(result, "a=one, b=two");
433    }
434
435    #[tokio::test]
436    async fn build_file_substitutes_empty_on_failure() {
437        let dir = tempfile::tempdir().unwrap();
438        let prompt_path = dir.path().join("AGENTS.md");
439        write(&prompt_path, "before !`exit 1` after").unwrap();
440
441        let prompt = Prompt::file(prompt_path, dir.path().to_path_buf());
442        let result = prompt.build().await.unwrap();
443        assert_eq!(result, "before  after");
444    }
445
446    #[tokio::test]
447    async fn build_file_trims_trailing_whitespace() {
448        let dir = tempfile::tempdir().unwrap();
449        let prompt_path = dir.path().join("AGENTS.md");
450        write(&prompt_path, "!`printf 'hi\\n\\n'`").unwrap();
451
452        let prompt = Prompt::file(prompt_path, dir.path().to_path_buf());
453        let result = prompt.build().await.unwrap();
454        assert_eq!(result, "hi");
455    }
456
457    #[test]
458    fn optional_file_source_skips_missing_file() {
459        let dir = tempfile::tempdir().unwrap();
460        write(dir.path().join("EXISTS.md"), "exists").unwrap();
461
462        let sources = vec![PromptSource::file("EXISTS.md"), PromptSource::file("MISSING.md").optional()];
463        let prompts = Prompt::from_sources(dir.path(), &sources).unwrap();
464        assert_eq!(prompts.len(), 1);
465    }
466
467    #[test]
468    fn optional_glob_source_skips_zero_matches() {
469        let dir = tempfile::tempdir().unwrap();
470        write(dir.path().join("EXISTS.md"), "exists").unwrap();
471
472        let sources = vec![PromptSource::file("EXISTS.md"), PromptSource::glob("nonexistent*.md").optional()];
473        let prompts = Prompt::from_sources(dir.path(), &sources).unwrap();
474        assert_eq!(prompts.len(), 1);
475    }
476
477    #[test]
478    fn required_glob_source_expands_to_one_prompt_per_match() {
479        let dir = tempfile::tempdir().unwrap();
480        let rules_dir = dir.path().join(".aether/rules");
481        create_dir_all(&rules_dir).unwrap();
482        write(rules_dir.join("a.md"), "a").unwrap();
483        write(rules_dir.join("b.md"), "b").unwrap();
484
485        let sources = vec![PromptSource::glob(".aether/rules/*.md")];
486        let prompts = Prompt::from_sources(dir.path(), &sources).unwrap();
487        assert_eq!(prompts.len(), 2);
488    }
489
490    #[test]
491    fn required_glob_source_with_no_matches_errors() {
492        let dir = tempfile::tempdir().unwrap();
493        let sources = vec![PromptSource::glob("nonexistent*.md")];
494        let err = Prompt::from_sources(dir.path(), &sources).unwrap_err();
495        assert!(matches!(err, PromptSourceError::ZeroMatch { .. }));
496    }
497
498    #[test]
499    fn optional_glob_source_still_errors_on_invalid_pattern() {
500        let dir = tempfile::tempdir().unwrap();
501        let sources = vec![PromptSource::glob("[invalid").optional()];
502        let err = Prompt::from_sources(dir.path(), &sources).unwrap_err();
503        assert!(matches!(err, PromptSourceError::InvalidGlobPattern { .. }));
504    }
505
506    #[test]
507    fn optional_file_source_skips_unresolved_variable() {
508        let dir = tempfile::tempdir().unwrap();
509        write(dir.path().join("EXISTS.md"), "exists").unwrap();
510
511        let sources = vec![
512            PromptSource::file("EXISTS.md"),
513            PromptSource::file("${DEFINITELY_NOT_SET_VAR_PROMPT_FILE}/foo.md").optional(),
514        ];
515        let prompts = Prompt::from_sources(dir.path(), &sources).unwrap();
516        assert_eq!(prompts.len(), 1);
517    }
518
519    #[test]
520    fn optional_glob_source_skips_unresolved_variable() {
521        let dir = tempfile::tempdir().unwrap();
522        write(dir.path().join("EXISTS.md"), "exists").unwrap();
523
524        let sources = vec![
525            PromptSource::file("EXISTS.md"),
526            PromptSource::glob("${DEFINITELY_NOT_SET_VAR_PROMPT_GLOB}/*.md").optional(),
527        ];
528        let prompts = Prompt::from_sources(dir.path(), &sources).unwrap();
529        assert_eq!(prompts.len(), 1);
530    }
531
532    #[test]
533    fn required_file_source_errors_on_unresolved_variable() {
534        let dir = tempfile::tempdir().unwrap();
535        let sources = vec![PromptSource::file("${DEFINITELY_NOT_SET_VAR_PROMPT_REQ}/foo.md")];
536        let err = Prompt::from_sources(dir.path(), &sources).unwrap_err();
537        assert!(matches!(err, PromptSourceError::UnresolvedVariable { .. }));
538    }
539
540    #[test]
541    fn prompt_source_string_shorthand_is_required_file() {
542        let source: PromptSource = serde_json::from_str(r#""SYSTEM.md""#).unwrap();
543        assert_eq!(source, PromptSource::file("SYSTEM.md"));
544    }
545
546    #[test]
547    fn optional_prompt_source_serializes_as_typed_object() {
548        let source = PromptSource::file("${WORKSPACE}/AGENTS.md").optional();
549        let value = serde_json::to_value(&source).unwrap();
550        assert_eq!(value, serde_json::json!({"type":"file","path":"${WORKSPACE}/AGENTS.md","optional":true}));
551
552        // Non-optional file stays as string shorthand
553        let source = PromptSource::file("SYSTEM.md");
554        let value = serde_json::to_value(&source).unwrap();
555        assert_eq!(value, serde_json::json!("SYSTEM.md"));
556    }
557
558    #[test]
559    fn optional_prompt_source_deserializes_from_typed_object() {
560        let source: PromptSource =
561            serde_json::from_str(r#"{"type":"file","path":"${WORKSPACE}/AGENTS.md","optional":true}"#).unwrap();
562        assert_eq!(source, PromptSource::file("${WORKSPACE}/AGENTS.md").optional());
563    }
564
565    #[test]
566    fn optional_glob_source_deserializes_from_typed_object() {
567        let source: PromptSource =
568            serde_json::from_str(r#"{"type":"glob","pattern":"${WORKSPACE}/.aether/rules/*.md","optional":true}"#)
569                .unwrap();
570        assert_eq!(source, PromptSource::glob("${WORKSPACE}/.aether/rules/*.md").optional());
571    }
572}
aether_core/core/prompt.rs

aether_core/core/
prompt.rs
