feat(ui): show live scan progress and ETA for cache refresh

Broadcast typed scan_progress messages over /ws/fetch-progress from the
manual refresh/rebuild paths of ModelScanner and RecipeScanner, and
render percent, processed/total, current file name and an EMA-smoothed
ETA in the loading overlay. Hardcoded refresh strings move to i18n
(common.scanProgress); WS connection failure falls back to the previous
static loading behavior.
This commit is contained in:
Will Miao
2026-09-03 11:38:27 +08:00
parent da71985c3e
commit 14da8a6f17
19 changed files with 1668 additions and 32 deletions
+186
View File
@@ -30,10 +30,14 @@ from py.utils.models import BaseModelMetadata
class RecordingWebSocketManager:
def __init__(self) -> None:
self.payloads: List[Dict[str, Any]] = []
self.broadcasts: List[Dict[str, Any]] = []
async def broadcast_init_progress(self, payload: Dict[str, Any]) -> None:
self.payloads.append(payload)
async def broadcast(self, payload: Dict[str, Any]) -> None:
self.broadcasts.append(payload)
def _normalize_path(path: Path) -> str:
return str(path).replace(os.sep, "/")
@@ -1395,3 +1399,185 @@ async def test_get_all_folders_invalidated_after_move(tmp_path: Path):
assert "new" in all_folders
assert "new/deep" in all_folders
assert set(cache.folders) <= set(all_folders)
@pytest.mark.asyncio
async def test_initialize_cache_broadcasts_scan_progress(tmp_path: Path, monkeypatch):
_create_files(tmp_path)
scanner = DummyScanner(tmp_path)
ws_stub = RecordingWebSocketManager()
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
await scanner._initialize_cache()
messages = ws_stub.broadcasts
assert messages, "expected scan_progress broadcasts"
started = messages[0]
assert started["type"] == "scan_progress"
assert started["status"] == "started"
assert started["stage"] == "scan_folders"
assert started["progress"] == 0
assert started["model_type"] == "dummy"
assert started["pageType"] == "dummy"
assert started["full_rebuild"] is True
count_messages = [m for m in messages if m["stage"] == "count_models"]
assert count_messages and count_messages[0]["total"] == 3
process_messages = [
m for m in messages
if m["stage"] == "process_models" and m["status"] == "processing"
]
assert process_messages, "expected at least one process_models update"
final_process = process_messages[-1]
assert final_process["processed"] == 3
assert final_process["total"] == 3
assert final_process["current_name"].endswith(".txt")
for message in process_messages:
assert 0 < message["progress"] <= 99
stages = [m["stage"] for m in messages]
assert "finalizing" in stages
completed = messages[-1]
assert completed["status"] == "completed"
assert completed["progress"] == 100
assert completed["elapsed_seconds"] >= 0
@pytest.mark.asyncio
async def test_initialize_cache_broadcasts_cancelled(tmp_path: Path, monkeypatch):
_create_files(tmp_path)
scanner = DummyScanner(tmp_path)
ws_stub = RecordingWebSocketManager()
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
original_process = DummyScanner._process_model_file
async def cancelling_process(self, file_path, root_path, **kwargs):
scanner.cancel_task()
return await original_process(self, file_path, root_path, **kwargs)
monkeypatch.setattr(DummyScanner, "_process_model_file", cancelling_process)
await scanner._initialize_cache()
messages = ws_stub.broadcasts
assert messages[0]["status"] == "started"
assert messages[-1]["status"] == "cancelled"
assert messages[-1]["elapsed_seconds"] >= 0
assert not any(m["status"] == "completed" for m in messages)
@pytest.mark.asyncio
async def test_initialize_cache_broadcasts_error(tmp_path: Path, monkeypatch):
scanner = DummyScanner(tmp_path)
ws_stub = RecordingWebSocketManager()
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
async def raising_gather(**_kwargs):
raise RuntimeError("boom")
monkeypatch.setattr(scanner, "_gather_model_data", raising_gather)
await scanner._initialize_cache()
messages = ws_stub.broadcasts
assert messages[0]["status"] == "started"
assert messages[-1]["status"] == "error"
assert messages[-1]["error"] == "boom"
@pytest.mark.asyncio
async def test_reconcile_cache_broadcasts_scan_progress(tmp_path: Path, monkeypatch):
_create_files(tmp_path)
scanner = DummyScanner(tmp_path)
await scanner._initialize_cache()
ws_stub = RecordingWebSocketManager()
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
new_file = tmp_path / "three.txt"
new_file.write_text("three", encoding="utf-8")
await scanner._reconcile_cache()
messages = ws_stub.broadcasts
assert messages, "expected scan_progress broadcasts"
started = messages[0]
assert started["type"] == "scan_progress"
assert started["status"] == "started"
assert started["stage"] == "reconcile_scan"
assert started["progress"] == 0
assert started["full_rebuild"] is False
process_messages = [
m for m in messages
if m["stage"] == "process_new" and m["status"] == "processing"
]
assert process_messages, "expected process_new progress updates"
assert process_messages[-1]["processed"] == 1
assert process_messages[-1]["total"] == 1
assert process_messages[-1]["current_name"] == "three.txt"
completed = messages[-1]
assert completed["status"] == "completed"
assert completed["progress"] == 100
assert completed["added"] == 1
assert completed["removed"] == 0
assert completed["elapsed_seconds"] >= 0
@pytest.mark.asyncio
async def test_reconcile_cache_broadcasts_cancelled(tmp_path: Path, monkeypatch):
_create_files(tmp_path)
scanner = DummyScanner(tmp_path)
await scanner._initialize_cache()
ws_stub = RecordingWebSocketManager()
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
new_file = tmp_path / "four.txt"
new_file.write_text("four", encoding="utf-8")
original_process = DummyScanner._process_model_file
async def cancelling_process(self, file_path, root_path, **kwargs):
scanner.cancel_task()
return await original_process(self, file_path, root_path, **kwargs)
monkeypatch.setattr(DummyScanner, "_process_model_file", cancelling_process)
await scanner._reconcile_cache()
messages = ws_stub.broadcasts
assert messages[0]["status"] == "started"
assert messages[-1]["status"] == "cancelled"
assert messages[-1]["elapsed_seconds"] >= 0
assert not any(m["status"] == "completed" for m in messages)
@pytest.mark.asyncio
async def test_reconcile_cache_broadcasts_error(tmp_path: Path, monkeypatch):
_create_files(tmp_path)
scanner = DummyScanner(tmp_path)
await scanner._initialize_cache()
ws_stub = RecordingWebSocketManager()
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
def raising_walk(*_args, **_kwargs):
raise RuntimeError("walk failed")
monkeypatch.setattr(model_scanner.os, "walk", raising_walk)
await scanner._reconcile_cache()
messages = ws_stub.broadcasts
assert messages[0]["status"] == "started"
assert messages[-1]["status"] == "error"
assert messages[-1]["error"] == "walk failed"
+131
View File
@@ -9,6 +9,7 @@ import pytest
from py.config import config
from py.services import model_scanner as model_scanner_module
from py.services import recipe_scanner as recipe_scanner_module
from py.services.model_cache import ModelCache
from py.services.model_hash_index import ModelHashIndex
from py.services.model_scanner import CacheBuildResult, ModelScanner
@@ -4965,3 +4966,133 @@ async def test_find_all_duplicate_recipes_include_prompt_missing_gen_params(reci
groups = await scanner.find_all_duplicate_recipes(include_prompt=True)
# Recipes without gen_params/prompt normalize to empty prompt and match
assert groups == {"abc:0.8\x1f": ["r1", "r2"]}
class RecordingRecipeWebSocketManager:
"""Minimal ws_manager stand-in that records broadcasts."""
def __init__(self) -> None:
self.payloads: list[Dict[str, Any]] = []
self.broadcasts: list[Dict[str, Any]] = []
async def broadcast_init_progress(self, payload: Dict[str, Any]) -> None:
self.payloads.append(payload)
async def broadcast(self, payload: Dict[str, Any]) -> None:
self.broadcasts.append(payload)
def _write_progress_recipe_files(recipes_dir: Path, count: int) -> None:
recipes_dir.mkdir(parents=True, exist_ok=True)
for idx in range(count):
recipe_path = recipes_dir / f"progress-recipe-{idx}.recipe.json"
recipe_path.write_text(
json.dumps(
{
"id": f"progress-recipe-{idx}",
"file_path": str(recipes_dir / f"img-{idx}.png"),
"title": f"Recipe {idx}",
"modified": 0.0,
"created_date": 0.0,
"loras": [],
}
),
encoding="utf-8",
)
@pytest.mark.asyncio
async def test_force_refresh_broadcasts_scan_progress(
tmp_path: Path, monkeypatch, recipe_scanner
):
scanner, _stub = recipe_scanner
recipes_dir = Path(config.loras_roots[0]) / "recipes"
_write_progress_recipe_files(recipes_dir, 3)
ws_stub = RecordingRecipeWebSocketManager()
monkeypatch.setattr(recipe_scanner_module, "ws_manager", ws_stub)
await scanner.get_cached_data(force_refresh=True)
# Wait for the FTS index build so no background task outlives the loop.
if scanner._fts_index_task:
await scanner._fts_index_task
messages = ws_stub.broadcasts
assert messages, "expected scan_progress broadcasts"
started = messages[0]
assert started["type"] == "scan_progress"
assert started["status"] == "started"
assert started["stage"] == "scan_folders"
assert started["progress"] == 0
assert started["model_type"] == "recipe"
assert started["pageType"] == "recipes"
assert started["full_rebuild"] is True
count_messages = [m for m in messages if m["stage"] == "count_models"]
assert count_messages and count_messages[0]["total"] == 3
process_messages = [
m
for m in messages
if m["stage"] == "process_models" and m["status"] == "processing"
]
assert process_messages, "expected at least one process_models update"
final_process = process_messages[-1]
assert final_process["processed"] == 3
assert final_process["total"] == 3
assert final_process["current_name"].endswith(".recipe.json")
for message in process_messages:
assert 0 < message["progress"] <= 99
completed = messages[-1]
assert completed["status"] == "completed"
assert completed["progress"] == 100
assert completed["elapsed_seconds"] >= 0
assert completed["total"] == 3
def test_sync_init_without_report_progress_does_not_broadcast(
tmp_path: Path, monkeypatch, recipe_scanner
):
"""Startup path (initialize_in_background) must not emit scan_progress."""
scanner, _stub = recipe_scanner
recipes_dir = Path(config.loras_roots[0]) / "recipes"
_write_progress_recipe_files(recipes_dir, 2)
ws_stub = RecordingRecipeWebSocketManager()
monkeypatch.setattr(recipe_scanner_module, "ws_manager", ws_stub)
# Invalidate the persistent cache so the sync path performs a full
# directory scan, exactly like a force refresh but without progress
# reporting (this is how initialize_in_background invokes it).
scanner._persistent_cache.save_cache([], {})
scanner._initialize_recipe_cache_sync()
assert ws_stub.broadcasts == []
def test_sync_init_reports_error_broadcast(
tmp_path: Path, monkeypatch, recipe_scanner
):
scanner, _stub = recipe_scanner
recipes_dir = Path(config.loras_roots[0]) / "recipes"
_write_progress_recipe_files(recipes_dir, 1)
ws_stub = RecordingRecipeWebSocketManager()
monkeypatch.setattr(recipe_scanner_module, "ws_manager", ws_stub)
scanner._persistent_cache.save_cache([], {})
def raising_scan(self, recipes_dir, progress_loop=None):
raise RuntimeError("boom")
monkeypatch.setattr(RecipeScanner, "_full_directory_scan_sync", raising_scan)
scanner._initialize_recipe_cache_sync(report_progress=True)
messages = ws_stub.broadcasts
assert messages[0]["status"] == "started"
assert messages[-1]["status"] == "error"
assert messages[-1]["error"] == "boom"