| 1 | //! FIM (Fill-in-the-Middle) edit tool. |
| 2 | //! |
| 3 | //! Reads a file, finds `prefix_anchor` and `suffix_anchor`, calls the active |
| 4 | //! route's `/beta/completions` FIM endpoint, and writes the generated middle |
| 5 | //! content back into the file. The URL is built from the session's own base URL |
| 6 | //! (`crates/tui/src/client.rs:3484`), so this works on any ChatCompletions |
| 7 | //! provider — it is not DeepSeek-specific. |
| 8 | |
| 9 | use std::fs; |
| 10 | |
| 11 | use async_trait::async_trait; |
| 12 | use serde_json::{Value, json}; |
| 13 | use thiserror::Error; |
| 14 | |
| 15 | use crate::client::DeepSeekClient; |
| 16 | |
| 17 | use super::spec::{ |
| 18 | ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, |
| 19 | optional_u64, required_str, |
| 20 | }; |
| 21 | |
| 22 | /// Result of a FIM edit operation |
| 23 | #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] |
| 24 | pub struct FimEditResult { |
| 25 | pub success: bool, |
| 26 | pub path: String, |
| 27 | pub generated_text: String, |
| 28 | pub prefix_end: usize, |
| 29 | pub suffix_start: usize, |
| 30 | pub message: String, |
| 31 | } |
| 32 | |
| 33 | /// Tool for performing Fill-in-the-Middle edits via the active route's FIM API. |
| 34 | /// (`DeepSeekClient` is the historical name of the shared provider client; it is |
| 35 | /// not a DeepSeek-only type.) |
| 36 | pub struct FimEditTool { |
| 37 | pub client: Option<DeepSeekClient>, |
| 38 | pub model: String, |
| 39 | } |
| 40 | |
| 41 | impl FimEditTool { |
| 42 | #[must_use] |
| 43 | pub fn new(client: Option<DeepSeekClient>, model: String) -> Self { |
| 44 | Self { client, model } |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | // === Errors === |
| 49 | |
| 50 | #[derive(Debug, Error)] |
| 51 | enum FimError { |
| 52 | #[error("Prefix anchor not found in file: '{0}'")] |
| 53 | PrefixNotFound(String), |
| 54 | #[error("Suffix anchor not found after prefix anchor: '{0}'")] |
| 55 | SuffixNotFound(String), |
| 56 | #[error("Prefix and suffix anchors overlap (suffix starts at {0}, prefix ends at {1})")] |
| 57 | AnchorsOverlap(usize, usize), |
| 58 | #[error("FIM API call failed: {0}")] |
| 59 | ApiFailed(String), |
| 60 | } |
| 61 | |
| 62 | #[async_trait] |
| 63 | impl ToolSpec for FimEditTool { |
| 64 | fn name(&self) -> &'static str { |
| 65 | "fim_edit" |
| 66 | } |
| 67 | |
| 68 | fn description(&self) -> &'static str { |
| 69 | "Edit a file using Fill-in-the-Middle (FIM) completion. Provide a file path, \ |
| 70 | prefix_anchor (text that appears before the section to replace), and \ |
| 71 | suffix_anchor (text that appears after the section to replace). The tool \ |
| 72 | calls the active route's fill-in-the-middle completion endpoint to \ |
| 73 | generate replacement content." |
| 74 | } |
| 75 | |
| 76 | fn input_schema(&self) -> Value { |
| 77 | json!({ |
| 78 | "type": "object", |
| 79 | "properties": { |
| 80 | "path": { |
| 81 | "type": "string", |
| 82 | "description": "Path to the file to edit (relative to workspace)" |
| 83 | }, |
| 84 | "prefix_anchor": { |
| 85 | "type": "string", |
| 86 | "description": "Text anchor marking the end of the prefix. Everything up to and including this anchor is kept as-is before the generated middle." |
| 87 | }, |
| 88 | "suffix_anchor": { |
| 89 | "type": "string", |
| 90 | "description": "Text anchor marking the start of the suffix. Everything from this anchor onward is kept as-is after the generated middle." |
| 91 | }, |
| 92 | "max_tokens": { |
| 93 | "type": "integer", |
| 94 | "description": "Maximum tokens to generate (default: 1024)" |
| 95 | } |
| 96 | }, |
| 97 | "required": ["path", "prefix_anchor", "suffix_anchor"] |
| 98 | }) |
| 99 | } |
| 100 | |
| 101 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 102 | vec![ |
| 103 | ToolCapability::WritesFiles, |
| 104 | ToolCapability::RequiresApproval, |
| 105 | ] |
| 106 | } |
| 107 | |
| 108 | fn approval_requirement(&self) -> ApprovalRequirement { |
| 109 | ApprovalRequirement::Suggest |
| 110 | } |
| 111 | |
| 112 | async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 113 | let path = required_str(&input, "path")?; |
| 114 | let prefix_anchor = required_str(&input, "prefix_anchor")?; |
| 115 | let suffix_anchor = required_str(&input, "suffix_anchor")?; |
| 116 | let max_tokens = optional_u64(&input, "max_tokens", 1024)?; |
| 117 | |
| 118 | // 1. Read the file |
| 119 | let resolved = context.resolve_path(path)?; |
| 120 | let content = fs::read_to_string(&resolved).map_err(|e| { |
| 121 | ToolError::execution_failed(format!("Failed to read {}: {}", resolved.display(), e)) |
| 122 | })?; |
| 123 | |
| 124 | // 2. Find prefix anchor |
| 125 | let prefix_pos = content.find(prefix_anchor).ok_or_else(|| { |
| 126 | ToolError::execution_failed( |
| 127 | FimError::PrefixNotFound(prefix_anchor.to_string()).to_string(), |
| 128 | ) |
| 129 | })?; |
| 130 | let prefix_end = prefix_pos + prefix_anchor.len(); |
| 131 | |
| 132 | // 3. Find suffix anchor (after prefix anchor) |
| 133 | let suffix_pos = content[prefix_end..].find(suffix_anchor).ok_or_else(|| { |
| 134 | ToolError::execution_failed( |
| 135 | FimError::SuffixNotFound(suffix_anchor.to_string()).to_string(), |
| 136 | ) |
| 137 | })?; |
| 138 | let suffix_start = prefix_end + suffix_pos; |
| 139 | |
| 140 | // 4. Validate anchors don't overlap |
| 141 | if suffix_start < prefix_end { |
| 142 | return Err(ToolError::execution_failed( |
| 143 | FimError::AnchorsOverlap(suffix_start, prefix_end).to_string(), |
| 144 | )); |
| 145 | } |
| 146 | |
| 147 | // 5. Extract prefix and suffix for the FIM API |
| 148 | let fim_prompt = content[..prefix_end].to_string(); |
| 149 | let fim_suffix = content[suffix_start..].to_string(); |
| 150 | |
| 151 | // 6. Call FIM API |
| 152 | let generated_text = match self.client.as_ref() { |
| 153 | Some(client) => client |
| 154 | .fim_completion(&self.model, &fim_prompt, &fim_suffix, max_tokens as u32) |
| 155 | .await |
| 156 | .map_err(|e| { |
| 157 | ToolError::execution_failed(FimError::ApiFailed(e.to_string()).to_string()) |
| 158 | })?, |
| 159 | None => { |
| 160 | return Err(ToolError::execution_failed( |
| 161 | "FIM API client not available".to_string(), |
| 162 | )); |
| 163 | } |
| 164 | }; |
| 165 | |
| 166 | // 7. Build the new content and write it back |
| 167 | let generated_len = generated_text.len(); |
| 168 | let new_content = format!("{fim_prompt}{generated_text}{fim_suffix}"); |
| 169 | crate::utils::write_atomic_workspace(&resolved, new_content.as_bytes()).map_err(|e| { |
| 170 | ToolError::execution_failed(format!("Failed to write {}: {}", resolved.display(), e)) |
| 171 | })?; |
| 172 | |
| 173 | let result = FimEditResult { |
| 174 | success: true, |
| 175 | path: path.to_string(), |
| 176 | generated_text, |
| 177 | prefix_end, |
| 178 | suffix_start, |
| 179 | message: format!( |
| 180 | "FIM edit applied to `{path}`. Generated {generated_len} chars between prefix_anchor end (byte {prefix_end}) and suffix_anchor start (byte {suffix_start}).", |
| 181 | ), |
| 182 | }; |
| 183 | |
| 184 | ToolResult::json(&result).map_err(|e| ToolError::execution_failed(e.to_string())) |
| 185 | } |
| 186 | } |
| 187 | |
| 188 | #[cfg(test)] |
| 189 | mod tests { |
| 190 | use super::*; |
| 191 | |
| 192 | #[test] |
| 193 | fn fim_edit_is_write_capable_but_not_read_only() { |
| 194 | let tool = FimEditTool::new(None, "fim-model".to_string()); |
| 195 | let capabilities = tool.capabilities(); |
| 196 | |
| 197 | assert!(capabilities.contains(&ToolCapability::WritesFiles)); |
| 198 | assert!(!capabilities.contains(&ToolCapability::ReadOnly)); |
| 199 | assert!(!tool.is_read_only()); |
| 200 | } |
| 201 | } |
| 202 |