feat: implement pending attachments feature in upload handling

migrations/0006_add_pending_attachments.sql

@@ -0,0 +1,16 @@
+-- Track attachments that are still uploading; moved to attachments on success.
+CREATE TABLE IF NOT EXISTS attachments_pending (
+    id TEXT PRIMARY KEY NOT NULL,
+    cipher_id TEXT NOT NULL,
+    file_name TEXT NOT NULL,
+    file_size INTEGER NOT NULL,
+    akey TEXT,
+    created_at TEXT NOT NULL,
+    updated_at TEXT NOT NULL,
+    organization_id TEXT,
+    FOREIGN KEY (cipher_id) REFERENCES ciphers(id) ON DELETE CASCADE
+);
+
+CREATE INDEX IF NOT EXISTS idx_attachments_pending_cipher ON attachments_pending(cipher_id);
+CREATE INDEX IF NOT EXISTS idx_attachments_pending_created_at ON attachments_pending(created_at);
+

sql/schema.sql

@@ -51,6 +51,21 @@ CREATE TABLE IF NOT EXISTS attachments (
 );
 CREATE INDEX IF NOT EXISTS idx_attachments_cipher ON attachments(cipher_id);
 
+-- Pending attachments table for in-flight uploads
+CREATE TABLE IF NOT EXISTS attachments_pending (
+    id TEXT PRIMARY KEY NOT NULL,
+    cipher_id TEXT NOT NULL,
+    file_name TEXT NOT NULL,
+    file_size INTEGER NOT NULL,
+    akey TEXT,
+    created_at TEXT NOT NULL,
+    updated_at TEXT NOT NULL,
+    organization_id TEXT,
+    FOREIGN KEY (cipher_id) REFERENCES ciphers(id) ON DELETE CASCADE
+);
+CREATE INDEX IF NOT EXISTS idx_attachments_pending_cipher ON attachments_pending(cipher_id);
+CREATE INDEX IF NOT EXISTS idx_attachments_pending_created_at ON attachments_pending(created_at);
+
 -- TwoFactor table for two-factor authentication
 -- Types: 0=Authenticator(TOTP), 1=Email, 5=Remember, 8=RecoveryCode
 CREATE TABLE IF NOT EXISTS twofactor (

src/entry.js

@@ -147,16 +147,32 @@ function getTotalLimitBytes(env) {
 // Get user's current attachment usage
 async function getUserAttachmentUsage(db, userId, excludeAttachmentId) {
   const query = excludeAttachmentId
-    ? `SELECT COALESCE(SUM(a.file_size), 0) as total
+    ? `SELECT COALESCE(SUM(file_size), 0) as total FROM (
-       FROM attachments a
+         SELECT a.file_size as file_size
-       JOIN ciphers c ON c.id = a.cipher_id
+         FROM attachments a
-       WHERE c.user_id = ?1 AND a.id != ?2`
+         JOIN ciphers c ON c.id = a.cipher_id
-    : `SELECT COALESCE(SUM(a.file_size), 0) as total
+         WHERE c.user_id = ?1 AND a.id != ?2
-       FROM attachments a
+         UNION ALL
-       JOIN ciphers c ON c.id = a.cipher_id
+         SELECT p.file_size as file_size
-       WHERE c.user_id = ?1`;
+         FROM attachments_pending p
+         JOIN ciphers c2 ON c2.id = p.cipher_id
+         WHERE c2.user_id = ?1 AND p.id != ?2
+       ) AS files`
+    : `SELECT COALESCE(SUM(file_size), 0) as total FROM (
+         SELECT a.file_size as file_size
+         FROM attachments a
+         JOIN ciphers c ON c.id = a.cipher_id
+         WHERE c.user_id = ?1
+         UNION ALL
+         SELECT p.file_size as file_size
+         FROM attachments_pending p
+         JOIN ciphers c2 ON c2.id = p.cipher_id
+         WHERE c2.user_id = ?1
+       ) AS files`;
 
-  const bindings = excludeAttachmentId ? [userId, excludeAttachmentId] : [userId];
+  const bindings = excludeAttachmentId
+    ? [userId, excludeAttachmentId]
+    : [userId];
 
   const result = await db
     .prepare(query)
@@ -259,20 +275,20 @@ async function handleAzureUpload(request, env, cipherId, attachmentId, token) {
     );
   }
 
-  // Fetch attachment record
+  // Fetch pending attachment record
-  const attachment = await db
+  const pending = await db
-    .prepare("SELECT * FROM attachments WHERE id = ?1")
+    .prepare("SELECT * FROM attachments_pending WHERE id = ?1")
     .bind(attachmentId)
     .first();
 
-  if (!attachment) {
+  if (!pending) {
-    return new Response(JSON.stringify({ error: "Attachment not found" }), {
+    return new Response(JSON.stringify({ error: "Attachment not found or already uploaded" }), {
       status: 404,
       headers: { "Content-Type": "application/json" },
     });
   }
 
-  if (attachment.cipher_id !== cipherId) {
+  if (pending.cipher_id !== cipherId) {
     return new Response(
       JSON.stringify({ error: "Attachment does not belong to cipher" }),
       { status: 400, headers: { "Content-Type": "application/json" } }
@@ -350,39 +366,29 @@ async function handleAzureUpload(request, env, cipherId, attachmentId, token) {
     );
   }
 
-  // Re-enforce limits with actual uploaded size (should be same, but be safe)
+  // Finalize upload: move pending -> attachments and touch revision timestamps
-  try {
-    await enforceLimits(db, env, userId, uploadedSize, attachmentId);
-  } catch (err) {
-    try {
-      await bucket.delete(r2Key);
-    } catch {
-      // Ignore cleanup errors
-    }
-    return new Response(JSON.stringify({ error: err.message }), {
-      status: 400,
-      headers: { "Content-Type": "application/json" },
-    });
-  }
-
-  // Update attachment record in database
   const now = nowString();
-  await db
+  await db.batch([
-    .prepare("UPDATE attachments SET file_size = ?1, updated_at = ?2 WHERE id = ?3")
+    db
-    .bind(uploadedSize, now, attachmentId)
+      .prepare(
-    .run();
+        "INSERT INTO attachments (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)"
-
+      )
-  // Touch cipher updated_at
+      .bind(
-  await db
+        attachmentId,
-    .prepare("UPDATE ciphers SET updated_at = ?1 WHERE id = ?2")
+        cipherId,
-    .bind(now, cipherId)
+        pending.file_name,
-    .run();
+        uploadedSize,
-
+        pending.akey,
-  // Touch user updated_at
+        pending.created_at || now,
-  await db
+        now,
-    .prepare("UPDATE users SET updated_at = ?1 WHERE id = ?2")
+        pending.organization_id || null
-    .bind(now, userId)
+      ),
-    .run();
+    db
+      .prepare("DELETE FROM attachments_pending WHERE id = ?1")
+      .bind(attachmentId),
+    db.prepare("UPDATE ciphers SET updated_at = ?1 WHERE id = ?2").bind(now, cipherId),
+    db.prepare("UPDATE users SET updated_at = ?1 WHERE id = ?2").bind(now, userId),
+  ]);
 
   return new Response(null, { status: 201 });
 }

src/handlers/attachments.rs

@@ -146,7 +146,7 @@ pub async fn create_attachment_v2(
 
     query!(
         &db,
-        "INSERT INTO attachments (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id)
+        "INSERT INTO attachments_pending (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id)
          VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6, ?7)",
         attachment_id,
         cipher.id,
@@ -165,8 +165,26 @@ pub async fn create_attachment_v2(
     let mut cipher_response: Cipher = cipher.into();
     hydrate_cipher_attachments(&db, &env, &mut cipher_response).await?;
 
-    touch_cipher_updated_at(&db, &cipher_id).await?;
+    // add pending attachment to response
-    db::touch_user_updated_at(&db, &claims.sub).await?;
+    let pending_attachment = AttachmentDB {
+        id: attachment_id.clone(),
+        cipher_id: cipher_id.clone(),
+        file_name,
+        file_size: declared_size,
+        akey: Some(key),
+        created_at: now.clone(),
+        updated_at: now,
+        organization_id: cipher_response.organization_id.clone(),
+    };
+
+    let pending_response = pending_attachment.to_response(None);
+    match &mut cipher_response.attachments {
+        Some(list) => list.push(pending_response),
+        None => cipher_response.attachments = Some(vec![pending_response]),
+    }
+
+    // no need to touch cipher updated_at and user updated_at here
+    // it will be touched in after upload
 
     Ok(Json(AttachmentUploadResponse {
         object: "attachment-fileUpload".to_string(),
@@ -190,8 +208,8 @@ pub async fn upload_attachment_v2_data(
 
     let _cipher = ensure_cipher_for_user(&db, &cipher_id, &claims.sub).await?;
 
-    let mut attachment = fetch_attachment(&db, &attachment_id).await?;
+    let mut pending = fetch_pending_attachment(&db, &attachment_id).await?;
-    if attachment.cipher_id != cipher_id {
+    if pending.cipher_id != cipher_id {
         return Err(AppError::BadRequest(
             "Attachment does not belong to cipher".to_string(),
         ));
@@ -202,45 +220,63 @@ pub async fn upload_attachment_v2_data(
     let actual_size = file_bytes.len() as i64;
 
     // Validate actual size against declared value deviation
-    if let Err(e) = validate_size_within_declared(&attachment, actual_size) {
+    if let Err(e) = validate_size_within_declared(&pending, actual_size) {
-        query!(&db, "DELETE FROM attachments WHERE id = ?1", attachment.id)
+        query!(
-            .map_err(|_| AppError::Database)?
+            &db,
-            .run()
+            "DELETE FROM attachments_pending WHERE id = ?1",
-            .await?;
+            pending.id
+        )
+        .map_err(|_| AppError::Database)?
+        .run()
+        .await?;
         return Err(e);
     }
 
     // Validate capacity limits (replace with actual size)
-    enforce_limits(&db, &env, &claims.sub, actual_size, Some(&attachment.id)).await?;
+    enforce_limits(&db, &env, &claims.sub, actual_size, Some(&pending.id)).await?;
 
     // Need a key
-    if attachment.akey.is_none() && key_override.is_none() {
+    if pending.akey.is_none() && key_override.is_none() {
         return Err(AppError::BadRequest(
             "No attachment key provided".to_string(),
         ));
     }
     if let Some(k) = key_override {
-        attachment.akey = Some(k);
+        pending.akey = Some(k);
     }
 
     // Save to R2
     upload_to_r2(
         &bucket,
-        &attachment.r2_key(),
+        &pending.r2_key(),
         content_type,
         file_bytes.to_vec(),
     )
     .await?;
 
-    // Update metadata
+    // Finalize: move pending -> attachments and touch timestamps
     let now = now_string();
     query!(
         &db,
-        "UPDATE attachments SET file_size = ?1, akey = COALESCE(?2, akey), updated_at = ?3 WHERE id = ?4",
+        "INSERT INTO attachments (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id)
+         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
+        pending.id,
+        pending.cipher_id,
+        pending.file_name,
         actual_size,
-        attachment.akey,
+        pending.akey,
+        pending.created_at,
         now,
-        attachment.id,
+        pending.organization_id,
+    )
+    .map_err(|_| AppError::Database)?
+    .run()
+    .await?;
+
+    query!(
+        &db,
+        "DELETE FROM attachments_pending WHERE id = ?1",
+        pending.id
     )
     .map_err(|_| AppError::Database)?
     .run()
@@ -312,7 +348,7 @@ pub async fn upload_attachment_legacy(
     touch_cipher_updated_at(&db, &cipher_id).await?;
     db::touch_user_updated_at(&db, &claims.sub).await?;
 
-    // 返回最新的 Cipher（含附件）
+    // reload cipher to return fresh updated_at and attachments state
     let mut cipher_response: Cipher = cipher.into();
     hydrate_cipher_attachments(&db, &env, &mut cipher_response).await?;
 
@@ -571,6 +607,18 @@ async fn fetch_attachment(db: &D1Database, attachment_id: &str) -> Result<Attach
         .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))
 }
 
+async fn fetch_pending_attachment(
+    db: &D1Database,
+    attachment_id: &str,
+) -> Result<AttachmentDB, AppError> {
+    db.prepare("SELECT * FROM attachments_pending WHERE id = ?1")
+        .bind(&[attachment_id.into()])?
+        .first(None)
+        .await
+        .map_err(|_| AppError::Database)?
+        .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))
+}
+
 async fn load_attachment_map_json(
     db: &D1Database,
     json_body: &str,
@@ -842,23 +890,40 @@ async fn user_attachment_usage(
     user_id: &str,
     exclude_attachment: Option<&str>,
 ) -> Result<i64, AppError> {
-    let query_str = if exclude_attachment.is_some() {
+    let (query_str, bindings): (String, Vec<JsValue>) = if let Some(id) = exclude_attachment {
-        "SELECT COALESCE(SUM(a.file_size), 0) as total
+        (
-         FROM attachments a
+            "SELECT COALESCE(SUM(file_size), 0) as total FROM (
-         JOIN ciphers c ON c.id = a.cipher_id
+                SELECT a.file_size AS file_size
-         WHERE c.user_id = ?1 AND a.id != ?2"
+                FROM attachments a
+                JOIN ciphers c ON c.id = a.cipher_id
+                WHERE c.user_id = ?1 AND a.id != ?2
+                UNION ALL
+                SELECT p.file_size AS file_size
+                FROM attachments_pending p
+                JOIN ciphers c2 ON c2.id = p.cipher_id
+                WHERE c2.user_id = ?1 AND p.id != ?2
+            ) AS files"
+                .to_string(),
+            vec![JsValue::from_str(user_id), JsValue::from_str(id)],
+        )
     } else {
-        "SELECT COALESCE(SUM(a.file_size), 0) as total
+        (
-         FROM attachments a
+            "SELECT COALESCE(SUM(file_size), 0) as total FROM (
-         JOIN ciphers c ON c.id = a.cipher_id
+                SELECT a.file_size AS file_size
-         WHERE c.user_id = ?1"
+                FROM attachments a
+                JOIN ciphers c ON c.id = a.cipher_id
+                WHERE c.user_id = ?1
+                UNION ALL
+                SELECT p.file_size AS file_size
+                FROM attachments_pending p
+                JOIN ciphers c2 ON c2.id = p.cipher_id
+                WHERE c2.user_id = ?1
+            ) AS files"
+                .to_string(),
+            vec![JsValue::from_str(user_id)],
+        )
     };
 
-    let mut bindings: Vec<JsValue> = vec![JsValue::from_str(user_id)];
-    if let Some(id) = exclude_attachment {
-        bindings.push(JsValue::from_str(id));
-    }
-
     let row: Option<Value> = db
         .prepare(query_str)
         .bind(&bindings)?

src/handlers/purge.rs

@@ -11,6 +11,8 @@ use worker::{query, D1Database, Env};
 
 /// Default number of days to keep soft-deleted items before purging
 const DEFAULT_PURGE_DAYS: i64 = 30;
+/// Retain pending attachments for at most this many days before cleanup
+const PENDING_RETENTION_DAYS: i64 = 1;
 
 /// Get the purge threshold days from environment variable or use default
 fn get_purge_days(env: &Env) -> i64 {
@@ -20,6 +22,43 @@ fn get_purge_days(env: &Env) -> i64 {
         .unwrap_or(DEFAULT_PURGE_DAYS)
 }
 
+/// Purge pending attachments older than the configured retention window.
+pub async fn purge_stale_pending_attachments(env: &Env) -> Result<u32, worker::Error> {
+    let db: D1Database = env.d1("vault1")?;
+    let now = Utc::now();
+    let pending_cutoff = now - Duration::days(PENDING_RETENTION_DAYS);
+    let pending_cutoff_str = pending_cutoff.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string();
+
+    let pending_count_result = query!(
+        &db,
+        "SELECT COUNT(*) as count FROM attachments_pending WHERE created_at < ?1",
+        pending_cutoff_str
+    )?
+    .first::<CountResult>(None)
+    .await?;
+
+    let pending_count = pending_count_result.map(|r| r.count).unwrap_or(0);
+
+    if pending_count > 0 {
+        query!(
+            &db,
+            "DELETE FROM attachments_pending WHERE created_at < ?1",
+            pending_cutoff_str
+        )?
+        .run()
+        .await?;
+        log::info!(
+            "Purged {} pending attachment(s) older than {} day(s)",
+            pending_count,
+            PENDING_RETENTION_DAYS
+        );
+    } else {
+        log::info!("No pending attachments to purge");
+    }
+
+    Ok(pending_count)
+}
+
 /// Purge soft-deleted ciphers that are older than the configured threshold.
 ///
 /// This function:

src/lib.rs

@@ -66,6 +66,19 @@ pub async fn scheduled(_event: ScheduledEvent, env: Env, _ctx: ScheduleContext)
     console_error_panic_hook::set_once();
     let _ = console_log::init_with_level(log::Level::Debug);
 
+    log::info!("Scheduled task triggered: purging stale pending attachments");
+    match handlers::purge::purge_stale_pending_attachments(&env).await {
+        Ok(count) => {
+            log::info!(
+                "Pending attachment purge completed: {} record(s) removed",
+                count
+            );
+        }
+        Err(e) => {
+            log::error!("Pending attachment purge failed: {:?}", e);
+        }
+    }
+
     log::info!("Scheduled task triggered: purging soft-deleted ciphers");
 
     match handlers::purge::purge_deleted_ciphers(&env).await {