use std::ops::Deref; use std::sync::Arc; use forge_domain::{Agent, *}; use serde_json::json; use tracing::debug; use crate::{AttachmentService, EnvironmentInfra, TemplateEngine, TerminalContextService}; /// Service responsible for setting user prompts in the conversation context #[derive(Clone)] pub struct UserPromptGenerator { services: Arc, agent: Agent, event: Event, current_time: chrono::DateTime, } impl> UserPromptGenerator { /// Creates a new UserPromptService pub fn new( service: Arc, agent: Agent, event: Event, current_time: chrono::DateTime, ) -> Self { Self { services: service, agent, event, current_time } } /// Sets the user prompt in the context based on agent configuration and /// event data pub async fn add_user_prompt( &self, conversation: Conversation, ) -> anyhow::Result { // Check if this is a resume BEFORE adding new messages let is_resume = conversation .context .as_ref() .map(|ctx| ctx.messages.iter().any(|msg| msg.has_role(Role::User))) .unwrap_or(false); let (conversation, content) = self.add_rendered_message(conversation).await?; let conversation = if is_resume { self.add_todos_on_resume(conversation)? } else { conversation }; let conversation = self.add_additional_context(conversation).await?; let conversation = if let Some(content) = content { self.add_attachments(conversation, &content).await? } else { conversation }; Ok(conversation) } /// Adds existing todos as a user message when resuming a conversation fn add_todos_on_resume(&self, mut conversation: Conversation) -> anyhow::Result { let mut context = conversation.context.take().unwrap_or_default(); // Load existing todos from session metrics let todos = conversation.metrics.todos.clone(); if !todos.is_empty() { // Format todos as markdown checklist let todo_content = self.format_todos_as_markdown(&todos); // Add as a droppable user message after the new task let todo_message = TextMessage { role: Role::User, content: todo_content, raw_content: None, tool_calls: None, thought_signature: None, reasoning_details: None, model: Some(self.agent.model.clone()), droppable: true, // Droppable so it can be removed during context compression phase: None, }; context = context.add_message(ContextMessage::Text(todo_message)); } Ok(conversation.context(context)) } /// Formats todos as a markdown checklist fn format_todos_as_markdown(&self, todos: &[Todo]) -> String { use std::fmt::Write; let mut content = String::from("**Current task list:**\n\n"); for todo in todos { let checkbox = match todo.status { TodoStatus::Completed => "[DONE]", TodoStatus::InProgress => "[IN_PROGRESS]", TodoStatus::Pending => "[PENDING]", TodoStatus::Cancelled => "[CANCELLED]", }; writeln!(content, "- {} {}", checkbox, todo.content) .expect("Writing to String should not fail"); } content } /// Adds additional context (piped input) as a droppable user message async fn add_additional_context( &self, mut conversation: Conversation, ) -> anyhow::Result { let mut context = conversation.context.take().unwrap_or_default(); if let Some(piped_input) = &self.event.additional_context { let piped_message = TextMessage { role: Role::User, content: piped_input.clone(), raw_content: None, tool_calls: None, thought_signature: None, reasoning_details: None, model: Some(self.agent.model.clone()), droppable: true, // Piped input is droppable phase: None, }; context = context.add_message(ContextMessage::Text(piped_message)); } Ok(conversation.context(context)) } /// Renders the user message content and adds it to the conversation /// Returns the conversation and the rendered content for attachment parsing async fn add_rendered_message( &self, mut conversation: Conversation, ) -> anyhow::Result<(Conversation, Option)> { let mut context = conversation.context.take().unwrap_or_default(); let event_value = self.event.value.clone(); let template_engine = TemplateEngine::default(); let content = if let Some(user_prompt) = &self.agent.user_prompt && self.event.value.is_some() { let user_input = self .event .value .as_ref() .and_then(|v| v.as_user_prompt().map(|u| u.as_str().to_string())) .unwrap_or_default(); let mut event_context = EventContext::new(EventContextValue::new(user_input)) .current_date(self.current_time.format("%Y-%m-%d").to_string()); // Check if context already contains user messages to determine if it's feedback let has_user_messages = context.messages.iter().any(|msg| msg.has_role(Role::User)); if has_user_messages { event_context = event_context.into_feedback(); } else { event_context = event_context.into_task(); } debug!(event_context = ?event_context, "Event context"); // Render the command first. let event_context = match self.event.value.as_ref().and_then(|v| v.as_command()) { Some(command) => { let rendered_prompt = template_engine.render_template( command.template.clone(), &json!({"parameters": command.parameters.join(" ")}), )?; event_context.event(EventContextValue::new(rendered_prompt)) } None => event_context, }; // Inject terminal context into the event context when available. let event_context = match TerminalContextService::new(self.services.clone()).get_terminal_context() { Some(ctx) => event_context.terminal_context(Some(ctx)), None => event_context, }; // Render the event value into agent's user prompt template. Some( template_engine.render_template( Template::new(user_prompt.template.as_str()), &event_context, )?, ) } else { // Use the raw event value as content if no user_prompt is provided event_value .as_ref() .and_then(|v| v.as_user_prompt().map(|p| p.deref().to_owned())) }; if let Some(content) = &content { // Create User Message let message = TextMessage { role: Role::User, content: content.clone(), raw_content: event_value, tool_calls: None, thought_signature: None, reasoning_details: None, model: Some(self.agent.model.clone()), droppable: false, phase: None, }; context = context.add_message(ContextMessage::Text(message)); } Ok((conversation.context(context), content)) } /// Parses and adds attachments to the conversation based on the provided /// content async fn add_attachments( &self, mut conversation: Conversation, content: &str, ) -> anyhow::Result { let mut context = conversation.context.take().unwrap_or_default(); // Parse Attachments (do NOT parse piped input for attachments) let attachments = self.services.attachments(content).await?; // Track file attachments as read operations in metrics let mut metrics = conversation.metrics.clone(); for attachment in &attachments { // Only track file content attachments (not images or directory listings). // Use the raw content_hash (computed before line-numbering) so that the // external-change detector, which hashes the raw file on disk, sees a // matching hash and does not raise a false "modified externally" warning. if let AttachmentContent::FileContent { info, .. } = &attachment.content { metrics = metrics.insert( attachment.path.clone(), FileOperation::new(ToolKind::Read) .content_hash(Some(info.content_hash.clone())), ); } } conversation.metrics = metrics; context = context.add_attachments(attachments, Some(self.agent.model.clone())); Ok(conversation.context(context)) } } #[cfg(test)] mod tests { use forge_domain::{ AgentId, AttachmentContent, Context, ContextMessage, ConversationId, FileInfo, ModelId, ProviderId, ToolKind, }; use pretty_assertions::assert_eq; use super::*; struct MockService; #[async_trait::async_trait] impl AttachmentService for MockService { async fn attachments(&self, _url: &str) -> anyhow::Result> { Ok(Vec::new()) } } impl crate::EnvironmentInfra for MockService { type Config = forge_config::ForgeConfig; fn get_environment(&self) -> forge_domain::Environment { use fake::{Fake, Faker}; Faker.fake() } fn get_config(&self) -> anyhow::Result { Ok(forge_config::ForgeConfig::default()) } async fn update_environment( &self, _ops: Vec, ) -> anyhow::Result<()> { Ok(()) } fn get_env_var(&self, _key: &str) -> Option { None } fn get_env_vars(&self) -> std::collections::BTreeMap { Default::default() } } fn fixture_agent_without_user_prompt() -> Agent { Agent::new( AgentId::from("test_agent"), ProviderId::OPENAI, ModelId::from("test-model"), ) } fn fixture_conversation() -> Conversation { Conversation::new(ConversationId::default()).context(Context::default()) } fn fixture_generator(agent: Agent, event: Event) -> UserPromptGenerator { UserPromptGenerator::new(Arc::new(MockService), agent, event, chrono::Local::now()) } #[tokio::test] async fn test_adds_context_as_droppable_message() { let agent = fixture_agent_without_user_prompt(); let event = Event::new("First Message").additional_context("Second Message"); let conversation = fixture_conversation(); let generator = fixture_generator(agent.clone(), event); let actual = generator.add_user_prompt(conversation).await.unwrap(); let messages = actual.context.unwrap().messages; assert_eq!( messages.len(), 2, "Should have context message and main message" ); // First message should be the context (droppable) let task_message = messages.first().unwrap(); assert_eq!(task_message.content().unwrap(), "First Message"); assert!( !task_message.is_droppable(), "Context message should be droppable" ); // Second message should not be droppable let context_message = messages.last().unwrap(); assert_eq!(context_message.content().unwrap(), "Second Message"); assert!( context_message.is_droppable(), "Main message should not be droppable" ); } #[tokio::test] async fn test_context_added_before_main_message() { let agent = fixture_agent_without_user_prompt(); let event = Event::new("First Message").additional_context("Second Message"); let conversation = fixture_conversation(); let generator = fixture_generator(agent.clone(), event); let actual = generator.add_user_prompt(conversation).await.unwrap(); let messages = actual.context.unwrap().messages; assert_eq!(messages.len(), 2); // Verify order: main message first, then additional context assert_eq!(messages[0].content().unwrap(), "First Message"); assert_eq!(messages[1].content().unwrap(), "Second Message"); } #[tokio::test] async fn test_no_context_only_main_message() { let agent = fixture_agent_without_user_prompt(); let event = Event::new("Simple task"); let conversation = fixture_conversation(); let generator = fixture_generator(agent.clone(), event); let actual = generator.add_user_prompt(conversation).await.unwrap(); let messages = actual.context.unwrap().messages; assert_eq!(messages.len(), 1, "Should only have the main message"); assert_eq!(messages[0].content().unwrap(), "Simple task"); } #[tokio::test] async fn test_empty_event_no_message_added() { let agent = fixture_agent_without_user_prompt(); let event = Event::empty(); let conversation = fixture_conversation(); let generator = fixture_generator(agent.clone(), event); let actual = generator.add_user_prompt(conversation).await.unwrap(); let messages = actual.context.unwrap().messages; assert_eq!( messages.len(), 0, "Should not add any message for empty event" ); } #[tokio::test] async fn test_raw_content_preserved_in_message() { let agent = fixture_agent_without_user_prompt(); let event = Event::new("Task text"); let conversation = fixture_conversation(); let generator = fixture_generator(agent.clone(), event); let actual = generator.add_user_prompt(conversation).await.unwrap(); let messages = actual.context.unwrap().messages; let message = messages.first().unwrap(); if let ContextMessage::Text(text_msg) = &**message { assert!( text_msg.raw_content.is_some(), "Raw content should be preserved" ); let raw = text_msg.raw_content.as_ref().unwrap(); assert_eq!(raw.as_user_prompt().unwrap().as_str(), "Task text"); } else { panic!("Expected TextMessage"); } } #[tokio::test] async fn test_attachments_tracked_as_read_operations() { // Setup - Create a service that returns file attachments struct MockServiceWithFiles; impl crate::EnvironmentInfra for MockServiceWithFiles { type Config = forge_config::ForgeConfig; fn get_environment(&self) -> forge_domain::Environment { use fake::{Fake, Faker}; Faker.fake() } fn get_config(&self) -> anyhow::Result { Ok(forge_config::ForgeConfig::default()) } async fn update_environment( &self, _ops: Vec, ) -> anyhow::Result<()> { Ok(()) } fn get_env_var(&self, _key: &str) -> Option { None } fn get_env_vars(&self) -> std::collections::BTreeMap { Default::default() } } #[async_trait::async_trait] impl AttachmentService for MockServiceWithFiles { async fn attachments(&self, _url: &str) -> anyhow::Result> { Ok(vec![ Attachment { path: "/test/file1.rs".to_string(), content: AttachmentContent::FileContent { content: "fn main() {}".to_string(), info: FileInfo::new(1, 1, 1, "hash1".to_string()), }, }, Attachment { path: "/test/file2.rs".to_string(), content: AttachmentContent::FileContent { content: "fn test() {}".to_string(), info: FileInfo::new(1, 1, 1, "hash2".to_string()), }, }, ]) } } let agent = fixture_agent_without_user_prompt(); let event = Event::new("Task with @[/test/file1.rs] and @[/test/file2.rs]"); let conversation = Conversation::new(ConversationId::default()); let generator = UserPromptGenerator::new( Arc::new(MockServiceWithFiles), agent.clone(), event, chrono::Local::now(), ); // Execute let actual = generator.add_user_prompt(conversation).await.unwrap(); // Assert - Both files should be tracked as read operations let file1_op = actual.metrics.file_operations.get("/test/file1.rs"); let file2_op = actual.metrics.file_operations.get("/test/file2.rs"); assert!(file1_op.is_some(), "file1.rs should be tracked in metrics"); assert!(file2_op.is_some(), "file2.rs should be tracked in metrics"); // Verify the operation is marked as Read let file1_metrics = file1_op.unwrap(); assert_eq!( file1_metrics.tool, ToolKind::Read, "file1.rs should be tracked as Read operation" ); assert!( file1_metrics.content_hash.is_some(), "file1.rs should have content hash" ); let file2_metrics = file2_op.unwrap(); assert_eq!( file2_metrics.tool, ToolKind::Read, "file2.rs should be tracked as Read operation" ); assert!( file2_metrics.content_hash.is_some(), "file2.rs should have content hash" ); // Verify both files are in files_accessed (since they are Read operations) assert!( actual.metrics.files_accessed.contains("/test/file1.rs"), "file1.rs should be in files_accessed" ); assert!( actual.metrics.files_accessed.contains("/test/file2.rs"), "file2.rs should be in files_accessed" ); } #[tokio::test] async fn test_todos_injected_on_resume() { // Setup - Simple mock that returns no attachments struct MockServiceWithTodos; impl crate::EnvironmentInfra for MockServiceWithTodos { type Config = forge_config::ForgeConfig; fn get_environment(&self) -> forge_domain::Environment { use fake::{Fake, Faker}; Faker.fake() } fn get_config(&self) -> anyhow::Result { Ok(forge_config::ForgeConfig::default()) } async fn update_environment( &self, _ops: Vec, ) -> anyhow::Result<()> { Ok(()) } fn get_env_var(&self, _key: &str) -> Option { None } fn get_env_vars(&self) -> std::collections::BTreeMap { Default::default() } } #[async_trait::async_trait] impl AttachmentService for MockServiceWithTodos { async fn attachments(&self, _url: &str) -> anyhow::Result> { Ok(Vec::new()) } } let agent = fixture_agent_without_user_prompt(); let event = Event::new("Continue working"); // Create a conversation with existing context (simulating resume) and todos // stored in metrics let conversation = Conversation::new(ConversationId::generate()) .context( Context::default() .add_message(ContextMessage::system("System message")) .add_message(ContextMessage::user("Previous task", None)), ) .metrics(Metrics::default().todos(vec![ Todo::new("Task 1").status(TodoStatus::Completed), Todo::new("Task 2").status(TodoStatus::InProgress), Todo::new("Task 3").status(TodoStatus::Pending), ])); let generator = UserPromptGenerator::new( Arc::new(MockServiceWithTodos), agent.clone(), event, chrono::Local::now(), ); // Execute let actual = generator.add_user_prompt(conversation).await.unwrap(); // Assert - Should have system, previous user, new user message, and todo list let messages = actual.context.unwrap().messages; assert_eq!(messages.len(), 4, "Should have 4 messages"); // First is system message assert_eq!(messages[0].content().unwrap(), "System message"); // Second is previous user task assert_eq!(messages[1].content().unwrap(), "Previous task"); // Third is the new user message assert_eq!(messages[2].content().unwrap(), "Continue working"); // Fourth should be the todo list (droppable) let todo_message = &messages[3]; assert!( todo_message.is_droppable(), "Todo message should be droppable" ); let todo_content = todo_message.content().unwrap(); assert!( todo_content.contains("Current task list:"), "Should contain task list header" ); assert!( todo_content.contains("[DONE] Task 1"), "Should contain completed task" ); assert!( todo_content.contains("[IN_PROGRESS] Task 2"), "Should contain in-progress task" ); assert!( todo_content.contains("[PENDING] Task 3"), "Should contain pending task" ); } #[tokio::test] async fn test_todos_not_injected_on_new_conversation() { // Setup - Simple mock with no attachments struct MockServiceNoTodos; impl crate::EnvironmentInfra for MockServiceNoTodos { type Config = forge_config::ForgeConfig; fn get_environment(&self) -> forge_domain::Environment { use fake::{Fake, Faker}; Faker.fake() } fn get_config(&self) -> anyhow::Result { Ok(forge_config::ForgeConfig::default()) } async fn update_environment( &self, _ops: Vec, ) -> anyhow::Result<()> { Ok(()) } fn get_env_var(&self, _key: &str) -> Option { None } fn get_env_vars(&self) -> std::collections::BTreeMap { Default::default() } } #[async_trait::async_trait] impl AttachmentService for MockServiceNoTodos { async fn attachments(&self, _url: &str) -> anyhow::Result> { Ok(Vec::new()) } } let agent = fixture_agent_without_user_prompt(); let event = Event::new("First task"); // Create a new conversation (no existing context, no todos) let conversation = Conversation::new(ConversationId::generate()); let generator = UserPromptGenerator::new( Arc::new(MockServiceNoTodos), agent.clone(), event, chrono::Local::now(), ); // Execute let actual = generator.add_user_prompt(conversation).await.unwrap(); // Assert - Should only have the user message, no todos let messages = actual.context.unwrap().messages; assert_eq!(messages.len(), 1, "Should only have user message"); assert_eq!(messages[0].content().unwrap(), "First task"); } }