From fd8c9fdb209c7a078db0209d4ba619ea4441cf75 Mon Sep 17 00:00:00 2001 From: wangfeng Date: Mon, 4 May 2026 18:11:56 -0700 Subject: [PATCH] feat(tools): FIM edit tool + strict tool mode for V4 endpoints (closes #662) --- crates/tui/src/client.rs | 10 +- crates/tui/src/config.rs | 1 + crates/tui/src/core/engine/tool_setup.rs | 1 + crates/tui/src/core/engine/turn_loop.rs | 6 +- crates/tui/src/main.rs | 1 + crates/tui/src/runtime_threads.rs | 1 + crates/tui/src/tools/fim.rs | 198 +++++++++++++++++++++-- crates/tui/src/tools/registry.rs | 7 + 8 files changed, 207 insertions(+), 18 deletions(-) diff --git a/crates/tui/src/client.rs b/crates/tui/src/client.rs index ce23b36c..c60b6abe 100644 --- a/crates/tui/src/client.rs +++ b/crates/tui/src/client.rs @@ -949,10 +949,11 @@ pub(super) fn parse_usage(usage: Option<&Value>) -> Usage { } } -/// Call the DeepSeek `/beta/completions` FIM endpoint. -/// -/// Returns the generated text (the "middle" between `prompt` and `suffix`). -pub async fn fim_completion( +impl DeepSeekClient { + /// Call the DeepSeek `/beta/completions` FIM endpoint. + /// + /// Returns the generated text (the "middle" between `prompt` and `suffix`). + pub async fn fim_completion( &self, model: &str, prompt: &str, @@ -982,6 +983,7 @@ pub async fn fim_completion( .and_then(serde_json::Value::as_str) .ok_or_else(|| anyhow::anyhow!("FIM response missing choices[0].text"))?; Ok(text.to_string()) + } } mod chat; diff --git a/crates/tui/src/config.rs b/crates/tui/src/config.rs index c7a0a6c4..749e48f5 100644 --- a/crates/tui/src/config.rs +++ b/crates/tui/src/config.rs @@ -2028,6 +2028,7 @@ fn merge_config(base: Config, override_cfg: Config) -> Config { per_model: override_cfg.context.per_model.or(base.context.per_model), }, subagents: override_cfg.subagents.or(base.subagents), + strict_tool_mode: override_cfg.strict_tool_mode.or(base.strict_tool_mode), runtime_api: override_cfg.runtime_api.or(base.runtime_api), } } diff --git a/crates/tui/src/core/engine/tool_setup.rs b/crates/tui/src/core/engine/tool_setup.rs index 0d33696e..67ebf236 100644 --- a/crates/tui/src/core/engine/tool_setup.rs +++ b/crates/tui/src/core/engine/tool_setup.rs @@ -33,6 +33,7 @@ impl Engine { builder = builder .with_review_tool(self.deepseek_client.clone(), self.session.model.clone()) .with_rlm_tool(self.deepseek_client.clone(), self.session.model.clone()) + .with_fim_tool(self.deepseek_client.clone(), self.session.model.clone()) .with_user_input_tool() .with_parallel_tool(); diff --git a/crates/tui/src/core/engine/turn_loop.rs b/crates/tui/src/core/engine/turn_loop.rs index 603309d0..5acd3d0a 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -235,7 +235,11 @@ impl Engine { system: self.session.system_prompt.clone(), tools: active_tools.clone(), tool_choice: if active_tools.is_some() { - Some(json!({ "type": "auto" })) + if self.config.strict_tool_mode { + Some(json!("required")) + } else { + Some(json!({ "type": "auto" })) + } } else { None }, diff --git a/crates/tui/src/main.rs b/crates/tui/src/main.rs index 40b664df..281d272b 100644 --- a/crates/tui/src/main.rs +++ b/crates/tui/src/main.rs @@ -3724,6 +3724,7 @@ async fn run_exec_agent( subagent_model_overrides: config.subagent_model_overrides(), memory_enabled: config.memory_enabled(), memory_path: config.memory_path(), + strict_tool_mode: config.strict_tool_mode.unwrap_or(false), goal_objective: None, }; diff --git a/crates/tui/src/runtime_threads.rs b/crates/tui/src/runtime_threads.rs index 33772552..f876ba69 100644 --- a/crates/tui/src/runtime_threads.rs +++ b/crates/tui/src/runtime_threads.rs @@ -1811,6 +1811,7 @@ impl RuntimeThreadManager { subagent_model_overrides: self.config.subagent_model_overrides(), memory_enabled: self.config.memory_enabled(), memory_path: self.config.memory_path(), + strict_tool_mode: self.config.strict_tool_mode.unwrap_or(false), goal_objective: None, }; diff --git a/crates/tui/src/tools/fim.rs b/crates/tui/src/tools/fim.rs index 05028359..b375a508 100644 --- a/crates/tui/src/tools/fim.rs +++ b/crates/tui/src/tools/fim.rs @@ -1,18 +1,190 @@ +//! FIM (Fill-in-the-Middle) edit tool. +//! +//! Reads a file, finds `prefix_anchor` and `suffix_anchor`, calls the +//! DeepSeek `/beta/completions` FIM endpoint, and writes the generated +//! middle content back into the file. + +use std::fs; + +use async_trait::async_trait; use serde_json::{Value, json}; -use crate::tools::spec::{ToolContext, ToolError, ToolResult, ToolSpec}; +use thiserror::Error; -pub struct FimEditTool; +use crate::client::DeepSeekClient; -#[async_trait::async_trait] -impl ToolSpec for FimEditTool { - fn name(&self) -> &'static str { "fim_edit" } - fn description(&self) -> &'static str { "Fill-in-the-middle edit via DeepSeek /beta FIM endpoint" } - fn input_schema(&self) -> Value { - json!({"type":"object","properties":{"path":{"type":"string"},"prefix_anchor":{"type":"string"},"suffix_anchor":{"type":"string"}},"required":["path","prefix_anchor","suffix_anchor"]}) - } - async fn execute(&self, input: Value, _ctx: &ToolContext) -> Result { - let path = input["path"].as_str().ok_or_else(|| ToolError::invalid_input("missing path"))?; - let _ = (path, &input["prefix_anchor"], &input["suffix_anchor"]); - Ok(ToolResult::text("FIM edit stub — wire to /beta endpoint in follow-up")) +use super::spec::{ + ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, + optional_u64, required_str, +}; + +/// Result of a FIM edit operation +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct FimEditResult { + pub success: bool, + pub path: String, + pub generated_text: String, + pub prefix_end: usize, + pub suffix_start: usize, + pub message: String, +} + +/// Tool for performing Fill-in-the-Middle edits via the DeepSeek FIM API. +pub struct FimEditTool { + pub client: Option, + pub model: String, +} + +impl FimEditTool { + #[must_use] + pub fn new(client: Option, model: String) -> Self { + Self { client, model } + } +} + +// === Errors === + +#[derive(Debug, Error)] +enum FimError { + #[error("Prefix anchor not found in file: '{0}'")] + PrefixNotFound(String), + #[error("Suffix anchor not found after prefix anchor: '{0}'")] + SuffixNotFound(String), + #[error("Prefix and suffix anchors overlap (suffix starts at {0}, prefix ends at {1})")] + AnchorsOverlap(usize, usize), + #[error("FIM API call failed: {0}")] + ApiFailed(String), +} + +#[async_trait] +impl ToolSpec for FimEditTool { + fn name(&self) -> &'static str { + "fim_edit" + } + + fn description(&self) -> &'static str { + "Edit a file using Fill-in-the-Middle (FIM) completion. Provide a file path, \ + prefix_anchor (text that appears before the section to replace), and \ + suffix_anchor (text that appears after the section to replace). The tool \ + calls DeepSeek's FIM endpoint to generate replacement content." + } + + fn input_schema(&self) -> Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file to edit (relative to workspace)" + }, + "prefix_anchor": { + "type": "string", + "description": "Text anchor marking the end of the prefix. Everything up to and including this anchor is kept as-is before the generated middle." + }, + "suffix_anchor": { + "type": "string", + "description": "Text anchor marking the start of the suffix. Everything from this anchor onward is kept as-is after the generated middle." + }, + "max_tokens": { + "type": "integer", + "description": "Maximum tokens to generate (default: 1024)" + } + }, + "required": ["path", "prefix_anchor", "suffix_anchor"] + }) + } + + fn capabilities(&self) -> Vec { + vec![ + ToolCapability::ReadOnly, + ToolCapability::WritesFiles, + ToolCapability::RequiresApproval, + ] + } + + fn approval_requirement(&self) -> ApprovalRequirement { + ApprovalRequirement::Suggest + } + + async fn execute(&self, input: Value, context: &ToolContext) -> Result { + let path = required_str(&input, "path")?; + let prefix_anchor = required_str(&input, "prefix_anchor")?; + let suffix_anchor = required_str(&input, "suffix_anchor")?; + let max_tokens = optional_u64(&input, "max_tokens", 1024); + + // 1. Read the file + let resolved = context.resolve_path(path)?; + let content = fs::read_to_string(&resolved).map_err(|e| { + ToolError::execution_failed(format!("Failed to read {}: {}", resolved.display(), e)) + })?; + + // 2. Find prefix anchor + let prefix_pos = content.find(prefix_anchor).ok_or_else(|| { + ToolError::execution_failed( + FimError::PrefixNotFound(prefix_anchor.to_string()).to_string(), + ) + })?; + let prefix_end = prefix_pos + prefix_anchor.len(); + + // 3. Find suffix anchor (after prefix anchor) + let suffix_pos = content[prefix_end..].find(suffix_anchor).ok_or_else(|| { + ToolError::execution_failed( + FimError::SuffixNotFound(suffix_anchor.to_string()).to_string(), + ) + })?; + let suffix_start = prefix_end + suffix_pos; + + // 4. Validate anchors don't overlap + if suffix_start < prefix_end { + return Err(ToolError::execution_failed( + FimError::AnchorsOverlap(suffix_start, prefix_end).to_string(), + )); + } + + // 5. Extract prefix and suffix for the FIM API + let fim_prompt = content[..prefix_end].to_string(); + let fim_suffix = content[suffix_start..].to_string(); + + // 6. Call FIM API + let generated_text = match self.client.as_ref() { + Some(client) => client + .fim_completion(&self.model, &fim_prompt, &fim_suffix, max_tokens as u32) + .await + .map_err(|e| { + ToolError::execution_failed(FimError::ApiFailed(e.to_string()).to_string()) + })?, + None => { + return Err(ToolError::execution_failed( + "FIM API client not available".to_string(), + )); + } + }; + + // 7. Build the new content and write it back + let generated_len = generated_text.len(); + let new_content = format!("{}{}{}", fim_prompt, generated_text, fim_suffix); + fs::write(&resolved, &new_content).map_err(|e| { + ToolError::execution_failed(format!( + "Failed to write {}: {}", + resolved.display(), + e + )) + })?; + + let result = FimEditResult { + success: true, + path: path.to_string(), + generated_text, + prefix_end, + suffix_start, + message: format!( + "FIM edit applied to `{}`. Generated {} chars between prefix_anchor end (byte {}) and suffix_anchor start (byte {}).", + path, + generated_len, + prefix_end, + suffix_start, + ), + }; + + ToolResult::json(&result).map_err(|e| ToolError::execution_failed(e.to_string())) } } diff --git a/crates/tui/src/tools/registry.rs b/crates/tui/src/tools/registry.rs index d85031bd..08a4cfe7 100644 --- a/crates/tui/src/tools/registry.rs +++ b/crates/tui/src/tools/registry.rs @@ -514,6 +514,13 @@ impl ToolRegistryBuilder { self.with_tool(Arc::new(NoteTool)) } + /// Include the FIM (Fill-in-the-Middle) edit tool. + #[must_use] + pub fn with_fim_tool(self, client: Option, model: String) -> Self { + use super::fim::FimEditTool; + self.with_tool(Arc::new(FimEditTool::new(client, model))) + } + /// Include the `remember` tool — model-callable bullet-add into the /// user memory file (#489). Only register when the user has opted /// in to the memory feature; without that, the tool would surface