diff --git a/.gitignore b/.gitignore index 7644ce20..0b89baef 100644 --- a/.gitignore +++ b/.gitignore @@ -49,6 +49,8 @@ tests/* !tests/test_scala_analyzer.py !tests/test_agent_limits.py !tests/test_module_scope.py +!tests/test_doc_write_guard.py +!tests/test_overview_generation.py # Jupyter *.ipynb diff --git a/codewiki/src/be/agent_tools/doc_guard.py b/codewiki/src/be/agent_tools/doc_guard.py new file mode 100644 index 00000000..5f0a11ea --- /dev/null +++ b/codewiki/src/be/agent_tools/doc_guard.py @@ -0,0 +1,191 @@ +"""Make a documentation agent actually write its page. + +A module agent can end its run with a text answer and never call +``str_replace_editor create`` (or call it only with arguments the tool +rejects). The run then "succeeds" but the page is missing, and the CLI only +notices at the very end; re-running repeats the same miss (issue #134). + +:func:`run_until_doc_written` checks the page after each run and, when it is +missing, continues the same conversation with a reminder naming the exact +file to create. Every miss is logged with the agent's final answer and a +summary of its tool calls, so a page that is still missing can be diagnosed. +""" + +from __future__ import annotations + +import logging +import os +from collections import Counter +from typing import Any + +from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + RetryPromptPart, + ToolCallPart, + ToolReturnPart, +) + +from codewiki.src.be.doc_layout import OVERVIEW_FILENAME, module_doc_file +from codewiki.src.be.module_naming import resolve_module_doc_path +from codewiki.src.be.overview_support import reusable_overview + +logger = logging.getLogger(__name__) + +# Follow-up turns sent to an agent that finished without writing its page. +DOC_WRITE_REMINDERS = 2 + +_EXCERPT_CHARS = 300 + + +class DocumentationWriteError(RuntimeError): + """An agent exhausted its attempts without creating its required page.""" + + +def doc_written( + docs_dir: str, module_name: str, module_path: list[str], module_tree: dict[str, Any] | None +) -> bool: + """Whether the page the agent for ``module_name`` must write exists.""" + if not module_path: + # The whole-repository agent writes {repo_name}.md, renamed to overview.md later. + return reusable_overview(os.path.join(docs_dir, f"{module_name}.md")) or reusable_overview( + os.path.join(docs_dir, OVERVIEW_FILENAME) + ) + path = resolve_module_doc_path(docs_dir, module_name, module_tree) + return path is not None and os.path.isfile(path) + + +def reminder_prompt(module_name: str, doc_file: str) -> str: + return ( + f"You finished without creating `{doc_file}`, so the documentation for " + f"`{module_name}` is missing. Create it now: call `str_replace_editor` with " + f'working_dir="docs", command="create", path="./{doc_file}" and the complete ' + f"Markdown page as file_text. If an earlier tool call returned an error, fix the " + f"arguments and try again. If the file already exists but is incomplete, replace " + f"its contents with str_replace_editor instead of create. " + f"Do not reply with the documentation as text." + ) + + +def _excerpt(text: Any) -> str: + text = " ".join(str(text or "").split()) + if len(text) > _EXCERPT_CHARS: + return text[:_EXCERPT_CHARS] + "..." + return text or "" + + +def _editor_command(part: ToolCallPart) -> str: + try: + args = part.args_as_dict() + except Exception: # noqa: BLE001 — malformed args are reported as such + return "" + return f"{args.get('command')} {args.get('path') or args.get('file')}" + + +def summarize_tool_activity(messages: list[Any]) -> str: + """One-line summary of an agent run: tool calls, editor writes and their results, + rejected tool calls and why the last response ended.""" + calls: Counter[str] = Counter() + writes: dict[str, str] = {} + results: dict[str, str] = {} + rejected: list[str] = [] + finish_reason = None + for message in messages: + if isinstance(message, ModelResponse): + finish_reason = getattr(message, "finish_reason", None) + for part in message.parts: + if isinstance(part, ToolCallPart): + calls[part.tool_name] += 1 + if part.tool_name == "str_replace_editor": + command = _editor_command(part) + if not command.startswith("view "): + writes[part.tool_call_id] = command + elif isinstance(message, ModelRequest): + for part in message.parts: + if isinstance(part, ToolReturnPart) and part.tool_call_id in writes: + lines = str(part.content).strip().splitlines() + results[part.tool_call_id] = _excerpt(lines[0] if lines else "") + elif isinstance(part, RetryPromptPart) and part.tool_name: + rejected.append(f"{part.tool_name}: {_excerpt(part.model_response())}") + + summary = [ + "tool calls: " + + (", ".join(f"{name} x{count}" for name, count in sorted(calls.items())) or "none") + ] + if writes: + summary.append( + "editor writes: " + + "; ".join( + f"{command} -> {results.get(call_id, 'no result')}" + for call_id, command in writes.items() + ) + ) + else: + summary.append("editor writes: none") + if rejected: + summary.append(f"rejected tool calls ({len(rejected)}): " + "; ".join(rejected[-3:])) + summary.append(f"last finish_reason: {finish_reason}") + return " | ".join(summary) + + +def run_usage(result: Any) -> Any: + """``usage`` is a property in pydantic-ai 2.x and a method in earlier releases.""" + usage = getattr(result, "usage", None) + return usage() if callable(usage) else usage + + +async def run_until_doc_written( + agent: Any, + user_prompt: str, + *, + deps: Any, + module_name: str, + module_path: list[str], + layout: str | None, + usage_limits: Any = None, +) -> tuple[Any, Any]: + """Run ``agent`` and remind it (up to :data:`DOC_WRITE_REMINDERS` times) to + write its page if it ended without one. + + Returns the last run's result and the token usage summed over all runs. + Each reminder is a separate run with its own ``usage_limits``, so a first + run that used most of its request budget still gets a chance to write. + Raises :class:`DocumentationWriteError` if all attempts miss the page. + """ + doc_file = module_doc_file(module_name, module_path, layout) + result = await agent.run(user_prompt, deps=deps, usage_limits=usage_limits) + usage = run_usage(result) + for attempt in range(1, DOC_WRITE_REMINDERS + 1): + if doc_written(deps.absolute_docs_path, module_name, module_path, deps.module_tree): + return result, usage + logger.warning( + "Agent for %s finished without writing %s (reminder %d/%d). Final answer: %s | %s", + module_name, + doc_file, + attempt, + DOC_WRITE_REMINDERS, + _excerpt(result.output), + summarize_tool_activity(result.new_messages()), + ) + result = await agent.run( + reminder_prompt(module_name, doc_file), + deps=deps, + usage_limits=usage_limits, + message_history=result.all_messages(), + ) + more = run_usage(result) + usage = more if usage is None else usage if more is None else usage + more + if not doc_written(deps.absolute_docs_path, module_name, module_path, deps.module_tree): + logger.error( + "Agent for %s still has not written %s after %d reminders. Final answer: %s | %s", + module_name, + doc_file, + DOC_WRITE_REMINDERS, + _excerpt(result.output), + summarize_tool_activity(result.new_messages()), + ) + raise DocumentationWriteError( + f"Agent for {module_name} has not written {doc_file} after " + f"{DOC_WRITE_REMINDERS} reminders" + ) + return result, usage diff --git a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py index 69a3d90b..29ff31d9 100644 --- a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py +++ b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py @@ -2,6 +2,7 @@ from pydantic_ai.usage import UsageLimits from codewiki.src.be.agent_tools.deps import CodeWikiDeps +from codewiki.src.be.agent_tools.doc_guard import run_until_doc_written from codewiki.src.be.doc_layout import config_layout, module_doc_file from codewiki.src.be.module_naming import skipped_report, plan_sub_module_specs, sub_module_report from codewiki.src.be.agent_tools.read_code_components import read_code_components_tool @@ -146,7 +147,8 @@ async def generate_sub_module_documentation( try: with usage_target(sub_module_name): - await sub_agent.run( + await run_until_doc_written( + sub_agent, format_user_prompt( module_name=deps.current_module_name, core_component_ids=core_component_ids, @@ -156,6 +158,9 @@ async def generate_sub_module_documentation( layout=layout, ), deps=ctx.deps, + module_name=sub_module_name, + module_path=list(deps.path_to_current_module), + layout=layout, usage_limits=UsageLimits(request_limit=ctx.deps.config.request_limit), ) finally: diff --git a/codewiki/src/be/caw_backend.py b/codewiki/src/be/caw_backend.py index 86d124d0..47075811 100644 --- a/codewiki/src/be/caw_backend.py +++ b/codewiki/src/be/caw_backend.py @@ -30,6 +30,7 @@ from caw import ToolGroup from codewiki.src.be.agent_tools.deps import CodeWikiDeps +from codewiki.src.be.agent_tools.doc_guard import DocumentationWriteError, doc_written from codewiki.src.be.backend import AgentReply, LLMBackend from codewiki.src.be.module_scope import ( complete_file_candidates, @@ -44,7 +45,8 @@ format_user_prompt, ) from codewiki.src.be.utils import is_complex_module, set_main_loop -from codewiki.src.be.doc_layout import config_layout, find_doc +from codewiki.src.be.doc_layout import config_layout, find_doc, module_doc_file +from codewiki.src.be.overview_support import reusable_overview, root_recovery_prompt from codewiki.src.config import MODULE_TREE_FILENAME, OVERVIEW_FILENAME, Config from codewiki.src.utils import file_manager @@ -393,13 +395,13 @@ def _run_module_agent_sync( # or a resume after a partial run silently skips every missing one. if not module_path: overview_docs_path = os.path.join(working_dir, OVERVIEW_FILENAME) - if os.path.exists(overview_docs_path): + if reusable_overview(overview_docs_path): logger.info("✓ Overview docs already exists at %s", overview_docs_path) return module_tree docs_path = find_doc(working_dir, module_name, module_tree) if module_path else None if docs_path is None: docs_path = os.path.join(working_dir, f"{module_name}.md") - if os.path.exists(docs_path): + if os.path.isfile(docs_path) and (module_path or reusable_overview(docs_path)): logger.info("✓ Module docs already exists at %s", docs_path) return module_tree @@ -463,6 +465,8 @@ def _run_module_agent_sync( module_path=list(module_path), layout=config_layout(config), ) + if not module_path: + user_prompt = root_recovery_prompt(user_prompt, working_dir, module_name) # caw forks claude / codex via subprocess.Popen without a cwd, so the # child CLI inherits Python's cwd — typically the repo root where the @@ -508,6 +512,7 @@ def _run_module_agent_sync( traj = agent.completion(user_prompt) finally: os.chdir(original_cwd) + file_manager.save_json(deps.module_tree, module_tree_path) # Nested sub-module agents are separate CLI processes that record # their own usage, so this covers only this module's own run. usage = self._record_usage( @@ -524,6 +529,18 @@ def _run_module_agent_sync( traj.total_tool_calls, (usage or {}).get("cost_usd", 0.0), ) + if not doc_written(working_dir, module_name, list(module_path), deps.module_tree): + doc_file = module_doc_file(module_name, list(module_path), config_layout(config)) + logger.error( + "Agent for %s finished without writing %s (turns=%d, tool_calls=%d). " + "Final answer: %s", + module_name, + doc_file, + traj.num_turns, + traj.total_tool_calls, + " ".join((getattr(traj, "result", None) or "").split())[:300] or "", + ) + raise DocumentationWriteError(f"Agent for {module_name} has not written {doc_file}") validate_module_file_coverage(deps.module_tree, components) file_manager.save_json(deps.module_tree, module_tree_path) return deps.module_tree diff --git a/codewiki/src/be/documentation_generator.py b/codewiki/src/be/documentation_generator.py index 94f7e9e2..84e0a6b8 100644 --- a/codewiki/src/be/documentation_generator.py +++ b/codewiki/src/be/documentation_generator.py @@ -1,6 +1,7 @@ import json import logging import os +from pathlib import Path import traceback from copy import deepcopy from typing import Any @@ -23,6 +24,14 @@ relative_link, ) from codewiki.src.be.utils import validate_mermaid_content +from codewiki.src.be.overview_support import ( + bounded_overview_prompt, + overview_content, + overview_error, + preserve_incomplete_overviews, + reusable_overview, + save_overview, +) from codewiki.src.be.module_naming import ( dedupe_module_tree_names, find_missing_module_docs, @@ -161,10 +170,9 @@ def build_overview_structure( """Build structure for overview generation with 1-depth children doc paths and target indicator. Children docs are referenced by absolute file path (``docs_path``) - rather than inlined, and ``components`` lists are stripped: inlining - the full tree plus docs blew past provider input caps (codex rejects - turns over 1,048,576 chars) on large repos. The overview agent reads - the referenced files itself. + rather than inlined, and ``components`` lists are stripped. The + completion prompt separately includes bounded child excerpts so API + backends need no file-reading tools. """ processed_module_tree = deepcopy(module_tree) @@ -185,9 +193,9 @@ def build_overview_structure( for child_name, child_info in module_info.items(): child_rel = module_doc_relpath(module_path + [child_name], layout) child_docs_path = self._resolve_child_docs_path(working_dir, child_name, module_tree) + child_info["link"] = relative_link(page_rel, child_rel) if child_docs_path is not None: child_info["docs_path"] = child_docs_path - child_info["link"] = relative_link(page_rel, child_rel) else: logger.warning(f"Module docs not found at {os.path.join(working_dir, child_rel)}") child_info["docs_path"] = None @@ -274,6 +282,7 @@ async def _generate_modules_in_order( # Process modules in dependency order final_module_tree = module_tree processed_modules = set() + failures = [] if len(module_tree) > 0: for module_path, module_name in processing_order: @@ -312,6 +321,7 @@ async def _generate_modules_in_order( processed_modules.add(module_key) except Exception as e: # noqa: BLE001 — one failed module must not abort the run + failures.append(f"{module_key}: {e}") logger.error(f"Failed to process module {module_key}: {e!s}") logger.error(f"Traceback: {traceback.format_exc()}") continue @@ -324,24 +334,20 @@ async def _generate_modules_in_order( else: logger.info("Processing whole repo because repo can fit in the context window") repo_name = os.path.basename(os.path.normpath(self.config.repo_path)) - final_module_tree = await self.backend.run_module_agent( - module_name=repo_name, - components=components, - core_component_ids=complete_file_candidates(leaf_nodes, components), - module_path=[], - working_dir=working_dir, - ) - - # save final_module_tree to module_tree.json - validate_module_file_coverage(final_module_tree, components) - file_manager.save_json( - final_module_tree, os.path.join(working_dir, MODULE_TREE_FILENAME) - ) - - # rename repo_name.md to overview.md repo_overview_path = os.path.join(working_dir, f"{repo_name}.md") - if os.path.exists(repo_overview_path): - os.rename(repo_overview_path, os.path.join(working_dir, OVERVIEW_FILENAME)) + overview_path = os.path.join(working_dir, OVERVIEW_FILENAME) + with preserve_incomplete_overviews([repo_overview_path, overview_path]): + final_module_tree = await self.backend.run_module_agent( + module_name=repo_name, + components=components, + core_component_ids=complete_file_candidates(leaf_nodes, components), + module_path=[], + working_dir=working_dir, + ) + validate_module_file_coverage(final_module_tree, components) + file_manager.save_json(final_module_tree, module_tree_path) + if os.path.exists(repo_overview_path): + os.replace(repo_overview_path, overview_path) # Also check resumed trees whose agents returned early because their # pages already existed. Do not silently bless an old partial tree. @@ -350,6 +356,9 @@ async def _generate_modules_in_order( # Move pages an agent saved in the wrong folder and fix links between pages organize_docs(working_dir, config_layout(self.config)) + if failures: + raise RuntimeError("Documentation failed for modules: " + "; ".join(failures)) + return working_dir async def generate_parent_module_docs( @@ -382,13 +391,13 @@ async def generate_parent_module_docs( existing_docs_path = ( resolve_module_doc_path(working_dir, module_name, module_tree) if module_path else None ) - if existing_docs_path is not None: + if existing_docs_path is not None and reusable_overview(existing_docs_path): logger.info(f"✓ Parent docs already exists at {existing_docs_path}") return module_tree - parent_docs_path = os.path.join( + parent_docs_path = existing_docs_path or os.path.join( working_dir, module_doc_relpath(module_path, config_layout(self.config)) ) - if os.path.exists(parent_docs_path): + if existing_docs_path is None and reusable_overview(parent_docs_path): logger.info(f"✓ Parent docs already exists at {parent_docs_path}") return module_tree @@ -412,52 +421,51 @@ async def generate_parent_module_docs( ) logger.debug(f"Overview prompt for {module_name}: {len(prompt)} chars") + child_info = repo_structure + for part in module_path: + child_info = child_info[part]["children"] + children = {} + for name, info in child_info.items(): + path = info.get("docs_path") + if path: + try: + children[name] = Path(path).read_text(encoding="utf-8") + except (OSError, UnicodeError) as exc: + logger.warning("Cannot read child documentation %s: %s", path, exc) + system_prompt = format_overview_system_prompt(self.config.get_prompt_addition()) + phase = "overview" if not module_path else "parent_modules" try: - request = prompt + correction = previous = "" for attempt in range(2): + request = bounded_overview_prompt( + prompt, children, system_prompt, correction=correction, previous=previous + ) with self.usage_ledger.in_phase(phase, target=module_name): parent_docs = self.backend.complete( request, - system_prompt=format_overview_system_prompt( - self.config.get_prompt_addition() - ), - ) - if not parent_docs: - raise RuntimeError( - f"LLM returned empty content for {module_name} overview " - f"(possible output truncation at max_tokens)" + system_prompt=system_prompt, ) - # Subscription backends may return raw Markdown instead of - # the requested wrapper. Validate either form before saving. - if "" in parent_docs and "" in parent_docs: - parent_content = ( - parent_docs.split("")[1].split("")[0].strip() - ) - else: - parent_content = parent_docs.strip() - report = await validate_mermaid_content(parent_content, repair=True) - if not report.errors: - if report.unavailable: - logger.warning(report.format(parent_docs_path)) - elif report.repairs: - logger.info(report.format(parent_docs_path)) - break + parent_content = overview_content(parent_docs) + correction = overview_error(parent_content) + previous = parent_content + if correction is None: + report = await validate_mermaid_content(parent_content, repair=True) + if not report.errors: + if report.unavailable: + logger.warning(report.format(parent_docs_path)) + elif report.repairs: + logger.info(report.format(parent_docs_path)) + save_overview(report.content, parent_docs_path) + break + correction = report.format(parent_docs_path) + previous = report.content if attempt == 1: raise RuntimeError( - f"Invalid Mermaid in {module_name} overview after one correction attempt:\n" - + report.format(parent_docs_path) + f"Invalid documentation for {module_name} at {parent_docs_path} " + f"after one correction attempt: {correction}" ) - request = ( - prompt - + "\n\nCorrect the Mermaid syntax in this previous draft. Return the complete " - "overview, preserving the other content.\n\n" - + report.content - + "\n\n" - + report.format(parent_docs_path) - ) - os.makedirs(os.path.dirname(parent_docs_path), exist_ok=True) - file_manager.save_text(report.content, parent_docs_path) + logger.warning("Retrying documentation for %s: %s", module_name, correction) logger.debug(f"Successfully generated parent documentation for: {module_name}") return module_tree diff --git a/codewiki/src/be/overview_support.py b/codewiki/src/be/overview_support.py new file mode 100644 index 00000000..6f69ba65 --- /dev/null +++ b/codewiki/src/be/overview_support.py @@ -0,0 +1,210 @@ +"""Bounded context, output checks and conservative recovery for overview pages.""" + +from __future__ import annotations + +import logging +import os +from contextlib import contextmanager +from pathlib import Path +import re +import tempfile + +from codewiki.src.be.prompt_template import MAX_USER_PROMPT_CHARS +from codewiki.src.be.utils import count_tokens, enc + +logger = logging.getLogger(__name__) + +OVERVIEW_INPUT_TOKENS = 32_000 +PREVIOUS_DRAFT_TOKENS = 8_000 +_TRUNCATED = "\n[Excerpt truncated; base the overview on the supplied content.]" +_PATH = re.compile(r"(? str | None: + """Recognize only empty pages and the specific preamble/path-only failure. + + Existing user prose is deliberately not subjected to new-output validation. + """ + if not content.strip(): + return "empty page" + match = _PREAMBLE.match(content.strip()) + if match: + rest = content.strip()[match.end() :] + if not _PATH.sub("", rest).strip(): + return "unfinished reading preamble with no documentation" + return None + + +def reusable_overview(path: str) -> bool: + if not os.path.isfile(path): + return False + content = Path(path).read_text(encoding="utf-8") + reason = incomplete_overview_reason(content) + if reason: + logger.warning("Regenerating incomplete documentation at %s: %s", path, reason) + return reason is None + + +def root_recovery_prompt(prompt: str, docs_dir: str, module_name: str) -> str: + paths = [Path(docs_dir, name) for name in (f"{module_name}.md", "overview.md")] + if any( + path.is_file() and incomplete_overview_reason(path.read_text(encoding="utf-8")) + for path in paths + ): + return prompt + ( + f"\n\nAn existing root page is incomplete. Write the finished page to {module_name}.md. " + "If that file already exists, replace its incomplete contents with str_replace_editor " + "instead of using create. Do not treat the incomplete overview.md as completed work." + ) + return prompt + + +@contextmanager +def preserve_incomplete_overviews(paths: list[str]): + """Keep old root pages available if an editing agent fails during recovery.""" + originals = {} + for path in paths: + if os.path.isfile(path): + content = Path(path).read_text(encoding="utf-8") + if incomplete_overview_reason(content): + originals[path] = content + try: + yield + except BaseException: + for path, content in originals.items(): + save_overview(content, path) + raise + + +def overview_content(response: str | None) -> str: + text = (response or "").strip() + if "" in text and "" in text: + text = text.split("", 1)[1].split("", 1)[0].strip() + return text + + +def overview_error(content: str) -> str | None: + reason = incomplete_overview_reason(content) + if reason: + return reason + heading = False + body_lines = [] + fence = None + for line in content.splitlines(): + stripped = line.strip() + marker = re.match(r"^(`{3,}|~{3,})", stripped) + if marker: + run = marker.group(1) + if fence is None: + fence = run + elif run[0] == fence[0] and len(run) >= len(fence): + fence = None + continue + if fence: + continue + if re.match(r"^#{1,6}\s+\S", stripped): + heading = True + else: + body_lines.append(stripped) + body = "\n".join(body_lines).strip() + if incomplete_overview_reason(body): + return "unfinished or empty explanatory body" + prose = _RELATIVE_PATH.sub("", _PATH.sub("", body)) + if not heading or not any(c.isalpha() for c in prose): + return "expected a Markdown heading and explanatory body outside code fences and file paths" + return None + + +def _clip(text: str, tokens: int) -> str: + encoded = enc.encode(text, disallowed_special=()) + if len(encoded) <= tokens: + return text + # Avoid decoding a partial UTF-8 character at the excerpt boundary. + prefix = enc.decode_bytes(encoded[: max(0, tokens)]).decode("utf-8", errors="ignore") + return prefix + _TRUNCATED + + +def bounded_overview_prompt( + base: str, + children: dict[str, str], + system: str | None, + *, + correction: str = "", + previous: str = "", +) -> str: + """Fairly distribute the remaining input budget over immediate child docs.""" + suffix = "" + if correction: + suffix = ( + "\n\nCorrect this previous draft and return the complete Markdown overview. " + "Preserve useful content. Validation failure: " + + correction + + "\n\n" + + _clip(previous, PREVIOUS_DRAFT_TOKENS) + + "\n" + ) + labels = {name: f"\n\n" for name in children} + + def render(allocations: dict[str, int]) -> str: + return ( + base + + "".join( + labels[name] + _clip(text, allocations[name]) + "\n" + for name, text in children.items() + ) + + suffix + ) + + sizes = {name: count_tokens(text) for name, text in children.items()} + # Empty excerpts still reserve labels and omission markers. + allocations = dict.fromkeys(children, 0) + mandatory = render(allocations) + remaining = OVERVIEW_INPUT_TOKENS - count_tokens((system or "") + mandatory) - 64 + if remaining < 0 or len((system or "") + mandatory) > MAX_USER_PROMPT_CHARS: + raise RuntimeError("Mandatory overview context exceeds the input budget") + pending = set(children) + while pending and remaining: + share = max(1, remaining // len(pending)) + for name in children: # stable ordering, including when the budget is very small + if name not in pending: + continue + take = min(share, sizes[name] - allocations[name], remaining) + allocations[name] += take + remaining -= take + if allocations[name] == sizes[name]: + pending.remove(name) + prompt = render(allocations) + # Token concatenation and long whitespace tokens can affect either cap. + while ( + count_tokens((system or "") + prompt) > OVERVIEW_INPUT_TOKENS + or len((system or "") + prompt) > MAX_USER_PROMPT_CHARS + ): + if not any(allocations.values()): + raise RuntimeError("Mandatory overview context exceeds the input budget") + allocations = {name: value // 2 for name, value in allocations.items()} + prompt = render(allocations) + return prompt + + +def save_overview(content: str, path: str) -> None: + """Replace a page only after validation, leaving the old file on failure.""" + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", dir=destination.parent, delete=False + ) as stream: + temporary = stream.name + stream.write(content) + mode = destination.stat().st_mode & 0o777 if destination.exists() else 0o644 + os.chmod(temporary, mode) + os.replace(temporary, destination) + finally: + if temporary and os.path.exists(temporary): + os.unlink(temporary) diff --git a/codewiki/src/be/prompt_template.py b/codewiki/src/be/prompt_template.py index 775e98d1..8e65fcb0 100644 --- a/codewiki/src/be/prompt_template.py +++ b/codewiki/src/be/prompt_template.py @@ -112,12 +112,13 @@ {repo_structure} -The core modules' documentation is NOT inlined above. Each top-level module carries a `docs_path` field with the absolute path to its documentation file — read those files with your file-reading tools before writing the overview (skip entries whose `docs_path` is null). When referencing a child's documentation, link it with the relative path in its `link` field, e.g. [alt text](). +Child documentation excerpts are supplied below as context. Some excerpts may be truncated or unavailable; use only the supplied evidence and do not invent missing details. No file-reading tools are available or needed. Link to child documentation using its relative `link` field, e.g. [alt text](). Please generate the overview of the `{repo_name}` repository in markdown format with the following structure: overview_content +Include a Markdown heading and explanatory prose. Return the finished page, without plans to read files or tool commands. """.strip() MODULE_OVERVIEW_PROMPT = """ @@ -133,12 +134,13 @@ {repo_structure} -The child modules' documentation is NOT inlined above. Each child of the target module carries a `docs_path` field with the absolute path to its documentation file — read those files with your file-reading tools before writing the overview (skip entries whose `docs_path` is null). When referencing a child's documentation, link it with the relative path in its `link` field, e.g. [alt text](). +Child documentation excerpts are supplied below as context. Some excerpts may be truncated or unavailable; use only the supplied evidence and do not invent missing details. No file-reading tools are available or needed. Link to child documentation using its relative `link` field, e.g. [alt text](). Please generate the overview of the `{module_name}` module in markdown format with the following structure: overview_content +Include a Markdown heading and explanatory prose. Return the finished page, without plans to read files or tool commands. """.strip() # These templates are formatted later; braces in literal Mermaid labels must diff --git a/codewiki/src/be/pydantic_ai_backend.py b/codewiki/src/be/pydantic_ai_backend.py index 3304c0b6..7d5bed95 100644 --- a/codewiki/src/be/pydantic_ai_backend.py +++ b/codewiki/src/be/pydantic_ai_backend.py @@ -18,6 +18,7 @@ from pydantic_ai.usage import UsageLimits from codewiki.src.be.agent_tools.deps import CodeWikiDeps +from codewiki.src.be.agent_tools.doc_guard import run_until_doc_written from codewiki.src.be.agent_tools.generate_sub_module_documentations import ( generate_sub_module_documentation_tool, ) @@ -28,6 +29,7 @@ from codewiki.src.be.dependency_analyzer.models.core import Node from codewiki.src.be.llm_services import call_llm, create_fallback_models, pop_last_usage from codewiki.src.be.module_scope import complete_file_candidates, validate_module_file_coverage +from codewiki.src.be.overview_support import reusable_overview, root_recovery_prompt from codewiki.src.be.prompt_template import ( format_leaf_system_prompt, format_system_prompt, @@ -135,13 +137,13 @@ async def run_module_agent( # or a resume after a partial run silently skips every missing one. if not module_path: overview_docs_path = os.path.join(working_dir, OVERVIEW_FILENAME) - if os.path.exists(overview_docs_path): + if reusable_overview(overview_docs_path): logger.info("✓ Overview docs already exists at %s", overview_docs_path) return module_tree docs_path = find_doc(working_dir, module_name, module_tree) if module_path else None if docs_path is None: docs_path = os.path.join(working_dir, f"{module_name}.md") - if os.path.exists(docs_path): + if os.path.isfile(docs_path) and (module_path or reusable_overview(docs_path)): logger.info("✓ Module docs already exists at %s", docs_path) return module_tree @@ -189,23 +191,33 @@ async def run_module_agent( scope_component_ids=core_component_ids, ) + user_prompt = format_user_prompt( + module_name=module_name, + core_component_ids=core_component_ids, + components=components, + module_tree=deps.module_tree, + module_path=module_path, + layout=config_layout(config), + ) + if not module_path: + user_prompt = root_recovery_prompt(user_prompt, working_dir, module_name) try: with track_usage(self.usage_ledger, config, "module_agent", module_name) as tracked: try: - await agent.run( - format_user_prompt( - module_name=module_name, - core_component_ids=core_component_ids, - components=components, - module_tree=deps.module_tree, - module_path=module_path, - layout=config_layout(config), - ), + await run_until_doc_written( + agent, + user_prompt, deps=deps, + module_name=module_name, + module_path=module_path, + layout=config_layout(config), usage_limits=UsageLimits(request_limit=config.request_limit), ) finally: self.last_usage = tracked.totals() + # Keep agent-inserted children resumable even if its own + # required page is still missing after reminders. + file_manager.save_json(deps.module_tree, module_tree_path) validate_module_file_coverage(deps.module_tree, components) file_manager.save_json(deps.module_tree, module_tree_path) return deps.module_tree diff --git a/tests/test_doc_layout.py b/tests/test_doc_layout.py index 0569373c..3dc34c4b 100644 --- a/tests/test_doc_layout.py +++ b/tests/test_doc_layout.py @@ -246,7 +246,7 @@ async def run_module_agent( return tree def complete(self, prompt, model=None, system_prompt=None): - return "[store](store.md)" + return "# Overview\n\nSee [store](store.md)." def test_generation_writes_nested_pages(tmp_path): @@ -267,7 +267,7 @@ def test_generation_writes_nested_pages(tmp_path): assert (tmp_path / "auth/session/store.md").read_text() == ( "# store\n\n[billing](../../billing.md)\n" ) - assert (tmp_path / "auth.md").read_text() == "[store](auth/session/store.md)" + assert (tmp_path / "auth.md").read_text() == "# Overview\n\nSee [store](auth/session/store.md)." assert find_missing_module_docs(TREE, str(tmp_path)) == [] # Resume: everything exists, nothing is regenerated diff --git a/tests/test_doc_write_guard.py b/tests/test_doc_write_guard.py new file mode 100644 index 00000000..77fb69b5 --- /dev/null +++ b/tests/test_doc_write_guard.py @@ -0,0 +1,322 @@ +"""An agent that ends without writing its page is reminded to write it (issue #134).""" + +import asyncio +import logging +import json +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from pydantic_ai import Agent +from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart +from pydantic_ai.models.function import AgentInfo, FunctionModel + +from codewiki.src.be.agent_tools.deps import CodeWikiDeps +from codewiki.src.be.agent_tools.doc_guard import ( + DOC_WRITE_REMINDERS, + DocumentationWriteError, + doc_written, + run_until_doc_written, +) +from codewiki.src.be.agent_tools.str_replace_editor import str_replace_editor_tool + + +def _deps(tmp_path, module_tree=None): + return CodeWikiDeps( + absolute_docs_path=str(tmp_path), + absolute_repo_path=str(tmp_path), + registry={}, + components={}, + path_to_current_module=[], + current_module_name="repo", + module_tree=module_tree or {}, + max_depth=2, + current_depth=1, + config=SimpleNamespace(), + ) + + +def _create(path): + return ToolCallPart( + "str_replace_editor", + {"working_dir": "docs", "command": "create", "path": path, "file_text": "# Page\n"}, + ) + + +def _scripted_agent(responses): + """Agent whose model replies with ``responses`` in order, one per request.""" + calls = [] + + def model(messages, info: AgentInfo): + # Older PydanticAI versions extend the same list after the callback. + calls.append(deepcopy(messages)) + return ModelResponse(parts=responses[len(calls) - 1]) + + agent = Agent(FunctionModel(model), deps_type=CodeWikiDeps, tools=[str_replace_editor_tool]) + return agent, calls + + +def _run(agent, deps, module_name, module_path): + return asyncio.run( + run_until_doc_written( + agent, + "document it", + deps=deps, + module_name=module_name, + module_path=module_path, + layout="flat", + ) + ) + + +def test_agent_reminded_until_page_written(tmp_path, caplog): + agent, calls = _scripted_agent( + [ + [TextPart("Here is the documentation: ...")], + [_create("./repo.md")], + [TextPart("done")], + ] + ) + with caplog.at_level(logging.WARNING): + result, usage = _run(agent, _deps(tmp_path), "repo", []) + + assert Path(tmp_path, "repo.md").read_text() == "# Page\n" + assert result.output == "done" + assert usage.requests == 3 + # The reminder continues the same conversation and names the exact file. + assert len(calls[1]) == 3 + assert "./repo.md" in calls[1][-1].parts[-1].content + assert "finished without writing repo.md" in caplog.text + assert "Here is the documentation" in caplog.text + assert "editor writes: none" in caplog.text + + +def test_written_page_needs_no_reminder(tmp_path): + agent, calls = _scripted_agent([[_create("./billing.md")], [TextPart("done")]]) + deps = _deps(tmp_path, {"billing": {"components": [], "children": {}}}) + _run(agent, deps, "billing", ["billing"]) + assert len(calls) == 2 + + +def test_gives_up_after_reminders_and_logs_tool_errors(tmp_path, caplog): + # The model keeps writing to the repo, which the tool rejects. + bad_write = [ + ToolCallPart( + "str_replace_editor", + {"working_dir": "repo", "command": "create", "path": "x.md", "file_text": "#"}, + ) + ] + responses = [] + for _ in range(DOC_WRITE_REMINDERS + 1): + responses += [bad_write, [TextPart("I wrote it")]] + agent, calls = _scripted_agent(responses) + with ( + caplog.at_level(logging.WARNING), + pytest.raises(DocumentationWriteError, match="2 reminders"), + ): + _run(agent, _deps(tmp_path), "repo", []) + + assert len(calls) == 2 * (DOC_WRITE_REMINDERS + 1) + assert not doc_written(str(tmp_path), "repo", [], {}) + assert "still has not written repo.md" in caplog.text + assert "The `view` command is the only allowed command" in caplog.text + + +def test_warning_reports_rejected_tool_calls(tmp_path, caplog): + # A call missing required arguments is rejected by validation before the tool runs. + agent, _ = _scripted_agent( + [ + [ToolCallPart("str_replace_editor", {"command": "create"})], + [TextPart("giving up")], + [_create("./repo.md")], + [TextPart("done")], + ] + ) + with caplog.at_level(logging.WARNING): + _run(agent, _deps(tmp_path), "repo", []) + assert Path(tmp_path, "repo.md").exists() + assert "rejected tool calls (1): str_replace_editor" in caplog.text + + +def test_doc_written_accepts_overview_for_whole_repo(tmp_path): + assert not doc_written(str(tmp_path), "repo", [], {}) + Path(tmp_path, "overview.md").write_text("# x\n") + assert doc_written(str(tmp_path), "repo", [], {}) + + +def test_doc_written_finds_nested_and_sanitized_pages(tmp_path): + tree = {"parent": {"children": {"Build Process": {"children": {}}}}} + (tmp_path / "build_process.md").write_text("# Build\n") + assert doc_written(str(tmp_path), "Build Process", ["parent", "Build Process"], tree) + (tmp_path / "parent").mkdir() + (tmp_path / "parent" / "child.md").write_text("# Child\n") + tree["parent"]["children"]["child"] = {"children": {}} + assert doc_written(str(tmp_path), "child", ["parent", "child"], tree) + + +def _backend(tmp_path, monkeypatch, kind, run): + from codewiki.src.be.backend import UsageLedger + from codewiki.src.config import Config + + config = Config.from_cli( + repo_path=str(tmp_path), + output_dir=str(tmp_path), + llm_api_key="test", + llm_base_url="https://example.invalid", + main_model="test", + cluster_model="test", + ) + if kind == "pydantic": + from codewiki.src.be import pydantic_ai_backend as mod + + class FakeAgent: + def __init__(self, *args, **kwargs): + pass + + async def run(self, prompt, deps, **kwargs): + run(prompt, deps) + return SimpleNamespace( + output="I'll write it", new_messages=lambda: [], all_messages=lambda: [] + ) + + monkeypatch.setattr(mod, "Agent", FakeAgent) + backend = mod.PydanticAIBackend.__new__(mod.PydanticAIBackend) + backend._fallback_models = None + backend._custom_instructions = "" + else: + pytest.importorskip("caw") + from codewiki.src.be import caw_backend as mod + + class FakeAgent: + def __init__(self, **kwargs): + self.deps = kwargs["tool_servers"][0]._deps + + def completion(self, prompt): + run(prompt, self.deps) + return SimpleNamespace( + result="I'll write it", total_usage=None, num_turns=1, total_tool_calls=0 + ) + + monkeypatch.setattr(mod, "CawAgent", FakeAgent) + backend = mod.CawBackend.__new__(mod.CawBackend) + backend._caw_provider = "codex" + backend._model = "test" + backend._repo_root = str(tmp_path) + backend._config = config + backend._usage_ledger = UsageLedger() + (tmp_path / "module_tree.json").write_text("{}") + return backend + + +@pytest.mark.parametrize("kind,expected_calls", [("pydantic", 3), ("caw", 1)]) +def test_backend_failure_is_explicit_and_preserves_reserved_children( + tmp_path, monkeypatch, kind, expected_calls +): + calls = [] + + def run(prompt, deps): + calls.append(prompt) + deps.module_tree["reserved"] = {"components": [], "children": {}} + + backend = _backend(tmp_path, monkeypatch, kind, run) + cwd = Path.cwd() + with pytest.raises(DocumentationWriteError, match="repo.md"): + asyncio.run(backend.run_module_agent("repo", {}, [], [], str(tmp_path))) + assert len(calls) == expected_calls + assert "reserved" in json.loads((tmp_path / "module_tree.json").read_text()) + assert Path.cwd() == cwd + + +@pytest.mark.parametrize("kind", ["pydantic", "caw"]) +@pytest.mark.parametrize("original,expected_calls", [("# Existing\n", 0), ("", 1)]) +def test_whole_repo_backend_skips_valid_overview_and_recovers_empty_one( + tmp_path, monkeypatch, kind, original, expected_calls +): + calls = [] + + def run(prompt, deps): + calls.append(prompt) + (tmp_path / "repo.md").write_text("# Repository\n\nComplete documentation.") + + backend = _backend(tmp_path, monkeypatch, kind, run) + (tmp_path / "overview.md").write_text(original) + asyncio.run(backend.run_module_agent("repo", {}, [], [], str(tmp_path))) + assert len(calls) == expected_calls + assert (tmp_path / "overview.md").read_text() == original + if calls: + assert "An existing root page is incomplete" in calls[0] + + +def test_subagent_failure_restores_dependencies(tmp_path, monkeypatch): + from codewiki.src.be.agent_tools import generate_sub_module_documentations as mod + from codewiki.src.be.dependency_analyzer.models.core import Node + + node = Node( + id="a.py::A", + name="A", + component_type="class", + file_path="a.py", + relative_path="a.py", + source_code="class A: pass", + ) + deps = _deps(tmp_path) + deps.components = {node.id: node} + deps.scope_component_ids = [node.id] + deps.config = SimpleNamespace( + agent_retries=1, request_limit=10, max_token_per_leaf_module=100, layout="flat" + ) + original_scope = deps.scope_component_ids + + class FakeAgent: + def __init__(self, *args, **kwargs): + pass + + async def fail(*args, **kwargs): + raise DocumentationWriteError("missing child.md") + + monkeypatch.setattr(mod, "Agent", FakeAgent) + monkeypatch.setattr(mod, "create_fallback_models", lambda config: None) + monkeypatch.setattr(mod, "run_until_doc_written", fail) + with pytest.raises(DocumentationWriteError): + asyncio.run( + mod.generate_sub_module_documentation(SimpleNamespace(deps=deps), {"child": [node.id]}) + ) + assert deps.current_module_name == "repo" and deps.current_depth == 1 + assert deps.path_to_current_module == [] + assert deps.scope_component_ids is original_scope + assert "child" in deps.module_tree + + +@pytest.mark.parametrize("kind", ["pydantic", "caw"]) +@pytest.mark.parametrize("fails", [False, True]) +def test_whole_repo_generation_replaces_or_restores_incomplete_overview( + tmp_path, monkeypatch, kind, fails +): + from codewiki.src.be.documentation_generator import DocumentationGenerator + + original = "I'll start by reading the child files.\n\n/opt/child.md" + page = tmp_path / "overview.md" + + def run(prompt, deps): + if fails: + page.write_text("partial edit") + raise RuntimeError("agent failed") + (tmp_path / "repo.md").write_text("# Repository\n\nComplete documentation.") + + backend = _backend(tmp_path, monkeypatch, kind, run) + backend._config.repo_path = str(tmp_path / "repo") + backend._config.docs_dir = str(tmp_path) + page.write_text(original) + gen = DocumentationGenerator.__new__(DocumentationGenerator) + gen.backend = backend + gen.config = backend._config + if fails: + with pytest.raises(RuntimeError, match="agent failed"): + asyncio.run(gen.generate_module_documentation({}, [])) + assert page.read_text() == original + else: + asyncio.run(gen.generate_module_documentation({}, [])) + assert page.read_text() == "# Repository\n\nComplete documentation." + assert not (tmp_path / "repo.md").exists() diff --git a/tests/test_mermaid_writers.py b/tests/test_mermaid_writers.py index 68de55a3..f6e5d2d7 100644 --- a/tests/test_mermaid_writers.py +++ b/tests/test_mermaid_writers.py @@ -220,18 +220,21 @@ def generator(tmp_path, replies): return gen +OVERVIEW_PROSE = "# Overview\n\nThe module coordinates the application components.\n\n" + + @pytest.mark.parametrize("module_path,filename", [([], "overview.md"), (["parent"], "parent.md")]) def test_overview_repairs_before_saving(tmp_path, parser, module_path, filename): - gen = generator(tmp_path, ["" + INVALID + ""]) + gen = generator(tmp_path, ["" + OVERVIEW_PROSE + INVALID + ""]) asyncio.run(gen.generate_parent_module_docs(module_path, str(tmp_path))) - assert (tmp_path / filename).read_text() == FIXED.strip() + assert (tmp_path / filename).read_text() == (OVERVIEW_PROSE + FIXED).strip() assert len(gen.backend.calls) == 1 def test_overview_retries_once_and_tracks_usage(tmp_path, parser): - gen = generator(tmp_path, [BROKEN, FIXED]) + gen = generator(tmp_path, [OVERVIEW_PROSE + BROKEN, OVERVIEW_PROSE + FIXED]) asyncio.run(gen.generate_parent_module_docs([], str(tmp_path))) - assert (tmp_path / "overview.md").read_text() == FIXED.strip() + assert (tmp_path / "overview.md").read_text() == (OVERVIEW_PROSE + FIXED).strip() assert len(gen.backend.calls) == 2 correction, system = gen.backend.calls[1] assert "" in correction and "Parse error" in correction @@ -241,7 +244,7 @@ def test_overview_retries_once_and_tracks_usage(tmp_path, parser): def test_invalid_overview_after_retry_is_not_saved(tmp_path, parser): - gen = generator(tmp_path, [BROKEN, BROKEN]) + gen = generator(tmp_path, [OVERVIEW_PROSE + BROKEN, OVERVIEW_PROSE + BROKEN]) with pytest.raises(RuntimeError, match="after one correction attempt"): asyncio.run(gen.generate_parent_module_docs([], str(tmp_path))) assert not (tmp_path / "overview.md").exists() @@ -253,9 +256,9 @@ async def unavailable(_): return DiagramResult("unavailable", "HTTP 503") monkeypatch.setattr(utils, "_validate_diagram", unavailable) - gen = generator(tmp_path, [INVALID]) + gen = generator(tmp_path, [OVERVIEW_PROSE + INVALID]) asyncio.run(gen.generate_parent_module_docs([], str(tmp_path))) - assert (tmp_path / "overview.md").read_text() == INVALID.strip() + assert (tmp_path / "overview.md").read_text() == (OVERVIEW_PROSE + INVALID).strip() assert len(gen.backend.calls) == 1 assert "validation unavailable" in caplog.text and "HTTP 503" in caplog.text diff --git a/tests/test_module_scope.py b/tests/test_module_scope.py index 4b648c46..71211118 100644 --- a/tests/test_module_scope.py +++ b/tests/test_module_scope.py @@ -373,6 +373,7 @@ def __init__(self, *args, **kwargs): async def run(self, prompt, deps, **kwargs): captured["ids"] = deps.module_component_ids() + (out / "tiny.md").write_text("# Tiny\n\nModule documentation.") return SimpleNamespace(output="ok") monkeypatch.setattr(mod, "Agent", FakeAgent) @@ -391,6 +392,7 @@ def __init__(self, **kwargs): captured["ids"] = toolkit._deps.module_component_ids() def completion(self, prompt): + (out / "tiny.md").write_text("# Tiny\n\nModule documentation.") return SimpleNamespace( total_usage=None, model="test", num_turns=1, total_tool_calls=0 ) @@ -441,6 +443,15 @@ def __init__(self, *args, **kwargs): async def run(self, prompt, deps, **kwargs): seen.extend(deps.module_component_ids()) + from codewiki.src.be.doc_layout import module_doc_file + + name = module_doc_file( + deps.current_module_name, deps.path_to_current_module, deps.config.layout + ) + path = Path(deps.absolute_docs_path, name) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("# Module\n\nModule documentation.") + return SimpleNamespace(output="ok") monkeypatch.setattr(tool, "Agent", FakeAgent) monkeypatch.setattr(tool, "create_fallback_models", lambda config: None) diff --git a/tests/test_overview_generation.py b/tests/test_overview_generation.py new file mode 100644 index 00000000..9090d5e7 --- /dev/null +++ b/tests/test_overview_generation.py @@ -0,0 +1,252 @@ +"""Regression coverage for missing context and incomplete pages in issue #134.""" + +import asyncio +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from codewiki.src.be.backend import UsageLedger +from codewiki.src.be.documentation_generator import DocumentationGenerator +from codewiki.src.be import overview_support as support +from codewiki.src.be.prompt_template import MAX_USER_PROMPT_CHARS +from codewiki.src.be.utils import count_tokens + +PREAMBLE = ( + "I'll start by reading the two child module documentation files.\n\n" + "/opt/runner/documentation/build_and_deployment.md " + "/opt/runner/documentation/bootstrap.md " + "/opt/runner/documentation/build_and_deployment.md" +) +VALID = "# Visão geral\n\nO módulo coordena a implantação e inicialização.\n" + + +class Backend: + def __init__(self, replies): + self.replies = iter(replies) + self.calls = [] + self.usage_ledger = UsageLedger() + + def complete(self, prompt, *, system_prompt=None): + self.calls.append((prompt, system_prompt)) + self.usage_ledger.add("complete", "", {"input_tokens": 10}) + return next(self.replies) + + +def generator(tmp_path, replies, layout="flat"): + tree = { + "parent": { + "components": [], + "children": { + "build_and_deployment": {"components": [], "children": {}}, + "bootstrap": {"components": [], "children": {}}, + }, + } + } + (tmp_path / "module_tree.json").write_text(json.dumps(tree)) + gen = DocumentationGenerator.__new__(DocumentationGenerator) + gen.config = SimpleNamespace( + layout=layout, repo_path="/tmp/example", get_prompt_addition=lambda: "Write in Portuguese." + ) + gen.backend = Backend(replies) + return gen + + +@pytest.mark.parametrize("layout", ["flat", "hierarchical"]) +def test_parent_receives_child_contents_and_relative_links(tmp_path, layout): + gen = generator(tmp_path, [VALID], layout) + child_dir = tmp_path if layout == "flat" else tmp_path / "parent" + child_dir.mkdir(exist_ok=True) + (child_dir / "build_and_deployment.md").write_text("# Build\n\nCHILD_CONTENT_SENTINEL") + (child_dir / "bootstrap.md").write_text("# Bootstrap\n\nInitialize the plugin.") + asyncio.run(gen.generate_parent_module_docs(["parent"], str(tmp_path))) + prompt, system = gen.backend.calls[0] + assert "CHILD_CONTENT_SENTINEL" in prompt + assert "Initialize the plugin" in prompt + assert '"link": "' + ("parent/" if layout == "hierarchical" else "") + "bootstrap.md" in prompt + assert "read those files with your" not in prompt + assert "Portuguese" in system + assert (tmp_path / "parent.md").read_text() == VALID.strip() + + +@pytest.mark.parametrize( + "response", + [ + PREAMBLE, + "", + None, + "# Only a title", + "```text\n# fake\nbody\n```", + "# Title\n/absolute/file.md", + "# Title\nrelative/file.md", + "# Title\n" + PREAMBLE, + ], +) +def test_invalid_output_is_corrected_once(tmp_path, response): + gen = generator(tmp_path, [response, "" + VALID + ""]) + asyncio.run(gen.generate_parent_module_docs(["parent"], str(tmp_path))) + assert (tmp_path / "parent.md").read_text() == VALID.strip() + assert len(gen.backend.calls) == 2 + assert "Validation failure" in gen.backend.calls[1][0] + assert gen.backend.usage_ledger.summary()["n_calls"] == 2 + + +@pytest.mark.parametrize("original", [None, "", PREAMBLE]) +def test_exhausted_correction_never_overwrites_old_page(tmp_path, original): + gen = generator(tmp_path, [PREAMBLE, PREAMBLE]) + path = tmp_path / "parent.md" + if original is not None: + path.write_text(original) + with pytest.raises(RuntimeError, match="after one correction attempt"): + asyncio.run(gen.generate_parent_module_docs(["parent"], str(tmp_path))) + assert len(gen.backend.calls) == 2 + assert path.read_text() == original if original is not None else not path.exists() + + +@pytest.mark.parametrize("original", ["", PREAMBLE, PREAMBLE.replace("I'll", "I’ll")]) +@pytest.mark.parametrize("module_path,filename", [([], "overview.md"), (["parent"], "parent.md")]) +def test_recognized_old_pages_are_repaired(tmp_path, original, module_path, filename): + gen = generator(tmp_path, [VALID]) + path = tmp_path / filename + path.write_text(original) + asyncio.run(gen.generate_parent_module_docs(module_path, str(tmp_path))) + assert path.read_text() == VALID.strip() + + +@pytest.mark.parametrize( + "original", + [ + VALID, + "A manually edited page without a heading.", + PREAMBLE + "\n\nUseful documentation follows.", + ], +) +def test_other_existing_prose_is_preserved(tmp_path, original): + gen = generator(tmp_path, []) + (tmp_path / "parent.md").write_text(original) + asyncio.run(gen.generate_parent_module_docs(["parent"], str(tmp_path))) + assert not gen.backend.calls + assert (tmp_path / "parent.md").read_text() == original + + +def test_missing_and_unreadable_children_warn_without_fabrication(tmp_path, monkeypatch, caplog): + gen = generator(tmp_path, [VALID]) + path = tmp_path / "bootstrap.md" + path.write_text("unreadable secret") + original = Path.read_text + + def read(path, *args, **kwargs): + if path.name == "bootstrap.md": + raise UnicodeError("invalid encoding") + return original(path, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", read) + asyncio.run(gen.generate_parent_module_docs(["parent"], str(tmp_path))) + assert "Cannot read child documentation" in caplog.text + assert "Module docs not found" in caplog.text + assert "unreadable secret" not in gen.backend.calls[0][0] + assert "bootstrap.md" in gen.backend.calls[0][0] + + +def test_sanitized_child_filename_is_read(tmp_path): + gen = generator(tmp_path, [VALID]) + path = tmp_path / "module_tree.json" + tree = json.loads(path.read_text()) + tree["parent"]["children"]["Build Deployment"] = tree["parent"]["children"].pop( + "build_and_deployment" + ) + path.write_text(json.dumps(tree)) + (tmp_path / "build_deployment.md").write_text("# Build\n\nSANITIZED_CONTENT") + asyncio.run(gen.generate_parent_module_docs(["parent"], str(tmp_path))) + assert "SANITIZED_CONTENT" in gen.backend.calls[0][0] + + +def test_context_budget_shares_space_and_redistributes_unused_allocations(monkeypatch): + monkeypatch.setattr(support, "OVERVIEW_INPUT_TOKENS", 700) + children = { + "small": "Small complete page.", + "large_a": "alpha " * 5000, + "large_b": "beta " * 5000, + } + prompt = support.bounded_overview_prompt("Structure", children, "System instructions") + assert count_tokens("System instructions" + prompt) <= 700 + assert "Small complete page." in prompt + assert prompt.count("alpha") > 150 and prompt.count("beta") > 150 + assert prompt.count("Excerpt truncated") == 2 + + +def test_large_unicode_context_and_correction_stay_within_caps(): + children = {"日本語": "日本語の説明。" * 20_000, "second": "source " * 50_000} + prompt = support.bounded_overview_prompt( + "Structure", children, "System", correction="invalid body", previous="draft " * 50_000 + ) + assert count_tokens("System" + prompt) <= support.OVERVIEW_INPUT_TOKENS + assert len("System" + prompt) <= MAX_USER_PROMPT_CHARS + assert "�" not in prompt + assert "日本語の説明" in prompt and "source" in prompt + + +def test_mandatory_context_overflow_is_explicit(monkeypatch): + monkeypatch.setattr(support, "OVERVIEW_INPUT_TOKENS", 100) + with pytest.raises(RuntimeError, match="Mandatory overview context"): + support.bounded_overview_prompt("structure " * 200, {}, "system") + + +def test_character_limit_also_bounds_whitespace_heavy_context(monkeypatch): + monkeypatch.setattr(support, "MAX_USER_PROMPT_CHARS", 1000) + prompt = support.bounded_overview_prompt("Structure", {"child": " " * 100_000}, "System") + assert len("System" + prompt) <= 1000 + + +def test_valid_prose_in_non_latin_script(): + assert support.overview_error("# 概要\n\nこれはシステムの説明です。") is None + + +def test_failed_atomic_replacement_preserves_existing_file(tmp_path, monkeypatch): + path = tmp_path / "overview.md" + path.write_text(PREAMBLE) + + def fail(*args): + raise OSError("disk unavailable") + + monkeypatch.setattr(support.os, "replace", fail) + with pytest.raises(OSError, match="disk unavailable"): + support.save_overview(VALID, str(path)) + assert path.read_text() == PREAMBLE + assert list(tmp_path.iterdir()) == [path] + + +def test_whole_repo_recovery_restores_old_page_on_failure(tmp_path): + path = tmp_path / "overview.md" + path.write_text(PREAMBLE) + with pytest.raises(RuntimeError), support.preserve_incomplete_overviews([str(path)]): + path.write_text("partial edit") + raise RuntimeError("agent failed") + assert path.read_text() == PREAMBLE + + +def test_parent_failure_does_not_skip_independent_modules_or_pass_silently(tmp_path): + gen = generator(tmp_path, []) + gen.config.docs_dir = str(tmp_path) + tree_path = tmp_path / "module_tree.json" + tree = json.loads(tree_path.read_text()) + tree["independent"] = {"components": [], "children": {}} + tree_path.write_text(json.dumps(tree)) + (tmp_path / "parent.md").write_text(PREAMBLE) + + async def module(module_name, **kwargs): + (tmp_path / f"{module_name}.md").write_text(VALID) + return tree + + async def overview(module_path, working_dir, components=None): + if module_path == ["parent"]: + raise RuntimeError("invalid parent response") + (tmp_path / "overview.md").write_text(VALID) + return tree + + gen.backend.run_module_agent = module + gen.generate_parent_module_docs = overview + with pytest.raises(RuntimeError, match="parent: invalid parent response"): + asyncio.run(gen.generate_module_documentation({}, [])) + assert (tmp_path / "independent.md").read_text() == VALID diff --git a/tests/test_processing_order_update.py b/tests/test_processing_order_update.py index 02b98ece..22fc8884 100644 --- a/tests/test_processing_order_update.py +++ b/tests/test_processing_order_update.py @@ -80,7 +80,7 @@ async def run_module_agent( def complete(self, prompt, model=None, system_prompt=None): self.complete_calls += 1 - return "overview" + return "# Overview\n\nThe repository coordinates its modules." def _generator(docs_dir: Path) -> tuple[DocumentationGenerator, FakeBackend]: diff --git a/tests/test_updater_orchestrator.py b/tests/test_updater_orchestrator.py index b4db3fb8..7db16ed4 100644 --- a/tests/test_updater_orchestrator.py +++ b/tests/test_updater_orchestrator.py @@ -46,7 +46,7 @@ def complete(self, prompt, *, model=None, system_prompt=None): else: decisions = [{"component_id": c, "action": "untracked"} for c in ids] return "```json\n" + json.dumps({"decisions": decisions}) + "\n```" - return "regenerated overview" + return "# Overview\n\nregenerated overview of the modules." async def run_module_agent( self, module_name, components, core_component_ids, module_path, working_dir