diff --git a/api/services/wiki/structure.py b/api/services/wiki/structure.py index af1fb70c7..ed0695172 100644 --- a/api/services/wiki/structure.py +++ b/api/services/wiki/structure.py @@ -139,6 +139,43 @@ def _parse_sections(root: ET.Element) -> tuple[list[WikiSection], list[str]]: return sections, root_sections +def _first_group(pattern: str, text: str) -> str: + """First capture group of `pattern` in `text`, or '' if no match.""" + m = re.search(pattern, text) + return m.group(1).strip() if m else "" + + +def _sections_via_regex(xml_text: str) -> tuple[list[WikiSection], list[str]]: + """Recover complete
...
blocks when strict XML parsing + fails (e.g. a truncated response). Mirrors _parse_sections.""" + sections: list[WikiSection] = [] + referenced: set[str] = set() + for i, block in enumerate(re.findall(r"", xml_text)): + sid = re.search(r'([\s\S]*?)", block) + page_refs = [ + m.strip() + for m in re.findall(r"([\s\S]*?)", block) + if m.strip() + ] + subs = [ + m.strip() + for m in re.findall(r"([\s\S]*?)", block) + if m.strip() + ] + sections.append( + WikiSection( + id=sid.group(1) if sid else f"section-{i + 1}", + title=title.group(1).strip() if title else "", + pages=page_refs, + subsections=subs or None, + ) + ) + referenced.update(subs) + root_sections = [s.id for s in sections if s.id not in referenced] + return sections, root_sections + + def parse_wiki_structure(text: str, comprehensive: bool) -> WikiStructureModel: """Parse the LLM's XML response into a WikiStructureModel. @@ -151,9 +188,21 @@ def parse_wiki_structure(text: str, comprehensive: bool) -> WikiStructureModel: text = re.sub(r"```\s*$", "", text) match = re.search(r"[\s\S]*?", text) - if not match: - raise ValueError("No valid XML found in response") - xml_text = match.group(0) + if match: + xml_text = match.group(0) + else: + # Truncated response: the model hit its output-token limit before + # emitting . Salvage from the opening tag to end-of-text + # (plus a synthetic close) so the regex fallbacks below can still recover + # the complete
/ blocks instead of failing the whole task. + open_match = re.search(r"[\s\S]*", text) + if not open_match: + raise ValueError("No valid XML found in response") + logger.warning( + "Response appears truncated (missing ); " + "salvaging complete blocks." + ) + xml_text = f"{open_match.group(0)}\n" # Strip control chars, then escape bare '&' that are not valid XML entities. xml_text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]", "", xml_text) @@ -167,22 +216,29 @@ def parse_wiki_structure(text: str, comprehensive: bool) -> WikiStructureModel: except ET.ParseError as e: logger.warning("Strict XML parse failed, using regex fallback: %s", e) - title = (root.findtext("title") if root is not None else None) or "" - description = (root.findtext("description") if root is not None else None) or "" + if root is not None: + title = root.findtext("title") or "" + description = root.findtext("description") or "" + pages = [_page_from_element(el, i) for i, el in enumerate(root.iter("page"))] + else: + # Strict parse failed (malformed / truncated): recover the header via + # regex. The wiki-level /<description> are emitted first, so the + # first match is the right one (page-level ones come later). + title = _first_group(r"<title>([\s\S]*?)", xml_text) + description = _first_group(r"([\s\S]*?)", xml_text) + pages = [] - pages = ( - [_page_from_element(el, i) for i, el in enumerate(root.iter("page"))] - if root is not None - else [] - ) if not pages: logger.warning("XML parsing yielded no pages; using regex fallback") pages = _pages_via_regex(xml_text) sections: list[WikiSection] = [] root_sections: list[str] = [] - if comprehensive and root is not None: - sections, root_sections = _parse_sections(root) + if comprehensive: + if root is not None: + sections, root_sections = _parse_sections(root) + else: + sections, root_sections = _sections_via_regex(xml_text) return WikiStructureModel( id="wiki", diff --git a/tests/backend/services/test_wiki_structure.py b/tests/backend/services/test_wiki_structure.py index 0b6a74bc3..7150990e3 100644 --- a/tests/backend/services/test_wiki_structure.py +++ b/tests/backend/services/test_wiki_structure.py @@ -106,3 +106,64 @@ def test_read_repo_file_tree(tmp_path, exclude_test_config): def test_detect_default_branch_non_git_dir(tmp_path): assert detect_default_branch(str(tmp_path)) == "main" + + +# A comprehensive response cut off mid-way (model hit its output-token limit): +# sections + page-1/page-2 are complete, page-3 is truncated, and there is no +# closing , , or . Mirrors the +# real failing log for AsyncFuncAI/deepwiki-open. +TRUNCATED_XML = """ + + DeepWiki-Open Wiki + An AI-powered documentation generator for repositories. + +
+ Overview + page-1 +
+
+ Extensibility and Customization + page-3 +
+ + + + Project Overview + high + README.md + page-2 + + + System Architecture + high + api/main.py + + + Deployment and Infrastructure + medium + + docker-compose.yml + Ollama-instruction.md""" + + +def test_parse_recovers_from_truncated_response(): + s = parse_wiki_structure(TRUNCATED_XML, comprehensive=True) + + # Header is recovered even though strict XML parsing fails on the truncation. + assert s.title == "DeepWiki-Open Wiki" + assert "AI-powered" in s.description + + # Only the COMPLETE blocks survive; the truncated page-3 is dropped + # (rather than failing the entire task, as it did before). + assert [p.id for p in s.pages] == ["page-1", "page-2"] + assert s.pages[0].filePaths == ["README.md"] + + # Sections were fully emitted before the cutoff -> recovered via regex. + assert {sec.id for sec in s.sections} == {"section-1", "section-2"} + assert set(s.rootSections) == {"section-1", "section-2"} + + +def test_parse_truncated_without_opening_tag_still_raises(): + # No at all -> genuinely unusable -> hard error stands. + with pytest.raises(ValueError): + parse_wiki_structure("some prose, no xml here at all", comprehensive=True)