diff --git a/src/youtube_extension/backend/code_generator.py b/src/youtube_extension/backend/code_generator.py index 62f1640a9..bc13b3543 100644 --- a/src/youtube_extension/backend/code_generator.py +++ b/src/youtube_extension/backend/code_generator.py @@ -10,17 +10,117 @@ template-based generation that still produces unique, video-specific output. """ +import asyncio +import contextlib import hashlib import json import logging import os +import shutil import tempfile from pathlib import Path -from typing import Any, Optional +from typing import Any, Callable, Optional, TypeVar from urllib.parse import parse_qs, urlparse logger = logging.getLogger(__name__) +_T = TypeVar("_T") + + +async def _run_offloop(func: Callable[..., _T], /, *args: Any, **kwargs: Any) -> _T: + """Run a blocking callable off the loop without abandoning it on cancellation. + + ``asyncio.to_thread`` cannot stop a worker thread that has already started, + so a cancellation delivered while the worker runs would return control to the + caller while the thread keeps mutating the scaffold on disk. Shield the worker + and wait for it to settle before re-raising ``CancelledError`` -- mirroring + ``_run_sync_rpc`` in ``services/cloud/cloud_tasks_queue.py`` -- so any + higher-level cleanup runs against a quiescent filesystem, never a live writer. + """ + task = asyncio.ensure_future(asyncio.to_thread(func, *args, **kwargs)) + cancelled = False + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + cancelled = True + except Exception: + if not cancelled: + raise + if cancelled: + # Surface the worker's own failure if it raised; otherwise honour the + # cancellation now that the thread has finished touching the disk. + with contextlib.suppress(Exception): + task.result() + raise asyncio.CancelledError + return task.result() + + +def _safe_rmtree(path: Path) -> None: + """Best-effort removal of a scaffold tree, tolerating absence and races.""" + shutil.rmtree(path, ignore_errors=True) + + +async def _make_scaffold_dir(prefix: str) -> Path: + """Create a temp scaffold directory off the loop, cleaning it up if the + caller is cancelled while ``mkdtemp`` is still running. + + ``asyncio.to_thread`` cannot interrupt ``mkdtemp`` once it has started, so a + cancellation delivered during the hop can still leave a directory on disk + whose path never reaches the caller -- the assignment that would hand it to + a higher-level cleanup scope never happens. This helper owns that window: it + waits for the worker to settle, removes any directory the worker created, + and only then propagates the ``CancelledError``. The removal is synchronous + because the directory is freshly created and empty, and cleanup on the + cancellation path must complete rather than risk a second cancellation. + """ + task = asyncio.ensure_future(asyncio.to_thread(tempfile.mkdtemp, prefix=prefix)) + cancelled = False + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + cancelled = True + except Exception: + if not cancelled: + raise + if cancelled: + created: Optional[str] = None + with contextlib.suppress(Exception): + created = task.result() + if created is not None: + _safe_rmtree(Path(created)) + raise asyncio.CancelledError + return Path(task.result()) + + +#: One scaffolding step. ``(path, None)`` creates a directory; ``(path, text)`` +#: writes a file. Steps are applied in list order. +WritePlan = list[tuple[Path, Optional[str]]] + + +def _apply_write_plan(plan: WritePlan) -> None: + """Apply an ordered scaffolding plan on the calling thread. + + This is the only place the generators touch the filesystem. It is written + as a plain synchronous function so callers can hand the whole batch to + ``asyncio.to_thread`` in a single hop, rather than paying a hop per file. + + Steps are applied strictly in order, so a directory step always lands + before the file steps that depend on it and the on-disk result matches + what the equivalent inline sequence produced. + + No exception is suppressed. A failing step raises exactly the error the + equivalent inline call would have raised, on the same step, leaving the + preceding steps applied -- identical to the previous behaviour. + """ + for path, content in plan: + if content is None: + path.mkdir(exist_ok=True) + else: + with open(path, "w") as handle: + handle.write(content) + def _extract_video_id(video_url: str) -> Optional[str]: """Extract the YouTube video ID from a URL, returning None if not found.""" @@ -149,31 +249,41 @@ async def generate_project(self, video_analysis: dict[str, Any], project_config: if build_plan: video_analysis["build_plan"] = build_plan - # Create temporary project directory - temp_dir = tempfile.mkdtemp(prefix="uvai_project_") - project_path = Path(temp_dir) - - # Generate project structure based on type - if project_type == "web": - result = await self._generate_web_project(project_path, video_analysis, technologies, features) - elif project_type == "api": - result = await self._generate_api_project(project_path, video_analysis, technologies, features) - elif project_type == "mobile": - result = await self._generate_mobile_project(project_path, video_analysis, technologies, features) - else: - result = await self._generate_web_project(project_path, video_analysis, technologies, features) - - result["project_path"] = str(project_path) - result["project_type"] = project_type - result["technologies"] = technologies - result["features"] = features - # Include the structured BuildPlan artifact in the result so that - # callers (e.g. API endpoints and tests) can inspect it. - if build_plan is not None: - result["build_plan"] = build_plan - - logger.info(f"✅ Project generated successfully at {project_path}") - return result + # Create temporary project directory. ``_make_scaffold_dir`` owns + # cleanup for the window where cancellation lands *during* mkdtemp: + # the directory can already exist on disk before its path reaches + # the cleanup scope below, so the helper removes it itself. + project_path = await _make_scaffold_dir("uvai_project_") + + # Past this point the path is known, so a cancelled or failed + # generation removes the scaffold here. The write hops are shielded + # (see ``_run_offloop``), so by the time cleanup runs the filesystem + # is settled and rmtree does not race a live writer. + try: + # Generate project structure based on type + if project_type == "web": + result = await self._generate_web_project(project_path, video_analysis, technologies, features) + elif project_type == "api": + result = await self._generate_api_project(project_path, video_analysis, technologies, features) + elif project_type == "mobile": + result = await self._generate_mobile_project(project_path, video_analysis, technologies, features) + else: + result = await self._generate_web_project(project_path, video_analysis, technologies, features) + + result["project_path"] = str(project_path) + result["project_type"] = project_type + result["technologies"] = technologies + result["features"] = features + # Include the structured BuildPlan artifact in the result so that + # callers (e.g. API endpoints and tests) can inspect it. + if build_plan is not None: + result["build_plan"] = build_plan + + logger.info(f"✅ Project generated successfully at {project_path}") + return result + except (Exception, asyncio.CancelledError): + await _run_offloop(_safe_rmtree, project_path) + raise except Exception as e: logger.error(f"❌ Project generation failed: {e}") @@ -234,25 +344,13 @@ async def _generate_react_project(self, project_path: Path, video_analysis: dict package_json["dependencies"]["tailwindcss"] = "^3.3.0" package_json["devDependencies"] = {"autoprefixer": "^10.4.14", "postcss": "^8.4.24"} - # Write package.json - with open(project_path / "package.json", "w") as f: - json.dump(package_json, f, indent=2) - - # Create src directory src_dir = project_path / "src" - src_dir.mkdir(exist_ok=True) - - # Create public directory public_dir = project_path / "public" - public_dir.mkdir(exist_ok=True) - # Generate index.html - index_html = self._generate_index_html(title) - with open(public_dir / "index.html", "w") as f: - f.write(index_html) - - # Generate main App component + # Build every artifact in memory first. This is pure string work and + # stays on the event loop; only the disk I/O below is offloaded. technologies = extracted_info.get("technologies", []) + index_html = self._generate_index_html(title) app_component = self._generate_react_app_component( title, technologies, @@ -261,23 +359,23 @@ async def _generate_react_project(self, project_path: Path, video_analysis: dict summary, key_concepts ) - with open(src_dir / "App.js", "w") as f: - f.write(app_component) - - # Generate index.js index_js = self._generate_react_index_js() - with open(src_dir / "index.js", "w") as f: - f.write(index_js) - - # Generate CSS app_css = self._generate_app_css(features) - with open(src_dir / "App.css", "w") as f: - f.write(app_css) - - # Generate README readme = self._generate_readme(title, "React", video_analysis) - with open(project_path / "README.md", "w") as f: - f.write(readme) + + # ``json.dumps`` produces exactly the bytes ``json.dump`` would have + # written for the same arguments; neither appends a trailing newline. + plan: WritePlan = [ + (project_path / "package.json", json.dumps(package_json, indent=2)), + (src_dir, None), + (public_dir, None), + (public_dir / "index.html", index_html), + (src_dir / "App.js", app_component), + (src_dir / "index.js", index_js), + (src_dir / "App.css", app_css), + (project_path / "README.md", readme), + ] + await _run_offloop(_apply_write_plan, plan) return { "framework": "react", @@ -313,25 +411,25 @@ async def _generate_vanilla_js_project(self, project_path: Path, video_analysis: summary, key_concepts ) - with open(project_path / "index.html", "w") as f: - f.write(index_html) # Generate main.js — NOW uses video-specific content main_js = self._generate_vanilla_main_js( title, tutorial_steps, features, key_concepts, fingerprint ) - with open(project_path / "main.js", "w") as f: - f.write(main_js) # Generate styles.css — NOW uses video-derived accent color styles_css = self._generate_vanilla_styles_css(title, features, fingerprint) - with open(project_path / "styles.css", "w") as f: - f.write(styles_css) # Generate README readme = self._generate_readme(title, "Vanilla JavaScript", video_analysis) - with open(project_path / "README.md", "w") as f: - f.write(readme) + + plan: WritePlan = [ + (project_path / "index.html", index_html), + (project_path / "main.js", main_js), + (project_path / "styles.css", styles_css), + (project_path / "README.md", readme), + ] + await _run_offloop(_apply_write_plan, plan) return { "framework": "vanilla", @@ -369,18 +467,19 @@ async def _generate_python_api(self, project_path: Path, video_analysis: dict, f requirements.append("sqlalchemy==2.0.23") if "authentication" in features: requirements.append("python-jose[cryptography]==3.3.0") - with open(project_path / "requirements.txt", "w") as f: - f.write("\n".join(requirements)) # Generate main.py main_py = self._generate_fastapi_main(title, features) - with open(project_path / "main.py", "w") as f: - f.write(main_py) # Generate README readme = self._generate_readme(title, "Python FastAPI", video_analysis) - with open(project_path / "README.md", "w") as f: - f.write(readme) + + plan: WritePlan = [ + (project_path / "requirements.txt", "\n".join(requirements)), + (project_path / "main.py", main_py), + (project_path / "README.md", readme), + ] + await _run_offloop(_apply_write_plan, plan) return { "framework": "fastapi", diff --git a/tests/unit/test_code_generator.py b/tests/unit/test_code_generator.py index 62d8c3d60..ce0541a2a 100644 --- a/tests/unit/test_code_generator.py +++ b/tests/unit/test_code_generator.py @@ -5,14 +5,18 @@ import asyncio import json import tempfile +import threading +import time from pathlib import Path import pytest from youtube_extension.backend.code_generator import ( ProjectCodeGenerator, + _apply_write_plan, _build_title, _extract_video_id, + _run_offloop, get_code_generator, ) @@ -800,3 +804,348 @@ def test_returns_same_instance_on_repeated_calls(self): gen2 = get_code_generator(use_ai_generation=False) assert gen1 is gen2 cg_module._code_generator = None # clean up + + +# =========================================================================== +# Scaffolding disk I/O runs off the event loop (issue #1250) +# =========================================================================== + + +class TestScaffoldingWritesOffLoop: + """Every filesystem call must land on a worker thread, not the loop. + + These assert *thread identity* rather than elapsed time: a wall-clock + threshold would be flaky under CI contention and would still pass if the + work ran on the loop but happened to be fast. + """ + + @staticmethod + def _recording_open(record: list[str]): + """Wrap ``builtins.open`` so each call records its executing thread.""" + import builtins + + real_open = builtins.open + + def _tracked(*args, **kwargs): + record.append(threading.current_thread().name) + return real_open(*args, **kwargs) + + return _tracked + + @pytest.mark.parametrize( + "generator_name", + [ + "_generate_react_project", + "_generate_vanilla_js_project", + "_generate_python_api", + ], + ) + async def test_generator_writes_never_touch_loop_thread( + self, monkeypatch, tmp_path, generator_name + ): + gen = ProjectCodeGenerator(use_ai_generation=False) + analysis = _build_video_analysis( + "Dashboard Tutorial", "auJzb1D-fag", "A summary.", ["state", "charts"] + ) + loop_thread = threading.current_thread().name + + seen: list[str] = [] + monkeypatch.setattr("builtins.open", self._recording_open(seen)) + + project = tmp_path / "project" + project.mkdir() + await getattr(gen, generator_name)(project, analysis, ["database"]) + + assert seen, "expected the generator to write at least one file" + offenders = [name for name in seen if name == loop_thread] + assert not offenders, ( + f"{generator_name} performed {len(offenders)} of {len(seen)} writes " + f"on the event loop thread ({loop_thread})" + ) + + async def test_mkdtemp_runs_off_loop(self, monkeypatch, tmp_path): + """The project directory itself is also created off-loop.""" + import youtube_extension.backend.code_generator as cg_module + + loop_thread = threading.current_thread().name + seen: list[str] = [] + real_mkdtemp = tempfile.mkdtemp + + def _tracked(*args, **kwargs): + seen.append(threading.current_thread().name) + return real_mkdtemp(*args, **kwargs) + + monkeypatch.setattr(cg_module.tempfile, "mkdtemp", _tracked) + + gen = ProjectCodeGenerator(use_ai_generation=False) + await gen.generate_project( + _build_video_analysis("T", "auJzb1D-fag", "S", []), + {"project_type": "web", "technologies": ["react"]}, + ) + + assert seen, "expected mkdtemp to be called" + assert loop_thread not in seen, ( + f"tempfile.mkdtemp ran on the event loop thread ({loop_thread})" + ) + + @pytest.mark.parametrize( + ("generator_name", "expected_files"), + [ + ("_generate_react_project", 6), + ("_generate_vanilla_js_project", 4), + ("_generate_python_api", 3), + ], + ) + async def test_batches_into_a_single_thread_hop( + self, monkeypatch, tmp_path, generator_name, expected_files + ): + """Cost is O(1) thread hops per generator, not O(files). + + Guards against a regression that offloads each write individually, + which would still pass the thread-identity tests above while paying + one context switch per file. + """ + import youtube_extension.backend.code_generator as cg_module + + real_to_thread = asyncio.to_thread + hops: list[str] = [] + + async def _counting(func, /, *args, **kwargs): + hops.append(getattr(func, "__name__", repr(func))) + return await real_to_thread(func, *args, **kwargs) + + monkeypatch.setattr(cg_module.asyncio, "to_thread", _counting) + + gen = ProjectCodeGenerator(use_ai_generation=False) + project = tmp_path / "project" + project.mkdir() + await getattr(gen, generator_name)( + project, + _build_video_analysis("T", "auJzb1D-fag", "S", []), + ["database"], + ) + + assert hops == ["_apply_write_plan"], ( + f"expected exactly one batched hop, got {hops}" + ) + written = [p for p in project.rglob("*") if p.is_file()] + assert len(written) == expected_files + + +class TestApplyWritePlan: + """Contract of the batched write helper itself.""" + + def test_applies_steps_in_order_so_dirs_precede_their_files(self, tmp_path): + nested = tmp_path / "src" + plan = [ + (nested, None), + (nested / "App.js", "console.log(1);"), + ] + _apply_write_plan(plan) + + assert nested.is_dir() + assert (nested / "App.js").read_text() == "console.log(1);" + + def test_directory_step_tolerates_an_existing_directory(self, tmp_path): + existing = tmp_path / "public" + existing.mkdir() + _apply_write_plan([(existing, None)]) # must not raise + assert existing.is_dir() + + def test_does_not_suppress_errors_and_leaves_earlier_steps_applied( + self, tmp_path + ): + """A failing step raises, exactly as the inline sequence did. + + Uses a real invalid path rather than a mock so the stdlib itself + produces the failure. + """ + good = tmp_path / "first.txt" + plan = [ + (good, "written"), + (tmp_path / "bad\x00name.txt", "never"), + (tmp_path / "third.txt", "unreached"), + ] + + with pytest.raises(ValueError, match="null"): + _apply_write_plan(plan) + + assert good.read_text() == "written" + assert not (tmp_path / "third.txt").exists() + + def test_writes_content_verbatim_without_adding_a_trailing_newline( + self, tmp_path + ): + target = tmp_path / "package.json" + payload = json.dumps({"name": "x", "version": "1.0.0"}, indent=2) + _apply_write_plan([(target, payload)]) + assert target.read_bytes() == payload.encode() + + +class TestRunOffloop: + """``_run_offloop`` must not abandon a worker thread on cancellation. + + ``asyncio.to_thread`` cannot interrupt a thread that has already started, so + the helper shields the worker and waits for it to settle before propagating + the cancellation — otherwise a caller's cleanup would race a live writer. + """ + + async def test_returns_worker_result_on_happy_path(self): + assert await _run_offloop(lambda a, b: a + b, 2, 3) == 5 + + async def test_propagates_worker_exception(self): + def _boom(): + raise ValueError("worker failed") + + with pytest.raises(ValueError, match="worker failed"): + await _run_offloop(_boom) + + async def test_cancellation_waits_for_worker_to_finish(self): + started = threading.Event() + finished = threading.Event() + + def _slow(): + started.set() + # Simulate a blocking write already in flight on the worker thread. + time.sleep(0.2) + finished.set() + + task = asyncio.ensure_future(_run_offloop(_slow)) + # Let the worker actually start before we cancel. + while not started.is_set(): + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + # The worker must have run to completion, not been abandoned mid-flight. + assert finished.is_set(), ( + "cancellation abandoned the worker before it finished" + ) + + +class TestScaffoldingCancellationSafety: + """A cancelled or failed generation must not leak its scaffold directory. + + Before the write hops were offloaded there was no cancellation point once + scaffolding began, so the directory always either fully materialised (and + its path was returned) or was never created. Offloading introduced ``await`` + points; ``generate_project`` therefore has to clean up a directory whose + path it will never hand back. + """ + + async def test_cancelled_generation_removes_orphan_scaffold( + self, monkeypatch, tmp_path + ): + import youtube_extension.backend.code_generator as cg_module + + project_dir = tmp_path / "uvai_project_cancel" + monkeypatch.setattr( + cg_module.tempfile, "mkdtemp", _tempdir_factory(project_dir) + ) + + started = threading.Event() + release = threading.Event() + real_apply = cg_module._apply_write_plan + + def _blocking_apply(plan): + # Park the worker mid-scaffold so we can cancel while it "writes". + started.set() + release.wait(5) + real_apply(plan) + + monkeypatch.setattr(cg_module, "_apply_write_plan", _blocking_apply) + + gen = ProjectCodeGenerator(use_ai_generation=False) + task = asyncio.ensure_future( + gen.generate_project( + _build_video_analysis("T", "auJzb1D-fag", "S", []), + {"project_type": "web", "technologies": ["react"]}, + ) + ) + + while not started.is_set(): + await asyncio.sleep(0.01) + task.cancel() + # Cancellation is now in flight; let the shielded worker finish so the + # filesystem is settled before cleanup runs. + release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert not project_dir.exists(), ( + "cancelled generation leaked its scaffold directory" + ) + + async def test_cancelled_during_mkdtemp_removes_created_dir( + self, monkeypatch, tmp_path + ): + """Cancellation *during* mkdtemp must not leak the created directory. + + The directory can materialise on the worker thread before its path is + assigned into the cleanup scope, so `_make_scaffold_dir` has to remove it + itself before propagating the cancellation. + """ + import youtube_extension.backend.code_generator as cg_module + + created = tmp_path / "uvai_project_mkdtemp_cancel" + started = threading.Event() + release = threading.Event() + + def _blocking_mkdtemp(prefix): + # Park inside mkdtemp so we can cancel before it returns the path. + started.set() + release.wait(5) + created.mkdir() + return str(created) + + monkeypatch.setattr(cg_module.tempfile, "mkdtemp", _blocking_mkdtemp) + + gen = ProjectCodeGenerator(use_ai_generation=False) + task = asyncio.ensure_future( + gen.generate_project( + _build_video_analysis("T", "auJzb1D-fag", "S", []), + {"project_type": "web", "technologies": ["react"]}, + ) + ) + + while not started.is_set(): + await asyncio.sleep(0.01) + task.cancel() + # Let the shielded mkdtemp worker finish creating the directory. + release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert not created.exists(), ( + "cancellation during mkdtemp leaked the scaffold directory" + ) + + async def test_failed_generation_removes_orphan_scaffold( + self, monkeypatch, tmp_path + ): + import youtube_extension.backend.code_generator as cg_module + + project_dir = tmp_path / "uvai_project_fail" + monkeypatch.setattr( + cg_module.tempfile, "mkdtemp", _tempdir_factory(project_dir) + ) + + def _exploding_apply(plan): + raise RuntimeError("disk exploded") + + monkeypatch.setattr(cg_module, "_apply_write_plan", _exploding_apply) + + gen = ProjectCodeGenerator(use_ai_generation=False) + with pytest.raises(RuntimeError, match="disk exploded"): + await gen.generate_project( + _build_video_analysis("T", "auJzb1D-fag", "S", []), + {"project_type": "web", "technologies": ["react"]}, + ) + + assert not project_dir.exists(), ( + "failed generation leaked its scaffold directory" + )