fix: reap expired skill uploads

daryllimyt · daryllimyt · commit 37e6998608c6 · 2026-04-23T11:42:11.000-04:00
diff --git a/tests/unit/test_skill_service.py b/tests/unit/test_skill_service.py
@@ -6,6 +6,7 @@
 import os
 import uuid
 from contextlib import asynccontextmanager
+from datetime import UTC, datetime, timedelta
 
 import pytest
 from dotenv import dotenv_values
@@ -790,6 +791,66 @@ async def test_create_draft_upload_rejects_invalid_normalized_content_type(
         ).scalars()
         assert upload_rows.all() == []
 
+    async def test_create_draft_upload_reaps_expired_incomplete_uploads(
+        self,
+        skill_service: SkillService,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Creating a new upload session should clean up older expired sessions."""
+
+        created = await skill_service.create_skill(SkillCreate(name="reap-uploads"))
+        stale_upload = await skill_service.create_draft_upload(
+            skill_id=created.id,
+            params=SkillUploadSessionCreate(
+                sha256=hashlib.sha256(b"stale upload").hexdigest(),
+                size_bytes=len(b"stale upload"),
+                content_type="text/plain; charset=utf-8",
+            ),
+        )
+        stale_upload_row = await skill_service.session.scalar(
+            select(SkillUploadModel).where(
+                SkillUploadModel.id == stale_upload.upload_id
+            )
+        )
+        assert stale_upload_row is not None
+        stale_upload_row.expires_at = datetime.now(UTC) - timedelta(minutes=1)
+        skill_service.session.add(stale_upload_row)
+        await skill_service.session.commit()
+
+        deleted: dict[str, str] = {}
+
+        async def fake_delete_file(*, key: str, bucket: str) -> None:
+            deleted["key"] = key
+            deleted["bucket"] = bucket
+
+        monkeypatch.setattr(
+            "tracecat.agent.skill.service.blob.delete_file",
+            fake_delete_file,
+        )
+
+        fresh_upload = await skill_service.create_draft_upload(
+            skill_id=created.id,
+            params=SkillUploadSessionCreate(
+                sha256=hashlib.sha256(b"fresh upload").hexdigest(),
+                size_bytes=len(b"fresh upload"),
+                content_type="text/plain; charset=utf-8",
+            ),
+        )
+
+        assert fresh_upload.upload_id != stale_upload.upload_id
+        assert deleted == {
+            "key": stale_upload.key,
+            "bucket": config.TRACECAT__BLOB_STORAGE_BUCKET_SKILLS,
+        }
+        assert (
+            await skill_service.session.scalar(
+                select(SkillUploadModel).where(
+                    SkillUploadModel.id == stale_upload.upload_id
+                )
+            )
+            is None
+        )
+
     async def test_attach_uploaded_blob_rejects_size_mismatch(
         self,
         skill_service: SkillService,
@@ -855,6 +916,67 @@ async def fake_open_download_stream(*, key: str, bucket: str):
             )
         assert iterated is False
 
+    async def test_attach_uploaded_blob_deletes_expired_staged_key(
+        self,
+        skill_service: SkillService,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Expired staged uploads should be deleted before the validation error returns."""
+
+        created = await skill_service.create_skill(SkillCreate(name="expired-upload"))
+        draft = await skill_service.get_draft(created.id)
+        assert draft is not None
+
+        upload = await skill_service.create_draft_upload(
+            skill_id=created.id,
+            params=SkillUploadSessionCreate(
+                sha256=hashlib.sha256(b"expired upload").hexdigest(),
+                size_bytes=len(b"expired upload"),
+                content_type="text/plain; charset=utf-8",
+            ),
+        )
+        upload_row = await skill_service.session.scalar(
+            select(SkillUploadModel).where(SkillUploadModel.id == upload.upload_id)
+        )
+        assert upload_row is not None
+        upload_row.expires_at = datetime.now(UTC) - timedelta(minutes=1)
+        skill_service.session.add(upload_row)
+        await skill_service.session.commit()
+
+        deleted: dict[str, str] = {}
+
+        async def fake_delete_file(*, key: str, bucket: str) -> None:
+            deleted["key"] = key
+            deleted["bucket"] = bucket
+
+        monkeypatch.setattr(
+            "tracecat.agent.skill.service.blob.delete_file",
+            fake_delete_file,
+        )
+
+        with pytest.raises(TracecatValidationError) as exc_info:
+            await skill_service.patch_draft(
+                skill_id=created.id,
+                params=SkillDraftPatch(
+                    base_revision=draft.draft_revision,
+                    operations=[
+                        SkillDraftAttachUploadedBlobOp(
+                            path="references/uploaded.txt",
+                            upload_id=upload.upload_id,
+                        )
+                    ],
+                ),
+            )
+
+        assert exc_info.value.detail == {
+            "code": "upload_expired",
+            "upload_id": str(upload.upload_id),
+        }
+        assert deleted == {
+            "key": upload.key,
+            "bucket": config.TRACECAT__BLOB_STORAGE_BUCKET_SKILLS,
+        }
+
     async def test_attach_uploaded_blob_stops_streaming_after_size_exceeded(
         self,
         skill_service: SkillService,
diff --git a/tracecat/agent/skill/service.py b/tracecat/agent/skill/service.py
@@ -179,6 +179,20 @@ def _staged_upload_key_for(self, *, upload_id: uuid.UUID, sha256: str) -> str:
         normalized_sha256 = self._normalize_sha256(sha256)
         return f"skills/{self.workspace_id}/uploads/{upload_id}/{normalized_sha256}"
 
+    def _staged_upload_prefix(self) -> str:
+        """Return the storage-prefix used for staged upload objects."""
+
+        return f"skills/{self.workspace_id}/uploads/"
+
+    def _is_staged_upload_object(self, upload: SkillUploadModel) -> bool:
+        """Return whether an upload still points at a temporary staged object."""
+
+        return (
+            upload.completed_at is None
+            and upload.blob_id is None
+            and upload.key.startswith(self._staged_upload_prefix())
+        )
+
     @staticmethod
     def _normalize_path(path: str) -> str:
         """Normalize and validate a relative POSIX draft path.
@@ -478,6 +492,11 @@ async def _materialize_uploaded_blob(self, upload: SkillUploadModel) -> SkillBlo
             return blob_row
 
         if upload.expires_at < datetime.now(UTC):
+            if self._is_staged_upload_object(upload):
+                await self._delete_staged_upload_object_best_effort(
+                    upload,
+                    reason="upload_expired",
+                )
             raise TracecatValidationError(
                 "Skill upload session has expired",
                 detail={"code": "upload_expired", "upload_id": str(upload.id)},
@@ -551,6 +570,51 @@ async def _materialize_uploaded_blob(self, upload: SkillUploadModel) -> SkillBlo
         await self.session.flush()
         return blob_row
 
+    async def _delete_staged_upload_object_best_effort(
+        self,
+        upload: SkillUploadModel,
+        *,
+        reason: str,
+    ) -> None:
+        """Delete a temporary staged upload object without failing the caller."""
+
+        try:
+            await blob.delete_file(key=upload.key, bucket=upload.bucket)
+        except Exception as exc:
+            self.logger.warning(
+                "Failed to delete staged skill upload object",
+                upload_id=str(upload.id),
+                key=upload.key,
+                bucket=upload.bucket,
+                reason=reason,
+                error=str(exc),
+            )
+
+    async def _reap_expired_incomplete_uploads(self) -> list[SkillUploadModel]:
+        """Delete expired incomplete upload rows and return staged objects to clean up."""
+
+        expired_stmt = (
+            select(SkillUploadModel)
+            .where(
+                SkillUploadModel.workspace_id == self.workspace_id,
+                SkillUploadModel.completed_at.is_(None),
+                SkillUploadModel.expires_at < datetime.now(UTC),
+            )
+            .with_for_update(skip_locked=True)
+        )
+        expired_uploads = (await self.session.execute(expired_stmt)).scalars().all()
+        if not expired_uploads:
+            return []
+
+        for upload in expired_uploads:
+            await self.session.delete(upload)
+        await self.session.flush()
+        return [
+            upload
+            for upload in expired_uploads
+            if self._is_staged_upload_object(upload)
+        ]
+
     async def _list_draft_rows(
         self, skill_id: uuid.UUID
     ) -> list[tuple[SkillDraftFile, SkillBlob]]:
@@ -1402,6 +1466,7 @@ async def create_draft_upload(
         skill = await self.get_skill(skill_id)
         if skill is None:
             raise TracecatNotFoundError(f"Skill '{skill_id}' not found")
+        expired_uploads = await self._reap_expired_incomplete_uploads()
 
         upload_id = uuid.uuid4()
         expires_at = datetime.now(UTC) + timedelta(seconds=DEFAULT_UPLOAD_TTL_SECONDS)
@@ -1424,6 +1489,11 @@ async def create_draft_upload(
         upload_row.id = upload_id
         self.session.add(upload_row)
         await self.session.commit()
+        for expired_upload in expired_uploads:
+            await self._delete_staged_upload_object_best_effort(
+                expired_upload,
+                reason="reap_expired_upload",
+            )
         return SkillUploadSessionRead(
             upload_id=upload_id,
             upload_url=await blob.generate_presigned_upload_url(
