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