| 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::ffi::OsString; |
| 13 | use std::io::Write as _; |
| 14 | use std::path::{Path, PathBuf}; |
| 15 | use std::time::Duration; |
| 16 | use tokio_util::sync::CancellationToken; |
| 17 | |
| 18 | use async_trait::async_trait; |
| 19 | use serde_json::{Value, json}; |
| 20 | |
| 21 | use super::spec::{ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, required_str}; |
| 22 | |
| 23 | /// Tool implementing `image_ocr`. Runs a local OCR backend and returns the |
| 24 | /// extracted text on success. |
| 25 | pub struct ImageOcrTool; |
| 26 | |
| 27 | #[async_trait] |
| 28 | impl ToolSpec for ImageOcrTool { |
| 29 | fn name(&self) -> &'static str { |
| 30 | "image_ocr" |
| 31 | } |
| 32 | |
| 33 | fn description(&self) -> &'static str { |
| 34 | "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." |
| 35 | } |
| 36 | |
| 37 | fn input_schema(&self) -> Value { |
| 38 | json!({ |
| 39 | "type": "object", |
| 40 | "properties": { |
| 41 | "path": { |
| 42 | "type": "string", |
| 43 | "description": "Path to the image file (relative to workspace or absolute). PNG / JPEG / TIFF supported." |
| 44 | } |
| 45 | }, |
| 46 | "required": ["path"] |
| 47 | }) |
| 48 | } |
| 49 | |
| 50 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 51 | vec![ToolCapability::ReadOnly, ToolCapability::Sandboxable] |
| 52 | } |
| 53 | |
| 54 | fn supports_parallel(&self) -> bool { |
| 55 | true |
| 56 | } |
| 57 | |
| 58 | async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 59 | let path_str = required_str(&input, "path")?; |
| 60 | // OCR text is file content: the same read guards as `read` apply. |
| 61 | let image_path = |
| 62 | crate::tools::file::resolve_guarded_read_path(context, path_str, "image_ocr")?; |
| 63 | let present = tokio::fs::try_exists(&image_path).await.unwrap_or(false); |
| 64 | if !present { |
| 65 | return Err(ToolError::execution_failed(format!( |
| 66 | "image_ocr: source path does not exist: {}", |
| 67 | image_path.display() |
| 68 | ))); |
| 69 | } |
| 70 | let text = ocr_image_path(&image_path, context).await?; |
| 71 | Ok(ToolResult::success(text)) |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | pub(crate) fn ocr_available() -> bool { |
| 76 | std::env::var_os("CODEWHALE_LOCAL_OCR_UNAVAILABLE").is_none() |
| 77 | && (crate::dependencies::resolve_tesseract().is_some() || native_ocr_available()) |
| 78 | } |
| 79 | |
| 80 | const OCR_TIMEOUT: Duration = Duration::from_secs(120); |
| 81 | const MAX_OCR_TEXT: usize = 16 * 1024 * 1024; |
| 82 | const MAX_OCR_DIAGNOSTIC: usize = 32 * 1024; |
| 83 | const NO_BACKEND: &str = "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."; |
| 84 | |
| 85 | /// Private bounded image snapshot staged before Host startup. Paths, image bytes, |
| 86 | /// full text and process diagnostics stay in Core. Native probing/resolution |
| 87 | /// happens only in the captured Core worker. |
| 88 | pub(crate) struct CapturedOcr { |
| 89 | path: PathBuf, |
| 90 | captured_sha256: String, |
| 91 | _staged: Option<tempfile::NamedTempFile>, |
| 92 | #[cfg(all(test, unix))] |
| 93 | overrides: Option<TestOverrides>, |
| 94 | } |
| 95 | impl CapturedOcr { |
| 96 | async fn capture( |
| 97 | path: &Path, |
| 98 | context: &ToolContext, |
| 99 | deadline: tokio::time::Instant, |
| 100 | ) -> Result<Self, ToolError> { |
| 101 | crate::tools::file::enforce_read_denylist(path, "image_ocr")?; |
| 102 | if crate::tools::file::is_codewhale_credential_path(path) { |
| 103 | return Err(ToolError::permission_denied( |
| 104 | "image_ocr cannot expose Codewhale configuration or credential-store files", |
| 105 | )); |
| 106 | } |
| 107 | if !path.is_absolute() { |
| 108 | return Err(ToolError::execution_failed( |
| 109 | "Image OCR capture requires an authorized absolute path", |
| 110 | )); |
| 111 | } |
| 112 | let path = path.to_path_buf(); |
| 113 | let cancel = context.cancel_token.clone(); |
| 114 | #[cfg(all(test, unix))] |
| 115 | let overrides = TEST_OVERRIDES.with(|slot| slot.borrow().clone()); |
| 116 | let worker = tokio::task::spawn_blocking(move || { |
| 117 | if cancel.as_ref().is_some_and(CancellationToken::is_cancelled) { |
| 118 | return Err(ToolError::cancelled("Image OCR was cancelled")); |
| 119 | } |
| 120 | let root = path.ancestors().last().expect("absolute source has a root"); |
| 121 | // Reuse the Fleet-anchored regular-file reader: no final or parent |
| 122 | // link can redirect the read after the path guards authorize it. |
| 123 | let mut source = crate::fs_confined::open_read(root, &path).map_err(|error| { |
| 124 | ToolError::execution_failed(format!( |
| 125 | "Failed to capture image {}: {error}", |
| 126 | path.display() |
| 127 | )) |
| 128 | })?; |
| 129 | let bytes = crate::tools::file::read_contract_source(&mut source, cancel.as_ref())?; |
| 130 | let captured_sha256 = crate::hashing::sha256_hex(&bytes); |
| 131 | // Reuse the fetched PDF staging contract. The temporary is private |
| 132 | // and retained by the input through Native and Tesseract work. |
| 133 | let mut staged = tempfile::NamedTempFile::new().map_err(|error| { |
| 134 | ToolError::execution_failed(format!("Failed to stage image OCR input: {error}")) |
| 135 | })?; |
| 136 | staged |
| 137 | .write_all(&bytes) |
| 138 | .and_then(|()| staged.flush()) |
| 139 | .map_err(|error| { |
| 140 | ToolError::execution_failed(format!("Failed to stage image OCR input: {error}")) |
| 141 | })?; |
| 142 | if cancel.as_ref().is_some_and(CancellationToken::is_cancelled) { |
| 143 | return Err(ToolError::cancelled("Image OCR was cancelled")); |
| 144 | } |
| 145 | Ok(Self { |
| 146 | path: staged.path().to_path_buf(), |
| 147 | captured_sha256, |
| 148 | _staged: Some(staged), |
| 149 | #[cfg(all(test, unix))] |
| 150 | overrides, |
| 151 | }) |
| 152 | }); |
| 153 | tokio::select! {biased; |
| 154 | ()=wait_cancel(context.cancel_token.as_ref())=>Err(ToolError::cancelled("Image OCR was cancelled")), |
| 155 | ()=tokio::time::sleep_until(deadline)=>Err(ToolError::Timeout {seconds:OCR_TIMEOUT.as_secs()}), |
| 156 | result=worker=>result.map_err(|error|ToolError::execution_failed(format!("Image OCR capture task: {error}")))?, |
| 157 | } |
| 158 | } |
| 159 | #[cfg(all(test, unix))] |
| 160 | pub(crate) fn for_test( |
| 161 | path: &Path, |
| 162 | native: TestNativeOcr, |
| 163 | tesseract: Option<OsString>, |
| 164 | ) -> Self { |
| 165 | Self { |
| 166 | path: path.to_path_buf(), |
| 167 | captured_sha256: crate::hashing::sha256_hex(path.as_os_str().as_encoded_bytes()), |
| 168 | _staged: None, // private broker test port; production always stages bytes |
| 169 | overrides: Some(TestOverrides { native, tesseract }), |
| 170 | } |
| 171 | } |
| 172 | pub(crate) fn digest(&self) -> String { |
| 173 | self.captured_sha256.clone() |
| 174 | } |
| 175 | pub(crate) fn native_step(self) -> OcrNativeStep { |
| 176 | #[cfg(all(test, unix))] |
| 177 | let result = self |
| 178 | .overrides |
| 179 | .as_ref() |
| 180 | .map(|value| (value.native)(&self.path)); |
| 181 | #[cfg(all(test, unix))] |
| 182 | let mut result = result.unwrap_or_else(|| try_native_ocr(&self.path)); |
| 183 | #[cfg(not(all(test, unix)))] |
| 184 | let mut result = try_native_ocr(&self.path); |
| 185 | if result |
| 186 | .as_ref() |
| 187 | .is_ok_and(|text| text.as_ref().is_some_and(|text| text.len() > MAX_OCR_TEXT)) |
| 188 | { |
| 189 | result = Err(ToolError::execution_failed( |
| 190 | "native OCR output exceeded the 16777216 byte safety limit", |
| 191 | )); |
| 192 | } |
| 193 | // Preserve the legacy native-first resolver order. A working native |
| 194 | // backend does not gain a new Tesseract probe or process launch. |
| 195 | let fallback = if matches!(result, Ok(Some(_))) { |
| 196 | None |
| 197 | } else { |
| 198 | #[cfg(all(test, unix))] |
| 199 | let supplied = self.overrides.as_ref().map(|value| value.tesseract.clone()); |
| 200 | #[cfg(all(test, unix))] |
| 201 | { |
| 202 | supplied |
| 203 | .unwrap_or_else(|| crate::dependencies::resolve_tesseract().map(OsString::from)) |
| 204 | } |
| 205 | #[cfg(not(all(test, unix)))] |
| 206 | { |
| 207 | crate::dependencies::resolve_tesseract().map(OsString::from) |
| 208 | } |
| 209 | }; |
| 210 | OcrNativeStep { |
| 211 | input: self, |
| 212 | fallback, |
| 213 | outcome: OcrOutcome::Native(result), |
| 214 | } |
| 215 | } |
| 216 | } |
| 217 | |
| 218 | pub(crate) struct OcrNativeStep { |
| 219 | input: CapturedOcr, |
| 220 | fallback: Option<OsString>, |
| 221 | outcome: OcrOutcome, |
| 222 | } |
| 223 | impl OcrNativeStep { |
| 224 | pub(crate) fn needs_tesseract(&self) -> bool { |
| 225 | self.fallback.is_some() |
| 226 | } |
| 227 | pub(crate) fn digest(&self) -> String { |
| 228 | let mut value = self.input.captured_sha256.as_bytes().to_vec(); |
| 229 | value.push(0); |
| 230 | if let Some(binary) = &self.fallback { |
| 231 | value.extend_from_slice(binary.as_encoded_bytes()); |
| 232 | } |
| 233 | crate::hashing::sha256_hex(&value) |
| 234 | } |
| 235 | pub(crate) fn projection(&self) -> Value { |
| 236 | let status = match &self.outcome { |
| 237 | OcrOutcome::Native(Ok(Some(_))) => "success", |
| 238 | OcrOutcome::Native(Ok(None)) => "unavailable", |
| 239 | _ => "error", |
| 240 | }; |
| 241 | json!({"kind":"ocr_process","state":"native","status":status,"can_fallback":self.needs_tesseract()}) |
| 242 | } |
| 243 | pub(crate) fn finish(self) -> OcrOutcome { |
| 244 | self.outcome |
| 245 | } |
| 246 | pub(crate) async fn tesseract( |
| 247 | self, |
| 248 | cancel: Option<&CancellationToken>, |
| 249 | deadline: tokio::time::Instant, |
| 250 | ) -> OcrOutcome { |
| 251 | let result=async { |
| 252 | if cancel.is_some_and(CancellationToken::is_cancelled) {return Err(ToolError::cancelled("Image OCR was cancelled"));} |
| 253 | if tokio::time::Instant::now()>=deadline {return Err(ToolError::Timeout {seconds:OCR_TIMEOUT.as_secs()});} |
| 254 | let binary=self.fallback.as_ref().ok_or_else(||ToolError::execution_failed("OCR fallback was not admitted"))?; |
| 255 | let mut command=tokio::process::Command::new(binary); |
| 256 | crate::utils::suppress_tokio_console_window(&mut command); |
| 257 | command.arg(&self.input.path).arg("-"); |
| 258 | crate::child_env::apply_to_tokio_command(&mut command,std::iter::empty::<(&str,&str)>()); |
| 259 | let stop=async { tokio::select! {biased;()=wait_cancel(cancel)=>{},()=tokio::time::sleep_until(deadline)=>{},} }; |
| 260 | let run=crate::process_tree::contained_output_with_input_bounded(&mut command,Vec::new(),MAX_OCR_TEXT,MAX_OCR_DIAGNOSTIC,stop).await |
| 261 | .map_err(|error|ToolError::execution_failed(format!("failed to launch tesseract: {error}")))?; |
| 262 | if cancel.is_some_and(CancellationToken::is_cancelled) {return Err(ToolError::cancelled("Image OCR was cancelled"));} |
| 263 | if run.stopped || tokio::time::Instant::now()>=deadline {return Err(ToolError::Timeout {seconds:OCR_TIMEOUT.as_secs()});} |
| 264 | Ok(run.output) |
| 265 | }.await; |
| 266 | OcrOutcome::Tesseract(result) |
| 267 | } |
| 268 | } |
| 269 | |
| 270 | /// Full OCR text/process output is retained privately, never copied to Host. |
| 271 | pub(crate) enum OcrOutcome { |
| 272 | Native(Result<Option<String>, ToolError>), |
| 273 | Tesseract(Result<std::process::Output, ToolError>), |
| 274 | } |
| 275 | impl OcrOutcome { |
| 276 | pub(crate) fn projection(&self) -> Value { |
| 277 | match self { |
| 278 | Self::Native(_) => unreachable!("native projection is stage-bound"), |
| 279 | Self::Tesseract(Err(_)) => { |
| 280 | json!({"kind":"ocr_process","state":"tesseract","status":"fault"}) |
| 281 | } |
| 282 | Self::Tesseract(Ok(output)) => { |
| 283 | json!({"kind":"ocr_process","state":"tesseract","status":"complete","success":output.status.success(),"exit_code":output.status.code()}) |
| 284 | } |
| 285 | } |
| 286 | } |
| 287 | fn into_rust_text(self) -> Result<String, ToolError> { |
| 288 | match self { |
| 289 | Self::Native(Ok(Some(text))) => Ok(text), |
| 290 | Self::Native(Ok(None)) => Err(ToolError::execution_failed(NO_BACKEND)), |
| 291 | Self::Native(Err(error)) | Self::Tesseract(Err(error)) => Err(error), |
| 292 | Self::Tesseract(Ok(output)) if output.status.success() => { |
| 293 | Ok(String::from_utf8_lossy(&output.stdout) |
| 294 | .trim_end() |
| 295 | .to_string()) |
| 296 | } |
| 297 | Self::Tesseract(Ok(output)) => Err(ToolError::execution_failed(format!( |
| 298 | "tesseract failed (exit {:?}): {}", |
| 299 | output.status.code(), |
| 300 | String::from_utf8_lossy(&output.stderr).trim() |
| 301 | ))), |
| 302 | } |
| 303 | } |
| 304 | fn into_host_text(self, result: ToolResult) -> Result<String, ToolError> { |
| 305 | #[derive(serde::Deserialize)] |
| 306 | #[serde(deny_unknown_fields)] |
| 307 | struct Decision { |
| 308 | kind: String, |
| 309 | code: String, |
| 310 | trim_end: bool, |
| 311 | message: Option<String>, |
| 312 | } |
| 313 | let invalid = || { |
| 314 | ToolError::execution_failed( |
| 315 | "OCR Host returned a malformed or inconsistent decision; no Rust fallback was attempted", |
| 316 | ) |
| 317 | }; |
| 318 | if !result.success { |
| 319 | return Err(invalid()); |
| 320 | } |
| 321 | let decision: Decision = |
| 322 | serde_json::from_value(result.metadata.ok_or_else(invalid)?).map_err(|_| invalid())?; |
| 323 | if decision.kind != "ocr_decision" { |
| 324 | return Err(invalid()); |
| 325 | } |
| 326 | match self { |
| 327 | Self::Native(Ok(Some(text))) |
| 328 | if decision.code == "native_success" |
| 329 | && !decision.trim_end |
| 330 | && decision.message.is_none() => |
| 331 | { |
| 332 | Ok(text) |
| 333 | } |
| 334 | Self::Native(Ok(None)) |
| 335 | if decision.code == "no_backend" |
| 336 | && !decision.trim_end |
| 337 | && decision.message.as_deref() == Some(NO_BACKEND) => |
| 338 | { |
| 339 | Err(ToolError::execution_failed(NO_BACKEND)) |
| 340 | } |
| 341 | Self::Native(Err(error)) |
| 342 | if decision.code == "native_error" |
| 343 | && !decision.trim_end |
| 344 | && decision.message.is_none() => |
| 345 | { |
| 346 | Err(error) |
| 347 | } |
| 348 | Self::Tesseract(Err(error)) |
| 349 | if decision.code == "fault" && !decision.trim_end && decision.message.is_none() => |
| 350 | { |
| 351 | Err(error) |
| 352 | } |
| 353 | Self::Tesseract(Ok(output)) |
| 354 | if output.status.success() |
| 355 | && decision.code == "tesseract_success" |
| 356 | && decision.trim_end |
| 357 | && decision.message.is_none() => |
| 358 | { |
| 359 | Ok(String::from_utf8_lossy(&output.stdout) |
| 360 | .trim_end() |
| 361 | .to_string()) |
| 362 | } |
| 363 | Self::Tesseract(Ok(output)) |
| 364 | if !output.status.success() |
| 365 | && decision.code == "execution" |
| 366 | && !decision.trim_end => |
| 367 | { |
| 368 | let prefix = format!("tesseract failed (exit {:?}): ", output.status.code()); |
| 369 | if decision.message.as_deref() != Some(prefix.as_str()) { |
| 370 | return Err(invalid()); |
| 371 | } |
| 372 | Err(ToolError::execution_failed(format!( |
| 373 | "{prefix}{}", |
| 374 | String::from_utf8_lossy(&output.stderr).trim() |
| 375 | ))) |
| 376 | } |
| 377 | _ => Err(invalid()), |
| 378 | } |
| 379 | } |
| 380 | } |
| 381 | |
| 382 | pub(crate) async fn ocr_image_path( |
| 383 | image_path: &Path, |
| 384 | context: &ToolContext, |
| 385 | ) -> Result<String, ToolError> { |
| 386 | let deadline = context |
| 387 | .turn_deadline |
| 388 | .unwrap_or_else(|| tokio::time::Instant::now() + OCR_TIMEOUT); |
| 389 | let input = CapturedOcr::capture(image_path, context, deadline).await?; |
| 390 | if context.features.enabled(crate::features::Feature::OcrHost) { |
| 391 | let mut captured_context = context.clone(); |
| 392 | captured_context.turn_deadline = Some(deadline); |
| 393 | let (outcome, decision) = crate::extension_host::manager() |
| 394 | .execute_ocr(input, &captured_context) |
| 395 | .await?; |
| 396 | return outcome.into_host_text(decision); |
| 397 | } |
| 398 | // Native Vision remains on the existing blocking pool. Cancellation cannot |
| 399 | // preempt framework FFI; an abandoned worker keeps its own captured input. |
| 400 | #[cfg(test)] |
| 401 | let scope = crate::test_support::env_scope_ticket(); |
| 402 | let native = tokio::task::spawn_blocking(move || { |
| 403 | #[cfg(test)] |
| 404 | let _scope = crate::test_support::join_env_scope(scope); |
| 405 | input.native_step() |
| 406 | }); |
| 407 | let step = tokio::select! {biased; |
| 408 | ()=wait_cancel(context.cancel_token.as_ref())=>return Err(ToolError::cancelled("Image OCR was cancelled")), |
| 409 | ()=tokio::time::sleep_until(deadline)=>return Err(ToolError::Timeout {seconds:OCR_TIMEOUT.as_secs()}), |
| 410 | result=native=>result.map_err(|error|ToolError::execution_failed(format!("Image OCR task: {error}")))?, |
| 411 | }; |
| 412 | let output = if step.needs_tesseract() { |
| 413 | step.tesseract(context.cancel_token.as_ref(), deadline) |
| 414 | .await |
| 415 | } else { |
| 416 | step.finish() |
| 417 | }; |
| 418 | if context |
| 419 | .cancel_token |
| 420 | .as_ref() |
| 421 | .is_some_and(CancellationToken::is_cancelled) |
| 422 | { |
| 423 | return Err(ToolError::cancelled("Image OCR was cancelled")); |
| 424 | } |
| 425 | output.into_rust_text() |
| 426 | } |
| 427 | async fn wait_cancel(cancel: Option<&CancellationToken>) { |
| 428 | match cancel { |
| 429 | Some(cancel) => cancel.cancelled().await, |
| 430 | None => std::future::pending::<()>().await, |
| 431 | } |
| 432 | } |
| 433 | |
| 434 | #[cfg(all(test, unix))] |
| 435 | type TestNativeOcr = |
| 436 | std::sync::Arc<dyn Fn(&Path) -> Result<Option<String>, ToolError> + Send + Sync>; |
| 437 | |
| 438 | #[cfg(all(test, unix))] |
| 439 | #[derive(Clone)] |
| 440 | struct TestOverrides { |
| 441 | native: TestNativeOcr, |
| 442 | tesseract: Option<OsString>, |
| 443 | } |
| 444 | #[cfg(all(test, unix))] |
| 445 | thread_local! {static TEST_OVERRIDES:std::cell::RefCell<Option<TestOverrides>>=const {std::cell::RefCell::new(None)};} |
| 446 | |
| 447 | #[cfg(target_os = "macos")] |
| 448 | fn native_ocr_available() -> bool { |
| 449 | // Classes can exist at link time while runtime Vision is unusable |
| 450 | // (restricted CI hosts); probe the ObjC class table once to match real use. |
| 451 | macos_vision::vision_runtime_available() |
| 452 | } |
| 453 | |
| 454 | #[cfg(not(target_os = "macos"))] |
| 455 | fn native_ocr_available() -> bool { |
| 456 | false |
| 457 | } |
| 458 | |
| 459 | #[cfg(not(target_os = "macos"))] |
| 460 | fn try_native_ocr(_image_path: &Path) -> Result<Option<String>, ToolError> { |
| 461 | Ok(None) |
| 462 | } |
| 463 | |
| 464 | #[cfg(target_os = "macos")] |
| 465 | #[link(name = "Vision", kind = "framework")] |
| 466 | unsafe extern "C" {} |
| 467 | |
| 468 | #[cfg(target_os = "macos")] |
| 469 | fn try_native_ocr(image_path: &Path) -> Result<Option<String>, ToolError> { |
| 470 | if !native_ocr_available() { |
| 471 | return Ok(None); |
| 472 | } |
| 473 | macos_vision::recognize_text(image_path).map(Some) |
| 474 | } |
| 475 | |
| 476 | #[cfg(target_os = "macos")] |
| 477 | mod macos_vision { |
| 478 | use super::*; |
| 479 | use objc2::msg_send; |
| 480 | use objc2::rc::{Retained, autoreleasepool}; |
| 481 | use objc2::runtime::{AnyClass, AnyObject}; |
| 482 | use objc2_foundation::{NSArray, NSDictionary, NSError, NSString, NSURL}; |
| 483 | use std::ptr; |
| 484 | |
| 485 | pub(super) fn recognize_text(image_path: &Path) -> Result<String, ToolError> { |
| 486 | autoreleasepool(|_| recognize_text_inner(image_path)) |
| 487 | } |
| 488 | |
| 489 | /// True when the Vision text-recognition classes resolve at runtime. |
| 490 | /// Does not attempt a full OCR round-trip (that needs an image and can |
| 491 | /// fail for image-specific reasons); class resolution is the cheap probe |
| 492 | /// used by `ocr_available` / tool registration. |
| 493 | pub(super) fn vision_runtime_available() -> bool { |
| 494 | use std::sync::OnceLock; |
| 495 | static AVAILABLE: OnceLock<bool> = OnceLock::new(); |
| 496 | *AVAILABLE.get_or_init(|| { |
| 497 | AnyClass::get(c"VNRecognizeTextRequest").is_some() |
| 498 | && AnyClass::get(c"VNImageRequestHandler").is_some() |
| 499 | }) |
| 500 | } |
| 501 | |
| 502 | fn recognize_text_inner(image_path: &Path) -> Result<String, ToolError> { |
| 503 | let url = NSURL::from_file_path(image_path).ok_or_else(|| { |
| 504 | ToolError::execution_failed(format!( |
| 505 | "image_ocr: failed to build file URL for {}", |
| 506 | image_path.display() |
| 507 | )) |
| 508 | })?; |
| 509 | |
| 510 | let request_class = AnyClass::get(c"VNRecognizeTextRequest").ok_or_else(|| { |
| 511 | ToolError::execution_failed("image_ocr: macOS Vision text request is unavailable") |
| 512 | })?; |
| 513 | let handler_class = AnyClass::get(c"VNImageRequestHandler").ok_or_else(|| { |
| 514 | ToolError::execution_failed("image_ocr: macOS Vision image handler is unavailable") |
| 515 | })?; |
| 516 | |
| 517 | let request = new_object(request_class, "VNRecognizeTextRequest")?; |
| 518 | // VNRequestTextRecognitionLevelAccurate is 0. Use accurate mode for |
| 519 | // screenshots and receipts; the tool is user-facing, not latency-critical. |
| 520 | // SAFETY: selectors and signatures match VNRecognizeTextRequest. |
| 521 | unsafe { |
| 522 | let _: () = msg_send![&*request, setRecognitionLevel: 0usize]; |
| 523 | let _: () = msg_send![&*request, setUsesLanguageCorrection: true]; |
| 524 | } |
| 525 | |
| 526 | let requests = NSArray::from_slice(&[&*request]); |
| 527 | let options: Retained<NSDictionary<NSString, AnyObject>> = NSDictionary::new(); |
| 528 | |
| 529 | let handler_alloc = alloc_object(handler_class, "VNImageRequestHandler")?; |
| 530 | // SAFETY: selector and signature match VNImageRequestHandler; consumes the alloc. |
| 531 | let handler_raw: *mut AnyObject = |
| 532 | unsafe { msg_send![handler_alloc, initWithURL: &*url, options: &*options] }; |
| 533 | // SAFETY: init returns +1; from_raw is null-checked. |
| 534 | let handler = unsafe { Retained::from_raw(handler_raw) }.ok_or_else(|| { |
| 535 | ToolError::execution_failed("image_ocr: failed to initialize Vision image handler") |
| 536 | })?; |
| 537 | |
| 538 | let mut error: *mut NSError = ptr::null_mut(); |
| 539 | // SAFETY: selector and signature match VNImageRequestHandler. |
| 540 | let ok: bool = |
| 541 | unsafe { msg_send![&*handler, performRequests: &*requests, error: &mut error] }; |
| 542 | if !ok { |
| 543 | return Err(ToolError::execution_failed(format!( |
| 544 | "image_ocr: macOS Vision failed{}", |
| 545 | vision_error_suffix(error) |
| 546 | ))); |
| 547 | } |
| 548 | |
| 549 | collect_recognized_text(&request) |
| 550 | } |
| 551 | |
| 552 | fn new_object(class: &AnyClass, label: &str) -> Result<Retained<AnyObject>, ToolError> { |
| 553 | // SAFETY: +1 or null; null handled by from_raw below. |
| 554 | let raw: *mut AnyObject = unsafe { msg_send![class, new] }; |
| 555 | // SAFETY: takes the +1 from `new`; null maps to Err. |
| 556 | unsafe { Retained::from_raw(raw) }.ok_or_else(|| { |
| 557 | ToolError::execution_failed(format!("image_ocr: failed to create {label}")) |
| 558 | }) |
| 559 | } |
| 560 | |
| 561 | fn alloc_object(class: &AnyClass, label: &str) -> Result<*mut AnyObject, ToolError> { |
| 562 | // SAFETY: +1 or null; null checked below. |
| 563 | let raw: *mut AnyObject = unsafe { msg_send![class, alloc] }; |
| 564 | if raw.is_null() { |
| 565 | Err(ToolError::execution_failed(format!( |
| 566 | "image_ocr: failed to allocate {label}" |
| 567 | ))) |
| 568 | } else { |
| 569 | Ok(raw) |
| 570 | } |
| 571 | } |
| 572 | |
| 573 | fn collect_recognized_text(request: &AnyObject) -> Result<String, ToolError> { |
| 574 | // SAFETY: autoreleased return; used synchronously, never stored. |
| 575 | let results: *mut AnyObject = unsafe { msg_send![request, results] }; |
| 576 | if results.is_null() { |
| 577 | return Ok(String::new()); |
| 578 | } |
| 579 | |
| 580 | // SAFETY: selector and signature match NSArray. |
| 581 | let count: usize = unsafe { msg_send![results, count] }; |
| 582 | let mut lines = Vec::new(); |
| 583 | for idx in 0..count { |
| 584 | // SAFETY: idx < count. |
| 585 | let observation: *mut AnyObject = unsafe { msg_send![results, objectAtIndex: idx] }; |
| 586 | if observation.is_null() { |
| 587 | continue; |
| 588 | } |
| 589 | // SAFETY: selector and signature match VNRecognizedTextObservation. |
| 590 | let candidates: *mut AnyObject = |
| 591 | unsafe { msg_send![observation, topCandidates: 1usize] }; |
| 592 | if candidates.is_null() { |
| 593 | continue; |
| 594 | } |
| 595 | // SAFETY: selector and signature match NSArray. |
| 596 | let candidate_count: usize = unsafe { msg_send![candidates, count] }; |
| 597 | if candidate_count == 0 { |
| 598 | continue; |
| 599 | } |
| 600 | // SAFETY: count > 0 checked above. |
| 601 | let candidate: *mut AnyObject = unsafe { msg_send![candidates, objectAtIndex: 0usize] }; |
| 602 | if candidate.is_null() { |
| 603 | continue; |
| 604 | } |
| 605 | // SAFETY: selector and signature match VNRecognizedText. |
| 606 | let text: *mut NSString = unsafe { msg_send![candidate, string] }; |
| 607 | if text.is_null() { |
| 608 | continue; |
| 609 | } |
| 610 | // SAFETY: `text` is non-null; used synchronously. |
| 611 | let line = unsafe { &*text }.to_string(); |
| 612 | let trimmed = line.trim(); |
| 613 | if !trimmed.is_empty() { |
| 614 | lines.push(trimmed.to_string()); |
| 615 | } |
| 616 | } |
| 617 | |
| 618 | Ok(lines.join("\n")) |
| 619 | } |
| 620 | |
| 621 | fn vision_error_suffix(error: *mut NSError) -> String { |
| 622 | if error.is_null() { |
| 623 | return String::new(); |
| 624 | } |
| 625 | // SAFETY: selector and signature match NSError. |
| 626 | let description: *mut NSString = unsafe { msg_send![error, localizedDescription] }; |
| 627 | if description.is_null() { |
| 628 | String::new() |
| 629 | } else { |
| 630 | // SAFETY: `description` is non-null; used synchronously. |
| 631 | format!(": {}", unsafe { &*description }) |
| 632 | } |
| 633 | } |
| 634 | } |
| 635 | |
| 636 | #[cfg(test)] |
| 637 | mod tests { |
| 638 | use super::*; |
| 639 | use std::fs; |
| 640 | use tempfile::tempdir; |
| 641 | |
| 642 | /// Resolve the checked-in OCR fixture path. The image lives at |
| 643 | /// `crates/tui/tests/fixtures/ocr_hello.png` (300x100 grayscale, |
| 644 | /// "HELLO OCR" rendered in Helvetica) and is committed for the |
| 645 | /// happy-path round-trip below. |
| 646 | fn ocr_fixture_path() -> std::path::PathBuf { |
| 647 | std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/ocr_hello.png") |
| 648 | } |
| 649 | |
| 650 | #[test] |
| 651 | fn tool_metadata_marks_image_ocr_read_only_and_parallel() { |
| 652 | let tool = ImageOcrTool; |
| 653 | assert_eq!(tool.name(), "image_ocr"); |
| 654 | assert!(tool.supports_parallel()); |
| 655 | let caps = tool.capabilities(); |
| 656 | assert!(caps.contains(&ToolCapability::ReadOnly)); |
| 657 | assert!(!caps.contains(&ToolCapability::WritesFiles)); |
| 658 | } |
| 659 | |
| 660 | #[tokio::test] |
| 661 | async fn image_ocr_rejects_missing_path() { |
| 662 | let tmp = tempdir().expect("tempdir"); |
| 663 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 664 | let err = ImageOcrTool |
| 665 | .execute(json!({"path": "definitely-not-here.png"}), &ctx) |
| 666 | .await |
| 667 | .expect_err("nonexistent path must reject before tesseract spawn"); |
| 668 | let msg = err.to_string(); |
| 669 | assert!( |
| 670 | msg.contains("does not exist"), |
| 671 | "error must call out missing path; got {msg}" |
| 672 | ); |
| 673 | } |
| 674 | |
| 675 | #[tokio::test] |
| 676 | async fn image_ocr_refuses_deny_listed_paths() { |
| 677 | // `.env` is on the default read deny-list, so no global guard setup |
| 678 | // is needed; the refusal precedes any OCR backend. |
| 679 | let tmp = tempdir().expect("tempdir"); |
| 680 | fs::copy(ocr_fixture_path(), tmp.path().join(".env")).unwrap(); |
| 681 | #[cfg(unix)] |
| 682 | std::os::unix::fs::symlink(tmp.path().join(".env"), tmp.path().join("pic.png")).unwrap(); |
| 683 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 684 | let mut paths = vec![".env"]; |
| 685 | if cfg!(unix) { |
| 686 | paths.push("pic.png"); |
| 687 | } |
| 688 | for path in paths { |
| 689 | let err = ImageOcrTool |
| 690 | .execute(json!({ "path": path }), &ctx) |
| 691 | .await |
| 692 | .expect_err("a deny-listed image must be refused"); |
| 693 | assert!( |
| 694 | matches!(err, ToolError::PermissionDenied { .. }), |
| 695 | "{path}: {err:?}" |
| 696 | ); |
| 697 | } |
| 698 | } |
| 699 | |
| 700 | #[tokio::test] |
| 701 | async fn image_ocr_recovers_hello_from_fixture_image() { |
| 702 | if !ocr_available() { |
| 703 | // Tool wouldn't be registered without a local OCR backend — mirror |
| 704 | // that here so the suite stays green on CI images that |
| 705 | // intentionally omit OCR tooling. |
| 706 | return; |
| 707 | } |
| 708 | let fixture = ocr_fixture_path(); |
| 709 | if !fixture.exists() { |
| 710 | // Fixture not committed (sparse / shallow checkout). Skip |
| 711 | // silently rather than failing the suite. |
| 712 | return; |
| 713 | } |
| 714 | let tmp = tempdir().expect("tempdir"); |
| 715 | // Stage the fixture under the workspace so the path resolver |
| 716 | // accepts the relative input — keeps the test independent of |
| 717 | // the workspace boundary check inside `resolve_path`. |
| 718 | let staged = tmp.path().join("ocr_hello.png"); |
| 719 | fs::copy(&fixture, &staged).unwrap(); |
| 720 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 721 | let result = match ImageOcrTool |
| 722 | .execute(json!({"path": "ocr_hello.png"}), &ctx) |
| 723 | .await |
| 724 | { |
| 725 | Ok(result) => result, |
| 726 | Err(err) => { |
| 727 | // Backend probe can still disagree with a live OCR run |
| 728 | // (restricted Vision, broken tesseract install, sandbox). |
| 729 | // Name promises coverage only when the backend works. |
| 730 | let msg = err.to_string(); |
| 731 | let _skip_reason = format!("OCR backend probe passed but execute failed: {msg}"); |
| 732 | let _ = &_skip_reason; |
| 733 | return; |
| 734 | } |
| 735 | }; |
| 736 | assert!(result.success); |
| 737 | // Tesseract reliably recovers "HELLO OCR" from the rendered |
| 738 | // PNG; allow either spacing variant. |
| 739 | let normalised = result.content.to_uppercase(); |
| 740 | assert!( |
| 741 | normalised.contains("HELLO") && normalised.contains("OCR"), |
| 742 | "expected OCR to recover HELLO OCR; got {:?}", |
| 743 | result.content |
| 744 | ); |
| 745 | } |
| 746 | } |
| 747 | |
| 748 | #[cfg(all(test, unix))] |
| 749 | #[path = "image_ocr/host_tests.rs"] |
| 750 | mod host_tests; |
| 751 |