From 583242dadfb6f6d6fe453cad8054b9cdec53f655 Mon Sep 17 00:00:00 2001 From: limityan Date: Wed, 15 Jul 2026 17:39:33 +0800 Subject: [PATCH] fix(review): require explicit follow-up delivery --- .../core/src/agentic/agents/registry/query.rs | 19 + .../core/src/agentic/agents/registry/tests.rs | 104 ++++ .../src/agentic/coordination/coordinator.rs | 559 +++++++++++++++++- .../tools/implementations/task/execution.rs | 5 + .../tools/implementations/task/input.rs | 36 +- .../task/launch_review_agent.rs | 8 +- .../agentic/tools/implementations/task/mod.rs | 12 +- .../tools/implementations/task/schema.rs | 10 + .../tools/implementations/task/tests.rs | 156 ++++- 9 files changed, 896 insertions(+), 13 deletions(-) diff --git a/src/crates/assembly/core/src/agentic/agents/registry/query.rs b/src/crates/assembly/core/src/agentic/agents/registry/query.rs index 2f9cf1ea8a..3b5ce45d33 100644 --- a/src/crates/assembly/core/src/agentic/agents/registry/query.rs +++ b/src/crates/assembly/core/src/agentic/agents/registry/query.rs @@ -160,6 +160,25 @@ impl AgentRegistry { None } + pub async fn get_subagent_is_review_for_workspace( + &self, + id: &str, + workspace_root: Option<&Path>, + ) -> Option { + self.ensure_user_custom_agents_loaded().await; + if let Some(workspace_root) = workspace_root { + let is_project_cache_loaded = + self.read_project_subagents().contains_key(workspace_root); + if !is_project_cache_loaded { + self.load_custom_agents(Some(workspace_root)).await; + } + } + + self.find_agent_entry(id, workspace_root) + .filter(|entry| entry.category == AgentCategory::SubAgent) + .map(|entry| is_review_agent_entry(&entry)) + } + fn entry_is_visible_for_query( entry: &AgentEntry, query: &SubagentQueryContext<'_>, diff --git a/src/crates/assembly/core/src/agentic/agents/registry/tests.rs b/src/crates/assembly/core/src/agentic/agents/registry/tests.rs index 1d6443f69e..6ca908a6e3 100644 --- a/src/crates/assembly/core/src/agentic/agents/registry/tests.rs +++ b/src/crates/assembly/core/src/agentic/agents/registry/tests.rs @@ -69,6 +69,30 @@ fn test_project_entry(id: &str, model: &str) -> AgentEntry { } } +fn test_project_custom_entry(id: &str, review: bool) -> AgentEntry { + let mut agent = CustomSubagent::new( + id.to_string(), + "Project custom subagent".to_string(), + vec!["Read".to_string()], + "prompt".to_string(), + review, + format!("{id}.md"), + CustomSubagentKind::Project, + ); + agent.data.review = review; + + AgentEntry { + category: AgentCategory::SubAgent, + source: AgentSource::Project, + subagent_source: Some(SubAgentSource::Project), + agent: Arc::new(agent), + visibility_policy: SubagentVisibilityPolicy::public(), + custom_config: Some(CustomSubagentConfig { + model: "fast".to_string(), + }), + } +} + fn insert_project_subagent(registry: &AgentRegistry, workspace: &Path, id: &str, model: &str) { let mut entries = HashMap::new(); entries.insert(id.to_string(), test_project_entry(id, model)); @@ -77,6 +101,67 @@ fn insert_project_subagent(registry: &AgentRegistry, workspace: &Path, id: &str, .insert(workspace.to_path_buf(), entries); } +#[tokio::test] +async fn review_lookup_is_scoped_to_the_requested_workspace() { + let registry = AgentRegistry::new(); + let review_workspace = PathBuf::from("review-workspace"); + let ordinary_workspace = PathBuf::from("ordinary-workspace"); + let agent_id = "SharedProjectAgent"; + + registry.write_project_subagents().insert( + review_workspace.clone(), + HashMap::from([( + agent_id.to_string(), + test_project_custom_entry(agent_id, true), + )]), + ); + registry.write_project_subagents().insert( + ordinary_workspace.clone(), + HashMap::from([( + agent_id.to_string(), + test_project_custom_entry(agent_id, false), + )]), + ); + + assert_eq!( + registry + .get_subagent_is_review_for_workspace(agent_id, Some(&review_workspace)) + .await, + Some(true) + ); + assert_eq!( + registry + .get_subagent_is_review_for_workspace(agent_id, Some(&ordinary_workspace)) + .await, + Some(false) + ); + assert_eq!( + registry + .get_subagent_is_review_for_workspace(agent_id, None) + .await, + None, + "a project agent must not leak into an unrelated workspace lookup" + ); +} + +#[tokio::test] +async fn review_lookup_cold_loads_the_requested_project_registry() { + let env = CustomAgentTestEnv::new("bitfun-project-review-lookup"); + let registry = AgentRegistry::new(); + let agent_id = "ProjectReviewer"; + write_project_custom_review_subagent( + &env.workspace_agents_dir.join("project-reviewer.md"), + agent_id, + ); + + assert_eq!( + registry + .get_subagent_is_review_for_workspace(agent_id, Some(&env.workspace_root)) + .await, + Some(true) + ); +} + #[test] fn top_level_modes_default_to_auto() { for agent_type in [ @@ -942,6 +1027,25 @@ fn write_project_custom_subagent(path: &Path, id: &str) { .expect("project subagent markdown should save"); } +fn write_project_custom_review_subagent(path: &Path, id: &str) { + let mut subagent = CustomSubagent::new_with_id( + id.to_string(), + id.to_string(), + "Project review subagent".to_string(), + vec!["Read".to_string()], + "Review the relevant files.".to_string(), + true, + path.to_string_lossy().to_string(), + CustomSubagentKind::Project, + "fast".to_string(), + UserContextPolicy::empty().with_workspace_instructions(), + ); + subagent.data.review = true; + subagent + .save_to_file(None) + .expect("project review subagent markdown should save"); +} + fn unique_suffix() -> String { use std::time::{SystemTime, UNIX_EPOCH}; diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 8375c57184..961ca7bd25 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -94,6 +94,25 @@ fn is_review_agent_type(agent_type: &str) -> bool { ) } +pub(crate) async fn validate_background_subagent_delivery( + agent_type: &str, + workspace_root: Option<&Path>, + allow_review_follow_up: bool, +) -> BitFunResult<()> { + let is_review = get_agent_registry() + .get_subagent_is_review_for_workspace(agent_type, workspace_root) + .await + .unwrap_or(false); + if !is_review || allow_review_follow_up { + return Ok(()); + } + + Err(BitFunError::Validation( + "Reviews wait for results by default so one final review can be returned. Retry without run_in_background. Only when the user explicitly asked not to wait, retry with run_in_background=true and allow_review_follow_up=true." + .to_string(), + )) +} + fn turn_review_manifest_for_agent( metadata: Option<&serde_json::Value>, agent_type: &str, @@ -5964,6 +5983,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet &self, target_session_id: &str, parent_session_id: &str, + background_allow_review_follow_up: Option, ) -> BitFunResult { let session = match self.session_manager.get_session(target_session_id) { Some(session) => session, @@ -5978,6 +5998,83 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet parent_session_id )) })?; + if let Some(allow_review_follow_up) = background_allow_review_follow_up { + let metadata = self + .session_manager + .load_session_metadata(&binding.session_storage_dir(), target_session_id) + .await? + .ok_or_else(|| { + BitFunError::NotFound(format!( + "Session metadata not found: {}", + target_session_id + )) + })?; + if !metadata.is_subagent() { + return Err(BitFunError::Validation(format!( + "Subagent execution target must be a subagent session: {}", + target_session_id + ))); + } + let created_by_marker = format!("session-{parent_session_id}"); + let owned_subagent = session_lineage_matches_parent( + metadata.relationship.as_ref(), + parent_session_id, + ) || metadata.created_by.as_deref() + == Some(created_by_marker.as_str()); + if !owned_subagent { + return Err(BitFunError::Validation(format!( + "Subagent session '{}' was not created by parent session '{}'", + target_session_id, parent_session_id + ))); + } + validate_background_subagent_delivery( + &metadata.agent_type, + metadata + .workspace_path + .as_deref() + .map(Path::new) + .or(Some(binding.root_path())), + allow_review_follow_up, + ) + .await?; + + let (session_view, _, _, _) = self + .session_manager + .restore_internal_session_view_from_storage_path_tail_timed( + &binding.session_storage_dir(), + target_session_id, + 0, + ) + .await?; + if session_view.kind != SessionKind::Subagent { + return Err(BitFunError::Validation(format!( + "Subagent execution target must be a subagent session: {}", + target_session_id + ))); + } + if !session_created_by_parent(&session_view, parent_session_id) + && !session_lineage_matches_parent( + metadata.relationship.as_ref(), + parent_session_id, + ) + { + return Err(BitFunError::Validation(format!( + "Subagent session '{}' was not created by parent session '{}'", + target_session_id, parent_session_id + ))); + } + validate_background_subagent_delivery( + &session_view.agent_type, + session_view + .config + .workspace_path + .as_deref() + .map(Path::new) + .or(Some(binding.root_path())), + allow_review_follow_up, + ) + .await?; + } self.session_manager .restore_internal_session_from_storage_path( &binding.session_storage_dir(), @@ -6011,6 +6108,15 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ))); } + if let Some(allow_review_follow_up) = background_allow_review_follow_up { + validate_background_subagent_delivery( + &session.agent_type, + session.config.workspace_path.as_deref().map(Path::new), + allow_review_follow_up, + ) + .await?; + } + Ok(session) } @@ -6089,6 +6195,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet async fn resolve_hidden_subagent_execution_request( &self, request: SubagentExecutionRequest, + background_allow_review_follow_up: Option, ) -> BitFunResult { let task_description = request.task_description.trim().to_string(); if task_description.is_empty() { @@ -6129,6 +6236,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .ensure_subagent_session_loaded_for_reuse( target_session_id, &parent_session_id, + background_allow_review_follow_up, ) .await?; if let Some(model_id) = model_id.as_deref() { @@ -6184,6 +6292,14 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .to_string(), ) })?; + if let Some(allow_review_follow_up) = background_allow_review_follow_up { + validate_background_subagent_delivery( + &agent_type, + Some(Path::new(&workspace_path)), + allow_review_follow_up, + ) + .await?; + } Ok(HiddenSubagentExecutionRequest { target_session_id: None, @@ -6229,6 +6345,14 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet let snapshot = self .capture_fork_agent_context_snapshot(&request.subagent_parent_info.session_id) .await?; + if let Some(allow_review_follow_up) = background_allow_review_follow_up { + validate_background_subagent_delivery( + &snapshot.parent_agent_type, + Some(Path::new(&snapshot.workspace_path)), + allow_review_follow_up, + ) + .await?; + } let mut session_config = snapshot.build_child_session_config(None); if let Some(model_id) = model_id { session_config.model_id = Some(model_id); @@ -6358,7 +6482,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet request: SubagentExecutionRequest, ) -> BitFunResult { let request = self - .resolve_hidden_subagent_execution_request(request) + .resolve_hidden_subagent_execution_request(request, None) .await?; self.prepare_hidden_subagent_execution_request(request) .await @@ -6423,7 +6547,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet parent_session_id: &str, subagent_session_id: &str, ) -> BitFunResult { - self.ensure_subagent_session_loaded_for_reuse(subagent_session_id, parent_session_id) + self.ensure_subagent_session_loaded_for_reuse(subagent_session_id, parent_session_id, None) .await?; let controls: Vec<(String, BackgroundSubagentTaskControl)> = self @@ -6557,9 +6681,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet &self, request: SubagentExecutionRequest, timeout_seconds: Option, + allow_review_follow_up: bool, ) -> BitFunResult { let request = self - .resolve_hidden_subagent_execution_request(request) + .resolve_hidden_subagent_execution_request(request, Some(allow_review_follow_up)) .await?; let request = self .prepare_hidden_subagent_execution_request(request) @@ -7663,11 +7788,12 @@ mod tests { use super::{ merge_prepended_messages_for_turn, normalize_subagent_max_concurrency, resolve_agent_session_create_created_by, resolve_agent_submission_turn_id, - should_require_tool_confirmation, turn_review_manifest_for_agent, ConversationCoordinator, - SubagentExecutionRequest, + should_require_tool_confirmation, turn_review_manifest_for_agent, + validate_background_subagent_delivery, ConversationCoordinator, SubagentExecutionRequest, }; + use crate::agentic::agents::{CustomSubagent, CustomSubagentKind, UserContextPolicy}; use crate::agentic::core::{ - InternalReminderKind, Message, MessageContent, MessageRole, MessageSemanticKind, + InternalReminderKind, Message, MessageContent, MessageRole, MessageSemanticKind, Session, SessionConfig, SessionKind, }; use crate::agentic::events::{EventQueue, EventQueueConfig, EventRouter}; @@ -7687,6 +7813,7 @@ mod tests { use crate::agentic::TurnSkillAgentSnapshot; use crate::infrastructure::PathManager; use crate::service::remote_ssh::workspace_state::init_remote_workspace_manager; + use crate::service::session::SessionMetadata; use bitfun_runtime_ports::{ AgentSessionCreateRequest, AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionSource, DelegationPolicy, DialogQueuePriority, DialogSubmissionPolicy, @@ -7785,6 +7912,426 @@ mod tests { assert_eq!(normalize_subagent_max_concurrency(usize::MAX), 64); } + #[tokio::test] + async fn review_background_delivery_requires_explicit_follow_up_permission() { + let error = validate_background_subagent_delivery("CodeReview", None, false) + .await + .expect_err("review background delivery should require explicit follow-up permission"); + + assert!(error.to_string().contains("one final review")); + assert!(error.to_string().contains("allow_review_follow_up=true")); + assert!( + validate_background_subagent_delivery("CodeReview", None, true) + .await + .is_ok() + ); + } + + #[tokio::test] + async fn explicit_review_follow_up_allows_fresh_background_resolution() { + let (coordinator, _) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-review-follow-up-fresh-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + let request = SubagentExecutionRequest { + task_description: "Review the current diff later".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: None, + subagent_type: Some("CodeReview".to_string()), + workspace_path: Some(workspace_path.to_string_lossy().into_owned()), + model_id: None, + subagent_parent_info: SubagentParentInfo { + session_id: "parent-session".to_string(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + }; + + let resolved = coordinator + .resolve_hidden_subagent_execution_request(request, Some(true)) + .await + .expect("explicit review follow-up should resolve a fresh background request"); + + assert_eq!(resolved.agent_type, "CodeReview"); + assert!(resolved.target_session_id.is_none()); + let _ = std::fs::remove_dir_all(workspace_path); + } + + #[tokio::test] + async fn explicit_review_follow_up_allows_reused_background_resolution() { + let (coordinator, _) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-review-follow-up-reuse-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + let review_session = coordinator + .create_hidden_agent_session( + None, + "Reusable review".to_string(), + "CodeReview".to_string(), + SessionConfig { + workspace_path: Some(workspace_path.to_string_lossy().into_owned()), + ..Default::default() + }, + Some("session-parent-session".to_string()), + SessionKind::Subagent, + ) + .await + .expect("review session should be created"); + let request = SubagentExecutionRequest { + task_description: "Continue the review later".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: Some(review_session.session_id.clone()), + subagent_type: None, + workspace_path: None, + model_id: None, + subagent_parent_info: SubagentParentInfo { + session_id: "parent-session".to_string(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + }; + + let resolved = coordinator + .resolve_hidden_subagent_execution_request(request, Some(true)) + .await + .expect("explicit review follow-up should resolve a reused background request"); + + assert_eq!( + resolved.target_session_id.as_deref(), + Some(review_session.session_id.as_str()) + ); + let _ = std::fs::remove_dir_all(workspace_path); + } + + #[tokio::test] + async fn non_review_background_delivery_preserves_existing_behavior() { + assert!( + validate_background_subagent_delivery("GeneralPurpose", None, false) + .await + .is_ok() + ); + assert!( + validate_background_subagent_delivery("ReviewFixer", None, false) + .await + .is_ok() + ); + } + + #[tokio::test] + async fn continued_review_session_cannot_run_in_background_implicitly() { + let (coordinator, session_manager) = test_coordinator(); + let review_session = coordinator + .create_hidden_agent_session( + None, + "Reusable review".to_string(), + "CodeReview".to_string(), + SessionConfig { + model_id: Some("primary".to_string()), + workspace_path: Some(std::env::temp_dir().to_string_lossy().into_owned()), + ..Default::default() + }, + Some("session-parent-session".to_string()), + SessionKind::Subagent, + ) + .await + .expect("review session should be created"); + let request = SubagentExecutionRequest { + task_description: "Continue reviewing the current diff".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: Some(review_session.session_id.clone()), + subagent_type: None, + workspace_path: None, + model_id: Some("fast".to_string()), + subagent_parent_info: SubagentParentInfo { + session_id: "parent-session".to_string(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + }; + + let error = coordinator + .start_background_subagent(request, None, false) + .await + .expect_err( + "continued review should not run in the background without explicit intent", + ); + + assert!(error.to_string().contains("one final review")); + assert!(error.to_string().contains("allow_review_follow_up=true")); + assert_eq!( + session_manager + .get_session(&review_session.session_id) + .expect("review session should remain available") + .config + .model_id + .as_deref(), + Some("primary"), + "rejected delivery must not mutate the review session model" + ); + } + + #[tokio::test] + async fn cold_review_session_is_rejected_before_full_restore() { + let (coordinator, session_manager) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-cold-review-background-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + let parent_session = session_manager + .create_session( + "Parent".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace_path.to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .expect("parent session should be created"); + let storage_binding = session_manager + .resolve_session_workspace_binding(&parent_session.session_id) + .await + .expect("parent workspace binding should resolve"); + let review_session_id = format!("session-cold-review-{}", uuid::Uuid::new_v4()); + let mut metadata = SessionMetadata::new( + review_session_id.clone(), + "Cold review".to_string(), + "CodeReview".to_string(), + "primary".to_string(), + ); + metadata.session_kind = SessionKind::Subagent; + metadata.created_by = Some(format!("session-{}", parent_session.session_id)); + metadata.workspace_path = Some(workspace_path.to_string_lossy().into_owned()); + metadata.relationship = Some(super::build_subagent_session_relationship( + Some(&SubagentParentInfo { + session_id: parent_session.session_id.clone(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }), + "CodeReview", + )); + session_manager + .save_session_metadata(&storage_binding.session_storage_dir(), &metadata) + .await + .expect("cold review metadata should be persisted"); + assert!(session_manager.get_session(&review_session_id).is_none()); + + let request = SubagentExecutionRequest { + task_description: "Continue the cold review".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: Some(review_session_id.clone()), + subagent_type: None, + workspace_path: None, + model_id: Some("fast".to_string()), + subagent_parent_info: SubagentParentInfo { + session_id: parent_session.session_id.clone(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + }; + + let error = coordinator + .resolve_hidden_subagent_execution_request(request, Some(false)) + .await + .expect_err("cold review should reject background delivery before restore"); + + assert!(error.to_string().contains("one final review")); + assert!(session_manager.get_session(&review_session_id).is_none()); + let persisted = session_manager + .load_session_metadata(&storage_binding.session_storage_dir(), &review_session_id) + .await + .expect("cold review metadata should remain readable") + .expect("cold review metadata should remain present"); + assert_eq!(persisted.model_name, "primary"); + let _ = std::fs::remove_dir_all(workspace_path); + } + + #[tokio::test] + async fn cold_subagent_owned_by_another_parent_is_rejected_before_restore() { + let (coordinator, session_manager) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-cold-foreign-subagent-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + let parent_session = session_manager + .create_session( + "Parent".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace_path.to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .expect("parent session should be created"); + let storage_binding = session_manager + .resolve_session_workspace_binding(&parent_session.session_id) + .await + .expect("parent workspace binding should resolve"); + let review_session_id = format!("session-foreign-review-{}", uuid::Uuid::new_v4()); + let mut metadata = SessionMetadata::new( + review_session_id.clone(), + "Foreign review".to_string(), + "CodeReview".to_string(), + "primary".to_string(), + ); + metadata.session_kind = SessionKind::Subagent; + metadata.created_by = Some("session-another-parent".to_string()); + metadata.workspace_path = Some(workspace_path.to_string_lossy().into_owned()); + session_manager + .save_session_metadata(&storage_binding.session_storage_dir(), &metadata) + .await + .expect("foreign review metadata should be persisted"); + + let request = SubagentExecutionRequest { + task_description: "Continue a foreign review".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: Some(review_session_id.clone()), + subagent_type: None, + workspace_path: None, + model_id: None, + subagent_parent_info: SubagentParentInfo { + session_id: parent_session.session_id.clone(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + }; + + let error = coordinator + .resolve_hidden_subagent_execution_request(request, Some(false)) + .await + .expect_err("foreign subagent should fail ownership preflight"); + + assert!(error + .to_string() + .contains("was not created by parent session")); + assert!(session_manager.get_session(&review_session_id).is_none()); + let _ = std::fs::remove_dir_all(workspace_path); + } + + #[tokio::test] + async fn cold_review_uses_read_only_session_workspace_before_restore() { + let (coordinator, session_manager) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-cold-review-workspace-preflight-test-{}", + uuid::Uuid::new_v4() + )); + let review_workspace = workspace_path.join("review-workspace"); + let reviewer_path = review_workspace + .join(".bitfun") + .join("agents") + .join("cold-reviewer.md"); + std::fs::create_dir_all( + reviewer_path + .parent() + .expect("reviewer path should have a parent"), + ) + .expect("review agent directory should exist"); + let mut reviewer = CustomSubagent::new_with_id( + "ColdWorkspaceReviewer".to_string(), + "ColdWorkspaceReviewer".to_string(), + "Project review subagent".to_string(), + vec!["Read".to_string()], + "Review the relevant files.".to_string(), + true, + reviewer_path.to_string_lossy().into_owned(), + CustomSubagentKind::Project, + "fast".to_string(), + UserContextPolicy::empty().with_workspace_instructions(), + ); + reviewer.data.review = true; + reviewer + .save_to_file(None) + .expect("project review subagent should save"); + + let parent_session = session_manager + .create_session( + "Parent".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace_path.to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .expect("parent session should be created"); + let storage_binding = session_manager + .resolve_session_workspace_binding(&parent_session.session_id) + .await + .expect("parent workspace binding should resolve"); + let review_session_id = format!("session-workspace-review-{}", uuid::Uuid::new_v4()); + let mut persisted_session = Session::new_with_id( + review_session_id.clone(), + "Workspace review".to_string(), + "ColdWorkspaceReviewer".to_string(), + SessionConfig { + workspace_path: Some(review_workspace.to_string_lossy().into_owned()), + ..Default::default() + }, + ); + persisted_session.kind = SessionKind::Subagent; + persisted_session.created_by = Some(format!("session-{}", parent_session.session_id)); + let persistence_manager = PersistenceManager::new(Arc::new( + PathManager::new().expect("path manager should initialize"), + )) + .expect("persistence manager should initialize"); + persistence_manager + .save_session(&storage_binding.session_storage_dir(), &persisted_session) + .await + .expect("cold review session should be persisted"); + let mut metadata = session_manager + .load_session_metadata(&storage_binding.session_storage_dir(), &review_session_id) + .await + .expect("cold review metadata should load") + .expect("cold review metadata should exist"); + metadata.workspace_path = Some(workspace_path.to_string_lossy().into_owned()); + session_manager + .save_session_metadata(&storage_binding.session_storage_dir(), &metadata) + .await + .expect("metadata workspace should be overridden for the mismatch fixture"); + + let request = SubagentExecutionRequest { + task_description: "Continue the workspace review".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: Some(review_session_id.clone()), + subagent_type: None, + workspace_path: None, + model_id: None, + subagent_parent_info: SubagentParentInfo { + session_id: parent_session.session_id.clone(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + }; + + let error = coordinator + .resolve_hidden_subagent_execution_request(request, Some(false)) + .await + .expect_err("stored review workspace should be checked before full restore"); + + assert!(error.to_string().contains("one final review")); + assert!(session_manager.get_session(&review_session_id).is_none()); + let _ = std::fs::remove_dir_all(workspace_path); + } + #[test] fn subagent_timeout_disable_clears_active_deadline() { use super::SubagentTimeoutAction; diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs index 6654921bd6..6a4114cdb0 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs @@ -57,6 +57,7 @@ struct BackgroundTaskStartRequest<'a> { subagent_context: Option>, prepared_prompt: String, timeout_seconds: Option, + allow_review_follow_up: bool, tool_call_id: String, session_id: String, dialog_turn_id: String, @@ -175,6 +176,7 @@ impl TaskTool { let model_id = invocation.model_id.clone(); let mut timeout_seconds = invocation.timeout_seconds; let run_in_background = invocation.run_in_background; + let allow_review_follow_up = invocation.allow_review_follow_up; let is_retry = invocation.is_retry; let requested_auto_retry = invocation.requested_auto_retry; let is_auto_retry = is_retry && requested_auto_retry; @@ -549,6 +551,7 @@ impl TaskTool { subagent_context, prepared_prompt, timeout_seconds, + allow_review_follow_up, tool_call_id, session_id, dialog_turn_id, @@ -598,6 +601,7 @@ impl TaskTool { subagent_context, prepared_prompt, timeout_seconds, + allow_review_follow_up, tool_call_id, session_id, dialog_turn_id, @@ -621,6 +625,7 @@ impl TaskTool { delegation_policy: context.delegation_policy().spawn_child(), }, timeout_seconds, + allow_review_follow_up, ) .await?; diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/input.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/input.rs index 17fdeff58c..10406173a1 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/input.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/input.rs @@ -74,6 +74,7 @@ pub(super) struct TaskInvocation { pub(super) model_id: Option, pub(super) timeout_seconds: Option, pub(super) run_in_background: bool, + pub(super) allow_review_follow_up: bool, pub(super) is_retry: bool, pub(super) requested_auto_retry: bool, } @@ -96,7 +97,12 @@ impl TaskTool { "action is not supported for DeepReview Task calls".to_string(), )); } - for field in ["fork_context", "session_id", "run_in_background"] { + for field in [ + "fork_context", + "session_id", + "run_in_background", + "allow_review_follow_up", + ] { if input.get(field).is_some() { return Err(BitFunError::tool(format!( "{field} is not allowed for DeepReview Task calls" @@ -114,6 +120,7 @@ impl TaskTool { model_id: Self::optional_trimmed_string(input, "model_id")?, timeout_seconds: Self::optional_timeout_seconds(input)?, run_in_background: false, + allow_review_follow_up: false, is_retry: input.get("retry").and_then(Value::as_bool).unwrap_or(false), requested_auto_retry: input .get("auto_retry") @@ -134,6 +141,13 @@ impl TaskTool { )); } let run_in_background = Self::optional_bool(input, "run_in_background")?.unwrap_or(false); + let allow_review_follow_up = + Self::optional_bool(input, "allow_review_follow_up")?.unwrap_or(false); + if action != TaskAction::Cancel && allow_review_follow_up && !run_in_background { + return Err(BitFunError::tool( + "allow_review_follow_up=true requires run_in_background=true".to_string(), + )); + } match action { TaskAction::Spawn => { @@ -179,6 +193,7 @@ impl TaskTool { model_id: Self::optional_trimmed_string(input, "model_id")?, timeout_seconds: None, run_in_background, + allow_review_follow_up, is_retry: false, requested_auto_retry: false, }) @@ -210,6 +225,7 @@ impl TaskTool { model_id: Self::optional_trimmed_string(input, "model_id")?, timeout_seconds: None, run_in_background, + allow_review_follow_up, is_retry: false, requested_auto_retry: false, }) @@ -225,6 +241,7 @@ impl TaskTool { "subagent_type", "model_id", "run_in_background", + "allow_review_follow_up", "retry", "auto_retry", "retry_coverage", @@ -242,6 +259,7 @@ impl TaskTool { model_id: None, timeout_seconds: None, run_in_background: false, + allow_review_follow_up: false, is_retry: false, requested_auto_retry: false, }) @@ -249,14 +267,28 @@ impl TaskTool { } } - pub(super) fn validate_invocation_input( + pub(super) async fn validate_invocation_input( input: &Value, is_deep_review_parent: bool, + workspace_root: Option<&std::path::Path>, ) -> ValidationResult { let invocation = match Self::parse_invocation(input, is_deep_review_parent) { Ok(invocation) => invocation, Err(error) => return Self::invalid_input(error.to_string()), }; + if invocation.action == TaskAction::Spawn && invocation.run_in_background { + if let Some(subagent_type) = invocation.subagent_type.as_deref() { + if let Err(error) = validate_background_subagent_delivery( + subagent_type, + workspace_root, + invocation.allow_review_follow_up, + ) + .await + { + return Self::invalid_input(error.to_string()); + } + } + } if invocation.action != TaskAction::Cancel { if let Some(result) = Self::validate_prompt_size(input) { return result; diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/launch_review_agent.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/launch_review_agent.rs index 1cce36d85d..38ba86d811 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/launch_review_agent.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/launch_review_agent.rs @@ -108,7 +108,13 @@ impl LaunchReviewAgentTool { } fn parse_invocation(input: &Value) -> BitFunResult { - for field in ["action", "fork_context", "session_id", "run_in_background"] { + for field in [ + "action", + "fork_context", + "session_id", + "run_in_background", + "allow_review_follow_up", + ] { if input.get(field).is_some() { return Err(BitFunError::tool(format!( "{field} is not supported for LaunchReviewAgent" diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs index 80dfa06867..e0c0c849a7 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs @@ -1,7 +1,9 @@ use crate::agentic::agents::{ get_agent_registry, AgentInfo, SubagentListScope, SubagentQueryContext, }; -use crate::agentic::coordination::{get_global_coordinator, SubagentExecutionRequest}; +use crate::agentic::coordination::{ + get_global_coordinator, validate_background_subagent_delivery, SubagentExecutionRequest, +}; use crate::agentic::deep_review::task_adapter::{ self as deep_review_task_adapter, DeepReviewLaunchBatchInfo, DeepReviewProviderQueueWaitOutcome, DeepReviewQueueWaitOutcome, DeepReviewQueueWaitSkipReason, @@ -187,8 +189,12 @@ impl Tool for TaskTool { input: &Value, context: Option<&ToolUseContext>, ) -> ValidationResult { - let _ = context; - Self::validate_invocation_input(input, false) + Self::validate_invocation_input( + input, + false, + context.and_then(ToolUseContext::workspace_root), + ) + .await } fn render_tool_use_message(&self, input: &Value, options: &ToolRenderOptions) -> String { diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/schema.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/schema.rs index ca65c316df..4b29986ab4 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/schema.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/schema.rs @@ -70,6 +70,13 @@ impl TaskTool { "description": "Optional for action='spawn' and action='send_input'. Defaults to false." }), ); + properties.insert( + "allow_review_follow_up".to_string(), + json!({ + "type": "boolean", + "description": "Optional for action='spawn' and action='send_input'. Use with run_in_background=true only when the user explicitly asked not to wait for a review result. This permits delivery in a later follow-up and does not change normal cancellation behavior." + }), + ); json!({ "type": "object", "properties": properties, @@ -111,6 +118,8 @@ The two modes are mutually exclusive: do not provide `subagent_type` when `fork_ `run_in_background` usage: - false: Wait for the agent to finish and return its result to you. - true: Run the agent in the background without blocking you. When the subagent finishes, its result will be delivered to you in a follow-up message. You can process remaining work before receiving the result. +- Review subagents are completion dependencies by default. Launch multiple review Task calls in one assistant message to run them concurrently, and wait for their results so you can merge one final review. +- If the user explicitly asks not to wait for a review result, set both `run_in_background=true` and `allow_review_follow_up=true`. This permits the result to arrive in a later follow-up; it does not change normal cancellation behavior. Never use it merely to improve parallelism. Usage notes: - Include a short description of what the agent will do for this round (for `spawn` and `send_input`). @@ -124,6 +133,7 @@ Usage notes: Examples (assume "example-reviewer" is present in the agent listing): - Start a new specialized subagent: `{ "action": "spawn", "description": "Inspect parser flow", "subagent_type": "example-reviewer", "prompt": "Inspect the parser flow in src/parser.rs and report risks, key functions, and any missing tests." }` +- Allow a review follow-up only when the user asked not to wait: `{ "action": "spawn", "description": "Review parser later", "subagent_type": "example-reviewer", "prompt": "Review the parser and report findings when finished.", "run_in_background": true, "allow_review_follow_up": true }` - Start by forking the current context: `{ "action": "spawn", "description": "Check migration impact", "fork_context": true, "prompt": "Using the current context, check whether the migration affects config loading. Stay read-only and report the answer with file references." }` - Continue an existing subagent with a specific model: `{ "action": "send_input", "description": "Continue parser review", "session_id": "subagent-session-123", "model_id": "fast", "prompt": "Continue from your prior parser review and focus on the error recovery paths." }` - Cancel a background subagent: `{ "action": "cancel", "session_id": "subagent-session-123" }` diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs index 29e74c5f62..ef7b82a651 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs @@ -127,6 +127,154 @@ fn task_schema_accepts_optional_model_id() { .any(|value| value.as_str() == Some("model_id"))); } +#[test] +fn task_schema_exposes_explicit_review_follow_up_control() { + let schema = TaskTool::new().input_schema(); + let follow_up = &schema["properties"]["allow_review_follow_up"]; + + assert_eq!(follow_up["type"], "boolean"); + let description = follow_up["description"] + .as_str() + .expect("allow_review_follow_up description should be a string"); + assert!(description.contains("explicitly")); + assert!(description.contains("review")); + assert!(description.contains("run_in_background=true")); + assert!(schema["properties"].get("detach_from_parent").is_none()); +} + +#[tokio::test] +async fn validate_input_rejects_review_background_without_explicit_follow_up_permission() { + let validation = TaskTool::new() + .validate_input( + &json!({ + "action": "spawn", + "description": "Review changes", + "prompt": "Review the current diff", + "subagent_type": "CodeReview", + "run_in_background": true + }), + None, + ) + .await; + + assert!(!validation.result); + let message = validation + .message + .as_deref() + .expect("validation should explain how review delivery works"); + assert!(message.contains("one final review")); + assert!(message.contains("allow_review_follow_up=true")); +} + +#[tokio::test] +async fn validate_input_accepts_explicit_review_follow_up() { + let validation = TaskTool::new() + .validate_input( + &json!({ + "action": "spawn", + "description": "Review later", + "prompt": "Review the current diff", + "subagent_type": "CodeReview", + "run_in_background": true, + "allow_review_follow_up": true + }), + None, + ) + .await; + + assert!(validation.result, "{:?}", validation.message); +} + +#[tokio::test] +async fn validate_input_rejects_review_follow_up_without_background_execution() { + let validation = TaskTool::new() + .validate_input( + &json!({ + "action": "spawn", + "description": "Review changes", + "prompt": "Review the current diff", + "subagent_type": "CodeReview", + "allow_review_follow_up": true + }), + None, + ) + .await; + + assert!(!validation.result); + assert!(validation.message.as_deref().is_some_and(|message| { + message.contains("allow_review_follow_up=true requires run_in_background=true") + })); +} + +#[test] +fn parse_input_preserves_review_follow_up_for_send_input() { + let invocation = TaskTool::parse_invocation( + &json!({ + "action": "send_input", + "description": "Continue review later", + "prompt": "Continue the review and report when finished", + "session_id": "review-session-1", + "run_in_background": true, + "allow_review_follow_up": true + }), + false, + ) + .expect("review follow-up send_input should parse"); + + assert!(invocation.run_in_background); + assert!(invocation.allow_review_follow_up); +} + +#[tokio::test] +async fn validate_input_preserves_non_review_background_tasks() { + let validation = TaskTool::new() + .validate_input( + &json!({ + "action": "spawn", + "description": "Investigate logs", + "prompt": "Inspect the logs and report later", + "subagent_type": "GeneralPurpose", + "run_in_background": true + }), + None, + ) + .await; + + assert!(validation.result, "{:?}", validation.message); +} + +#[tokio::test] +async fn validate_input_rejects_review_follow_up_for_cancel() { + let validation = TaskTool::new() + .validate_input( + &json!({ + "action": "cancel", + "session_id": "subagent-session-1", + "allow_review_follow_up": true + }), + None, + ) + .await; + + assert!(!validation.result); + assert!(validation + .message + .as_deref() + .is_some_and(|message| message.contains("allow_review_follow_up is not allowed"))); +} + +#[test] +fn joined_review_tasks_remain_concurrency_safe() { + let input = json!({ + "action": "spawn", + "description": "Review changes", + "prompt": "Review the current diff", + "subagent_type": "CodeReview" + }); + + assert!(TaskTool::new().is_concurrency_safe(Some(&input))); +} + #[test] fn task_schema_describes_spawn_context_modes_as_exclusive() { let description = TaskTool::new().render_description(); @@ -457,7 +605,13 @@ async fn validate_input_rejects_timeout_for_regular_parent() { #[tokio::test] async fn launch_review_agent_rejects_task_context_controls() { let context = test_tool_context("DeepReview"); - for field in ["action", "fork_context", "session_id", "run_in_background"] { + for field in [ + "action", + "fork_context", + "session_id", + "run_in_background", + "allow_review_follow_up", + ] { let mut input = json!({ "description": "delegate", "prompt": "Review security-sensitive files",