diff --git a/packages/pickled-core/src/pickled_core/mine/features_stage.py b/packages/pickled-core/src/pickled_core/mine/features_stage.py index 49f1ab1..1354622 100644 --- a/packages/pickled-core/src/pickled_core/mine/features_stage.py +++ b/packages/pickled-core/src/pickled_core/mine/features_stage.py @@ -47,7 +47,7 @@ def _filter_story_paths( return filtered -async def _draft_one( +def _draft_one_sync( story_path: Path, feature_path: Path, llm: LLMClient, @@ -55,6 +55,15 @@ async def _draft_one( overwrite: bool, interactive: bool, ) -> FeatureResult: + """Draft one feature synchronously (used by both serial and parallel paths). + + This function performs blocking I/O (LLM HTTP call + file writes). The + parallel path schedules it on a worker thread via :func:`asyncio.to_thread` + so that ``--max-parallel`` actually overlaps LLM calls. Calling the LLM + directly inside an ``async def`` would block the event loop and serialise + every draft, defeating the flag (and producing N×slower mining runs for + workspaces with many surfaces). + """ surface_id = _story_surface_id(story_path) if feature_path.is_file() and not overwrite: return FeatureResult( @@ -122,7 +131,8 @@ async def _run_parallel( async def _one(story_path: Path, feature_path: Path) -> None: async with sem: results.append( - await _draft_one( + await asyncio.to_thread( + _draft_one_sync, story_path, feature_path, llm, @@ -178,19 +188,16 @@ def run_features( _run_parallel(jobs, llm, overwrite=overwrite, max_parallel=max_parallel) ) else: - results = [] - for story_path, feature_path in jobs: - results.append( - asyncio.run( - _draft_one( - story_path, - feature_path, - llm, - overwrite=overwrite, - interactive=True, - ) - ) + results = [ + _draft_one_sync( + story_path, + feature_path, + llm, + overwrite=overwrite, + interactive=True, ) + for story_path, feature_path in jobs + ] return FeaturesStageResult(output_dir=paths.root, results=results, skipped_entire_stage=False) diff --git a/packages/pickled-core/tests/test_mine_features_stage.py b/packages/pickled-core/tests/test_mine_features_stage.py index f2a2f7a..6ad0f12 100644 --- a/packages/pickled-core/tests/test_mine_features_stage.py +++ b/packages/pickled-core/tests/test_mine_features_stage.py @@ -2,6 +2,8 @@ from __future__ import annotations +import threading +import time from pathlib import Path from typing import Any @@ -92,3 +94,70 @@ class _Client: assert not result.skipped_entire_stage assert len(result.results) == 1 assert result.results[0].surface_id == "tiny_target_greet" + + +def test_max_parallel_actually_overlaps_llm_calls( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """`--max-parallel N` must let up to N LLM drafts run concurrently. + + Regression for a bug where ``_draft_one`` was an ``async def`` that + invoked the synchronous ``FeatureDrafter.draft_from_story`` directly: + every coroutine blocked the event loop for the entire LLM call, so + the asyncio.Semaphore never released to a waiting task and the + advertised ``--max-parallel`` flag silently degraded to 1-way + serial execution. With 50 surfaces × 5s LLM calls, that turns a + ~63 s mining run with ``--max-parallel 8`` into ~250 s. + """ + stories = tmp_path / "stories" + stories.mkdir() + for index in range(8): + (stories / f"surf_{index}.story.md").write_text("# story\n", encoding="utf-8") + + state: dict[str, int] = {"current": 0, "peak": 0} + lock = threading.Lock() + + class _SlowDrafter: + def __init__(self, llm: Any) -> None: + _ = llm + + def draft_from_story(self, story: str) -> DraftResult: + _ = story + with lock: + state["current"] += 1 + if state["current"] > state["peak"]: + state["peak"] = state["current"] + try: + time.sleep(0.2) + finally: + with lock: + state["current"] -= 1 + return DraftResult( + text="Feature: x\n\n Scenario: y\n Given z\n", + rationale="", + warnings=(), + ) + + monkeypatch.setattr("pickled_core.mine.features_stage.FeatureDrafter", _SlowDrafter) + + class _Client: + pass + + started = time.time() + result = run_features( + tmp_path, + llm=_Client(), # type: ignore[arg-type] + quick=True, + overwrite=True, + max_parallel=4, + ) + elapsed = time.time() - started + + assert len(result.results) == 8 + # 4-way parallelism over 8 jobs of 0.2s each: ideal ~0.4s, accept up to 1.5s + # for slow CI jitter; serial would take ~1.6s and we want a clear gap. + assert state["peak"] >= 2, ( + f"max-parallel=4 should overlap LLM calls; observed peak={state['peak']}" + ) + assert elapsed < 1.5, f"max-parallel=4 should run in well under serial time; got {elapsed:.2f}s"