From 85119a38ec47f6ea25d86f9dbc95d0d2caf59e3a Mon Sep 17 00:00:00 2001 From: widat Date: Sat, 25 Apr 2026 11:17:51 +0700 Subject: [PATCH 1/3] feat(execute_command): gate RTK command-line rewrite via models.json use_rtk --- crates/code_assistant/src/agent/runner.rs | 5 + crates/code_assistant/src/mcp/handler.rs | 1 + crates/code_assistant/src/tests/mocks.rs | 2 + crates/code_assistant/src/tools/core/tool.rs | 4 + .../src/tools/impls/execute_command.rs | 149 ++++++++++++++++-- 5 files changed, 151 insertions(+), 10 deletions(-) diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index 2814767b..e07c5ac0 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -1050,6 +1050,7 @@ impl Agent { ui: Some(ui), tool_id: Some(tool_id.to_string()), session_id: None, + model_name: None, permission_handler: None, // Will be handled by sub-agent runner sub_agent_runner, }; @@ -2240,6 +2241,10 @@ impl Agent { ui: Some(self.ui.as_ref()), tool_id: Some(tool_request.id.clone()), session_id: self.session_id.clone(), + model_name: self + .session_model_config + .as_ref() + .map(|config| config.model_name.clone()), permission_handler: self.permission_handler.as_deref(), sub_agent_runner: self.sub_agent_runner.as_deref(), diff --git a/crates/code_assistant/src/mcp/handler.rs b/crates/code_assistant/src/mcp/handler.rs index 09c8bf47..d2146d0c 100644 --- a/crates/code_assistant/src/mcp/handler.rs +++ b/crates/code_assistant/src/mcp/handler.rs @@ -269,6 +269,7 @@ impl MessageHandler { ui: None, tool_id: None, session_id: None, + model_name: None, permission_handler: None, sub_agent_runner: None, }; diff --git a/crates/code_assistant/src/tests/mocks.rs b/crates/code_assistant/src/tests/mocks.rs index fefded22..7d409837 100644 --- a/crates/code_assistant/src/tests/mocks.rs +++ b/crates/code_assistant/src/tests/mocks.rs @@ -231,6 +231,7 @@ pub fn create_test_tool_context<'a>( ui, tool_id, session_id: None, + model_name: None, permission_handler: None, sub_agent_runner: None, } @@ -1137,6 +1138,7 @@ impl ToolTestFixture { ui: self.ui.as_ref().map(|ui| ui as &dyn UserInterface), tool_id: self.tool_id.clone(), session_id: None, + model_name: None, permission_handler: self.permission_handler.as_deref(), sub_agent_runner: None, } diff --git a/crates/code_assistant/src/tools/core/tool.rs b/crates/code_assistant/src/tools/core/tool.rs index fc4fe96d..9f648ff9 100644 --- a/crates/code_assistant/src/tools/core/tool.rs +++ b/crates/code_assistant/src/tools/core/tool.rs @@ -25,6 +25,9 @@ pub struct ToolContext<'a> { /// tool behavior; leave `None` when the diag log is not relevant /// (MCP, tests). pub session_id: Option, + /// Optional active model display name (from models.json) for model-specific + /// tool behavior flags. + pub model_name: Option, /// Optional permission handler for potentially sensitive operations pub permission_handler: Option<&'a dyn PermissionMediator>, @@ -45,6 +48,7 @@ impl<'a> ToolContext<'a> { ui: None, tool_id: None, session_id: None, + model_name: None, permission_handler: None, sub_agent_runner: None, } diff --git a/crates/code_assistant/src/tools/impls/execute_command.rs b/crates/code_assistant/src/tools/impls/execute_command.rs index cf68e4ba..b718bf57 100644 --- a/crates/code_assistant/src/tools/impls/execute_command.rs +++ b/crates/code_assistant/src/tools/impls/execute_command.rs @@ -6,9 +6,12 @@ use crate::ui::streaming::DisplayFragment; use crate::ui::UserInterface; use anyhow::{anyhow, Result}; use command_executor::{SandboxCommandRequest, StreamingCallback}; +use llm::provider_config::ConfigurationSystem; use serde::{Deserialize, Serialize}; use serde_json::json; use std::path::PathBuf; +use std::process::Stdio; +use tracing::{debug, info, warn}; // Input type for the execute_command tool #[derive(Deserialize, Serialize)] @@ -33,6 +36,111 @@ pub struct ExecuteCommandOutput { pub success: bool, } +fn parse_rtk_response_optimized_command(stdout: &[u8]) -> Option { + let response_optimized_command = String::from_utf8_lossy(stdout).trim().to_string(); + if response_optimized_command.is_empty() { + None + } else { + Some(response_optimized_command) + } +} + +fn should_use_rtk_for_model(model_name: Option<&str>) -> bool { + let Some(model_name) = model_name else { + debug!("No active model name in tool context; use_rtk defaults to false"); + return false; + }; + + let config_system = match ConfigurationSystem::load() { + Ok(config) => config, + Err(err) => { + warn!("Failed to load configuration system for use_rtk lookup: {err}"); + return false; + } + }; + + let Some(model_config) = config_system.get_model(model_name) else { + warn!( + "Model '{}' not found in models.json; use_rtk defaults to false", + model_name + ); + return false; + }; + + model_config + .config + .get("use_rtk") + .and_then(|v| v.as_bool()) + .unwrap_or(false) +} + +async fn rewrite_command_line_with_rtk( + original_command_line: &str, + use_rtk: bool, +) -> String { + if !use_rtk { + info!("RTK command-line rewrite disabled, using original command"); + return original_command_line.to_string(); + } + + info!("RTK command-line rewrite enabled, attempting rewrite"); + let rtk_process_output = tokio::process::Command::new("rtk") + .arg("rewrite") + .arg(original_command_line) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await; + + let rtk_process_output = match rtk_process_output { + Ok(output) => output, + Err(err) => { + warn!("RTK command-line rewrite is enabled but failed to execute `rtk`: {err}"); + return original_command_line.to_string(); + } + }; + + match parse_rtk_response_optimized_command(&rtk_process_output.stdout) { + Some(rewritten_command_line) => { + if !rtk_process_output.status.success() { + let stderr = String::from_utf8_lossy(&rtk_process_output.stderr); + info!( + "RTK returned non-zero status ({:?}) but produced a rewritten command line; using it. stderr: {}", + rtk_process_output.status.code(), + stderr.trim() + ); + } + + if rewritten_command_line == original_command_line { + info!("RTK command-line rewrite returned unchanged command"); + return rewritten_command_line; + } + + info!("RTK produced a rewritten command line"); + debug!( + "RTK command-line rewrite details: original={:?} rewritten={:?}", + original_command_line, rewritten_command_line + ); + rewritten_command_line + } + None => { + if rtk_process_output.status.success() { + warn!( + "RTK command-line rewrite returned empty output with success status, using original command" + ); + } else { + let stderr = String::from_utf8_lossy(&rtk_process_output.stderr); + debug!( + "RTK command-line rewrite returned no output (status: {:?}); using original command. stderr: {}", + rtk_process_output.status.code(), + stderr.trim() + ); + } + original_command_line.to_string() + } + } +} + // Render implementation for output formatting impl Render for ExecuteCommandOutput { fn status(&self) -> String { @@ -174,6 +282,9 @@ impl Tool for ExecuteCommandTool { context: &mut ToolContext<'a>, input: &mut Self::Input, ) -> Result { + let use_rtk = should_use_rtk_for_model(context.model_name.as_deref()); + let rewritten_command_line = rewrite_command_line_with_rtk(&input.command_line, use_rtk).await; + // Diag logging: snapshot session_id up-front so every log line in this // call can correlate with the .diag.log file for the session. let diag_session = context.session_id.clone(); @@ -182,8 +293,8 @@ impl Tool for ExecuteCommandTool { crate::session::diag::log( sid, format_args!( - "ExecuteCommandTool::execute: entered tool_id={:?} project={} cmd={:?}", - diag_tool_id, input.project, input.command_line + "ExecuteCommandTool::execute: entered tool_id={:?} project={} cmd={:?} rewritten={:?}", + diag_tool_id, input.project, input.command_line, rewritten_command_line ), ); } @@ -233,7 +344,7 @@ impl Tool for ExecuteCommandTool { tool_id: context.tool_id.as_deref(), tool_name: "execute_command", reason: PermissionRequestReason::ExecuteCommand { - command_line: &input.command_line, + command_line: &rewritten_command_line, working_dir: Some(effective_working_dir.as_path()), }, }) @@ -278,7 +389,7 @@ impl Tool for ExecuteCommandTool { context .command_executor .execute_streaming( - &input.command_line, + &rewritten_command_line, Some(&effective_working_dir), Some(&callback), Some(&sandbox_request), @@ -290,7 +401,7 @@ impl Tool for ExecuteCommandTool { context .command_executor .execute_streaming( - &input.command_line, + &rewritten_command_line, Some(&effective_working_dir), None, Some(&sandbox_request), @@ -327,7 +438,7 @@ impl Tool for ExecuteCommandTool { Ok(ExecuteCommandOutput { project: input.project.clone(), - command_line: input.command_line.clone(), + command_line: rewritten_command_line, working_dir: working_dir_path, output: result.output, success: result.success, @@ -437,14 +548,18 @@ mod tests { let result = tool.execute(&mut context, &mut input).await?; // Verify result - assert_eq!(result.command_line, "ls -la"); + assert!( + result.command_line == "ls -la" || result.command_line.ends_with("ls -la"), + "command should be original or RTK-rewritten variant, got: {}", + result.command_line + ); assert_eq!(result.output, "Command output"); // Match expected output from mock assert!(result.success); // Verify command was executed with correct parameters let commands = fixture.command_executor().get_captured_commands(); assert_eq!(commands.len(), 1); - assert_eq!(commands[0].command_line, "ls -la"); + assert_eq!(commands[0].command_line, result.command_line); assert_eq!(commands[0].working_dir, Some(PathBuf::from("./root/src"))); Ok(()) @@ -474,14 +589,19 @@ mod tests { let result = tool.execute(&mut context, &mut input).await?; // Verify result shows failure - assert_eq!(result.command_line, "rm -rf /tmp/nonexistent"); + assert!( + result.command_line == "rm -rf /tmp/nonexistent" + || result.command_line.ends_with("rm -rf /tmp/nonexistent"), + "command should be original or RTK-rewritten variant, got: {}", + result.command_line + ); assert_eq!(result.output, "Command failed: permission denied"); assert!(!result.success); // Verify command was executed let commands = fixture.command_executor().get_captured_commands(); assert_eq!(commands.len(), 1); - assert_eq!(commands[0].command_line, "rm -rf /tmp/nonexistent"); + assert_eq!(commands[0].command_line, result.command_line); assert_eq!(commands[0].working_dir, Some(PathBuf::from("./root"))); Ok(()) @@ -620,4 +740,13 @@ mod tests { Ok(()) } + + #[test] + fn test_parse_rtk_response_optimized_command() { + assert_eq!( + parse_rtk_response_optimized_command(b"rg -n \"foo\" src\n"), + Some("rg -n \"foo\" src".to_string()) + ); + assert_eq!(parse_rtk_response_optimized_command(b"\n \t"), None); + } } From b81431c76c7aef016e1afde7b913bdb3050c096f Mon Sep 17 00:00:00 2001 From: widat Date: Sat, 25 Apr 2026 17:04:27 +0700 Subject: [PATCH 2/3] feat : move use_rtk above config --- .../src/tools/impls/execute_command.rs | 14 ++++---------- crates/llm/src/provider_config.rs | 4 ++++ 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/crates/code_assistant/src/tools/impls/execute_command.rs b/crates/code_assistant/src/tools/impls/execute_command.rs index b718bf57..c5be74b9 100644 --- a/crates/code_assistant/src/tools/impls/execute_command.rs +++ b/crates/code_assistant/src/tools/impls/execute_command.rs @@ -67,17 +67,10 @@ fn should_use_rtk_for_model(model_name: Option<&str>) -> bool { return false; }; - model_config - .config - .get("use_rtk") - .and_then(|v| v.as_bool()) - .unwrap_or(false) + model_config.use_rtk } -async fn rewrite_command_line_with_rtk( - original_command_line: &str, - use_rtk: bool, -) -> String { +async fn rewrite_command_line_with_rtk(original_command_line: &str, use_rtk: bool) -> String { if !use_rtk { info!("RTK command-line rewrite disabled, using original command"); return original_command_line.to_string(); @@ -283,7 +276,8 @@ impl Tool for ExecuteCommandTool { input: &mut Self::Input, ) -> Result { let use_rtk = should_use_rtk_for_model(context.model_name.as_deref()); - let rewritten_command_line = rewrite_command_line_with_rtk(&input.command_line, use_rtk).await; + let rewritten_command_line = + rewrite_command_line_with_rtk(&input.command_line, use_rtk).await; // Diag logging: snapshot session_id up-front so every log line in this // call can correlate with the .diag.log file for the session. diff --git a/crates/llm/src/provider_config.rs b/crates/llm/src/provider_config.rs index 2bc96f75..bc5f264b 100644 --- a/crates/llm/src/provider_config.rs +++ b/crates/llm/src/provider_config.rs @@ -25,6 +25,10 @@ pub struct ModelConfig { pub provider: String, /// Model ID within the provider pub id: String, + /// Whether command output should be optimized through RTK for this model. + /// Defaults to false when omitted in models.json. + #[serde(default)] + pub use_rtk: bool, /// Model-specific configuration pub config: serde_json::Value, /// Maximum context window supported by the model (token count) From 1876293ad0ec1d47f0c1111ebdbd95cfcda5466f Mon Sep 17 00:00:00 2001 From: widat Date: Mon, 27 Apr 2026 21:16:27 +0700 Subject: [PATCH 3/3] fix: delete imperative function --- .../src/tools/impls/execute_command.rs | 159 ++++++++---------- 1 file changed, 67 insertions(+), 92 deletions(-) diff --git a/crates/code_assistant/src/tools/impls/execute_command.rs b/crates/code_assistant/src/tools/impls/execute_command.rs index f327c810..465aa6d9 100644 --- a/crates/code_assistant/src/tools/impls/execute_command.rs +++ b/crates/code_assistant/src/tools/impls/execute_command.rs @@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize}; use serde_json::json; use std::path::PathBuf; use std::process::Stdio; -use tracing::{debug, info, warn}; +use tracing::{debug, warn}; // Input type for the execute_command tool #[derive(Deserialize, Serialize)] @@ -70,70 +70,6 @@ fn should_use_rtk_for_model(model_name: Option<&str>) -> bool { model_config.use_rtk } -async fn rewrite_command_line_with_rtk(original_command_line: &str, use_rtk: bool) -> String { - if !use_rtk { - info!("RTK command-line rewrite disabled, using original command"); - return original_command_line.to_string(); - } - - info!("RTK command-line rewrite enabled, attempting rewrite"); - let rtk_process_output = tokio::process::Command::new("rtk") - .arg("rewrite") - .arg(original_command_line) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .output() - .await; - - let rtk_process_output = match rtk_process_output { - Ok(output) => output, - Err(err) => { - warn!("RTK command-line rewrite is enabled but failed to execute `rtk`: {err}"); - return original_command_line.to_string(); - } - }; - - match parse_rtk_response_optimized_command(&rtk_process_output.stdout) { - Some(rewritten_command_line) => { - if !rtk_process_output.status.success() { - let stderr = String::from_utf8_lossy(&rtk_process_output.stderr); - info!( - "RTK returned non-zero status ({:?}) but produced a rewritten command line; using it. stderr: {}", - rtk_process_output.status.code(), - stderr.trim() - ); - } - - if rewritten_command_line == original_command_line { - info!("RTK command-line rewrite returned unchanged command"); - return rewritten_command_line; - } - - info!("RTK produced a rewritten command line"); - debug!( - "RTK command-line rewrite details: original={:?} rewritten={:?}", - original_command_line, rewritten_command_line - ); - rewritten_command_line - } - None => { - if rtk_process_output.status.success() { - warn!( - "RTK command-line rewrite returned empty output with success status, using original command" - ); - } else { - let stderr = String::from_utf8_lossy(&rtk_process_output.stderr); - debug!( - "RTK command-line rewrite returned no output (status: {:?}); using original command. stderr: {}", - rtk_process_output.status.code(), - stderr.trim() - ); - } - original_command_line.to_string() - } - } -} - // Render implementation for output formatting impl Render for ExecuteCommandOutput { fn status(&self) -> String { @@ -275,23 +211,62 @@ impl Tool for ExecuteCommandTool { context: &mut ToolContext<'a>, input: &mut Self::Input, ) -> Result { + let use_rtk = should_use_rtk_for_model(context.model_name.as_deref()); - let rewritten_command_line = - rewrite_command_line_with_rtk(&input.command_line, use_rtk).await; - - // Diag logging: snapshot session_id up-front so every log line in this - // call can correlate with the .diag.log file for the session. - let diag_session = context.session_id.clone(); - let diag_tool_id = context.tool_id.clone(); - if let Some(sid) = diag_session.as_deref() { - crate::session::diag::log( - sid, - format_args!( - "ExecuteCommandTool::execute: entered tool_id={:?} project={} cmd={:?} rewritten={:?}", - diag_tool_id, input.project, input.command_line, rewritten_command_line - ), - ); - } + let original_command_line = input.command_line.clone(); + let effective_command_line = if !use_rtk { + original_command_line.clone() + } else { + match tokio::process::Command::new("rtk") + .arg("rewrite") + .arg(&original_command_line) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await + { + Ok(rtk_process_output) => { + if let Some(output) = + parse_rtk_response_optimized_command(&rtk_process_output.stdout) + { + if !rtk_process_output.status.success() { + let stderr = String::from_utf8_lossy(&rtk_process_output.stderr); + debug!( + "RTK returned non-zero status ({:?}) but produced a rewritten command line; using it. stderr: {}", + rtk_process_output.status.code(), + stderr.trim() + ); + } + + if output != original_command_line { + debug!( + "RTK command-line rewrite details: original={:?} rewritten={:?}", + original_command_line, output + ); + } + output + } else { + if rtk_process_output.status.success() { + warn!( + "RTK command-line rewrite returned empty output with success status, using original command" + ); + } else { + let stderr = String::from_utf8_lossy(&rtk_process_output.stderr); + debug!( + "RTK command-line rewrite returned no output (status: {:?}); using original command. stderr: {}", + rtk_process_output.status.code(), + stderr.trim() + ); + } + original_command_line.clone() + } + } + Err(err) => { + warn!("RTK command-line rewrite is enabled but failed to execute `rtk`: {err}"); + original_command_line.clone() + } + } + }; // Get explorer for the specified project let explorer = context @@ -334,14 +309,14 @@ impl Tool for ExecuteCommandTool { })?; let decision = handler - .request_permission(PermissionRequest { - tool_id: context.tool_id.as_deref(), - tool_name: "execute_command", - reason: PermissionRequestReason::ExecuteCommand { - command_line: &rewritten_command_line, - working_dir: Some(effective_working_dir.as_path()), - }, - }) + .request_permission(PermissionRequest { + tool_id: context.tool_id.as_deref(), + tool_name: "execute_command", + reason: PermissionRequestReason::ExecuteCommand { + command_line: &effective_command_line, + working_dir: Some(effective_working_dir.as_path()), + }, + }) .await?; match decision { @@ -372,7 +347,7 @@ impl Tool for ExecuteCommandTool { context .command_executor .execute_streaming( - &rewritten_command_line, + &effective_command_line, Some(&effective_working_dir), Some(&callback), Some(&sandbox_request), @@ -384,7 +359,7 @@ impl Tool for ExecuteCommandTool { context .command_executor .execute_streaming( - &rewritten_command_line, + &effective_command_line, Some(&effective_working_dir), None, Some(&sandbox_request), @@ -395,7 +370,7 @@ impl Tool for ExecuteCommandTool { Ok(ExecuteCommandOutput { project: input.project.clone(), - command_line: rewritten_command_line, + command_line: effective_command_line, working_dir: working_dir_path, output: result.output, success: result.success,