Skip to main content

agent_diva_agent/
context.rs

1//! Context builder for assembling prompts
2
3use crate::skills::SkillsLoader;
4use agent_diva_core::soul::SoulStateStore;
5use agent_diva_providers::Message;
6use agent_diva_tools::sanitize::truncate_tool_result;
7use std::path::Path;
8use std::path::PathBuf;
9
10const DEFAULT_AGENT_NAME: &str = "agent-diva";
11const DEFAULT_AGENT_EMOJI: &str = "🐈";
12const DEFAULT_AGENT_ROLE: &str = "helpful AI assistant";
13
14/// Runtime controls for soul prompt injection.
15#[derive(Debug, Clone)]
16pub struct SoulContextSettings {
17    pub enabled: bool,
18    pub max_chars: usize,
19    pub bootstrap_once: bool,
20}
21
22impl Default for SoulContextSettings {
23    fn default() -> Self {
24        Self {
25            enabled: true,
26            max_chars: 4000,
27            bootstrap_once: true,
28        }
29    }
30}
31
32/// Builds the context for LLM requests
33pub struct ContextBuilder {
34    workspace: PathBuf,
35    skills_loader: SkillsLoader,
36    soul_settings: SoulContextSettings,
37}
38
39impl ContextBuilder {
40    /// Create a new context builder
41    pub fn new(workspace: PathBuf) -> Self {
42        let skills_loader = SkillsLoader::new(&workspace, None);
43        Self {
44            workspace,
45            skills_loader,
46            soul_settings: SoulContextSettings::default(),
47        }
48    }
49
50    /// Create a new context builder with skills
51    pub fn with_skills(workspace: PathBuf, builtin_skills_dir: Option<PathBuf>) -> Self {
52        let skills_loader = SkillsLoader::new(&workspace, builtin_skills_dir);
53        Self {
54            workspace,
55            skills_loader,
56            soul_settings: SoulContextSettings::default(),
57        }
58    }
59
60    /// Override soul context settings.
61    pub fn set_soul_settings(&mut self, settings: SoulContextSettings) {
62        self.soul_settings = settings;
63    }
64
65    /// Build system prompt from workspace files and memory
66    pub fn build_system_prompt(&self) -> String {
67        let workspace_path = self.workspace.display();
68        let now = chrono::Local::now().format("%Y-%m-%d %H:%M (%A)");
69        let identity_header = self.load_identity_header();
70
71        let mut prompt = format!(
72            r#"{identity_header}
73
74You have access to tools that allow you to:
75- Read, write, and edit files
76- Execute shell commands
77- Search the web and fetch web pages
78- Send messages to users on chat channels
79- Schedule reminders and recurring jobs (cron)
80
81## Current Time
82{now}
83
84## Workspace
85Your workspace is at: {workspace_path}
86- Memory files: {workspace_path}/memory/MEMORY.md
87- Memory history log: {workspace_path}/memory/HISTORY.md"#
88        );
89
90        if self.soul_settings.enabled {
91            self.append_soul_sections(&mut prompt);
92        }
93
94        // Skills - progressive loading
95        // 1) Always-loaded skills (full content)
96        let always_skills = self.skills_loader.get_always_skills();
97        if !always_skills.is_empty() {
98            let always_content = self.skills_loader.load_skills_for_context(&always_skills);
99            if !always_content.is_empty() {
100                prompt.push_str("\n\n## Active Skills\n");
101                prompt.push_str(&always_content);
102            }
103        }
104
105        // 2) Available skills summary
106        let skills_summary = self.skills_loader.build_skills_summary();
107        if !skills_summary.is_empty() {
108            prompt.push_str("\n\n## Skills\n");
109            prompt.push_str(
110                "The following skills extend your capabilities. To use a skill, read its SKILL.md file using the read_file tool.\n",
111            );
112            prompt
113                .push_str("Skills with available=\"false\" need dependencies installed first.\n\n");
114            prompt.push_str(&skills_summary);
115        }
116
117        prompt.push_str(
118            r#"
119
120IMPORTANT: When responding to direct questions or conversations, reply directly with your text response.
121Only use the 'message' tool when you need to send a message to a specific chat channel.
122For normal conversation, just respond with text - do not call the message tool.
123When a user asks to create a reminder, timer, or recurring schedule, use the 'cron' tool instead of saying the feature is unavailable.
124When the user asks about prior work, project status, recent conclusions, commitments, or user preferences, prefer memory-backed answers over guesses.
125If the system already injected recalled memory for the current turn, use that context first.
126Use 'memory_recall' for compact answer-oriented recall, 'memory_search' for broader discovery, and 'memory_get' to fetch a full record or source fragment.
127Use 'diary_list' and 'diary_read' only to expand diary dates or specific diary details.
128Do not treat missing or unmatched memory as established fact.
129
130Always be helpful, accurate, and concise. When using tools, explain what you're doing."#,
131        );
132
133        prompt.push_str(&format!(
134            "\nWhen remembering something, write to {}/memory/MEMORY.md",
135            workspace_path
136        ));
137
138        prompt
139    }
140
141    fn append_soul_sections(&self, prompt: &mut String) {
142        let sections = [
143            ("AGENTS.md", "Agent Rules"),
144            ("SOUL.md", "Soul"),
145            ("IDENTITY.md", "Identity"),
146            ("USER.md", "User Profile"),
147        ];
148
149        for (rel, title) in sections {
150            if let Some(content) = self.read_soul_file(rel) {
151                self.append_section(prompt, title, &content);
152            }
153        }
154
155        if self.should_include_bootstrap() {
156            if let Some(content) = self.read_soul_file("BOOTSTRAP.md") {
157                let _ = SoulStateStore::new(&self.workspace).mark_bootstrap_seeded();
158                self.append_section(prompt, "Bootstrap", &content);
159            }
160        }
161    }
162
163    fn should_include_bootstrap(&self) -> bool {
164        if !self.soul_settings.bootstrap_once {
165            return true;
166        }
167        let store = SoulStateStore::new(&self.workspace);
168        !store.is_bootstrap_completed()
169    }
170
171    fn read_soul_file(&self, rel: &str) -> Option<String> {
172        let path = self.workspace.join(rel);
173        read_trimmed_markdown(&path, self.soul_settings.max_chars)
174    }
175
176    fn append_section(&self, prompt: &mut String, title: &str, content: &str) {
177        prompt.push_str("\n\n## ");
178        prompt.push_str(title);
179        prompt.push('\n');
180        prompt.push_str(content);
181    }
182
183    fn load_identity_header(&self) -> String {
184        let Some(content) = self.read_soul_file("IDENTITY.md") else {
185            return default_identity_header();
186        };
187
188        let name = parse_identity_field(&content, &["name", "agent", "assistant"])
189            .unwrap_or_else(|| DEFAULT_AGENT_NAME.to_string());
190        let emoji = parse_identity_field(&content, &["emoji", "icon", "signature"])
191            .unwrap_or_else(|| DEFAULT_AGENT_EMOJI.to_string());
192        let role = parse_identity_field(&content, &["role", "nature", "type"])
193            .unwrap_or_else(|| DEFAULT_AGENT_ROLE.to_string());
194        let voice = parse_identity_field(&content, &["voice", "style", "vibe"]);
195
196        let mut header = format!("# {} {}\n\nYou are {}, a {}.", name, emoji, name, role);
197        if let Some(voice) = voice {
198            header.push_str(" Preferred communication style: ");
199            header.push_str(&voice);
200            header.push('.');
201        }
202        header
203    }
204
205    /// Build the complete message list for an LLM call
206    pub fn build_messages(
207        &self,
208        history: Vec<agent_diva_core::session::ChatMessage>,
209        current_message: String,
210        channel: Option<&str>,
211        chat_id: Option<&str>,
212    ) -> Vec<Message> {
213        let mut messages = Vec::new();
214
215        // System prompt
216        let mut system_prompt = self.build_system_prompt();
217        if let (Some(ch), Some(id)) = (channel, chat_id) {
218            system_prompt.push_str(&format!(
219                "\n\n## Current Session\nChannel: {}\nChat ID: {}",
220                ch, id
221            ));
222        }
223        messages.push(Message::system(system_prompt));
224
225        // History - convert from ChatMessage to Message
226        for msg in history {
227            let message = match msg.role.as_str() {
228                "user" => Message::user(&msg.content),
229                "assistant" => {
230                    let mut m = Message::assistant(&msg.content);
231                    // Restore tool_calls from session history
232                    if let Some(ref tc_values) = msg.tool_calls {
233                        if let Ok(calls) =
234                            serde_json::from_value::<Vec<agent_diva_providers::ToolCallRequest>>(
235                                serde_json::Value::Array(tc_values.clone()),
236                            )
237                        {
238                            m.tool_calls = Some(calls);
239                        }
240                    }
241                    if let Some(reasoning) = msg.reasoning_content {
242                        m.reasoning_content = Some(reasoning);
243                    }
244                    if let Some(thinking_blocks) = msg.thinking_blocks {
245                        m.thinking_blocks = Some(thinking_blocks);
246                    }
247                    m
248                }
249                "tool" => {
250                    let tool_call_id = msg.tool_call_id.unwrap_or_default();
251                    let mut m = Message::tool(msg.content, tool_call_id);
252                    m.name = msg.name;
253                    m
254                }
255                _ => continue,
256            };
257            messages.push(message);
258        }
259
260        // Current message
261        messages.push(Message::user(current_message));
262
263        messages
264    }
265
266    /// Add a tool result to the message list
267    ///
268    /// Large tool results are truncated to prevent oversized API requests
269    /// that could cause 400 errors from LLM providers.
270    pub fn add_tool_result(
271        &self,
272        messages: &mut Vec<Message>,
273        tool_call_id: String,
274        _tool_name: String,
275        result: String,
276    ) {
277        // Truncate large tool results to prevent API errors
278        let truncated_result = truncate_tool_result(&result);
279        messages.push(Message::tool(truncated_result, tool_call_id));
280    }
281
282    /// Add an assistant message with optional tool calls
283    pub fn add_assistant_message(
284        &self,
285        messages: &mut Vec<Message>,
286        content: Option<String>,
287        tool_calls: Option<Vec<agent_diva_providers::ToolCallRequest>>,
288        reasoning_content: Option<String>,
289        thinking_blocks: Option<Vec<serde_json::Value>>,
290    ) {
291        let mut msg = Message::assistant(content.unwrap_or_default());
292        if let Some(calls) = tool_calls {
293            msg.tool_calls = Some(calls);
294        }
295        if let Some(reasoning) = reasoning_content {
296            msg.reasoning_content = Some(reasoning);
297        }
298        if let Some(blocks) = thinking_blocks {
299            msg.thinking_blocks = Some(blocks);
300        }
301        messages.push(msg);
302    }
303}
304
305impl Default for ContextBuilder {
306    fn default() -> Self {
307        Self::new(PathBuf::from("."))
308    }
309}
310
311fn read_trimmed_markdown(path: &Path, max_chars: usize) -> Option<String> {
312    let content = std::fs::read_to_string(path).ok()?;
313    let trimmed = content.trim();
314    if trimmed.is_empty() {
315        return None;
316    }
317
318    if trimmed.chars().count() <= max_chars {
319        return Some(trimmed.to_string());
320    }
321
322    let mut out = String::new();
323    for (idx, ch) in trimmed.chars().enumerate() {
324        if idx >= max_chars.saturating_sub(3) {
325            break;
326        }
327        out.push(ch);
328    }
329    out.push_str("...");
330    Some(out)
331}
332
333fn parse_identity_field(content: &str, keys: &[&str]) -> Option<String> {
334    for line in content.lines() {
335        let line = line.trim().trim_start_matches(&['-', '*'][..]).trim();
336        if line.is_empty() {
337            continue;
338        }
339
340        // Split on ASCII or full-width colon to avoid manual byte indices.
341        let (prefix, value_part) = match line.split_once(':').or_else(|| line.split_once('īŧš')) {
342            Some((p, v)) => (p.trim(), v.trim()),
343            None => continue,
344        };
345
346        for key in keys {
347            if prefix.eq_ignore_ascii_case(key) && !value_part.is_empty() {
348                return Some(value_part.to_string());
349            }
350        }
351    }
352    None
353}
354
355fn default_identity_header() -> String {
356    format!(
357        "# {} {}\n\nYou are {}, a {}.",
358        DEFAULT_AGENT_NAME, DEFAULT_AGENT_EMOJI, DEFAULT_AGENT_NAME, DEFAULT_AGENT_ROLE
359    )
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365    use std::fs;
366    use tempfile::TempDir;
367
368    #[test]
369    fn test_build_system_prompt() {
370        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
371        let prompt = builder.build_system_prompt();
372        assert!(prompt.contains("agent-diva"));
373        assert!(prompt.contains("/tmp/test"));
374        assert!(prompt.contains("prefer memory-backed answers over guesses"));
375        assert!(prompt.contains("'memory_search' for broader discovery"));
376        assert!(!prompt.contains("## Long-term Memory"));
377    }
378
379    #[test]
380    fn test_build_messages() {
381        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
382        let messages =
383            builder.build_messages(vec![], "Hello".to_string(), Some("cli"), Some("test"));
384        assert_eq!(messages.len(), 2); // system + user
385        assert_eq!(messages[0].role, "system");
386        assert_eq!(messages[1].role, "user");
387        assert_eq!(messages[1].content, "Hello");
388    }
389
390    #[test]
391    fn test_build_system_prompt_includes_skills_sections() {
392        let workspace = TempDir::new().unwrap();
393        let skills_dir = workspace.path().join("skills");
394        fs::create_dir_all(skills_dir.join("always-skill")).unwrap();
395        fs::write(
396            skills_dir.join("always-skill").join("SKILL.md"),
397            "---\nname: always-skill\ndescription: Always loaded\nmetadata: '{\"nanobot\":{\"always\":true}}'\n---\n\n# Always skill body\n",
398        )
399        .unwrap();
400
401        let builder = ContextBuilder::with_skills(workspace.path().to_path_buf(), None);
402        let prompt = builder.build_system_prompt();
403
404        assert!(prompt.contains("## Active Skills"));
405        assert!(prompt.contains("## Skills"));
406        assert!(prompt.contains("<skills>"));
407    }
408
409    #[test]
410    fn test_add_tool_result() {
411        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
412        let mut messages = vec![Message::user("test")];
413        builder.add_tool_result(
414            &mut messages,
415            "call_123".to_string(),
416            "read_file".to_string(),
417            "file content".to_string(),
418        );
419        assert_eq!(messages.len(), 2);
420        assert_eq!(messages[1].role, "tool");
421    }
422
423    #[test]
424    fn test_add_assistant_message() {
425        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
426        let mut messages = vec![Message::user("test")];
427
428        // Test with reasoning content
429        builder.add_assistant_message(
430            &mut messages,
431            Some("response".to_string()),
432            None,
433            Some("reasoning".to_string()),
434            None,
435        );
436
437        assert_eq!(messages.len(), 2);
438        assert_eq!(messages[1].role, "assistant");
439        assert_eq!(messages[1].content, "response");
440        assert_eq!(messages[1].reasoning_content, Some("reasoning".to_string()));
441    }
442
443    #[test]
444    fn test_build_system_prompt_includes_soul_sections_in_order() {
445        let workspace = TempDir::new().unwrap();
446        fs::write(workspace.path().join("AGENTS.md"), "# Repo Rules").unwrap();
447        fs::write(workspace.path().join("SOUL.md"), "# Core Traits").unwrap();
448        fs::write(workspace.path().join("IDENTITY.md"), "# Identity").unwrap();
449        fs::write(workspace.path().join("USER.md"), "# Preferences").unwrap();
450        fs::write(workspace.path().join("BOOTSTRAP.md"), "# Bootstrap Steps").unwrap();
451
452        let builder = ContextBuilder::new(workspace.path().to_path_buf());
453        let prompt = builder.build_system_prompt();
454
455        let idx_agents = prompt.find("## Agent Rules").unwrap();
456        let idx_soul = prompt.find("## Soul").unwrap();
457        let idx_identity = prompt.find("## Identity").unwrap();
458        let idx_user = prompt.find("## User Profile").unwrap();
459        let idx_bootstrap = prompt.find("## Bootstrap").unwrap();
460
461        assert!(idx_agents < idx_soul);
462        assert!(idx_soul < idx_identity);
463        assert!(idx_identity < idx_user);
464        assert!(idx_user < idx_bootstrap);
465    }
466
467    #[test]
468    fn test_build_system_prompt_skips_bootstrap_when_completed() {
469        let workspace = TempDir::new().unwrap();
470        fs::write(workspace.path().join("BOOTSTRAP.md"), "# Bootstrap Steps").unwrap();
471        let store = SoulStateStore::new(workspace.path());
472        let state = agent_diva_core::soul::SoulState {
473            bootstrap_completed_at: Some(chrono::Utc::now()),
474            ..Default::default()
475        };
476        store.save(&state).unwrap();
477
478        let builder = ContextBuilder::new(workspace.path().to_path_buf());
479        let prompt = builder.build_system_prompt();
480        assert!(!prompt.contains("## Bootstrap"));
481    }
482
483    #[test]
484    fn test_read_trimmed_markdown_respects_char_limit() {
485        let temp = TempDir::new().unwrap();
486        let path = temp.path().join("SOUL.md");
487        fs::write(&path, "abcdefghij").unwrap();
488
489        let got = read_trimmed_markdown(&path, 6).unwrap();
490        assert_eq!(got, "abc...");
491        assert!(got.chars().count() <= 6);
492    }
493
494    #[test]
495    fn test_build_system_prompt_uses_identity_file_for_header() {
496        let workspace = TempDir::new().unwrap();
497        fs::write(
498            workspace.path().join("IDENTITY.md"),
499            "# Identity\n- Name: Nova\n- Emoji: ✨\n- Role: strategic coding partner\n- Style: concise and direct\n",
500        )
501        .unwrap();
502        let builder = ContextBuilder::new(workspace.path().to_path_buf());
503        let prompt = builder.build_system_prompt();
504        assert!(prompt.contains("# Nova ✨"));
505        assert!(prompt.contains("You are Nova, a strategic coding partner."));
506        assert!(prompt.contains("Preferred communication style: concise and direct."));
507    }
508
509    #[test]
510    fn test_build_system_prompt_identity_header_falls_back_to_default() {
511        let workspace = TempDir::new().unwrap();
512        let builder = ContextBuilder::new(workspace.path().to_path_buf());
513        let prompt = builder.build_system_prompt();
514        assert!(prompt.contains("# agent-diva 🐈"));
515        assert!(prompt.contains("You are agent-diva, a helpful AI assistant."));
516    }
517
518    #[test]
519    fn test_build_system_prompt_empty_identity_falls_back_to_default() {
520        let workspace = TempDir::new().unwrap();
521        fs::write(workspace.path().join("IDENTITY.md"), "   \n").unwrap();
522        let builder = ContextBuilder::new(workspace.path().to_path_buf());
523        let prompt = builder.build_system_prompt();
524        assert!(prompt.contains("# agent-diva 🐈"));
525    }
526
527    #[test]
528    fn test_build_system_prompt_long_identity_is_trimmed_by_max_chars() {
529        let workspace = TempDir::new().unwrap();
530        let long_name = "N".repeat(6000);
531        fs::write(
532            workspace.path().join("IDENTITY.md"),
533            format!("- Name: {}\n- Role: helper", long_name),
534        )
535        .unwrap();
536        let mut builder = ContextBuilder::new(workspace.path().to_path_buf());
537        builder.set_soul_settings(SoulContextSettings {
538            enabled: true,
539            max_chars: 120,
540            bootstrap_once: true,
541        });
542        let prompt = builder.build_system_prompt();
543        assert!(prompt.contains("You are "));
544        assert!(prompt.chars().count() > 120);
545    }
546
547    #[test]
548    fn test_parse_identity_field_handles_markdown_list() {
549        let raw = "- Name: Diva\n- Style: pragmatic";
550        assert_eq!(
551            parse_identity_field(raw, &["name"]).as_deref(),
552            Some("Diva")
553        );
554        assert_eq!(
555            parse_identity_field(raw, &["style"]).as_deref(),
556            Some("pragmatic")
557        );
558    }
559
560    #[test]
561    fn test_parse_identity_field_supports_chinese_voice_line() {
562        let raw = "- Voice: įŽ€æ´ã€åŽžį”¨ã€åäŊœ";
563        assert_eq!(
564            parse_identity_field(raw, &["voice"]).as_deref(),
565            Some("įŽ€æ´ã€åŽžį”¨ã€åäŊœ")
566        );
567    }
568}