diff --git a/src/apps/desktop/src/computer_use/windows_ax_ui.rs b/src/apps/desktop/src/computer_use/windows_ax_ui.rs index c102ccd584..bbcb411f84 100644 --- a/src/apps/desktop/src/computer_use/windows_ax_ui.rs +++ b/src/apps/desktop/src/computer_use/windows_ax_ui.rs @@ -305,9 +305,9 @@ fn read_cached_is_offscreen(element: &IUIAutomationElement) -> bool { /// Read bounding rect as `(center_x, center_y, Some((l, t, r, b)))`. Returns /// `rect=None` when the element has no meaningful `BoundingRectangle`. -fn read_cached_bounding_rect_full( - element: &IUIAutomationElement, -) -> (i32, i32, Option<(i32, i32, i32, i32)>) { +type CachedBoundingRect = (i32, i32, Option<(i32, i32, i32, i32)>); + +fn read_cached_bounding_rect_full(element: &IUIAutomationElement) -> CachedBoundingRect { unsafe { match element.CachedBoundingRectangle() { Ok(r) if r.right > r.left && r.bottom > r.top => ( diff --git a/src/apps/desktop/src/computer_use/windows_bg_input.rs b/src/apps/desktop/src/computer_use/windows_bg_input.rs index 11a6322680..9c07c860c5 100644 --- a/src/apps/desktop/src/computer_use/windows_bg_input.rs +++ b/src/apps/desktop/src/computer_use/windows_bg_input.rs @@ -132,7 +132,7 @@ fn il_name(rid: u32) -> &'static str { #[repr(C)] #[derive(Clone, Copy)] -struct KEYBDINPUT { +struct KeybdInput { wVk: u16, wScan: u16, dwFlags: u32, @@ -142,7 +142,7 @@ struct KEYBDINPUT { #[repr(C)] #[derive(Clone, Copy)] -struct MOUSEINPUT { +struct MouseInput { dx: i32, dy: i32, mouseData: u32, @@ -153,7 +153,7 @@ struct MOUSEINPUT { #[repr(C)] #[derive(Clone, Copy)] -struct HARDWAREINPUT { +struct HardwareInput { uMsg: u32, wParamL: u16, wParamH: u16, @@ -164,14 +164,14 @@ struct HARDWAREINPUT { #[repr(C)] #[derive(Clone, Copy)] union INPUT_0 { - ki: KEYBDINPUT, - mi: MOUSEINPUT, - hi: HARDWAREINPUT, + ki: KeybdInput, + mi: MouseInput, + hi: HardwareInput, } #[repr(C)] #[derive(Clone, Copy)] -struct INPUT { +struct Input { r#type: u32, Anonymous: INPUT_0, } @@ -190,7 +190,7 @@ struct TOKEN_MANDATORY_LABEL { #[link(name = "user32")] extern "system" { - fn SendInput(c_inputs: u32, p_inputs: *const INPUT, cb_size: i32) -> u32; + fn SendInput(c_inputs: u32, p_inputs: *const Input, cb_size: i32) -> u32; fn AttachThreadInput(id_attach: u32, id_attach_to: u32, f_attach: i32) -> i32; fn MapVirtualKeyW(code: u32, map_type: u32) -> u32; /// `VkKeyScanW` — translate a Unicode char to a virtual-key code + shift @@ -846,15 +846,15 @@ fn make_key_lparam(scan: u32, down: bool) -> LPARAM { /// One `SendInput` keyboard event carrying a Unicode code unit (`KEYEVENTF_ /// UNICODE`). `up` adds `KEYEVENTF_KEYUP`. -fn unicode_event(unit: u16, up: bool) -> INPUT { +fn unicode_event(unit: u16, up: bool) -> Input { let mut flags = KEYEVENTF_UNICODE; if up { flags |= KEYEVENTF_KEYUP; } - INPUT { + Input { r#type: INPUT_KEYBOARD, Anonymous: INPUT_0 { - ki: KEYBDINPUT { + ki: KeybdInput { wVk: 0, wScan: unit, dwFlags: flags, @@ -867,12 +867,12 @@ fn unicode_event(unit: u16, up: bool) -> INPUT { /// One `SendInput` keyboard event for a virtual-key code. `up` adds /// `KEYEVENTF_KEYUP`. -fn vk_event(vk: u16, scan: u32, up: bool) -> INPUT { +fn vk_event(vk: u16, scan: u32, up: bool) -> Input { let flags = if up { KEYEVENTF_KEYUP } else { 0 }; - INPUT { + Input { r#type: INPUT_KEYBOARD, Anonymous: INPUT_0 { - ki: KEYBDINPUT { + ki: KeybdInput { wVk: vk, wScan: scan as u16, dwFlags: flags, @@ -889,7 +889,7 @@ fn vk_event(vk: u16, scan: u32, up: bool) -> INPUT { /// `SendInput` reads `ev.len()` `INPUT` records from `ev.as_ptr()`; every /// record is fully initialized above. `cbSize` is the true `size_of::`. unsafe fn send_unicode(text: &str) -> BitFunResult<()> { - let mut ev: Vec = Vec::with_capacity(text.len() * 2); + let mut ev: Vec = Vec::with_capacity(text.len() * 2); for u in text.encode_utf16() { ev.push(unicode_event(u, false)); ev.push(unicode_event(u, true)); @@ -900,7 +900,7 @@ unsafe fn send_unicode(text: &str) -> BitFunResult<()> { let sent = SendInput( ev.len() as u32, ev.as_ptr(), - std::mem::size_of::() as i32, + std::mem::size_of::() as i32, ); if sent as usize != ev.len() { return Err(BitFunError::service(format!( @@ -917,7 +917,7 @@ unsafe fn send_unicode(text: &str) -> BitFunResult<()> { /// # Safety /// `SendInput` reads a fully-initialized `INPUT` array; `cbSize` is correct. unsafe fn send_key_combo(keycode: u16, modifiers: &[u16]) -> BitFunResult<()> { - let mut ev: Vec = Vec::with_capacity(modifiers.len() * 2 + 2); + let mut ev: Vec = Vec::with_capacity(modifiers.len() * 2 + 2); for &m in modifiers { let m_scan = MapVirtualKeyW(m as u32, MAPVK_VK_TO_VSC); ev.push(vk_event(m, m_scan, false)); @@ -935,7 +935,7 @@ unsafe fn send_key_combo(keycode: u16, modifiers: &[u16]) -> BitFunResult<()> { let sent = SendInput( ev.len() as u32, ev.as_ptr(), - std::mem::size_of::() as i32, + std::mem::size_of::() as i32, ); if sent as usize != ev.len() { return Err(BitFunError::service(format!( @@ -1268,3 +1268,55 @@ pub(super) fn parse_key_chord(keys: &[String]) -> BitFunResult<(Vec, u16)> }; Ok((modifiers, keycode)) } + +#[cfg(test)] +mod tests { + use super::{HardwareInput, Input, KeybdInput, MouseInput, INPUT_0}; + use std::mem::{align_of, offset_of, size_of}; + + #[test] + fn send_input_ffi_layout_matches_winuser() { + assert_eq!(size_of::(), 8); + assert_eq!(align_of::(), 4); + assert_eq!(offset_of!(HardwareInput, uMsg), 0); + assert_eq!(offset_of!(HardwareInput, wParamL), 4); + assert_eq!(offset_of!(HardwareInput, wParamH), 6); + + if cfg!(target_pointer_width = "64") { + assert_eq!(size_of::(), 24); + assert_eq!(align_of::(), 8); + assert_eq!(offset_of!(KeybdInput, dwExtraInfo), 16); + assert_eq!(size_of::(), 32); + assert_eq!(align_of::(), 8); + assert_eq!(offset_of!(MouseInput, dwExtraInfo), 24); + assert_eq!(size_of::(), 32); + assert_eq!(align_of::(), 8); + assert_eq!(size_of::(), 40); + assert_eq!(align_of::(), 8); + assert_eq!(offset_of!(Input, Anonymous), 8); + } else { + assert_eq!(size_of::(), 16); + assert_eq!(align_of::(), 4); + assert_eq!(offset_of!(KeybdInput, dwExtraInfo), 12); + assert_eq!(size_of::(), 24); + assert_eq!(align_of::(), 4); + assert_eq!(offset_of!(MouseInput, dwExtraInfo), 20); + assert_eq!(size_of::(), 24); + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 28); + assert_eq!(align_of::(), 4); + assert_eq!(offset_of!(Input, Anonymous), 4); + } + + assert_eq!(offset_of!(KeybdInput, wVk), 0); + assert_eq!(offset_of!(KeybdInput, wScan), 2); + assert_eq!(offset_of!(KeybdInput, dwFlags), 4); + assert_eq!(offset_of!(KeybdInput, time), 8); + assert_eq!(offset_of!(MouseInput, dx), 0); + assert_eq!(offset_of!(MouseInput, dy), 4); + assert_eq!(offset_of!(MouseInput, mouseData), 8); + assert_eq!(offset_of!(MouseInput, dwFlags), 12); + assert_eq!(offset_of!(MouseInput, time), 16); + assert_eq!(offset_of!(Input, r#type), 0); + } +} diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index 1a182510a1..ff5fce6ae7 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -289,6 +289,45 @@ struct CompressionTriggerBudget { safety_reserve_tokens: usize, } +// Fields are declared in reverse parameter order so dropping an unconsumed +// input preserves the previous function-parameter drop order. Call sites keep +// struct literal fields in the original evaluation order. +struct TurnPromptScaffoldInput<'a> { + stage: &'a str, + runtime_context_needs: RuntimeContextNeeds, + tool_listing_sections: ToolListingSections, + supports_image_understanding: bool, + model_name: &'a str, + current_agent: &'a dyn crate::agentic::agents::Agent, + context: &'a ExecutionContext, +} + +struct FinalizeRoundInput<'a> { + context_window: usize, + tool_definitions: Option>, + reminder_text: &'a str, + messages: &'a [Message], + prepended_reminders: &'a [&'a str], + primary_model_facts: &'a PrimaryModelFacts, + execution_context_vars: &'a HashMap, + round_group_id: Option, + round_number: usize, + agent_type: String, + context: &'a ExecutionContext, + ai_client: Arc, +} + +struct CompressionModelSummaryInput<'a> { + trace_config: Option, + primary_supports_image_understanding: bool, + prepended_prompt_reminders: &'a PrependedPromptReminders, + tool_definitions: &'a Option>, + workspace: Option<&'a WorkspaceBinding>, + dialog_turn_id: &'a str, + runtime_messages: &'a [Message], + ai_client: Arc, +} + /// Execution engine pub struct ExecutionEngine { round_executor: Arc, @@ -1079,51 +1118,45 @@ impl ExecutionEngine { async fn resolve_turn_prompt_scaffold( &self, - context: &ExecutionContext, - current_agent: &dyn crate::agentic::agents::Agent, - model_name: &str, - supports_image_understanding: bool, - tool_listing_sections: ToolListingSections, - runtime_context_needs: RuntimeContextNeeds, - stage: &str, + input: TurnPromptScaffoldInput<'_>, ) -> BitFunResult { debug!( "Resolving turn prompt scaffold: session_id={}, turn_id={}, stage={}, agent={}, model={}", - context.session_id, - context.dialog_turn_id, - stage, - current_agent.name(), - model_name + input.context.session_id, + input.context.dialog_turn_id, + input.stage, + input.current_agent.name(), + input.model_name ); let prompt_context = Self::build_prompt_context( - context, - model_name, - supports_image_understanding, - tool_listing_sections, - runtime_context_needs, + input.context, + input.model_name, + input.supports_image_understanding, + input.tool_listing_sections, + input.runtime_context_needs, ) .await; let prepended_prompt_reminders = self .build_cached_prepended_prompt_reminders( - &context.session_id, - current_agent, + &input.context.session_id, + input.current_agent, prompt_context.as_ref(), - &context.context, + &input.context.context, ) .await; let system_prompt = self .resolve_cached_system_prompt( - &context.session_id, - current_agent, + &input.context.session_id, + input.current_agent, prompt_context.as_ref(), ) .await?; Self::log_turn_prompt_scaffold( - &context.session_id, - &context.dialog_turn_id, - stage, + &input.context.session_id, + &input.context.dialog_turn_id, + input.stage, system_prompt.len(), &prepended_prompt_reminders, ); @@ -1297,81 +1330,68 @@ impl ExecutionEngine { .collect() } - async fn run_finalize_round( - &self, - ai_client: Arc, - context: &ExecutionContext, - agent_type: String, - round_number: usize, - round_group_id: Option, - execution_context_vars: &HashMap, - primary_model_facts: &PrimaryModelFacts, - prepended_reminders: &[&str], - messages: &[Message], - reminder_text: &str, - tool_definitions: Option>, - context_window: usize, - ) -> BitFunResult { + async fn run_finalize_round(&self, input: FinalizeRoundInput<'_>) -> BitFunResult { // Keep the original tool definitions attached to the finalize request // even though finalize forbids tool execution at runtime. Dropping the // tools here would change the provider request shape, which breaks // prompt/prefix cache reuse and turns the finalize round into a cache // miss for providers that key caching on the full request schema. - let finalize_tool_names = Self::finalize_tool_names(tool_definitions.as_deref()); + let finalize_tool_names = Self::finalize_tool_names(input.tool_definitions.as_deref()); let finalize_runtime_tool_restrictions = - Self::finalize_runtime_tool_restrictions(context, &finalize_tool_names); + Self::finalize_runtime_tool_restrictions(input.context, &finalize_tool_names); let mut final_ai_messages = Self::build_ai_messages_for_send( - messages, - &ai_client.config.format, - context + input.messages, + &input.ai_client.config.format, + input + .context .workspace .as_ref() .map(|workspace| workspace.root_path()), - &context.dialog_turn_id, - primary_model_facts.supports_image_inputs, - prepended_reminders, + &input.context.dialog_turn_id, + input.primary_model_facts.supports_image_inputs, + input.prepended_reminders, ) .await?; - final_ai_messages.push(AIMessage::user(render_system_reminder(reminder_text))); + final_ai_messages.push(AIMessage::user(render_system_reminder(input.reminder_text))); final_ai_messages.push(AIMessage::user(Self::FINALIZE_USER_FOLLOWUP.to_string())); let model_exchange_trace_dir = self .session_manager - .persistent_model_exchange_trace_dir(&context.session_id) + .persistent_model_exchange_trace_dir(&input.context.session_id) .await; let round_context = RoundContext { - session_id: context.session_id.clone(), - subagent_parent_info: context.subagent_parent_info.clone(), - dialog_turn_id: context.dialog_turn_id.clone(), - turn_index: context.turn_index, - round_number, - round_group_id, - workspace: context.workspace.clone(), + session_id: input.context.session_id.clone(), + subagent_parent_info: input.context.subagent_parent_info.clone(), + dialog_turn_id: input.context.dialog_turn_id.clone(), + turn_index: input.context.turn_index, + round_number: input.round_number, + round_group_id: input.round_group_id, + workspace: input.context.workspace.clone(), model_exchange_trace_dir, available_tools: finalize_tool_names, collapsed_tools: Vec::new(), unlocked_collapsed_tools: Vec::new(), - model_name: ai_client.config.model.clone(), - primary_model_facts: primary_model_facts.clone(), - agent_type, - context_vars: execution_context_vars.clone(), - delegation_policy: context.delegation_policy, + model_name: input.ai_client.config.model.clone(), + primary_model_facts: input.primary_model_facts.clone(), + agent_type: input.agent_type, + context_vars: input.execution_context_vars.clone(), + delegation_policy: input.context.delegation_policy, runtime_tool_restrictions: finalize_runtime_tool_restrictions, steering_interrupt: None, cancellation_token: CancellationToken::new(), - workspace_services: context.workspace_services.clone(), - terminal_port: context.terminal_port.clone(), - remote_exec_port: context.remote_exec_port.clone(), - recover_partial_on_cancel: context.recover_partial_on_cancel, + workspace_services: input.context.workspace_services.clone(), + terminal_port: input.context.terminal_port.clone(), + remote_exec_port: input.context.remote_exec_port.clone(), + recover_partial_on_cancel: input.context.recover_partial_on_cancel, }; self.round_executor .execute_round( - ai_client, + input.ai_client, round_context, final_ai_messages, - tool_definitions, - Some(context_window), + input.tool_definitions, + Some(input.context_window), ) .await } @@ -1679,32 +1699,25 @@ impl ExecutionEngine { async fn generate_compression_model_summary( &self, - ai_client: Arc, - runtime_messages: &[Message], - dialog_turn_id: &str, - workspace: Option<&WorkspaceBinding>, - tool_definitions: &Option>, - prepended_prompt_reminders: &PrependedPromptReminders, - primary_supports_image_understanding: bool, - trace_config: Option, + input: CompressionModelSummaryInput<'_>, ) -> BitFunResult> { let request_messages = self .build_compression_request_messages( - runtime_messages, - dialog_turn_id, - workspace, - &ai_client.config.format, - primary_supports_image_understanding, - prepended_prompt_reminders, + input.runtime_messages, + input.dialog_turn_id, + input.workspace, + &input.ai_client.config.format, + input.primary_supports_image_understanding, + input.prepended_prompt_reminders, ) .await?; let raw_summary = self .request_compression_summary_with_retry( - ai_client, + input.ai_client, request_messages, - tool_definitions.clone(), - trace_config, + input.tool_definitions.clone(), + input.trace_config, 2, ) .await?; @@ -1851,15 +1864,15 @@ impl ExecutionEngine { let tool_definitions = tool_manifest.map(|manifest| manifest.tool_definitions); let turn_prompt_scaffold = self - .resolve_turn_prompt_scaffold( + .resolve_turn_prompt_scaffold(TurnPromptScaffoldInput { context, - current_agent.as_ref(), - &ai_client.config.model, - primary_supports_image_understanding, + current_agent: current_agent.as_ref(), + model_name: &ai_client.config.model, + supports_image_understanding: primary_supports_image_understanding, tool_listing_sections, runtime_context_needs, - "compression_scaffold", - ) + stage: "compression_scaffold", + }) .await?; Ok(CompressionRuntimeScaffold { @@ -1944,16 +1957,16 @@ impl ExecutionEngine { ) .await; let model_summary = match self - .generate_compression_model_summary( + .generate_compression_model_summary(CompressionModelSummaryInput { ai_client, - &runtime_messages, + runtime_messages: &runtime_messages, dialog_turn_id, workspace, tool_definitions, prepended_prompt_reminders, primary_supports_image_understanding, trace_config, - ) + }) .await { Ok(summary) => summary, @@ -2229,16 +2242,16 @@ impl ExecutionEngine { ) .await; let model_summary = match self - .generate_compression_model_summary( - scaffold.ai_client.clone(), - &runtime_messages, - &dialog_turn_id, - context.workspace.as_ref(), - &scaffold.tool_definitions, - &scaffold.prepended_prompt_reminders, - scaffold.primary_supports_image_understanding, + .generate_compression_model_summary(CompressionModelSummaryInput { + ai_client: scaffold.ai_client.clone(), + runtime_messages: &runtime_messages, + dialog_turn_id: &dialog_turn_id, + workspace: context.workspace.as_ref(), + tool_definitions: &scaffold.tool_definitions, + prepended_prompt_reminders: &scaffold.prepended_prompt_reminders, + primary_supports_image_understanding: scaffold.primary_supports_image_understanding, trace_config, - ) + }) .await { Ok(summary) => summary, @@ -2648,15 +2661,15 @@ impl ExecutionEngine { // It is refreshed after successful context compression so the first // post-compaction request builds the new provider-side prefix cache. let mut turn_prompt_scaffold = self - .resolve_turn_prompt_scaffold( - &context, - current_agent.as_ref(), - &ai_client.config.model, - primary_supports_image_understanding, - tool_listing_sections.clone(), + .resolve_turn_prompt_scaffold(TurnPromptScaffoldInput { + context: &context, + current_agent: current_agent.as_ref(), + model_name: &ai_client.config.model, + supports_image_understanding: primary_supports_image_understanding, + tool_listing_sections: tool_listing_sections.clone(), runtime_context_needs, - "turn_start", - ) + stage: "turn_start", + }) .await?; // Add System Prompt to the beginning of message list (only for this execution, not persisted) @@ -2916,15 +2929,15 @@ impl ExecutionEngine { messages = compressed_messages; turn_prompt_scaffold = self - .resolve_turn_prompt_scaffold( - &context, - current_agent.as_ref(), - &ai_client.config.model, - primary_supports_image_understanding, - tool_listing_sections.clone(), + .resolve_turn_prompt_scaffold(TurnPromptScaffoldInput { + context: &context, + current_agent: current_agent.as_ref(), + model_name: &ai_client.config.model, + supports_image_understanding: primary_supports_image_understanding, + tool_listing_sections: tool_listing_sections.clone(), runtime_context_needs, - "after_context_compression", - ) + stage: "after_context_compression", + }) .await?; Self::apply_turn_prompt_scaffold_to_messages( &mut messages, @@ -3551,20 +3564,20 @@ impl ExecutionEngine { .prepended_prompt_reminders .ordered_reminders(); let final_round_result = self - .run_finalize_round( - ai_client.clone(), - &context, - agent_type.clone(), - completed_rounds, - finalize_round_group_id.clone(), - &execution_context_vars, - &primary_model_facts, - &finalize_prepended_reminders, - &messages, - finalize_reminder, - tool_definitions.clone(), + .run_finalize_round(FinalizeRoundInput { + ai_client: ai_client.clone(), + context: &context, + agent_type: agent_type.clone(), + round_number: completed_rounds, + round_group_id: finalize_round_group_id.clone(), + execution_context_vars: &execution_context_vars, + primary_model_facts: &primary_model_facts, + prepended_reminders: &finalize_prepended_reminders, + messages: &messages, + reminder_text: finalize_reminder, + tool_definitions: tool_definitions.clone(), context_window, - ) + }) .await?; let mut accepted = final_round_result.had_assistant_text @@ -3581,20 +3594,20 @@ impl ExecutionEngine { context.session_id, context.dialog_turn_id ); let retry_result = self - .run_finalize_round( - ai_client.clone(), - &context, - agent_type.clone(), - completed_rounds, - finalize_round_group_id.clone(), - &execution_context_vars, - &primary_model_facts, - &finalize_prepended_reminders, - &messages, - finalize_reminder, - tool_definitions.clone(), + .run_finalize_round(FinalizeRoundInput { + ai_client: ai_client.clone(), + context: &context, + agent_type: agent_type.clone(), + round_number: completed_rounds, + round_group_id: finalize_round_group_id.clone(), + execution_context_vars: &execution_context_vars, + primary_model_facts: &primary_model_facts, + prepended_reminders: &finalize_prepended_reminders, + messages: &messages, + reminder_text: finalize_reminder, + tool_definitions: tool_definitions.clone(), context_window, - ) + }) .await?; if !retry_result.had_assistant_text || Self::assistant_has_tool_calls(&retry_result.assistant_message) diff --git a/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs b/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs index 8b834cb6ba..2b20872c7d 100644 --- a/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs +++ b/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs @@ -110,30 +110,34 @@ struct WorkspaceModelExchangeTraceSink { trace_paths: DashMap, } +// Reverse declaration order preserves the previous function-parameter drop +// order; struct literals remain in the original evaluation order. +struct WorkspaceModelExchangeTraceInput { + model_id: String, + api_format: String, + provider: String, + operation_trigger: Option, + operation_id: String, + operation_kind: String, + turn_id: String, + session_id: String, + policy: ModelExchangeTracePolicy, + trace_session_dir: PathBuf, +} + impl WorkspaceModelExchangeTraceSink { - fn new( - trace_session_dir: PathBuf, - policy: ModelExchangeTracePolicy, - session_id: String, - turn_id: String, - operation_kind: String, - operation_id: String, - operation_trigger: Option, - provider: String, - api_format: String, - model_id: String, - ) -> Self { + fn new(input: WorkspaceModelExchangeTraceInput) -> Self { Self { - trace_session_dir, - policy, - session_id, - turn_id, - operation_kind, - operation_id, - operation_trigger, - provider, - api_format, - model_id, + trace_session_dir: input.trace_session_dir, + policy: input.policy, + session_id: input.session_id, + turn_id: input.turn_id, + operation_kind: input.operation_kind, + operation_id: input.operation_id, + operation_trigger: input.operation_trigger, + provider: input.provider, + api_format: input.api_format, + model_id: input.model_id, trace_paths: DashMap::new(), } } @@ -402,16 +406,18 @@ pub(super) async fn prepare_model_exchange_trace_for_workspace( Some(ModelExchangeTraceConfig { sink: Arc::new(WorkspaceModelExchangeTraceSink::new( - trace_session_dir, - policy, - session_id.to_string(), - turn_id.to_string(), - operation.kind.to_string(), - operation.id.to_string(), - operation.trigger.map(str::to_string), - ai_client.config.format.clone(), - ai_client.config.format.clone(), - ai_client.config.model.clone(), + WorkspaceModelExchangeTraceInput { + trace_session_dir, + policy, + session_id: session_id.to_string(), + turn_id: turn_id.to_string(), + operation_kind: operation.kind.to_string(), + operation_id: operation.id.to_string(), + operation_trigger: operation.trigger.map(str::to_string), + provider: ai_client.config.format.clone(), + api_format: ai_client.config.format.clone(), + model_id: ai_client.config.model.clone(), + }, )), capture_request_body: policy.capture_request_body, }) @@ -477,3 +483,76 @@ fn sequence_allocators() -> &'static DashMap>>> { static ALLOCATORS: OnceLock>>>> = OnceLock::new(); ALLOCATORS.get_or_init(DashMap::new) } + +#[cfg(test)] +mod tests { + use super::*; + + struct SequenceAllocatorCleanup(String); + + impl Drop for SequenceAllocatorCleanup { + fn drop(&mut self) { + sequence_allocators().remove(&self.0); + } + } + + #[tokio::test] + async fn trace_request_preserves_operation_and_model_identity() { + let directory = tempfile::tempdir().expect("trace directory should be created"); + let _allocator_cleanup = + SequenceAllocatorCleanup(directory.path().to_string_lossy().to_string()); + let sink = WorkspaceModelExchangeTraceSink::new(WorkspaceModelExchangeTraceInput { + trace_session_dir: directory.path().to_path_buf(), + policy: ModelExchangeTracePolicy { + mode: ModelExchangeTracingMode::Full, + capture_request_body: true, + capture_response_text: true, + capture_reasoning: true, + capture_tool_calls: true, + capture_usage: true, + capture_provider_metadata: true, + }, + session_id: "session-identity".to_string(), + turn_id: "turn-identity".to_string(), + operation_kind: "context_compression".to_string(), + operation_id: "compression-identity".to_string(), + operation_trigger: Some("manual".to_string()), + provider: "provider-format".to_string(), + api_format: "api-format".to_string(), + model_id: "model-identity".to_string(), + }); + + let handle = sink + .request_attempt_started(&ModelExchangeRequestAttempt { + request_url: "https://example.invalid/model".to_string(), + request_body: Some(serde_json::json!({"request": "body"})), + attempt_number: 2, + }) + .await + .expect("trace request should be recorded"); + let path = sink + .trace_paths + .get(&handle.trace_id) + .expect("trace path should be registered") + .value() + .clone(); + let record = sink + .read_record(&path) + .await + .expect("trace record should be readable"); + + assert_eq!(record.session_id, "session-identity"); + assert_eq!(record.turn_id, "turn-identity"); + assert_eq!(record.operation_kind, "context_compression"); + assert_eq!(record.operation_id, "compression-identity"); + assert_eq!(record.operation_trigger.as_deref(), Some("manual")); + assert_eq!(record.request.provider, "provider-format"); + assert_eq!(record.request.api_format, "api-format"); + assert_eq!(record.request.model_id, "model-identity"); + assert_eq!(record.request.attempt_number, 2); + assert_eq!( + record.request.body, + Some(serde_json::json!({"request": "body"})) + ); + } +} diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs index 065384acba..69be58e989 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs @@ -50,7 +50,7 @@ fn json_object_metadata(value: Value) -> serde_json::Map { } } -async fn deliver_background_bash_result( +struct BackgroundBashResultDelivery { parent_session_id: String, parent_agent_type: String, parent_workspace_path: Option, @@ -61,7 +61,21 @@ async fn deliver_background_bash_result( metadata: serde_json::Map, terminal_session_id: String, failure_context: &'static str, -) { +} + +async fn deliver_background_bash_result(delivery: BackgroundBashResultDelivery) { + let BackgroundBashResultDelivery { + parent_session_id, + parent_agent_type, + parent_workspace_path, + parent_remote_connection_id, + parent_remote_ssh_host, + delivery_text, + display_text, + metadata, + terminal_session_id, + failure_context, + } = delivery; let runtime = match CoreServiceAgentRuntime::global_agent_runtime_with_lifecycle_delivery() { Ok(runtime) => runtime, Err(error) => { @@ -1240,18 +1254,18 @@ impl BashTool { "outputFile": output_file_reference_for_task.clone(), }); - deliver_background_bash_result( - parent_session_id.clone(), - parent_agent_type.clone(), - parent_workspace_path.clone(), - parent_remote_connection_id.clone(), - parent_remote_ssh_host.clone(), + deliver_background_bash_result(BackgroundBashResultDelivery { + parent_session_id: parent_session_id.clone(), + parent_agent_type: parent_agent_type.clone(), + parent_workspace_path: parent_workspace_path.clone(), + parent_remote_connection_id: parent_remote_connection_id.clone(), + parent_remote_ssh_host: parent_remote_ssh_host.clone(), delivery_text, display_text, - json_object_metadata(metadata), - terminal_session_id.clone(), - "result", - ) + metadata: json_object_metadata(metadata), + terminal_session_id: terminal_session_id.clone(), + failure_context: "result", + }) .await; delivery_sent = true; break; @@ -1280,18 +1294,18 @@ impl BashTool { "error": message.clone(), }); - deliver_background_bash_result( - parent_session_id.clone(), - parent_agent_type.clone(), - parent_workspace_path.clone(), - parent_remote_connection_id.clone(), - parent_remote_ssh_host.clone(), + deliver_background_bash_result(BackgroundBashResultDelivery { + parent_session_id: parent_session_id.clone(), + parent_agent_type: parent_agent_type.clone(), + parent_workspace_path: parent_workspace_path.clone(), + parent_remote_connection_id: parent_remote_connection_id.clone(), + parent_remote_ssh_host: parent_remote_ssh_host.clone(), delivery_text, display_text, - json_object_metadata(metadata), - terminal_session_id.clone(), - "error result", - ) + metadata: json_object_metadata(metadata), + terminal_session_id: terminal_session_id.clone(), + failure_context: "error result", + }) .await; delivery_sent = true; break; @@ -1322,18 +1336,18 @@ impl BashTool { "error": "stream_ended_without_completion", }); - deliver_background_bash_result( - parent_session_id.clone(), - parent_agent_type.clone(), - parent_workspace_path.clone(), - parent_remote_connection_id.clone(), - parent_remote_ssh_host.clone(), + deliver_background_bash_result(BackgroundBashResultDelivery { + parent_session_id: parent_session_id.clone(), + parent_agent_type: parent_agent_type.clone(), + parent_workspace_path: parent_workspace_path.clone(), + parent_remote_connection_id: parent_remote_connection_id.clone(), + parent_remote_ssh_host: parent_remote_ssh_host.clone(), delivery_text, display_text, - json_object_metadata(metadata), - terminal_session_id.clone(), - "stream-end result", - ) + metadata: json_object_metadata(metadata), + terminal_session_id: terminal_session_id.clone(), + failure_context: "stream-end result", + }) .await; } }); diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/computer_use_actions.rs b/src/crates/assembly/core/src/agentic/tools/implementations/computer_use_actions.rs index 37bae081e5..7c9c963327 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/computer_use_actions.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/computer_use_actions.rs @@ -24,9 +24,10 @@ use super::control_hub::{coded_tool_error, err_response, ControlHubError, ErrorC /// row the dispatcher injects an `app_state.loop_warning` so the model is /// forced off the failing path on its **next** turn (`/Screenshot policy/ /// Mandatory screenshot moments` in `claw_mode.md`). -static APP_LOOP_TRACKER: std::sync::OnceLock< - std::sync::Mutex>, -> = std::sync::OnceLock::new(); +type AppLoopTracker = + std::sync::OnceLock>>; + +static APP_LOOP_TRACKER: AppLoopTracker = std::sync::OnceLock::new(); fn loop_tracker_observe( pid: Option, diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/get_file_diff_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/get_file_diff_tool.rs index 9d7060991f..511b54af32 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/get_file_diff_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/get_file_diff_tool.rs @@ -33,6 +33,19 @@ use std::path::Path; /// 3. Return full file content pub struct GetFileDiffTool; +type ExactReviewTarget = (String, String, Vec, Option); + +struct ExactReviewDiffRequest<'a> { + workspace_root: &'a Path, + logical_path: &'a str, + base_revision: &'a str, + head_revision: &'a str, + paths: &'a [String], + fingerprint: Option<&'a str>, + diff_offset: usize, + cursor_binding: &'a str, +} + #[derive(Debug, Clone, PartialEq, Eq)] enum ProviderFileDiffRoute { Identity { @@ -279,7 +292,7 @@ impl GetFileDiffTool { fn exact_review_target( relative_path: &str, context: &ToolUseContext, - ) -> BitFunResult, Option)>> { + ) -> BitFunResult> { let Some(evidence) = Self::target_evidence(context)? else { return Ok(None); }; @@ -308,17 +321,17 @@ impl GetFileDiffTool { ))) } - async fn exact_review_diff( - &self, - workspace_root: &Path, - logical_path: &str, - base_revision: &str, - head_revision: &str, - paths: &[String], - fingerprint: Option<&str>, - diff_offset: usize, - cursor_binding: &str, - ) -> BitFunResult { + async fn exact_review_diff(&self, request: ExactReviewDiffRequest<'_>) -> BitFunResult { + let ExactReviewDiffRequest { + workspace_root, + logical_path, + base_revision, + head_revision, + paths, + fingerprint, + diff_offset, + cursor_binding, + } = request; let diff_content = GitService::get_review_diff(workspace_root, base_revision, head_revision, paths) .await @@ -1607,19 +1620,19 @@ Usage: BitFunError::tool("Workspace root is required for Review target diff".to_string()) })?; let data = self - .exact_review_diff( + .exact_review_diff(ExactReviewDiffRequest { workspace_root, logical_path, - &base_revision, - &head_revision, - &paths, - fingerprint.as_deref(), + base_revision: &base_revision, + head_revision: &head_revision, + paths: &paths, + fingerprint: fingerprint.as_deref(), diff_offset, - prepared_evidence + cursor_binding: prepared_evidence .as_ref() .map(ReviewTargetEvidence::fingerprint) .unwrap_or_default(), - ) + }) .await?; let data = Self::apply_review_diff_budget( data, 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 1b5abccd8a..6758aa2df2 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 @@ -29,6 +29,22 @@ fn build_deep_review_subagent_context( values } +struct BackgroundTaskStartRequest<'a> { + coordinator: &'a std::sync::Arc, + context: &'a ToolUseContext, + context_mode: SubagentContextMode, + target_session_id: Option, + subagent_type: Option, + effective_workspace_path: Option, + model_id: Option, + subagent_context: Option>, + prepared_prompt: String, + timeout_seconds: Option, + tool_call_id: String, + session_id: String, + dialog_turn_id: String, +} + impl TaskTool { pub(super) async fn load_configured_tool_execution_timeout() -> Option { let service = GlobalConfigManager::get_service().await.ok()?; @@ -501,8 +517,8 @@ impl TaskTool { }); let prepared_prompt = prompt; if run_in_background { - return Self::start_background_task( - &coordinator, + return Self::start_background_task(BackgroundTaskStartRequest { + coordinator: &coordinator, context, context_mode, target_session_id, @@ -515,7 +531,7 @@ impl TaskTool { tool_call_id, session_id, dialog_turn_id, - ) + }) .await; } @@ -548,20 +564,23 @@ impl TaskTool { } async fn start_background_task( - coordinator: &std::sync::Arc, - context: &ToolUseContext, - context_mode: SubagentContextMode, - target_session_id: Option, - subagent_type: Option, - effective_workspace_path: Option, - model_id: Option, - subagent_context: Option>, - prepared_prompt: String, - timeout_seconds: Option, - tool_call_id: String, - session_id: String, - dialog_turn_id: String, + request: BackgroundTaskStartRequest<'_>, ) -> BitFunResult> { + let BackgroundTaskStartRequest { + coordinator, + context, + context_mode, + target_session_id, + subagent_type, + effective_workspace_path, + model_id, + subagent_context, + prepared_prompt, + timeout_seconds, + tool_call_id, + session_id, + dialog_turn_id, + } = request; let parent_info = SubagentParentInfo { tool_call_id, session_id, diff --git a/src/crates/assembly/core/src/service/config/mode_config_canonicalizer.rs b/src/crates/assembly/core/src/service/config/mode_config_canonicalizer.rs index 188657b995..6b6549f546 100644 --- a/src/crates/assembly/core/src/service/config/mode_config_canonicalizer.rs +++ b/src/crates/assembly/core/src/service/config/mode_config_canonicalizer.rs @@ -144,28 +144,42 @@ fn stored_agent_profile_from_tool_selection( } } - stored_agent_profile_from_overrides( + stored_agent_profile_from_overrides(StoredAgentProfileOverrides { agent_id, added_tools, removed_tools, disabled_user_skills, enabled_user_skills, subagent_overrides, - &default_tools, + default_tools: &default_tools, valid_tools, - ) + }) } -fn stored_agent_profile_from_overrides( - agent_id: &str, +struct StoredAgentProfileOverrides<'a> { + agent_id: &'a str, added_tools: Vec, removed_tools: Vec, disabled_user_skills: Vec, enabled_user_skills: Vec, subagent_overrides: ParentSubagentOverrideConfig, - default_tools: &[String], - valid_tools: &HashSet, + default_tools: &'a [String], + valid_tools: &'a HashSet, +} + +fn stored_agent_profile_from_overrides( + overrides: StoredAgentProfileOverrides<'_>, ) -> Option { + let StoredAgentProfileOverrides { + agent_id, + added_tools, + removed_tools, + disabled_user_skills, + enabled_user_skills, + subagent_overrides, + default_tools, + valid_tools, + } = overrides; let profile_id = resolve_profile_id(agent_id); let default_set: HashSet = default_tools.iter().cloned().collect(); let mut added_tools = normalize_tools(added_tools, valid_tools); @@ -250,14 +264,16 @@ fn canonicalize_agent_profile( } Ok(stored_agent_profile_from_overrides( - profile_id, - stored.added_tools, - stored.removed_tools, - stored.disabled_user_skills, - stored.enabled_user_skills, - stored.subagent_overrides, - default_tools, - valid_tools, + StoredAgentProfileOverrides { + agent_id: profile_id, + added_tools: stored.added_tools, + removed_tools: stored.removed_tools, + disabled_user_skills: stored.disabled_user_skills, + enabled_user_skills: stored.enabled_user_skills, + subagent_overrides: stored.subagent_overrides, + default_tools, + valid_tools, + }, )) } @@ -531,6 +547,7 @@ mod tests { use super::{ agent_profile_member_mode_ids_for, canonicalize_agent_profile, normalize_skill_override_lists, stored_agent_profile_from_overrides, + StoredAgentProfileOverrides, }; use crate::service::config::types::AgentSubagentOverrideState; use serde_json::Value; @@ -557,16 +574,16 @@ mod tests { #[test] fn stored_agent_profile_from_overrides_keeps_enabled_user_skills() { let valid_tools = HashSet::new(); - let stored = stored_agent_profile_from_overrides( - "agentic", - Vec::new(), - Vec::new(), - Vec::new(), - vec!["user::bitfun-system::pdf".to_string()], - Default::default(), - &[], - &valid_tools, - ) + let stored = stored_agent_profile_from_overrides(StoredAgentProfileOverrides { + agent_id: "agentic", + added_tools: Vec::new(), + removed_tools: Vec::new(), + disabled_user_skills: Vec::new(), + enabled_user_skills: vec!["user::bitfun-system::pdf".to_string()], + subagent_overrides: Default::default(), + default_tools: &[], + valid_tools: &valid_tools, + }) .expect("mode config should be retained when skill overrides exist"); assert_eq!(stored.profile_id, "coding_shared"); @@ -585,16 +602,16 @@ mod tests { "builtin::builtin::Explore".to_string(), AgentSubagentOverrideState::Disabled, ); - let stored = stored_agent_profile_from_overrides( - "debug", - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - subagent_overrides.clone(), - &[], - &valid_tools, - ) + let stored = stored_agent_profile_from_overrides(StoredAgentProfileOverrides { + agent_id: "debug", + added_tools: Vec::new(), + removed_tools: Vec::new(), + disabled_user_skills: Vec::new(), + enabled_user_skills: Vec::new(), + subagent_overrides: subagent_overrides.clone(), + default_tools: &[], + valid_tools: &valid_tools, + }) .expect("mode config should be retained when subagent overrides exist"); assert_eq!(stored.profile_id, "coding_shared"); diff --git a/src/crates/execution/agent-runtime/src/deep_review/team_definition.rs b/src/crates/execution/agent-runtime/src/deep_review/team_definition.rs index dca6af743d..2bad10fe33 100644 --- a/src/crates/execution/agent-runtime/src/deep_review/team_definition.rs +++ b/src/crates/execution/agent-runtime/src/deep_review/team_definition.rs @@ -63,16 +63,28 @@ pub struct ReviewTeamDefinition { pub hidden_agent_ids: Vec, } -fn review_role( - key: &str, - subagent_id: &str, - fun_name: &str, - role_name: &str, - description: &str, - responsibilities: &[&str], - accent_color: &str, - conditional: bool, -) -> ReviewTeamRoleDefinition { +struct ReviewRoleInput<'a>( + &'a str, + &'a str, + &'a str, + &'a str, + &'a str, + &'a [&'a str], + &'a str, + bool, +); + +fn review_role(input: ReviewRoleInput<'_>) -> ReviewTeamRoleDefinition { + let ReviewRoleInput( + key, + subagent_id, + fun_name, + role_name, + description, + responsibilities, + accent_color, + conditional, + ) = input; ReviewTeamRoleDefinition { key: key.to_string(), subagent_id: subagent_id.to_string(), @@ -95,16 +107,28 @@ fn role_directives(entries: &[(&str, &str)]) -> BTreeMap { .collect() } -fn strategy_profile( - level: &str, - label: &str, - summary: &str, - token_impact: &str, - runtime_impact: &str, - default_model_slot: &str, - prompt_directive: &str, - directives: &[(&str, &str)], -) -> ReviewStrategyManifestProfile { +struct StrategyProfileInput<'a>( + &'a str, + &'a str, + &'a str, + &'a str, + &'a str, + &'a str, + &'a str, + &'a [(&'a str, &'a str)], +); + +fn strategy_profile(input: StrategyProfileInput<'_>) -> ReviewStrategyManifestProfile { + let StrategyProfileInput( + level, + label, + summary, + token_impact, + runtime_impact, + default_model_slot, + prompt_directive, + directives, + ) = input; ReviewStrategyManifestProfile { level: level.to_string(), label: label.to_string(), @@ -119,7 +143,7 @@ fn strategy_profile( pub fn default_review_team_definition() -> ReviewTeamDefinition { let core_roles = vec![ - review_role( + review_role(ReviewRoleInput( "businessLogic", REVIEWER_BUSINESS_LOGIC_AGENT_TYPE, "Logic Reviewer", @@ -132,8 +156,8 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { ], "#2563eb", false, - ), - review_role( + )), + review_role(ReviewRoleInput( "performance", REVIEWER_PERFORMANCE_AGENT_TYPE, "Performance Reviewer", @@ -146,8 +170,8 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { ], "#d97706", false, - ), - review_role( + )), + review_role(ReviewRoleInput( "security", REVIEWER_SECURITY_AGENT_TYPE, "Security Reviewer", @@ -160,8 +184,8 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { ], "#dc2626", false, - ), - review_role( + )), + review_role(ReviewRoleInput( "architecture", REVIEWER_ARCHITECTURE_AGENT_TYPE, "Architecture Reviewer", @@ -174,8 +198,8 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { ], "#0891b2", false, - ), - review_role( + )), + review_role(ReviewRoleInput( "frontend", REVIEWER_FRONTEND_AGENT_TYPE, "Frontend Reviewer", @@ -188,8 +212,8 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { ], "#059669", true, - ), - review_role( + )), + review_role(ReviewRoleInput( "judge", REVIEW_JUDGE_AGENT_TYPE, "Review Arbiter", @@ -203,13 +227,13 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { ], "#7c3aed", false, - ), + )), ]; let strategy_profiles = BTreeMap::from([ ( "quick".to_string(), - strategy_profile( + strategy_profile(StrategyProfileInput( "quick", "Quick", "Quick keeps built-in target-matched reviewers, skips user-added specialists, and reports reduced coverage.", @@ -243,11 +267,11 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { "This was a quick review. Focus on confirming or rejecting each finding efficiently. If a finding's evidence is thin, reject it rather than spending time verifying.", ), ], - ), + )), ), ( "normal".to_string(), - strategy_profile( + strategy_profile(StrategyProfileInput( "normal", "Normal", "Normal stays practical for slower models, limits optional expansion, and uses summary-first on large changes.", @@ -281,11 +305,11 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { "Validate each finding's logical consistency and evidence quality. Spot-check code only when a claim needs verification.", ), ], - ), + )), ), ( "deep".to_string(), - strategy_profile( + strategy_profile(StrategyProfileInput( "deep", "Deep", "Thorough multi-pass review with the longest budget for risky or release-sensitive changes.", @@ -319,7 +343,7 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { "This was a deep review with potentially complex findings. Cross-validate findings across reviewers for consistency. For each finding, verify the evidence supports the conclusion and the suggested fix is safe. Pay extra attention to overlapping findings across reviewers or same-role instances.", ), ], - ), + )), ), ]); @@ -361,3 +385,65 @@ pub fn default_review_team_definition() -> ReviewTeamDefinition { hidden_agent_ids, } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_team_preserves_role_and_strategy_profile_values() { + let definition = default_review_team_definition(); + assert_eq!( + definition + .core_roles + .iter() + .map(|role| role.key.as_str()) + .collect::>(), + [ + "businessLogic", + "performance", + "security", + "architecture", + "frontend", + "judge", + ] + ); + assert_eq!( + definition + .core_roles + .iter() + .map(|role| ( + role.subagent_id.as_str(), + role.accent_color.as_str(), + role.conditional + )) + .collect::>(), + [ + (REVIEWER_BUSINESS_LOGIC_AGENT_TYPE, "#2563eb", false), + (REVIEWER_PERFORMANCE_AGENT_TYPE, "#d97706", false), + (REVIEWER_SECURITY_AGENT_TYPE, "#dc2626", false), + (REVIEWER_ARCHITECTURE_AGENT_TYPE, "#0891b2", false), + (REVIEWER_FRONTEND_AGENT_TYPE, "#059669", true), + (REVIEW_JUDGE_AGENT_TYPE, "#7c3aed", false), + ] + ); + assert_eq!( + definition + .strategy_profiles + .iter() + .map(|(key, profile)| { + ( + key.as_str(), + profile.default_model_slot.as_str(), + profile.role_directives.len(), + ) + }) + .collect::>(), + [ + ("deep", "primary", 6), + ("normal", "fast", 6), + ("quick", "fast", 6) + ] + ); + } +} diff --git a/src/crates/execution/agent-runtime/tests/custom_agent_mode_contracts.rs b/src/crates/execution/agent-runtime/tests/custom_agent_mode_contracts.rs index 77323e2667..f2ea543f29 100644 --- a/src/crates/execution/agent-runtime/tests/custom_agent_mode_contracts.rs +++ b/src/crates/execution/agent-runtime/tests/custom_agent_mode_contracts.rs @@ -11,16 +11,30 @@ use bitfun_agent_runtime::prompt::UserContextPolicy; use std::fs; use std::path::{Path, PathBuf}; +struct BuildModeDefinitionInput<'a>( + Option<&'a str>, + Option<&'a str>, + Option<&'a str>, + Option>, + Option, + Option<&'a str>, + Option, + CustomAgentLevel, +); + fn build_mode_definition( - id: Option<&str>, - name: Option<&str>, - description: Option<&str>, - tools: Option>, - readonly: Option, - model: Option<&str>, - user_context_policy: Option, - level: CustomAgentLevel, + input: BuildModeDefinitionInput<'_>, ) -> Result { + let BuildModeDefinitionInput( + id, + name, + description, + tools, + readonly, + model, + user_context_policy, + level, + ) = input; CustomAgentDefinition::from_front_matter_fields( id, name, @@ -38,7 +52,7 @@ fn build_mode_definition( #[test] fn custom_mode_defaults_generate_id_and_default_policy() { - let parsed = build_mode_definition( + let parsed = build_mode_definition(BuildModeDefinitionInput( None, Some("PlannerPlus"), Some("Custom planning mode"), @@ -47,7 +61,7 @@ fn custom_mode_defaults_generate_id_and_default_policy() { None, None, CustomAgentLevel::User, - ) + )) .expect("mode definition should be valid"); assert_eq!(parsed.definition.id, "PlannerPlus"); @@ -97,7 +111,7 @@ fn custom_mode_rejects_review_flag() { fn custom_mode_markdown_save_omits_default_fields() { let dir = TestTempDir::new("bitfun-runtime-custom-mode-defaults"); let path = dir.join("planner.md"); - let definition = build_mode_definition( + let definition = build_mode_definition(BuildModeDefinitionInput( Some("PlannerPlus"), Some("PlannerPlus"), Some("Custom planning mode"), @@ -106,7 +120,7 @@ fn custom_mode_markdown_save_omits_default_fields() { None, None, CustomAgentLevel::User, - ) + )) .expect("mode definition should be valid") .definition; @@ -132,7 +146,7 @@ fn custom_mode_markdown_save_round_trips_custom_policy_and_model() { let dir = TestTempDir::new("bitfun-runtime-custom-mode-custom"); let path = dir.join("planner.md"); let policy = UserContextPolicy::empty().with_workspace_instructions(); - let definition = build_mode_definition( + let definition = build_mode_definition(BuildModeDefinitionInput( Some("PlannerPlus"), Some("PlannerPlus"), Some("Custom planning mode"), @@ -141,7 +155,7 @@ fn custom_mode_markdown_save_round_trips_custom_policy_and_model() { Some("primary"), Some(policy.clone()), CustomAgentLevel::User, - ) + )) .expect("mode definition should be valid") .definition; @@ -166,7 +180,7 @@ fn custom_mode_markdown_save_round_trips_empty_custom_policy() { let dir = TestTempDir::new("bitfun-runtime-custom-mode-empty-policy"); let path = dir.join("planner.md"); let policy = UserContextPolicy::empty(); - let definition = build_mode_definition( + let definition = build_mode_definition(BuildModeDefinitionInput( Some("PlannerPlus"), Some("PlannerPlus"), Some("Custom planning mode"), @@ -175,7 +189,7 @@ fn custom_mode_markdown_save_round_trips_empty_custom_policy() { None, Some(policy.clone()), CustomAgentLevel::User, - ) + )) .expect("mode definition should be valid") .definition; @@ -264,7 +278,7 @@ fn custom_mode_discovery_rejects_project_scoped_modes_without_dropping_valid_age #[test] fn custom_agent_validation_filters_invalid_tools_and_falls_back_model() { - let mut definition = build_mode_definition( + let mut definition = build_mode_definition(BuildModeDefinitionInput( Some("PlannerPlus"), Some("PlannerPlus"), Some("Custom planning mode"), @@ -277,7 +291,7 @@ fn custom_agent_validation_filters_invalid_tools_and_falls_back_model() { Some("missing-model"), None, CustomAgentLevel::User, - ) + )) .expect("mode definition should be valid") .definition; @@ -435,7 +449,7 @@ impl Drop for TestTempDir { } fn write_mode(path: &Path, id: &str, level: CustomAgentLevel) { - let definition = build_mode_definition( + let definition = build_mode_definition(BuildModeDefinitionInput( Some(id), Some(id), Some("Custom planning mode"), @@ -444,7 +458,7 @@ fn write_mode(path: &Path, id: &str, level: CustomAgentLevel) { None, None, level, - ) + )) .expect("mode definition should be valid") .definition; custom_agent_save_markdown_file(path, &definition).expect("mode markdown should save"); diff --git a/src/crates/execution/agent-runtime/tests/custom_subagent_contracts.rs b/src/crates/execution/agent-runtime/tests/custom_subagent_contracts.rs index ef898d2e04..eee729b5ae 100644 --- a/src/crates/execution/agent-runtime/tests/custom_subagent_contracts.rs +++ b/src/crates/execution/agent-runtime/tests/custom_subagent_contracts.rs @@ -14,16 +14,21 @@ use bitfun_agent_runtime::custom_subagent::{ use std::fs; use std::path::PathBuf; +struct BuildDefinitionInput<'a>( + Option<&'a str>, + Option<&'a str>, + Option<&'a str>, + Option>, + Option, + Option, + Option<&'a str>, + CustomSubagentKind, +); + fn build_definition( - id: Option<&str>, - name: Option<&str>, - description: Option<&str>, - tools: Option>, - readonly: Option, - review: Option, - model: Option<&str>, - level: CustomSubagentKind, + input: BuildDefinitionInput<'_>, ) -> Result { + let BuildDefinitionInput(id, name, description, tools, readonly, review, model, level) = input; CustomSubagentDefinition::from_front_matter_fields( id, name, @@ -97,7 +102,7 @@ fn custom_subagent_kind_remains_project_or_user() { #[test] fn custom_subagent_definition_from_front_matter_preserves_schema_and_defaults() { - let definition = build_definition( + let definition = build_definition(BuildDefinitionInput( Some("ReviewExtra"), Some("Additional code reviewer"), Some("Review agent for changed files"), @@ -106,7 +111,7 @@ fn custom_subagent_definition_from_front_matter_preserves_schema_and_defaults() Some(true), Some("deepseek-reasoner"), CustomSubagentKind::User, - ) + )) .expect("front matter fields should build a definition"); assert_eq!(definition.id, "ReviewExtra"); @@ -126,7 +131,7 @@ fn custom_subagent_definition_from_front_matter_preserves_schema_and_defaults() #[test] fn custom_subagent_definition_reports_legacy_missing_field_errors() { - let missing_name = build_definition( + let missing_name = build_definition(BuildDefinitionInput( None, None, Some("Additional code reviewer"), @@ -135,12 +140,12 @@ fn custom_subagent_definition_reports_legacy_missing_field_errors() { None, None, CustomSubagentKind::Project, - ) + )) .expect_err("missing name should fail"); assert_eq!(missing_name, CustomSubagentDefinitionError::MissingName); assert_eq!(missing_name.message(), "Missing name field"); - let missing_description = build_definition( + let missing_description = build_definition(BuildDefinitionInput( Some("ReviewExtra"), Some("Additional code reviewer"), None, @@ -149,7 +154,7 @@ fn custom_subagent_definition_reports_legacy_missing_field_errors() { None, None, CustomSubagentKind::Project, - ) + )) .expect_err("missing description should fail"); assert_eq!( missing_description, @@ -162,7 +167,7 @@ fn custom_subagent_definition_reports_legacy_missing_field_errors() { fn custom_subagent_markdown_io_writes_canonical_front_matter() { let dir = TestTempDir::new("bitfun-agent-runtime-subagent"); let path = dir.join("reviewer.md"); - let definition = build_definition( + let definition = build_definition(BuildDefinitionInput( Some("Reviewer"), Some("Review changed code"), Some("Review changed files and report findings"), @@ -171,7 +176,7 @@ fn custom_subagent_markdown_io_writes_canonical_front_matter() { Some(true), Some("deepseek-reasoner"), CustomSubagentKind::Project, - ) + )) .expect("definition should be valid"); custom_subagent_save_markdown_file(&path, &definition).expect("definition should save"); diff --git a/src/crates/execution/agent-stream/src/lib.rs b/src/crates/execution/agent-stream/src/lib.rs index 2b1e818411..b4497a2178 100644 --- a/src/crates/execution/agent-stream/src/lib.rs +++ b/src/crates/execution/agent-stream/src/lib.rs @@ -468,6 +468,16 @@ pub struct StreamProcessor { event_sink: Arc, } +struct GracefulShutdownInput { + session_id: String, + turn_id: String, + round_id: String, + attempt_id: String, + attempt_index: u32, + tool_calls: Vec, + reason: String, +} + impl StreamProcessor { const WATCHDOG_GRACE_SECS: u64 = 2; @@ -553,29 +563,29 @@ impl StreamProcessor { /// Execute graceful shutdown from context async fn graceful_shutdown_from_ctx(&self, ctx: &mut StreamContext, reason: String) { ctx.force_finish_pending_tool_calls(); - self.graceful_shutdown( - ctx.session_id.clone(), - ctx.dialog_turn_id.clone(), - ctx.round_id.clone(), - ctx.attempt_id.clone(), - ctx.attempt_index, - ctx.tool_calls.clone(), + self.graceful_shutdown(GracefulShutdownInput { + session_id: ctx.session_id.clone(), + turn_id: ctx.dialog_turn_id.clone(), + round_id: ctx.round_id.clone(), + attempt_id: ctx.attempt_id.clone(), + attempt_index: ctx.attempt_index, + tool_calls: ctx.tool_calls.clone(), reason, - ) + }) .await; } /// Graceful shutdown: cleanup all unfinished tool states and notify frontend - async fn graceful_shutdown( - &self, - session_id: String, - turn_id: String, - round_id: String, - attempt_id: String, - attempt_index: u32, - tool_calls: Vec, - reason: String, - ) { + async fn graceful_shutdown(&self, input: GracefulShutdownInput) { + let GracefulShutdownInput { + session_id, + turn_id, + round_id, + attempt_id, + attempt_index, + tool_calls, + reason, + } = input; debug!( "Starting graceful shutdown: session_id={}, reason={}", session_id, reason @@ -1162,11 +1172,11 @@ impl StreamProcessor { #[cfg(test)] mod tests { use super::{ - HiddenTextTag, SseLogCollector, SseLogConfig, StreamEventSink, StreamProcessOptions, - StreamProcessor, + GracefulShutdownInput, HiddenTextTag, SseLogCollector, SseLogConfig, StreamEventSink, + StreamProcessOptions, StreamProcessor, ToolCall, }; use super::{UnifiedResponse, UnifiedTokenUsage, UnifiedToolCall}; - use bitfun_events::{AgenticEvent, AgenticEventPriority as EventPriority}; + use bitfun_events::{AgenticEvent, AgenticEventPriority as EventPriority, ToolEventData}; use futures::StreamExt; use serde_json::json; use std::sync::Arc; @@ -1197,6 +1207,46 @@ mod tests { StreamProcessor::new(Arc::new(NoopEventSink)) } + #[tokio::test] + async fn graceful_shutdown_emits_tool_cleanup_before_turn_cancellation() { + let sink = Arc::new(RecordingEventSink::default()); + let processor = StreamProcessor::new(sink.clone()); + + processor + .graceful_shutdown(GracefulShutdownInput { + session_id: "session_1".to_string(), + turn_id: "turn_1".to_string(), + round_id: "round_1".to_string(), + attempt_id: "attempt_1".to_string(), + attempt_index: 1, + tool_calls: vec![ToolCall { + tool_id: "tool_1".to_string(), + tool_name: "Read".to_string(), + arguments: json!({}), + raw_arguments: None, + is_error: false, + recovered_from_truncation: false, + }], + reason: "User cancelled stream processing".to_string(), + }) + .await; + + let events = sink.events.lock().await; + assert_eq!(events.len(), 2); + assert!(matches!( + &events[0], + AgenticEvent::ToolEvent { + tool_event: ToolEventData::Cancelled { tool_id, .. }, + .. + } if tool_id == "tool_1" + )); + assert!(matches!( + &events[1], + AgenticEvent::DialogTurnCancelled { session_id, turn_id } + if session_id == "session_1" && turn_id == "turn_1" + )); + } + #[test] fn derives_watchdog_timeout_from_stream_idle_timeout() { assert_eq!(StreamProcessor::derive_watchdog_timeout(None), None); diff --git a/src/crates/services/services-core/src/filesystem/tree.rs b/src/crates/services/services-core/src/filesystem/tree.rs index d9a8873d19..00923daa67 100644 --- a/src/crates/services/services-core/src/filesystem/tree.rs +++ b/src/crates/services/services-core/src/filesystem/tree.rs @@ -338,6 +338,18 @@ impl Default for FileTreeService { } } +struct SearchFileContentInput<'a> { + path: &'a Path, + file_name: &'a str, + matcher: &'a Regex, + results: &'a Arc>>, + max_results: usize, + should_stop: &'a Arc, + limit_reached: &'a Arc, + cancel_flag: Option<&'a Arc>, + progress_sink: Option<&'a Arc>, +} + impl FileTreeService { pub fn new(options: FileTreeOptions) -> Self { Self { options } @@ -1238,17 +1250,17 @@ impl FileTreeService { } } - if let Err(error) = Self::search_file_content_lines( + if let Err(error) = Self::search_file_content_lines(SearchFileContentInput { path, - &file_name, - matcher.as_ref(), - &results, + file_name: &file_name, + matcher: matcher.as_ref(), + results: &results, max_results, - &should_stop, - &limit_reached, - cancel_flag.as_ref(), - progress_sink.as_ref(), - ) { + should_stop: &should_stop, + limit_reached: &limit_reached, + cancel_flag: cancel_flag.as_ref(), + progress_sink: progress_sink.as_ref(), + }) { warn!( "Failed to search file content {}: {}", path.display(), @@ -1447,17 +1459,18 @@ impl FileTreeService { !should_stop.load(Ordering::Relaxed) } - fn search_file_content_lines( - path: &Path, - file_name: &str, - matcher: &Regex, - results: &Arc>>, - max_results: usize, - should_stop: &Arc, - limit_reached: &Arc, - cancel_flag: Option<&Arc>, - progress_sink: Option<&Arc>, - ) -> FileSystemResult<()> { + fn search_file_content_lines(input: SearchFileContentInput<'_>) -> FileSystemResult<()> { + let SearchFileContentInput { + path, + file_name, + matcher, + results, + max_results, + should_stop, + limit_reached, + cancel_flag, + progress_sink, + } = input; if should_stop.load(Ordering::Relaxed) || cancellation_requested(cancel_flag) { should_stop.store(true, Ordering::Relaxed); return Ok(()); @@ -1578,3 +1591,48 @@ pub enum SearchMatchType { FileName, Content, } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn content_search_preserves_matching_lines_and_preview_ranges() { + let root = tempfile::tempdir().expect("create temp search directory"); + std::fs::write( + root.path().join("sample.txt"), + "first line\r\nbefore needle after\r\n", + ) + .expect("write search fixture"); + + let outcome = FileTreeService::default() + .search_file_contents( + root.path().to_str().expect("utf-8 temp path"), + "needle", + FileContentSearchOptions { + case_sensitive: true, + use_regex: false, + whole_word: false, + max_results: 10, + max_file_size_bytes: 1024, + cancel_flag: None, + }, + ) + .await + .expect("content search"); + + assert!(!outcome.truncated); + assert_eq!(outcome.results.len(), 1); + assert_eq!(outcome.results[0].line_number, Some(2)); + assert_eq!( + outcome.results[0].matched_content.as_deref(), + Some("before needle after") + ); + assert_eq!( + outcome.results[0].preview_before.as_deref(), + Some("before ") + ); + assert_eq!(outcome.results[0].preview_inside.as_deref(), Some("needle")); + assert_eq!(outcome.results[0].preview_after.as_deref(), Some(" after")); + } +} diff --git a/src/crates/services/services-integrations/src/mcp/protocol/transport_remote.rs b/src/crates/services/services-integrations/src/mcp/protocol/transport_remote.rs index 77e984a27e..faa0a13bec 100644 --- a/src/crates/services/services-integrations/src/mcp/protocol/transport_remote.rs +++ b/src/crates/services/services-integrations/src/mcp/protocol/transport_remote.rs @@ -202,7 +202,7 @@ impl StreamableHttpClient for BitFunStreamableHttpClient { } } - let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed(); + let event_stream = SseStream::from_bytes_stream(response.bytes_stream()).boxed(); Ok(event_stream) } @@ -303,7 +303,7 @@ impl StreamableHttpClient for BitFunStreamableHttpClient { match content_type.as_deref() { Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => { - let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed(); + let event_stream = SseStream::from_bytes_stream(response.bytes_stream()).boxed(); Ok(StreamableHttpPostResponse::Sse(event_stream, session_id)) } Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {