1use std::collections::{HashMap, HashSet};
2use std::path::Path;
3
4use anyhow::{Context, Result, bail};
5use sqlx::{Row, SqliteConnection};
6
7use super::wire::{AttachmentAddPayload, BlobUploadContract, ChangeWire};
8use crate::attachments::decode::validate_image;
9use crate::attachments::lifecycle::{
10 ByteCount, LifecyclePolicy, PruneSummary, SystemClock, acquire_lease, ensure_local_capacity,
11 prune, release_lease, release_reservation, reserve_upload,
12};
13use crate::attachments::storage::{
14 blob_inventory_row, object_path, sha256_hex, store_validated_blob,
15};
16use crate::change_log::op_type;
17use crate::db::Database;
18
19#[derive(Debug, Clone)]
20pub(super) struct MissingLocalBlob {
21 pub sha256: String,
22 pub byte_size: i64,
23 pub media_type: String,
24 pub width: Option<i64>,
25 pub height: Option<i64>,
26}
27
28#[derive(Debug)]
29pub(super) struct MissingBlobPage {
30 pub blobs: Vec<MissingLocalBlob>,
31 pub has_more: bool,
32}
33
34#[derive(Debug)]
35pub(super) struct PreparedBlobUpload {
36 pub bytes: Vec<u8>,
37 pub lease_id: String,
38}
39
40#[derive(Debug)]
41pub struct ServerBlobDownload {
42 pub bytes: Vec<u8>,
43}
44
45pub(super) fn attachment_blob_contract(change: &ChangeWire) -> Result<Option<BlobUploadContract>> {
46 if change.op_type != op_type::ATTACHMENT_ADD {
47 return Ok(None);
48 }
49 let payload = AttachmentAddPayload::from_change(change)?;
50 Ok(Some(BlobUploadContract {
51 workspace_id: payload.workspace_id,
52 sha256: payload.sha256,
53 byte_size: payload.byte_size,
54 media_type: payload.media_type,
55 width: payload
56 .width
57 .context("validated attachment payload missing width")?,
58 height: payload
59 .height
60 .context("validated attachment payload missing height")?,
61 }))
62}
63
64pub(super) fn attachment_blob_contracts(changes: &[ChangeWire]) -> Result<Vec<BlobUploadContract>> {
65 let mut contracts = Vec::new();
66 for change in changes {
67 if let Some(contract) = attachment_blob_contract(change)? {
68 contracts.push(contract);
69 }
70 }
71 Ok(contracts)
72}
73
74fn same_blob_content_contract(left: &BlobUploadContract, right: &BlobUploadContract) -> bool {
75 left.byte_size == right.byte_size
76 && left.media_type == right.media_type
77 && left.width == right.width
78 && left.height == right.height
79}
80
81pub(super) fn unique_blob_content_contracts(
82 contracts: &[BlobUploadContract],
83) -> Result<Vec<BlobUploadContract>> {
84 let mut seen = HashMap::new();
85 let mut unique = Vec::new();
86 for contract in contracts {
87 if let Some(existing) = seen.get(&contract.sha256) {
88 if !same_blob_content_contract(existing, contract) {
89 bail!("error blob-inventory-metadata-mismatch");
90 }
91 continue;
92 }
93 seen.insert(contract.sha256.clone(), contract.clone());
94 unique.push(contract.clone());
95 }
96 Ok(unique)
97}
98
99pub(super) fn unique_blob_admission_contracts(
100 contracts: &[BlobUploadContract],
101) -> Vec<BlobUploadContract> {
102 let mut seen = HashSet::new();
103 contracts
104 .iter()
105 .filter(|contract| seen.insert((contract.workspace_id.as_str(), contract.sha256.as_str())))
106 .cloned()
107 .collect()
108}
109
110pub(super) fn unique_blob_contracts(changes: &[ChangeWire]) -> Result<Vec<BlobUploadContract>> {
111 let contracts = attachment_blob_contracts(changes)?;
112 unique_blob_content_contracts(&contracts)?;
113 Ok(unique_blob_admission_contracts(&contracts))
114}
115
116pub(super) fn contract_by_hash(
117 contracts: &[BlobUploadContract],
118 sha256: &str,
119) -> Result<BlobUploadContract> {
120 contracts
121 .iter()
122 .find(|contract| contract.sha256 == sha256)
123 .cloned()
124 .context("payload missing attachment blob contract")
125}
126
127impl Database {
128 pub(super) async fn prepare_blob_upload(
129 &self,
130 blob_dir: &Path,
131 contract: &BlobUploadContract,
132 ) -> Result<PreparedBlobUpload> {
133 let (row, lease_id) = {
134 let mut conn = self.acquire_writer().await?;
135 let row = blob_inventory_row(&mut conn, &contract.sha256)
136 .await?
137 .filter(|row| row.available)
138 .context("error attachment-blob-local-missing")?;
139 let lease_id =
140 acquire_lease(&mut conn, &contract.sha256, "transfer", &SystemClock).await?;
141 (row, lease_id)
142 };
143 let result = async {
144 let bytes = tokio::fs::read(object_path(blob_dir, &contract.sha256)?)
145 .await
146 .context("error attachment-blob-local-missing")?;
147 if sha256_hex(&bytes) != contract.sha256
148 || row.byte_size
149 != i64::try_from(bytes.len()).context("attachment bytes exceed i64")?
150 {
151 bail!("error attachment-blob-local-invalid");
152 }
153 let validated = validate_image(bytes.clone(), Some(row.media_type.clone())).await?;
154 if (validated.facts.width, validated.facts.height) != (contract.width, contract.height)
155 || validated.facts.media_type != contract.media_type
156 || contract.byte_size != row.byte_size
157 {
158 bail!("error attachment-blob-local-invalid");
159 }
160 Ok(PreparedBlobUpload {
161 bytes,
162 lease_id: lease_id.clone(),
163 })
164 }
165 .await;
166 if result.is_err() {
167 let _ = self.finish_blob_upload(&lease_id).await;
168 }
169 result
170 }
171
172 pub(super) async fn finish_blob_upload(&self, lease_id: &str) -> Result<()> {
173 let mut conn = self.acquire_writer().await?;
174 release_lease(&mut conn, lease_id).await
175 }
176
177 pub(super) async fn missing_local_blob_page(&self, limit: usize) -> Result<MissingBlobPage> {
178 let mut conn = self.acquire_reader().await?;
179 missing_local_blob_page(&mut conn, limit).await
180 }
181
182 pub(super) async fn missing_local_blob_counts(&self) -> Result<ByteCount> {
183 let mut conn = self.acquire_reader().await?;
184 missing_local_blob_counts(&mut conn).await
185 }
186
187 pub async fn missing_sync_attachment_counts(&self) -> Result<ByteCount> {
188 self.missing_local_blob_counts().await
189 }
190
191 pub(super) async fn store_downloaded_blob(
192 &self,
193 blob_dir: &Path,
194 policy: LifecyclePolicy,
195 blob: &MissingLocalBlob,
196 bytes: Vec<u8>,
197 ) -> Result<()> {
198 let expected = usize::try_from(blob.byte_size).context("attachment bytes exceed usize")?;
199 if bytes.len() != expected || sha256_hex(&bytes) != blob.sha256 {
200 bail!("error attachment-blob-remote-invalid");
201 }
202 let validated = validate_image(bytes, Some(blob.media_type.clone()))
203 .await
204 .context("error attachment-blob-remote-invalid")?;
205 if (blob.width, blob.height) != (Some(validated.facts.width), Some(validated.facts.height))
206 {
207 bail!("error attachment-blob-remote-invalid");
208 }
209 let mut conn = self.acquire_writer().await?;
210 let reservation = ensure_local_capacity(
211 &mut conn,
212 blob_dir,
213 &blob.sha256,
214 blob.byte_size,
215 policy,
216 &SystemClock,
217 )
218 .await?;
219 let result = store_validated_blob(&mut conn, blob_dir, validated).await;
220 if let Some(reservation) = reservation {
221 release_reservation(&mut conn, &reservation).await?;
222 }
223 result.map(|_| ())
224 }
225
226 pub async fn prepare_server_blob_uploads(
227 &self,
228 blob_dir: &Path,
229 policy: LifecyclePolicy,
230 blobs: &[BlobUploadContract],
231 ) -> Result<Vec<String>> {
232 super::wire::validate_blob_contracts(blobs)?;
233 let mut conn = self.acquire_writer().await?;
234 let mut missing = Vec::new();
235 let mut missing_hashes = HashSet::new();
236 for blob in blobs {
237 if blob_available(&mut conn, blob_dir, &blob.sha256).await? {
238 reserve_upload(
239 &mut conn,
240 &blob.workspace_id,
241 &blob.sha256,
242 blob.byte_size,
243 policy.quota_bytes,
244 &SystemClock,
245 )
246 .await?;
247 } else if missing_hashes.insert(blob.sha256.as_str()) {
248 missing.push(blob.sha256.clone());
249 }
250 }
251 Ok(missing)
252 }
253
254 pub async fn store_server_blob(
255 &self,
256 blob_dir: &Path,
257 policy: LifecyclePolicy,
258 contract: &BlobUploadContract,
259 bytes: Vec<u8>,
260 ) -> Result<()> {
261 super::wire::validate_blob_contracts(std::slice::from_ref(contract))?;
262 if sha256_hex(&bytes) != contract.sha256
263 || usize::try_from(contract.byte_size).ok() != Some(bytes.len())
264 {
265 bail!("error blob-hash-or-size-mismatch");
266 }
267 let validated = validate_image(bytes, Some(contract.media_type.clone()))
268 .await
269 .context("error blob-validation-failed")?;
270 if (validated.facts.width, validated.facts.height) != (contract.width, contract.height) {
271 bail!("error blob-validation-failed");
272 }
273 let mut conn = self.acquire_writer().await?;
274 let reservation = reserve_upload(
275 &mut conn,
276 &contract.workspace_id,
277 &contract.sha256,
278 contract.byte_size,
279 policy.quota_bytes,
280 &SystemClock,
281 )
282 .await?;
283 let result = store_validated_blob(&mut conn, blob_dir, validated).await;
284 if result.is_err()
285 && let Some(reservation) = reservation
286 {
287 release_reservation(&mut conn, &reservation).await?;
288 }
289 result.map(|_| ())
290 }
291
292 pub async fn read_server_blob(
293 &self,
294 blob_dir: &Path,
295 sha256: &str,
296 ) -> Result<Option<ServerBlobDownload>> {
297 super::wire::validate_blob_hashes(&[sha256.to_string()])?;
298 let lease = {
299 let mut conn = self.acquire_writer().await?;
300 if !blob_available(&mut conn, blob_dir, sha256).await? {
301 return Ok(None);
302 }
303 acquire_lease(&mut conn, sha256, "transfer", &SystemClock).await?
304 };
305 let result = tokio::fs::read(object_path(blob_dir, sha256)?).await;
306 let mut conn = self.acquire_writer().await?;
307 release_lease(&mut conn, &lease).await?;
308 Ok(Some(ServerBlobDownload {
309 bytes: result.context("error blob-read-failed")?,
310 }))
311 }
312
313 pub async fn maintain_server_blobs(
314 &self,
315 blob_dir: &Path,
316 policy: LifecyclePolicy,
317 ) -> Result<PruneSummary> {
318 let mut conn = self.acquire_writer().await?;
319 prune(&mut conn, blob_dir, policy, true, &SystemClock).await
320 }
321}
322
323async fn missing_local_blob_page(
324 conn: &mut SqliteConnection,
325 limit: usize,
326) -> Result<MissingBlobPage> {
327 let rows = sqlx::query(
328 "SELECT ta.sha256, MAX(ta.byte_size) AS byte_size,
329 MAX(ta.media_type) AS media_type, MAX(ta.width) AS width,
330 MAX(ta.height) AS height
331 FROM task_attachments AS ta INDEXED BY idx_task_attachments_live_sha256
332 JOIN tasks t ON t.workspace_id = ta.workspace_id AND t.id = ta.task_id
333 JOIN changes c ON c.change_id = ta.created_by_change_id
334 AND c.op_type = 'attachment_add' AND c.server_seq IS NOT NULL
335 LEFT JOIN blob_inventory bi ON bi.sha256 = ta.sha256
336 WHERE ta.deleted = 0 AND t.deleted = 0
337 GROUP BY ta.sha256
338 HAVING MAX(COALESCE(bi.available, 0)) = 0
339 ORDER BY ta.sha256 LIMIT ?",
340 )
341 .bind(i64::try_from(limit.saturating_add(1))?)
342 .fetch_all(&mut *conn)
343 .await?;
344 let has_more = rows.len() > limit;
345 let blobs = rows
346 .into_iter()
347 .take(limit)
348 .map(|row| MissingLocalBlob {
349 sha256: row.get("sha256"),
350 byte_size: row.get("byte_size"),
351 media_type: row.get("media_type"),
352 width: row.get("width"),
353 height: row.get("height"),
354 })
355 .collect();
356 Ok(MissingBlobPage { blobs, has_more })
357}
358
359pub(super) async fn missing_local_blob_counts(conn: &mut SqliteConnection) -> Result<ByteCount> {
360 let (count, bytes): (i64, i64) = sqlx::query_as(
361 "SELECT COUNT(*), COALESCE(SUM(byte_size), 0) FROM (
362 SELECT ta.sha256, MAX(ta.byte_size) AS byte_size
363 FROM task_attachments AS ta INDEXED BY idx_task_attachments_live_sha256
364 JOIN tasks t ON t.workspace_id = ta.workspace_id AND t.id = ta.task_id
365 JOIN changes c ON c.change_id = ta.created_by_change_id
366 AND c.op_type = 'attachment_add' AND c.server_seq IS NOT NULL
367 LEFT JOIN blob_inventory bi ON bi.sha256 = ta.sha256
368 WHERE ta.deleted = 0 AND t.deleted = 0
369 GROUP BY ta.sha256
370 HAVING MAX(COALESCE(bi.available, 0)) = 0
371 )",
372 )
373 .fetch_one(&mut *conn)
374 .await?;
375 Ok(ByteCount {
376 count: u64::try_from(count)?,
377 bytes: u64::try_from(bytes)?,
378 })
379}
380
381pub(super) async fn blob_available(
382 conn: &mut SqliteConnection,
383 blob_dir: &Path,
384 sha256: &str,
385) -> Result<bool> {
386 let Some(row) = blob_inventory_row(conn, sha256).await? else {
387 return Ok(false);
388 };
389 Ok(row.available && object_path(blob_dir, sha256)?.exists())
390}
391
392#[cfg(test)]
393mod tests {
394 use serde_json::json;
395
396 use super::*;
397
398 fn attachment_change(op_type: &str) -> ChangeWire {
399 ChangeWire {
400 change_id: "AAAAAAAAAAAAAAA0".to_string(),
401 client_id: "client".to_string(),
402 local_seq: 1,
403 entity_type: "task".to_string(),
404 entity_id: "BBBBBBBBBBBBBBBB".to_string(),
405 field: Some("attachments".to_string()),
406 op_type: op_type.to_string(),
407 payload: json!({
408 "workspace_id": "0000000000000000",
409 "workspace_key": "default",
410 "attachment_id": "7KQ9A1X4MV2P8D6R",
411 "sha256": "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789",
412 "byte_size": 12,
413 "media_type": "image/png",
414 "filename": null,
415 "alt_text": null,
416 "width": 320,
417 "height": 240,
418 "created_at": "2026-06-01T00:00:00Z"
419 }),
420 base_version: None,
421 created_at: "2026-06-01T00:00:00Z".to_string(),
422 server_seq: None,
423 }
424 }
425
426 #[test]
427 fn attachment_blob_contract_uses_typed_add_payload() {
428 let add = attachment_change(op_type::ATTACHMENT_ADD);
429 let contract = attachment_blob_contract(&add).unwrap().unwrap();
430 assert_eq!(contract.workspace_id, "0000000000000000");
431 assert_eq!(
432 contract.sha256,
433 "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789"
434 );
435 assert_eq!(contract.byte_size, 12);
436 assert_eq!(contract.media_type, "image/png");
437 assert_eq!((contract.width, contract.height), (320, 240));
438
439 let mut malformed = add;
440 malformed.payload["sha256"] = json!("invalid");
441 assert!(
442 attachment_blob_contract(&malformed)
443 .unwrap_err()
444 .to_string()
445 .contains("invalid-sha256")
446 );
447
448 assert!(
449 attachment_blob_contract(&attachment_change(op_type::ATTACHMENT_DELETE))
450 .unwrap()
451 .is_none()
452 );
453 assert!(
454 attachment_blob_contract(&attachment_change(op_type::NOTE_ADD))
455 .unwrap()
456 .is_none()
457 );
458 }
459
460 #[test]
461 fn blob_contracts_deduplicate_content_and_admission_separately() {
462 let contract = attachment_blob_contract(&attachment_change(op_type::ATTACHMENT_ADD))
463 .unwrap()
464 .unwrap();
465 let mut other_workspace = contract.clone();
466 other_workspace.workspace_id = "0000000000000001".to_string();
467 let contracts = vec![contract.clone(), contract.clone(), other_workspace.clone()];
468
469 assert_eq!(unique_blob_content_contracts(&contracts).unwrap().len(), 1);
470 assert_eq!(unique_blob_admission_contracts(&contracts).len(), 2);
471
472 let mut mismatch = contract.clone();
473 mismatch.height += 1;
474 let error = unique_blob_content_contracts(&[contract, mismatch]).unwrap_err();
475 assert!(
476 error
477 .to_string()
478 .contains("blob-inventory-metadata-mismatch")
479 );
480 }
481
482 async fn seed_missing_attachments(database: &Database, count: usize) {
483 let mut conn = database.acquire_writer().await.unwrap();
484 for index in 0..count {
485 let task_id = format!("{index:016X}");
486 let attachment_id = format!("{:016X}", index + 100);
487 let change_id = format!("{:016X}", index + 200);
488 let sha256 = format!("{index:064x}");
489 sqlx::query(
490 "INSERT INTO tasks(
491 workspace_id, id, title, description, project_id, status, priority,
492 created_at, updated_at, queue_activity_at
493 ) VALUES ('0000000000000000', ?, 'task', '', 'project', 'inbox', 'none',
494 '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z',
495 '2026-01-01T00:00:00Z')",
496 )
497 .bind(&task_id)
498 .execute(&mut *conn)
499 .await
500 .unwrap();
501 sqlx::query(
502 "INSERT INTO changes(
503 change_id, client_id, local_seq, entity_type, entity_id, field, op_type,
504 payload, created_at, server_seq
505 ) VALUES (?, 'remote', ?, 'task', ?, 'attachments', 'attachment_add', '{}',
506 '2026-01-01T00:00:00Z', ?)",
507 )
508 .bind(&change_id)
509 .bind(i64::try_from(index + 1).unwrap())
510 .bind(&task_id)
511 .bind(i64::try_from(index + 1).unwrap())
512 .execute(&mut *conn)
513 .await
514 .unwrap();
515 sqlx::query(
516 "INSERT INTO task_attachments(
517 workspace_id, attachment_id, task_id, sha256, byte_size, media_type,
518 width, height, created_at, created_by_change_id
519 ) VALUES ('0000000000000000', ?, ?, ?, 4, 'image/png', 1, 1,
520 '2026-01-01T00:00:00Z', ?)",
521 )
522 .bind(attachment_id)
523 .bind(task_id)
524 .bind(sha256)
525 .bind(change_id)
526 .execute(&mut *conn)
527 .await
528 .unwrap();
529 }
530 }
531
532 #[tokio::test]
533 async fn imported_local_only_attachment_is_not_scheduled_for_download() {
534 let temp = tempfile::tempdir().unwrap();
535 let database = Database::open(&temp.path().join("replica.sqlite"))
536 .await
537 .unwrap();
538 let mut conn = database.acquire_writer().await.unwrap();
539 sqlx::query(
540 "INSERT INTO tasks(
541 workspace_id, id, title, description, project_id, status, priority,
542 created_at, updated_at, queue_activity_at
543 ) VALUES ('0000000000000000', '0000000000000001', 'task', '', 'project',
544 'inbox', 'none', '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z',
545 '2026-01-01T00:00:00Z')",
546 )
547 .execute(&mut *conn)
548 .await
549 .unwrap();
550 sqlx::query(
551 "INSERT INTO task_attachments(
552 workspace_id, attachment_id, task_id, sha256, byte_size, media_type,
553 width, height, created_at, created_by_change_id
554 ) VALUES ('0000000000000000', '0000000000000002', '0000000000000001',
555 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa',
556 4, 'image/png', 1, 1, '2026-01-01T00:00:00Z', NULL)",
557 )
558 .execute(&mut *conn)
559 .await
560 .unwrap();
561 drop(conn);
562
563 let page = database.missing_local_blob_page(2).await.unwrap();
564 assert!(page.blobs.is_empty());
565 assert_eq!(
566 database.missing_sync_attachment_counts().await.unwrap(),
567 ByteCount { count: 0, bytes: 0 }
568 );
569 }
570
571 #[tokio::test]
572 async fn missing_blob_page_is_bounded_and_reports_more() {
573 let temp = tempfile::tempdir().unwrap();
574 let database = Database::open(&temp.path().join("replica.sqlite"))
575 .await
576 .unwrap();
577 seed_missing_attachments(&database, 3).await;
578 {
579 let mut conn = database.acquire_writer().await.unwrap();
580 crate::attachments::storage::upsert_inventory_available(
581 &mut conn,
582 &format!("{:064x}", 0),
583 4,
584 "image/png",
585 )
586 .await
587 .unwrap();
588 }
589
590 let page_before_maintenance = database.missing_local_blob_page(2).await.unwrap();
591 assert_eq!(page_before_maintenance.blobs.len(), 2);
592 assert!(!page_before_maintenance.has_more);
593
594 let blob_dir = temp.path().join("blobs");
595 let mut conn = database.acquire_writer().await.unwrap();
596 let missing = crate::attachments::lifecycle::reconcile_missing_objects(
597 &mut conn,
598 &blob_dir,
599 &SystemClock,
600 )
601 .await
602 .unwrap();
603 assert_eq!(missing.count, 1);
604 drop(conn);
605
606 let page = database.missing_local_blob_page(2).await.unwrap();
607 assert_eq!(page.blobs.len(), 2);
608 assert!(page.has_more);
609 assert_eq!(
610 database.missing_sync_attachment_counts().await.unwrap(),
611 ByteCount {
612 count: 3,
613 bytes: 12
614 }
615 );
616 }
617}