aether_core/core/
prompt.rs

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

aether_core/core/
prompt.rs
