perf: async file saving in web generation endpoint

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     "transformers>=4.36.0",
     "uvicorn>=0.24.0",
     "safetensors>=0.4.0",
+    "aiofiles>=23.2.1",
 ]
 
 [project.optional-dependencies]

src/controllers/web_controller.py

@@ -5,6 +5,7 @@
 import html
 from urllib.parse import urlparse
 from typing import Optional
+import aiofiles
 
 from fastapi import APIRouter, Request, Form, HTTPException, Depends
 from fastapi.responses import HTMLResponse, JSONResponse
@@ -115,20 +116,20 @@ def _generate_task():
         filepath = os.path.join(OUTPUT_DIR, filename)
         filepath_1_7 = os.path.join(OUTPUT_DIR, f"{normalized_id.replace('/', '_')}_ai_sbom_1_7.json")
 
-        def _save_task():
+        def _format_task():
             # Generate Formatted JSON strings
             json_1_6 = export_aibom(aibom, bom_type="cyclonedx", spec_version="1.6")
             json_1_7 = export_aibom(aibom, bom_type="cyclonedx", spec_version="1.7")
-
-            os.makedirs(OUTPUT_DIR, exist_ok=True)
-            with open(filepath, "w", encoding="utf-8") as f:
-                f.write(json_1_6)
-            with open(filepath_1_7, "w", encoding="utf-8") as f:
-                f.write(json_1_7)
-            log_sbom_generation(sanitized_model_id)
             return json_1_6, json_1_7
 
-        json_1_6, json_1_7 = await loop.run_in_executor(None, _save_task)
+        json_1_6, json_1_7 = await loop.run_in_executor(None, _format_task)
+
+        os.makedirs(OUTPUT_DIR, exist_ok=True)
+        async with aiofiles.open(filepath, "w", encoding="utf-8") as f:
+            await f.write(json_1_6)
+        async with aiofiles.open(filepath_1_7, "w", encoding="utf-8") as f:
+            await f.write(json_1_7)
+        log_sbom_generation(sanitized_model_id)
 
         # Extract score
         completeness_score = None

tests/test_web_controller.py

@@ -152,6 +152,31 @@ def test_generate_route_writes_utf8_json_to_patched_output_dir(self):
         finally:
             shutil.rmtree(temp_root, ignore_errors=True)
 
+    @patch("src.controllers.web_controller.aiofiles.open")
+    def test_generate_uses_aiofiles_async(self, mock_aiofiles_open):
+        from unittest.mock import AsyncMock
+        mock_file = AsyncMock()
+        mock_aiofiles_open.return_value.__aenter__.return_value = mock_file
+
+        app = FastAPI()
+        app.include_router(web_controller.router)
+
+        with patch.object(web_controller, "OUTPUT_DIR", "/tmp/dummy"), \
+             patch.object(web_controller, "HfApi", _FakeHfApi), \
+             patch.object(web_controller, "AIBOMService", _FakeService), \
+             patch.object(web_controller, "log_sbom_generation", lambda model_id: None), \
+             patch.object(web_controller, "get_sbom_count", lambda: "0"):
+
+            client = TestClient(app)
+            response = client.post(
+                "/generate",
+                data={"model_id": "owner/local-web-model"},
+            )
+
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(mock_aiofiles_open.call_count, 2)
+        # Verify that await f.write(...) was called via the context manager
+        self.assertEqual(mock_file.write.call_count, 2)
 
 if __name__ == "__main__":
     unittest.main()