fix: try recovery from truncated xml string

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 <section>...</section> 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"<section\b[\s\S]*?</section>", xml_text)):
+        sid = re.search(r'<section\s+id="([^"]+)"', block)
+        title = re.search(r"<title>([\s\S]*?)</title>", block)
+        page_refs = [
+            m.strip()
+            for m in re.findall(r"<page_ref>([\s\S]*?)</page_ref>", block)
+            if m.strip()
+        ]
+        subs = [
+            m.strip()
+            for m in re.findall(r"<section_ref>([\s\S]*?)</section_ref>", 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"<wiki_structure>[\s\S]*?</wiki_structure>", text)
-    if not match:
-        raise ValueError("No valid <wiki_structure> 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 </wiki_structure>. Salvage from the opening tag to end-of-text
+        # (plus a synthetic close) so the regex fallbacks below can still recover
+        # the complete <section>/<page> blocks instead of failing the whole task.
+        open_match = re.search(r"<wiki_structure>[\s\S]*", text)
+        if not open_match:
+            raise ValueError("No valid <wiki_structure> XML found in response")
+        logger.warning(
+            "Response appears truncated (missing </wiki_structure>); "
+            "salvaging complete blocks."
+        )
+        xml_text = f"{open_match.group(0)}\n</wiki_structure>"
 
     # 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 <title>/<description> are emitted first, so the
+        # first match is the right one (page-level ones come later).
+        title = _first_group(r"<title>([\s\S]*?)</title>", xml_text)
+        description = _first_group(r"<description>([\s\S]*?)</description>", 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",

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 </relevant_files>, </page>, </pages> or </wiki_structure>. Mirrors the
+# real failing log for AsyncFuncAI/deepwiki-open.
+TRUNCATED_XML = """
+<wiki_structure>
+  <title>DeepWiki-Open Wiki</title>
+  <description>An AI-powered documentation generator for repositories.</description>
+  <sections>
+    <section id="section-1">
+      <title>Overview</title>
+      <pages><page_ref>page-1</page_ref></pages>
+    </section>
+    <section id="section-2">
+      <title>Extensibility and Customization</title>
+      <pages><page_ref>page-3</page_ref></pages>
+    </section>
+  </sections>
+  <pages>
+    <page id="page-1">
+      <title>Project Overview</title>
+      <importance>high</importance>
+      <relevant_files><file_path>README.md</file_path></relevant_files>
+      <related_pages><related>page-2</related></related_pages>
+    </page>
+    <page id="page-2">
+      <title>System Architecture</title>
+      <importance>high</importance>
+      <relevant_files><file_path>api/main.py</file_path></relevant_files>
+    </page>
+    <page id="page-3">
+      <title>Deployment and Infrastructure</title>
+      <importance>medium</importance>
+      <relevant_files>
+        <file_path>docker-compose.yml</file_path>
+        <file_path>Ollama-instruction.md</file_path>"""
+
+
+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 <page> 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 <wiki_structure> at all -> genuinely unusable -> hard error stands.
+    with pytest.raises(ValueError):
+        parse_wiki_structure("some prose, no xml here at all", comprehensive=True)