Coverage for src/lilbee/server/handlers/wiki.py: 100%

1"""Wiki build, synthesis, and single-page generation handlers (SSE-streamed).""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import threading 

7from collections.abc import AsyncGenerator, Callable 

8from functools import partial 

9from typing import Any 

10 

11from lilbee.app import services as svc_mod 

12from lilbee.core.config import cfg 

13from lilbee.runtime.progress import DetailedProgressCallback 

14from lilbee.server.handlers.sse import SseStream 

15from lilbee.server.models import WikiGenerateResult 

16from lilbee.wiki import run_full_build, run_full_synthesize 

17 

18_Summary = dict[str, Any] 

19 

20 

21async def _wiki_run_stream( 

22 run: Callable[[DetailedProgressCallback, threading.Event], _Summary], label: str 

23) -> AsyncGenerator[str, None]: 

24 """Run a wiki job off the event loop, yielding its progress as SSE. 

25 

26 Emits wiki_phase and wiki_page events while the job runs, then a done event 

27 carrying the run summary. The run gets the stream's cancel event, so a 

28 client disconnect stops it at the next source boundary; without that the 

29 worker keeps building the whole corpus and holds the wiki build mutex, 

30 blocking every other surface. 

31 """ 

32 sse = SseStream() 

33 

34 async def _run() -> _Summary: 

35 try: 

36 return await asyncio.to_thread(run, sse.callback, sse.cancel) 

37 finally: 

38 sse.queue.put_nowait(None) 

39 

40 task = asyncio.create_task(_run()) 

41 async for event in sse.drain(task, label): 

42 yield event 

43 frame = sse.terminal_frame(task, dict) 

44 if frame is not None: 

45 yield frame 

46 

47 

48async def wiki_build_stream() -> AsyncGenerator[str, None]: 

49 """Build the concept and entity wiki, streaming progress.""" 

50 async for event in _wiki_run_stream( 

51 lambda on_progress, cancel: dict(run_full_build(cfg, on_progress, cancel)), 

52 "Wiki build stream", 

53 ): 

54 yield event 

55 

56 

57async def wiki_synthesize_stream() -> AsyncGenerator[str, None]: 

58 """Generate synthesis pages for cross-source clusters, streaming progress.""" 

59 async for event in _wiki_run_stream( 

60 lambda on_progress, cancel: dict(run_full_synthesize(cfg, on_progress, cancel)), 

61 "Wiki synthesize stream", 

62 ): 

63 yield event 

64 

65 

66async def wiki_generate_stream(slug: str) -> AsyncGenerator[str, None]: 

67 """Generate one indexed page, streaming progress. 

68 

69 The done event carries the written page's read slug and path; a stale 

70 index entry surfaces as an error event. 

71 """ 

72 async for event in _wiki_run_stream(partial(_generate_one_page, slug), "Wiki generate stream"): 

73 yield event 

74 

75 

76def _generate_one_page( 

77 slug: str, on_progress: DetailedProgressCallback, cancel: threading.Event 

78) -> _Summary: 

79 """Write one indexed page and shape the done payload for the wire.""" 

80 from lilbee.wiki.browse import page_slug 

81 from lilbee.wiki.lazy import generate_stub_page 

82 

83 path = generate_stub_page( 

84 slug, svc_mod.get_services().store, on_progress=on_progress, cancel=cancel 

85 ) 

86 if path is None: 

87 raise RuntimeError(f"index entry for {slug} is stale; its sources are gone") 

88 result = WikiGenerateResult( 

89 slug=page_slug(path, cfg.data_root / cfg.wiki_dir), path=path.as_posix() 

90 ) 

91 return result.model_dump()