From eb2ff9dd5c201ad39dc06e65630f72b93605b040 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sat, 7 Mar 2026 22:56:26 -0500 Subject: [PATCH] Add project_memory attribute for prompt nodes Prompt nodes now discover project docs (AGENTS.md, CLAUDE.md, etc.) and pass them as a system prompt to one_shot LLM calls. The project_memory attribute defaults to true and can be set to false to disable this behavior. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/arc-workflows/src/cli/backend.rs | 9 +- crates/arc-workflows/src/cli/cli_backend.rs | 5 +- crates/arc-workflows/src/graph/types.rs | 14 ++ crates/arc-workflows/src/handler/agent.rs | 1 + crates/arc-workflows/src/handler/prompt.rs | 172 +++++++++++++++++++- 5 files changed, 197 insertions(+), 4 deletions(-) diff --git a/crates/arc-workflows/src/cli/backend.rs b/crates/arc-workflows/src/cli/backend.rs index dded4371c..6a27d5c46 100644 --- a/crates/arc-workflows/src/cli/backend.rs +++ b/crates/arc-workflows/src/cli/backend.rs @@ -199,6 +199,7 @@ impl CodergenBackend for AgentApiBackend { &self, node: &Node, prompt: &str, + system_prompt: Option<&str>, stage_dir: &std::path::Path, ) -> Result { let client = Client::from_env() @@ -215,9 +216,15 @@ impl CodergenBackend for AgentApiBackend { .max_tokens() .or_else(|| arc_llm::catalog::get_model_info(model).and_then(|m| m.limits.max_output)); + let mut messages = Vec::new(); + if let Some(sys) = system_prompt { + messages.push(arc_llm::types::Message::system(sys)); + } + messages.push(arc_llm::types::Message::user(prompt)); + let request = arc_llm::types::Request { model: model.to_string(), - messages: vec![arc_llm::types::Message::user(prompt)], + messages, provider, reasoning_effort: Some(node.reasoning_effort().to_string()), tools: None, diff --git a/crates/arc-workflows/src/cli/cli_backend.rs b/crates/arc-workflows/src/cli/cli_backend.rs index 78597a84b..bdd4ec9ab 100644 --- a/crates/arc-workflows/src/cli/cli_backend.rs +++ b/crates/arc-workflows/src/cli/cli_backend.rs @@ -696,10 +696,13 @@ impl CodergenBackend for BackendRouter { &self, node: &Node, prompt: &str, + system_prompt: Option<&str>, stage_dir: &Path, ) -> Result { // CLI backend doesn't support one_shot, always route to API - self.api_backend.one_shot(node, prompt, stage_dir).await + self.api_backend + .one_shot(node, prompt, system_prompt, stage_dir) + .await } } diff --git a/crates/arc-workflows/src/graph/types.rs b/crates/arc-workflows/src/graph/types.rs index 3d7442bcf..e6230aa57 100644 --- a/crates/arc-workflows/src/graph/types.rs +++ b/crates/arc-workflows/src/graph/types.rs @@ -217,6 +217,11 @@ impl Node { self.bool_attr("allow_partial").unwrap_or(false) } + #[must_use] + pub fn project_memory(&self) -> bool { + self.bool_attr("project_memory").unwrap_or(true) + } + #[must_use] pub fn retry_policy(&self) -> Option<&str> { self.str_attr("retry_policy") @@ -525,6 +530,15 @@ mod tests { assert!(!node.allow_partial()); assert_eq!(node.retry_policy(), None); assert_eq!(node.max_visits(), None); + assert!(node.project_memory()); + } + + #[test] + fn node_project_memory_false_overrides_default() { + let mut node = Node::new("x"); + node.attrs + .insert("project_memory".to_string(), AttrValue::Boolean(false)); + assert!(!node.project_memory()); } #[test] diff --git a/crates/arc-workflows/src/handler/agent.rs b/crates/arc-workflows/src/handler/agent.rs index f6f28e8f0..a83d0115e 100644 --- a/crates/arc-workflows/src/handler/agent.rs +++ b/crates/arc-workflows/src/handler/agent.rs @@ -45,6 +45,7 @@ pub trait CodergenBackend: Send + Sync { &self, _node: &Node, _prompt: &str, + _system_prompt: Option<&str>, _stage_dir: &Path, ) -> Result { Err(ArcError::Validation( diff --git a/crates/arc-workflows/src/handler/prompt.rs b/crates/arc-workflows/src/handler/prompt.rs index bb8e39d73..7c3f07c0f 100644 --- a/crates/arc-workflows/src/handler/prompt.rs +++ b/crates/arc-workflows/src/handler/prompt.rs @@ -2,6 +2,8 @@ use std::path::Path; use async_trait::async_trait; +use arc_llm::provider::Provider; + use crate::context::keys; use crate::context::Context; use crate::error::ArcError; @@ -33,7 +35,7 @@ impl Handler for PromptHandler { context: &Context, graph: &Graph, logs_root: &Path, - _services: &EngineServices, + services: &EngineServices, ) -> Result { // 1. Build prompt (prepend fidelity preamble if present) let raw_prompt = node @@ -48,6 +50,25 @@ impl Handler for PromptHandler { format!("{preamble}\n\n{expanded}") }; + // 1b. Discover project docs for system prompt when project_memory is enabled + let system_prompt = if node.project_memory() { + let working_dir = services.sandbox.working_directory(); + let provider = node + .llm_provider() + .and_then(|s| s.parse::().ok()) + .unwrap_or(Provider::Anthropic); + let docs = + arc_agent::discover_project_docs(&*services.sandbox, working_dir, working_dir, provider).await; + tracing::debug!(node = %node.id, doc_count = docs.len(), "Project docs discovered for prompt node"); + if docs.is_empty() { + None + } else { + Some(docs.join("\n\n")) + } + } else { + None + }; + // 2. Write prompt to logs let visit = crate::engine::visit_from_context(context); let stage_dir = crate::engine::node_dir(logs_root, &node.id, visit); @@ -57,7 +78,9 @@ impl Handler for PromptHandler { // 3. Call LLM backend (one_shot) let (response_text, stage_usage, backend_files_touched) = if let Some(backend) = &self.backend { - let result = backend.one_shot(node, &prompt, &stage_dir).await; + let result = backend + .one_shot(node, &prompt, system_prompt.as_deref(), &stage_dir) + .await; match result { Ok(CodergenResult::Full(outcome)) => { let status_json = serde_json::to_string_pretty(&outcome) @@ -191,6 +214,7 @@ mod tests { &self, _node: &Node, _prompt: &str, + _system_prompt: Option<&str>, _stage_dir: &Path, ) -> Result { Ok(CodergenResult::Text { @@ -235,6 +259,7 @@ mod tests { struct OneShotCapturingBackend { captured_prompt: Arc>>, + captured_system_prompt: Arc>>>, } #[async_trait] @@ -256,9 +281,12 @@ mod tests { &self, _node: &Node, prompt: &str, + system_prompt: Option<&str>, _stage_dir: &Path, ) -> Result { *self.captured_prompt.lock().unwrap() = Some(prompt.to_string()); + *self.captured_system_prompt.lock().unwrap() = + Some(system_prompt.map(String::from)); Ok(CodergenResult::Text { text: "classified".to_string(), usage: None, @@ -270,6 +298,7 @@ mod tests { let captured = Arc::new(Mutex::new(None)); let backend = OneShotCapturingBackend { captured_prompt: captured.clone(), + captured_system_prompt: Arc::new(Mutex::new(None)), }; let handler = PromptHandler::new(Some(Box::new(backend))); @@ -295,4 +324,143 @@ mod tests { ); assert!(prompt.ends_with("Classify this")); } + + #[tokio::test] + async fn prompt_handler_passes_system_prompt_when_project_memory_enabled() { + use std::sync::Mutex; + + use arc_agent::Sandbox; + + struct CapturingBackend { + captured_system_prompt: Arc>>>, + } + + #[async_trait] + impl CodergenBackend for CapturingBackend { + async fn run( + &self, + _node: &Node, + _prompt: &str, + _context: &Context, + _thread_id: Option<&str>, + _emitter: &Arc, + _stage_dir: &Path, + _sandbox: &Arc, + ) -> Result { + panic!("run() should not be called for prompt handler"); + } + + async fn one_shot( + &self, + _node: &Node, + _prompt: &str, + system_prompt: Option<&str>, + _stage_dir: &Path, + ) -> Result { + *self.captured_system_prompt.lock().unwrap() = + Some(system_prompt.map(String::from)); + Ok(CodergenResult::Text { + text: "ok".to_string(), + usage: None, + files_touched: Vec::new(), + }) + } + } + + let captured_sys = Arc::new(Mutex::new(None)); + let backend = CapturingBackend { + captured_system_prompt: captured_sys.clone(), + }; + let handler = PromptHandler::new(Some(Box::new(backend))); + + // project_memory defaults to true; sandbox working_directory points to cwd + // which likely has no AGENTS.md/CLAUDE.md, so system_prompt should be None + let mut node = Node::new("classify"); + node.attrs.insert( + "prompt".to_string(), + AttrValue::String("Classify this".to_string()), + ); + let context = Context::new(); + let graph = Graph::new("test"); + let tmp = TempDir::new().unwrap(); + + handler + .execute(&node, &context, &graph, tmp.path(), &make_services()) + .await + .unwrap(); + + // With project_memory=true (default), one_shot is called (system_prompt captured) + let sys = captured_sys.lock().unwrap().clone(); + assert!(sys.is_some(), "one_shot should have been called"); + } + + #[tokio::test] + async fn prompt_handler_passes_none_system_prompt_when_project_memory_false() { + use std::sync::Mutex; + + use arc_agent::Sandbox; + + struct CapturingBackend { + captured_system_prompt: Arc>>>, + } + + #[async_trait] + impl CodergenBackend for CapturingBackend { + async fn run( + &self, + _node: &Node, + _prompt: &str, + _context: &Context, + _thread_id: Option<&str>, + _emitter: &Arc, + _stage_dir: &Path, + _sandbox: &Arc, + ) -> Result { + panic!("run() should not be called for prompt handler"); + } + + async fn one_shot( + &self, + _node: &Node, + _prompt: &str, + system_prompt: Option<&str>, + _stage_dir: &Path, + ) -> Result { + *self.captured_system_prompt.lock().unwrap() = + Some(system_prompt.map(String::from)); + Ok(CodergenResult::Text { + text: "ok".to_string(), + usage: None, + files_touched: Vec::new(), + }) + } + } + + let captured_sys = Arc::new(Mutex::new(None)); + let backend = CapturingBackend { + captured_system_prompt: captured_sys.clone(), + }; + let handler = PromptHandler::new(Some(Box::new(backend))); + + let mut node = Node::new("classify"); + node.attrs.insert( + "prompt".to_string(), + AttrValue::String("Classify this".to_string()), + ); + node.attrs.insert( + "project_memory".to_string(), + AttrValue::Boolean(false), + ); + let context = Context::new(); + let graph = Graph::new("test"); + let tmp = TempDir::new().unwrap(); + + handler + .execute(&node, &context, &graph, tmp.path(), &make_services()) + .await + .unwrap(); + + let sys = captured_sys.lock().unwrap().clone(); + assert_eq!(sys, Some(None), "system_prompt should be None when project_memory=false"); + } }