| 1 | //! Publish the rich tool image already admitted to model history through the |
| 2 | //! existing immutable session evidence store for authenticated UI retrieval. |
| 3 | |
| 4 | use std::io; |
| 5 | use std::path::PathBuf; |
| 6 | |
| 7 | use base64::Engine as _; |
| 8 | use base64::engine::general_purpose::STANDARD; |
| 9 | use codewhale_tools::ToolResultContentBlock; |
| 10 | use serde_json::{Value, json}; |
| 11 | |
| 12 | use crate::tools::large_output_router::{ |
| 13 | EVIDENCE_RETENTION_SECS, EvidenceArtifact, EvidenceRetentionState, evidence_is_expired, |
| 14 | publish_evidence_metadata, read_evidence_metadata, unix_millis_now, |
| 15 | }; |
| 16 | use crate::tools::spec::RichToolResult; |
| 17 | |
| 18 | /// No plugin-supplied descriptor may cross this authority boundary. Publication |
| 19 | /// failure only omits the UI image; the validated model content remains intact. |
| 20 | pub(super) async fn project( |
| 21 | mut rich: RichToolResult, |
| 22 | session_id: &str, |
| 23 | call_id: &str, |
| 24 | tool_name: &str, |
| 25 | ) -> RichToolResult { |
| 26 | if let Some(metadata) = rich.result.metadata.as_mut().and_then(Value::as_object_mut) { |
| 27 | metadata.remove("tool_media"); |
| 28 | } |
| 29 | if rich.content_blocks.is_empty() { |
| 30 | return rich; |
| 31 | } |
| 32 | let fallback = rich.result.clone(); |
| 33 | let session_id = session_id.to_owned(); |
| 34 | let call_id = call_id.to_owned(); |
| 35 | let tool_name = tool_name.to_owned(); |
| 36 | // Image decoding and confined disk publication must not block the reactor. |
| 37 | tokio::task::spawn_blocking(move || { |
| 38 | let mut rich = crate::image_attach::bound_rich_tool_result(rich); |
| 39 | if rich.content_blocks.is_empty() { |
| 40 | return rich; |
| 41 | } |
| 42 | match publish_image(&rich.content_blocks[0], &session_id, &call_id, &tool_name) { |
| 43 | Ok(descriptor) => { |
| 44 | let metadata = rich.result.metadata.get_or_insert_with(|| json!({})); |
| 45 | if !metadata.is_object() { |
| 46 | *metadata = json!({}); |
| 47 | } |
| 48 | metadata["tool_media"] = json!([descriptor]); |
| 49 | } |
| 50 | Err(_) => rich.result.content.push_str( |
| 51 | "\n[Tool image preview unavailable: session evidence could not be published.]", |
| 52 | ), |
| 53 | } |
| 54 | rich |
| 55 | }) |
| 56 | .await |
| 57 | .unwrap_or_else(|_| { |
| 58 | let mut rich = RichToolResult::plain(fallback); |
| 59 | rich.result |
| 60 | .content |
| 61 | .push_str("\n[Tool image omitted: image preparation failed.]"); |
| 62 | rich |
| 63 | }) |
| 64 | } |
| 65 | |
| 66 | fn publish_image( |
| 67 | block: &ToolResultContentBlock, |
| 68 | session_id: &str, |
| 69 | call_id: &str, |
| 70 | tool_name: &str, |
| 71 | ) -> io::Result<Value> { |
| 72 | let ToolResultContentBlock::Image { mime_type, data } = block; |
| 73 | let invalid = || io::Error::new(io::ErrorKind::InvalidData, "invalid tool image evidence"); |
| 74 | let bytes = STANDARD.decode(data).map_err(|_| invalid())?; |
| 75 | // The caller passed the shared rich-image validator; retain the shared decode |
| 76 | // guard when deriving dimensions rather than trusting metadata or headers. |
| 77 | let (_, width, height) = |
| 78 | crate::image_attach::decode_and_guard_image(&bytes).map_err(|_| invalid())?; |
| 79 | let identity = serde_json::to_vec(&(session_id, call_id, 0u8)).map_err(|_| invalid())?; |
| 80 | let handle = format!("art_image_{}", crate::hashing::sha256_hex(&identity)); |
| 81 | let storage_path = |
| 82 | PathBuf::from(crate::artifacts::ARTIFACTS_DIR_NAME).join(format!("{handle}.image")); |
| 83 | let digest = crate::hashing::sha256_hex(&bytes); |
| 84 | let now = unix_millis_now(); |
| 85 | let proposed = EvidenceArtifact { |
| 86 | handle: handle.clone(), |
| 87 | digest: digest.clone(), |
| 88 | size_bytes: bytes.len() as u64, |
| 89 | content_type: mime_type.clone(), |
| 90 | tool_name: tool_name.to_owned(), |
| 91 | call_id: call_id.to_owned(), |
| 92 | origin_session: session_id.to_owned(), |
| 93 | generation: 1, |
| 94 | redacted: false, |
| 95 | encoding: "binary".to_owned(), |
| 96 | retention_state: EvidenceRetentionState::Live, |
| 97 | created_at_unix_ms: now, |
| 98 | retain_until_unix_ms: now.saturating_add(EVIDENCE_RETENTION_SECS * 1_000), |
| 99 | storage_path: storage_path.clone(), |
| 100 | }; |
| 101 | let matches_image = |existing: &EvidenceArtifact| { |
| 102 | existing.digest == digest |
| 103 | && existing.size_bytes == bytes.len() as u64 |
| 104 | && existing.content_type == *mime_type |
| 105 | && existing.call_id == call_id |
| 106 | && existing.tool_name == tool_name |
| 107 | && existing.origin_session == session_id |
| 108 | && existing.handle == handle |
| 109 | && existing.storage_path == storage_path |
| 110 | && existing.generation == 1 |
| 111 | && existing.encoding == "binary" |
| 112 | && !existing.redacted |
| 113 | && !evidence_is_expired(existing, unix_millis_now()) |
| 114 | }; |
| 115 | let artifact = match read_evidence_metadata(session_id, &handle) { |
| 116 | Ok(existing) if matches_image(&existing) => existing, |
| 117 | Ok(_) => return Err(invalid()), |
| 118 | Err(error) if error.kind() == io::ErrorKind::NotFound => proposed, |
| 119 | Err(error) => return Err(error), |
| 120 | }; |
| 121 | crate::artifacts::write_session_relative_immutable(session_id, &storage_path, &bytes)?; |
| 122 | if let Err(error) = publish_evidence_metadata(session_id, &artifact) { |
| 123 | // Concurrent identical replays can propose different timestamps. Keep |
| 124 | // the winner's immutable lifetime, never replace or extend it. |
| 125 | if error.kind() != io::ErrorKind::AlreadyExists |
| 126 | || !matches_image(&read_evidence_metadata(session_id, &handle)?) |
| 127 | { |
| 128 | return Err(error); |
| 129 | } |
| 130 | } |
| 131 | Ok(json!({ |
| 132 | "version": 1, "session_id": session_id, "artifact_id": handle, |
| 133 | "tool_call_id": call_id, "media_type": mime_type, "byte_size": bytes.len(), |
| 134 | "sha256": digest, "width": width, "height": height, |
| 135 | })) |
| 136 | } |
| 137 | |
| 138 | #[cfg(test)] |
| 139 | mod tests { |
| 140 | use super::*; |
| 141 | use crate::tools::spec::ToolResult; |
| 142 | |
| 143 | struct ArtifactRoot(Option<PathBuf>); |
| 144 | impl Drop for ArtifactRoot { |
| 145 | fn drop(&mut self) { |
| 146 | crate::artifacts::set_test_artifact_sessions_root(self.0.take()); |
| 147 | } |
| 148 | } |
| 149 | |
| 150 | fn image_result(pixel: u8, mime: &str) -> RichToolResult { |
| 151 | let mut bytes = std::io::Cursor::new(Vec::new()); |
| 152 | image::DynamicImage::ImageRgba8(image::RgbaImage::from_pixel( |
| 153 | 2, |
| 154 | 1, |
| 155 | image::Rgba([pixel, 0, 0, 255]), |
| 156 | )) |
| 157 | .write_to(&mut bytes, image::ImageFormat::Png) |
| 158 | .unwrap(); |
| 159 | let mut result = ToolResult::success("Captured observation".to_owned()); |
| 160 | result.metadata = Some(json!({"tool_media":[{"artifact_id":"forged"}],"keep":true})); |
| 161 | RichToolResult::with_content_blocks( |
| 162 | result, |
| 163 | vec![ToolResultContentBlock::Image { |
| 164 | mime_type: mime.to_owned(), |
| 165 | data: STANDARD.encode(bytes.into_inner()), |
| 166 | }], |
| 167 | ) |
| 168 | } |
| 169 | |
| 170 | #[test] |
| 171 | fn tool_media_publication_is_immutable_and_replay_preserves_manifest() { |
| 172 | // Serialize the shared disk fixture before entering the async runtime. |
| 173 | let _guard = crate::artifacts::TEST_ARTIFACT_SESSIONS_GUARD |
| 174 | .lock() |
| 175 | .unwrap(); |
| 176 | tokio::runtime::Builder::new_current_thread() |
| 177 | .enable_all() |
| 178 | .build() |
| 179 | .unwrap() |
| 180 | .block_on(check_tool_media_publication_is_immutable_and_replay_preserves_manifest()); |
| 181 | } |
| 182 | |
| 183 | async fn check_tool_media_publication_is_immutable_and_replay_preserves_manifest() { |
| 184 | let temp = tempfile::tempdir().unwrap(); |
| 185 | let _root = ArtifactRoot(crate::artifacts::set_test_artifact_sessions_root(Some( |
| 186 | temp.path().to_owned(), |
| 187 | ))); |
| 188 | let first = project( |
| 189 | image_result(3, "image/png"), |
| 190 | "media_owner", |
| 191 | "call/a", |
| 192 | "screenshot", |
| 193 | ) |
| 194 | .await; |
| 195 | let metadata = first.result.metadata.as_ref().unwrap(); |
| 196 | let descriptor = &metadata["tool_media"][0]; |
| 197 | assert_eq!(descriptor["version"], 1); |
| 198 | assert_eq!(descriptor["width"], 2); |
| 199 | assert_eq!(descriptor["height"], 1); |
| 200 | assert_eq!(metadata["keep"], true); |
| 201 | assert_eq!(first.content_blocks.len(), 1); |
| 202 | let handle = descriptor["artifact_id"].as_str().unwrap(); |
| 203 | let manifest_before = read_evidence_metadata("media_owner", handle).unwrap(); |
| 204 | let replay = project( |
| 205 | image_result(3, "image/png"), |
| 206 | "media_owner", |
| 207 | "call/a", |
| 208 | "screenshot", |
| 209 | ) |
| 210 | .await; |
| 211 | assert_eq!(replay.result.metadata, first.result.metadata); |
| 212 | assert_eq!( |
| 213 | serde_json::to_value(read_evidence_metadata("media_owner", handle).unwrap()).unwrap(), |
| 214 | serde_json::to_value(&manifest_before).unwrap() |
| 215 | ); |
| 216 | let (parallel_a, parallel_b) = tokio::join!( |
| 217 | project( |
| 218 | image_result(9, "image/png"), |
| 219 | "media_owner", |
| 220 | "concurrent", |
| 221 | "screenshot" |
| 222 | ), |
| 223 | project( |
| 224 | image_result(9, "image/png"), |
| 225 | "media_owner", |
| 226 | "concurrent", |
| 227 | "screenshot" |
| 228 | ), |
| 229 | ); |
| 230 | assert!( |
| 231 | parallel_a |
| 232 | .result |
| 233 | .metadata |
| 234 | .as_ref() |
| 235 | .unwrap() |
| 236 | .get("tool_media") |
| 237 | .is_some() |
| 238 | ); |
| 239 | assert_eq!(parallel_a.result.metadata, parallel_b.result.metadata); |
| 240 | let conflict = project( |
| 241 | image_result(4, "image/png"), |
| 242 | "media_owner", |
| 243 | "call/a", |
| 244 | "screenshot", |
| 245 | ) |
| 246 | .await; |
| 247 | assert!( |
| 248 | conflict |
| 249 | .result |
| 250 | .metadata |
| 251 | .as_ref() |
| 252 | .unwrap() |
| 253 | .get("tool_media") |
| 254 | .is_none() |
| 255 | ); |
| 256 | assert_eq!( |
| 257 | conflict.content_blocks.len(), |
| 258 | 1, |
| 259 | "UI publication conflict must retain validated model image" |
| 260 | ); |
| 261 | assert!(conflict.result.content.contains("preview unavailable")); |
| 262 | let other_call = project( |
| 263 | image_result(4, "image/png"), |
| 264 | "media_owner", |
| 265 | "call:a", |
| 266 | "screenshot", |
| 267 | ) |
| 268 | .await; |
| 269 | assert_ne!( |
| 270 | other_call.result.metadata.unwrap()["tool_media"][0]["artifact_id"], |
| 271 | descriptor["artifact_id"], |
| 272 | "raw call IDs must not collide after sanitization" |
| 273 | ); |
| 274 | let other_session = project( |
| 275 | image_result(3, "image/png"), |
| 276 | "other_owner", |
| 277 | "call/a", |
| 278 | "screenshot", |
| 279 | ) |
| 280 | .await; |
| 281 | assert_ne!( |
| 282 | other_session.result.metadata.unwrap()["tool_media"][0]["artifact_id"], |
| 283 | descriptor["artifact_id"] |
| 284 | ); |
| 285 | assert!( |
| 286 | !temp.path().join("media_owner.json").exists(), |
| 287 | "publication must not create a second session state authority" |
| 288 | ); |
| 289 | } |
| 290 | |
| 291 | #[test] |
| 292 | fn tool_media_rejects_spoofed_mime_bytes_and_plugin_descriptors() { |
| 293 | // Serialize the shared disk fixture before entering the async runtime. |
| 294 | let _guard = crate::artifacts::TEST_ARTIFACT_SESSIONS_GUARD |
| 295 | .lock() |
| 296 | .unwrap(); |
| 297 | tokio::runtime::Builder::new_current_thread() |
| 298 | .enable_all() |
| 299 | .build() |
| 300 | .unwrap() |
| 301 | .block_on(check_tool_media_rejects_spoofed_mime_bytes_and_plugin_descriptors()); |
| 302 | } |
| 303 | |
| 304 | async fn check_tool_media_rejects_spoofed_mime_bytes_and_plugin_descriptors() { |
| 305 | let temp = tempfile::tempdir().unwrap(); |
| 306 | let _root = ArtifactRoot(crate::artifacts::set_test_artifact_sessions_root(Some( |
| 307 | temp.path().to_owned(), |
| 308 | ))); |
| 309 | let mime = project( |
| 310 | image_result(0, "image/jpeg"), |
| 311 | "media_owner", |
| 312 | "bad-mime", |
| 313 | "screenshot", |
| 314 | ) |
| 315 | .await; |
| 316 | assert!(mime.content_blocks.is_empty()); |
| 317 | assert!(mime.result.metadata.unwrap().get("tool_media").is_none()); |
| 318 | let mut bad = image_result(0, "image/png"); |
| 319 | bad.content_blocks = vec![ToolResultContentBlock::Image { |
| 320 | mime_type: "image/png".to_owned(), |
| 321 | data: STANDARD.encode(b"not an image"), |
| 322 | }]; |
| 323 | let bytes = project(bad, "media_owner", "bad-bytes", "screenshot").await; |
| 324 | assert!(bytes.content_blocks.is_empty()); |
| 325 | assert!(bytes.result.metadata.unwrap().get("tool_media").is_none()); |
| 326 | let mut plain = image_result(0, "image/png"); |
| 327 | plain.content_blocks.clear(); |
| 328 | let plain = project(plain, "media_owner", "no-image", "screenshot").await; |
| 329 | assert!(plain.result.metadata.unwrap().get("tool_media").is_none()); |
| 330 | assert!(!temp.path().join("media_owner").exists()); |
| 331 | } |
| 332 | } |
| 333 |