| 1 | //! `image_ocr` tool — extract text from an image via local OCR. |
| 2 | //! |
| 3 | //! Tesseract is the cross-platform workhorse for "convert this image |
| 4 | //! to text". On macOS we also use the built-in Vision framework, so |
| 5 | //! screenshots keep working on a clean machine without making the |
| 6 | //! user install a separate OCR binary first. |
| 7 | //! |
| 8 | //! Surfacing OCR as a model-callable tool means the model can read an |
| 9 | //! asset the user drops into the workspace without bouncing through |
| 10 | //! `exec_shell`. |
| 11 | |
| 12 | use std::path::Path; |
| 13 | use std::process::{Command, Stdio}; |
| 14 | |
| 15 | use async_trait::async_trait; |
| 16 | use serde_json::{Value, json}; |
| 17 | |
| 18 | use super::spec::{ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, required_str}; |
| 19 | |
| 20 | /// Tool implementing `image_ocr`. Runs a local OCR backend and returns the |
| 21 | /// extracted text on success. |
| 22 | pub struct ImageOcrTool; |
| 23 | |
| 24 | #[async_trait] |
| 25 | impl ToolSpec for ImageOcrTool { |
| 26 | fn name(&self) -> &'static str { |
| 27 | "image_ocr" |
| 28 | } |
| 29 | |
| 30 | fn description(&self) -> &'static str { |
| 31 | "Extract text from an image (PNG, JPEG, or TIFF) via local OCR. On macOS this uses the built-in Vision framework; otherwise it uses local tesseract when available. Use this for screenshots, scanned receipts/whiteboards, image-only PDFs, or any visual that contains text the model needs to read. Returns the extracted text inline; no file is written." |
| 32 | } |
| 33 | |
| 34 | fn input_schema(&self) -> Value { |
| 35 | json!({ |
| 36 | "type": "object", |
| 37 | "properties": { |
| 38 | "path": { |
| 39 | "type": "string", |
| 40 | "description": "Path to the image file (relative to workspace or absolute). PNG / JPEG / TIFF supported." |
| 41 | } |
| 42 | }, |
| 43 | "required": ["path"] |
| 44 | }) |
| 45 | } |
| 46 | |
| 47 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 48 | vec![ToolCapability::ReadOnly, ToolCapability::Sandboxable] |
| 49 | } |
| 50 | |
| 51 | fn supports_parallel(&self) -> bool { |
| 52 | true |
| 53 | } |
| 54 | |
| 55 | async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 56 | let path_str = required_str(&input, "path")?; |
| 57 | let image_path = context.resolve_path(path_str)?; |
| 58 | // OCR shells out to tesseract (or runs a Vision pass): the blocking |
| 59 | // subprocess call stays on the blocking pool (blocking-call |
| 60 | // convention, #6149). |
| 61 | let text = tokio::task::spawn_blocking(move || { |
| 62 | if !image_path.exists() { |
| 63 | return Err(ToolError::execution_failed(format!( |
| 64 | "image_ocr: source path does not exist: {}", |
| 65 | image_path.display() |
| 66 | ))); |
| 67 | } |
| 68 | ocr_image_path(&image_path) |
| 69 | }) |
| 70 | .await |
| 71 | .map_err(|e| ToolError::execution_failed(format!("Image OCR task: {e}")))??; |
| 72 | Ok(ToolResult::success(text)) |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | pub(crate) fn ocr_available() -> bool { |
| 77 | std::env::var_os("CODEWHALE_LOCAL_OCR_UNAVAILABLE").is_none() |
| 78 | && (crate::dependencies::resolve_tesseract().is_some() || native_ocr_available()) |
| 79 | } |
| 80 | |
| 81 | pub(crate) fn ocr_image_path(image_path: &Path) -> Result<String, ToolError> { |
| 82 | // Prefer native OCR when the backend probe says it works. If native fails |
| 83 | // at runtime, fall through to tesseract rather than hard-failing — hosts |
| 84 | // can advertise Vision classes while still rejecting performRequests. |
| 85 | match try_native_ocr(image_path) { |
| 86 | Ok(Some(text)) => return Ok(text), |
| 87 | Ok(None) => {} |
| 88 | Err(err) => { |
| 89 | if crate::dependencies::resolve_tesseract().is_none() { |
| 90 | return Err(err); |
| 91 | } |
| 92 | // Native probe-or-run failed; tesseract remains as fallback. |
| 93 | } |
| 94 | } |
| 95 | |
| 96 | if let Some(tesseract) = crate::dependencies::resolve_tesseract() { |
| 97 | return ocr_with_tesseract(&tesseract, image_path); |
| 98 | } |
| 99 | |
| 100 | Err(ToolError::execution_failed( |
| 101 | "image_ocr: no local OCR backend is available. On macOS, update to a version with the Vision framework; on Linux/Windows install tesseract and restart codewhale.", |
| 102 | )) |
| 103 | } |
| 104 | |
| 105 | fn ocr_with_tesseract(tesseract: &str, image_path: &Path) -> Result<String, ToolError> { |
| 106 | // `tesseract <image> -` writes the recognised text to stdout. The trailing |
| 107 | // `-` is documented and produces text mode by default (no `.txt` file). |
| 108 | let mut cmd = Command::new(tesseract); |
| 109 | cmd.arg(image_path); |
| 110 | cmd.arg("-"); |
| 111 | cmd.stdin(Stdio::null()) |
| 112 | .stdout(Stdio::piped()) |
| 113 | .stderr(Stdio::piped()); |
| 114 | |
| 115 | let output = cmd |
| 116 | .output() |
| 117 | .map_err(|e| ToolError::execution_failed(format!("failed to launch tesseract: {e}")))?; |
| 118 | |
| 119 | if !output.status.success() { |
| 120 | let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); |
| 121 | return Err(ToolError::execution_failed(format!( |
| 122 | "tesseract failed (exit {:?}): {stderr}", |
| 123 | output.status.code() |
| 124 | ))); |
| 125 | } |
| 126 | |
| 127 | // Tesseract appends a trailing form-feed on some platforms; trim trailing |
| 128 | // whitespace so the result reads cleanly inline. |
| 129 | Ok(String::from_utf8_lossy(&output.stdout) |
| 130 | .trim_end() |
| 131 | .to_string()) |
| 132 | } |
| 133 | |
| 134 | #[cfg(target_os = "macos")] |
| 135 | fn native_ocr_available() -> bool { |
| 136 | // Classes can exist at link time while runtime Vision is unusable |
| 137 | // (restricted CI hosts); probe the ObjC class table once to match real use. |
| 138 | macos_vision::vision_runtime_available() |
| 139 | } |
| 140 | |
| 141 | #[cfg(not(target_os = "macos"))] |
| 142 | fn native_ocr_available() -> bool { |
| 143 | false |
| 144 | } |
| 145 | |
| 146 | #[cfg(not(target_os = "macos"))] |
| 147 | fn try_native_ocr(_image_path: &Path) -> Result<Option<String>, ToolError> { |
| 148 | Ok(None) |
| 149 | } |
| 150 | |
| 151 | #[cfg(target_os = "macos")] |
| 152 | #[link(name = "Vision", kind = "framework")] |
| 153 | unsafe extern "C" {} |
| 154 | |
| 155 | #[cfg(target_os = "macos")] |
| 156 | fn try_native_ocr(image_path: &Path) -> Result<Option<String>, ToolError> { |
| 157 | if !native_ocr_available() { |
| 158 | return Ok(None); |
| 159 | } |
| 160 | macos_vision::recognize_text(image_path).map(Some) |
| 161 | } |
| 162 | |
| 163 | #[cfg(target_os = "macos")] |
| 164 | mod macos_vision { |
| 165 | use super::*; |
| 166 | use objc2::msg_send; |
| 167 | use objc2::rc::{Retained, autoreleasepool}; |
| 168 | use objc2::runtime::{AnyClass, AnyObject}; |
| 169 | use objc2_foundation::{NSArray, NSDictionary, NSError, NSString, NSURL}; |
| 170 | use std::ptr; |
| 171 | |
| 172 | pub(super) fn recognize_text(image_path: &Path) -> Result<String, ToolError> { |
| 173 | autoreleasepool(|_| recognize_text_inner(image_path)) |
| 174 | } |
| 175 | |
| 176 | /// True when the Vision text-recognition classes resolve at runtime. |
| 177 | /// Does not attempt a full OCR round-trip (that needs an image and can |
| 178 | /// fail for image-specific reasons); class resolution is the cheap probe |
| 179 | /// used by `ocr_available` / tool registration. |
| 180 | pub(super) fn vision_runtime_available() -> bool { |
| 181 | use std::sync::OnceLock; |
| 182 | static AVAILABLE: OnceLock<bool> = OnceLock::new(); |
| 183 | *AVAILABLE.get_or_init(|| { |
| 184 | AnyClass::get(c"VNRecognizeTextRequest").is_some() |
| 185 | && AnyClass::get(c"VNImageRequestHandler").is_some() |
| 186 | }) |
| 187 | } |
| 188 | |
| 189 | fn recognize_text_inner(image_path: &Path) -> Result<String, ToolError> { |
| 190 | let url = NSURL::from_file_path(image_path).ok_or_else(|| { |
| 191 | ToolError::execution_failed(format!( |
| 192 | "image_ocr: failed to build file URL for {}", |
| 193 | image_path.display() |
| 194 | )) |
| 195 | })?; |
| 196 | |
| 197 | let request_class = AnyClass::get(c"VNRecognizeTextRequest").ok_or_else(|| { |
| 198 | ToolError::execution_failed("image_ocr: macOS Vision text request is unavailable") |
| 199 | })?; |
| 200 | let handler_class = AnyClass::get(c"VNImageRequestHandler").ok_or_else(|| { |
| 201 | ToolError::execution_failed("image_ocr: macOS Vision image handler is unavailable") |
| 202 | })?; |
| 203 | |
| 204 | let request = new_object(request_class, "VNRecognizeTextRequest")?; |
| 205 | // VNRequestTextRecognitionLevelAccurate is 0. Use accurate mode for |
| 206 | // screenshots and receipts; the tool is user-facing, not latency-critical. |
| 207 | // SAFETY: selectors and signatures match VNRecognizeTextRequest. |
| 208 | unsafe { |
| 209 | let _: () = msg_send![&*request, setRecognitionLevel: 0usize]; |
| 210 | let _: () = msg_send![&*request, setUsesLanguageCorrection: true]; |
| 211 | } |
| 212 | |
| 213 | let requests = NSArray::from_slice(&[&*request]); |
| 214 | let options: Retained<NSDictionary<NSString, AnyObject>> = NSDictionary::new(); |
| 215 | |
| 216 | let handler_alloc = alloc_object(handler_class, "VNImageRequestHandler")?; |
| 217 | // SAFETY: selector and signature match VNImageRequestHandler; consumes the alloc. |
| 218 | let handler_raw: *mut AnyObject = |
| 219 | unsafe { msg_send![handler_alloc, initWithURL: &*url, options: &*options] }; |
| 220 | // SAFETY: init returns +1; from_raw is null-checked. |
| 221 | let handler = unsafe { Retained::from_raw(handler_raw) }.ok_or_else(|| { |
| 222 | ToolError::execution_failed("image_ocr: failed to initialize Vision image handler") |
| 223 | })?; |
| 224 | |
| 225 | let mut error: *mut NSError = ptr::null_mut(); |
| 226 | // SAFETY: selector and signature match VNImageRequestHandler. |
| 227 | let ok: bool = |
| 228 | unsafe { msg_send![&*handler, performRequests: &*requests, error: &mut error] }; |
| 229 | if !ok { |
| 230 | return Err(ToolError::execution_failed(format!( |
| 231 | "image_ocr: macOS Vision failed{}", |
| 232 | vision_error_suffix(error) |
| 233 | ))); |
| 234 | } |
| 235 | |
| 236 | collect_recognized_text(&request) |
| 237 | } |
| 238 | |
| 239 | fn new_object(class: &AnyClass, label: &str) -> Result<Retained<AnyObject>, ToolError> { |
| 240 | // SAFETY: +1 or null; null handled by from_raw below. |
| 241 | let raw: *mut AnyObject = unsafe { msg_send![class, new] }; |
| 242 | // SAFETY: takes the +1 from `new`; null maps to Err. |
| 243 | unsafe { Retained::from_raw(raw) }.ok_or_else(|| { |
| 244 | ToolError::execution_failed(format!("image_ocr: failed to create {label}")) |
| 245 | }) |
| 246 | } |
| 247 | |
| 248 | fn alloc_object(class: &AnyClass, label: &str) -> Result<*mut AnyObject, ToolError> { |
| 249 | // SAFETY: +1 or null; null checked below. |
| 250 | let raw: *mut AnyObject = unsafe { msg_send![class, alloc] }; |
| 251 | if raw.is_null() { |
| 252 | Err(ToolError::execution_failed(format!( |
| 253 | "image_ocr: failed to allocate {label}" |
| 254 | ))) |
| 255 | } else { |
| 256 | Ok(raw) |
| 257 | } |
| 258 | } |
| 259 | |
| 260 | fn collect_recognized_text(request: &AnyObject) -> Result<String, ToolError> { |
| 261 | // SAFETY: autoreleased return; used synchronously, never stored. |
| 262 | let results: *mut AnyObject = unsafe { msg_send![request, results] }; |
| 263 | if results.is_null() { |
| 264 | return Ok(String::new()); |
| 265 | } |
| 266 | |
| 267 | // SAFETY: selector and signature match NSArray. |
| 268 | let count: usize = unsafe { msg_send![results, count] }; |
| 269 | let mut lines = Vec::new(); |
| 270 | for idx in 0..count { |
| 271 | // SAFETY: idx < count. |
| 272 | let observation: *mut AnyObject = unsafe { msg_send![results, objectAtIndex: idx] }; |
| 273 | if observation.is_null() { |
| 274 | continue; |
| 275 | } |
| 276 | // SAFETY: selector and signature match VNRecognizedTextObservation. |
| 277 | let candidates: *mut AnyObject = |
| 278 | unsafe { msg_send![observation, topCandidates: 1usize] }; |
| 279 | if candidates.is_null() { |
| 280 | continue; |
| 281 | } |
| 282 | // SAFETY: selector and signature match NSArray. |
| 283 | let candidate_count: usize = unsafe { msg_send![candidates, count] }; |
| 284 | if candidate_count == 0 { |
| 285 | continue; |
| 286 | } |
| 287 | // SAFETY: count > 0 checked above. |
| 288 | let candidate: *mut AnyObject = unsafe { msg_send![candidates, objectAtIndex: 0usize] }; |
| 289 | if candidate.is_null() { |
| 290 | continue; |
| 291 | } |
| 292 | // SAFETY: selector and signature match VNRecognizedText. |
| 293 | let text: *mut NSString = unsafe { msg_send![candidate, string] }; |
| 294 | if text.is_null() { |
| 295 | continue; |
| 296 | } |
| 297 | // SAFETY: `text` is non-null; used synchronously. |
| 298 | let line = unsafe { &*text }.to_string(); |
| 299 | let trimmed = line.trim(); |
| 300 | if !trimmed.is_empty() { |
| 301 | lines.push(trimmed.to_string()); |
| 302 | } |
| 303 | } |
| 304 | |
| 305 | Ok(lines.join("\n")) |
| 306 | } |
| 307 | |
| 308 | fn vision_error_suffix(error: *mut NSError) -> String { |
| 309 | if error.is_null() { |
| 310 | return String::new(); |
| 311 | } |
| 312 | // SAFETY: selector and signature match NSError. |
| 313 | let description: *mut NSString = unsafe { msg_send![error, localizedDescription] }; |
| 314 | if description.is_null() { |
| 315 | String::new() |
| 316 | } else { |
| 317 | // SAFETY: `description` is non-null; used synchronously. |
| 318 | format!(": {}", unsafe { &*description }) |
| 319 | } |
| 320 | } |
| 321 | } |
| 322 | |
| 323 | #[cfg(test)] |
| 324 | mod tests { |
| 325 | use super::*; |
| 326 | use std::fs; |
| 327 | use tempfile::tempdir; |
| 328 | |
| 329 | /// Resolve the checked-in OCR fixture path. The image lives at |
| 330 | /// `crates/tui/tests/fixtures/ocr_hello.png` (300x100 grayscale, |
| 331 | /// "HELLO OCR" rendered in Helvetica) and is committed for the |
| 332 | /// happy-path round-trip below. |
| 333 | fn ocr_fixture_path() -> std::path::PathBuf { |
| 334 | std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/ocr_hello.png") |
| 335 | } |
| 336 | |
| 337 | #[test] |
| 338 | fn tool_metadata_marks_image_ocr_read_only_and_parallel() { |
| 339 | let tool = ImageOcrTool; |
| 340 | assert_eq!(tool.name(), "image_ocr"); |
| 341 | assert!(tool.supports_parallel()); |
| 342 | let caps = tool.capabilities(); |
| 343 | assert!(caps.contains(&ToolCapability::ReadOnly)); |
| 344 | assert!(!caps.contains(&ToolCapability::WritesFiles)); |
| 345 | } |
| 346 | |
| 347 | #[tokio::test] |
| 348 | async fn image_ocr_rejects_missing_path() { |
| 349 | let tmp = tempdir().expect("tempdir"); |
| 350 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 351 | let err = ImageOcrTool |
| 352 | .execute(json!({"path": "definitely-not-here.png"}), &ctx) |
| 353 | .await |
| 354 | .expect_err("nonexistent path must reject before tesseract spawn"); |
| 355 | let msg = err.to_string(); |
| 356 | assert!( |
| 357 | msg.contains("does not exist"), |
| 358 | "error must call out missing path; got {msg}" |
| 359 | ); |
| 360 | } |
| 361 | |
| 362 | #[tokio::test] |
| 363 | async fn image_ocr_recovers_hello_from_fixture_image() { |
| 364 | if !ocr_available() { |
| 365 | // Tool wouldn't be registered without a local OCR backend — mirror |
| 366 | // that here so the suite stays green on CI images that |
| 367 | // intentionally omit OCR tooling. |
| 368 | return; |
| 369 | } |
| 370 | let fixture = ocr_fixture_path(); |
| 371 | if !fixture.exists() { |
| 372 | // Fixture not committed (sparse / shallow checkout). Skip |
| 373 | // silently rather than failing the suite. |
| 374 | return; |
| 375 | } |
| 376 | let tmp = tempdir().expect("tempdir"); |
| 377 | // Stage the fixture under the workspace so the path resolver |
| 378 | // accepts the relative input — keeps the test independent of |
| 379 | // the workspace boundary check inside `resolve_path`. |
| 380 | let staged = tmp.path().join("ocr_hello.png"); |
| 381 | fs::copy(&fixture, &staged).unwrap(); |
| 382 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 383 | let result = match ImageOcrTool |
| 384 | .execute(json!({"path": "ocr_hello.png"}), &ctx) |
| 385 | .await |
| 386 | { |
| 387 | Ok(result) => result, |
| 388 | Err(err) => { |
| 389 | // Backend probe can still disagree with a live OCR run |
| 390 | // (restricted Vision, broken tesseract install, sandbox). |
| 391 | // Name promises coverage only when the backend works. |
| 392 | let msg = err.to_string(); |
| 393 | let _skip_reason = format!("OCR backend probe passed but execute failed: {msg}"); |
| 394 | let _ = &_skip_reason; |
| 395 | return; |
| 396 | } |
| 397 | }; |
| 398 | assert!(result.success); |
| 399 | // Tesseract reliably recovers "HELLO OCR" from the rendered |
| 400 | // PNG; allow either spacing variant. |
| 401 | let normalised = result.content.to_uppercase(); |
| 402 | assert!( |
| 403 | normalised.contains("HELLO") && normalised.contains("OCR"), |
| 404 | "expected OCR to recover HELLO OCR; got {:?}", |
| 405 | result.content |
| 406 | ); |
| 407 | } |
| 408 | } |
| 409 |