From 956147d864312214329fe2b246a6c9c2ed782e7b Mon Sep 17 00:00:00 2001 From: mountain Date: Tue, 7 Jul 2026 10:12:51 +0800 Subject: [PATCH] fix: resolve five P1 defects from the SDK review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - AgentRunner.run: offload to a worker-thread event loop when called from inside a running loop (Jupyter, FastAPI handlers) — mirrors pipeline._run_async; Runner.run_sync raised RuntimeError there. - SQLiteStorage: create connections with check_same_thread=False so close() can actually close connections created by worker threads. Each thread still gets its own connection via threading.local; with the default True those closes raised ProgrammingError (silently swallowed) and leaked every worker connection. - CloudBackend.query: non-streaming chat completions now use a 300s timeout and a single attempt. The default 30s ReadTimeout fired before generation finished and the retry loop re-billed the full server-side retrieval + generation up to three times. _request gains retries/timeout overrides; the exhausted-retry path also no longer sleeps before raising. - MarkdownParser: content before the first heading (abstract/preamble) becomes a node instead of being silently dropped and unretrievable; a file with no headings at all yields a single document node instead of zero nodes (which pushed an empty page list into the pipeline). - LegacyCloudAPI.is_retrieval_ready: API failures (revoked key, network down) now propagate as PageIndexAPIError instead of reading as "not ready", which turned polling loops into infinite loops. Adds regression tests for each fix. Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS --- pageindex/agent.py | 18 +++++++++++++++-- pageindex/backend/cloud.py | 22 +++++++++++++++------ pageindex/cloud_api.py | 14 ++++++++----- pageindex/parser/markdown.py | 33 +++++++++++++++++++++++++++++-- pageindex/storage/sqlite.py | 7 ++++++- tests/test_agent.py | 27 +++++++++++++++++++++++++ tests/test_cloud_backend.py | 17 ++++++++++++++++ tests/test_legacy_sdk_contract.py | 11 +++++++++++ tests/test_markdown_parser.py | 18 +++++++++++++++++ tests/test_sqlite_storage.py | 20 +++++++++++++++++++ 10 files changed, 171 insertions(+), 16 deletions(-) diff --git a/pageindex/agent.py b/pageindex/agent.py index 677c0ea..a186fa4 100644 --- a/pageindex/agent.py +++ b/pageindex/agent.py @@ -135,7 +135,14 @@ class AgentRunner: self._instructions = instructions or OPEN_SYSTEM_PROMPT def run(self, question: str) -> str: - """Sync non-streaming query. Returns answer string.""" + """Sync non-streaming query. Returns answer string. + + Safe to call from within a running event loop (Jupyter, FastAPI + handlers): the agent then runs on a private loop in a worker thread, + mirroring pipeline._run_async — Runner.run_sync would otherwise raise + RuntimeError in that situation. + """ + import asyncio from agents import Agent, Runner from agents.model_settings import ModelSettings agent = Agent( @@ -146,5 +153,12 @@ class AgentRunner: model=self._model, model_settings=ModelSettings(parallel_tool_calls=False), ) - result = Runner.run_sync(agent, question) + try: + asyncio.get_running_loop() + except RuntimeError: + result = Runner.run_sync(agent, question) + else: + import concurrent.futures + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + result = pool.submit(asyncio.run, Runner.run(agent, question)).result() return result.final_output diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py index b479e81..095b9d9 100644 --- a/pageindex/backend/cloud.py +++ b/pageindex/backend/cloud.py @@ -49,10 +49,14 @@ class CloudBackend: ) self._folder_warning_shown = True - def _request(self, method: str, path: str, **kwargs) -> dict: + def _request(self, method: str, path: str, retries: int = 3, **kwargs) -> dict: + """HTTP helper. ``retries`` caps total attempts — pass 1 for + non-idempotent, expensive calls (e.g. chat completions) where a + retry would redo the full server-side work.""" url = f"{API_BASE}{path}" + kwargs.setdefault("timeout", 30) last_status: int | None = None - for attempt in range(3): + for attempt in range(retries): if attempt and "files" in kwargs: # Rewind file objects before a retry — the previous attempt # consumed them, and re-sending without seek(0) would upload @@ -62,10 +66,12 @@ class CloudBackend: if hasattr(fobj, "seek"): fobj.seek(0) try: - resp = requests.request(method, url, headers=self._headers, timeout=30, **kwargs) + resp = requests.request(method, url, headers=self._headers, **kwargs) if resp.status_code in (429, 500, 502, 503): - logger.warning("Cloud API %s %s returned %d, retrying...", method, path, resp.status_code) last_status = resp.status_code + if attempt == retries - 1: + break + logger.warning("Cloud API %s %s returned %d, retrying...", method, path, resp.status_code) time.sleep(2 ** attempt) continue if resp.status_code != 200: @@ -74,7 +80,7 @@ class CloudBackend: status_code=resp.status_code) return resp.json() if resp.content else {} except requests.RequestException as e: - if attempt == 2: + if attempt == retries - 1: raise CloudAPIError(f"Cloud API request failed: {e}") from e time.sleep(2 ** attempt) raise CloudAPIError(f"Cloud API {method} {path} failed after retries" @@ -301,7 +307,11 @@ class CloudBackend: "doc_ids cannot be empty; pass None to query the whole collection" ) doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection) - resp = self._request("POST", "/chat/completions/", json={ + # A non-streaming completion returns nothing until generation + # finishes, so it needs far more than the default 30s. retries=1: + # retrying this non-idempotent call would redo the full server-side + # retrieval + generation (and bill it) on every attempt. + resp = self._request("POST", "/chat/completions/", retries=1, timeout=300, json={ "messages": [{"role": "user", "content": question}], "doc_id": doc_id, "stream": False, diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index 303702a..f04182f 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -85,11 +85,15 @@ class LegacyCloudAPI: return response.json() def is_retrieval_ready(self, doc_id: str) -> bool: - try: - result = self.get_tree(doc_id) - return result.get("retrieval_ready", False) - except PageIndexAPIError: - return False + """Return whether retrieval is ready for ``doc_id``. + + API failures (revoked key, network down, unknown doc) propagate as + PageIndexAPIError instead of reading as "not ready" — swallowing them + turned ``while not is_retrieval_ready(...)`` polling loops into + infinite loops. + """ + result = self.get_tree(doc_id) + return result.get("retrieval_ready", False) def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> dict[str, Any]: payload = { diff --git a/pageindex/parser/markdown.py b/pageindex/parser/markdown.py index f62013c..0eba5bb 100644 --- a/pageindex/parser/markdown.py +++ b/pageindex/parser/markdown.py @@ -17,7 +17,7 @@ class MarkdownParser: lines = content.split("\n") headers = self._extract_headers(lines) - nodes = self._build_nodes(headers, lines, model) + nodes = self._build_nodes(headers, lines, model, doc_title=path.stem) return ParsedDocument(doc_name=path.stem, nodes=nodes) @@ -42,8 +42,37 @@ class MarkdownParser: }) return headers - def _build_nodes(self, headers: list[dict], lines: list[str], model: str | None) -> list[ContentNode]: + def _build_nodes(self, headers: list[dict], lines: list[str], model: str | None, + doc_title: str = "Document") -> list[ContentNode]: nodes = [] + + # A file with no headings at all still has content — index it as a + # single node instead of producing zero nodes (which would push an + # empty page list into the LLM pipeline). + if not headers: + text = "\n".join(lines).strip() + if text: + nodes.append(ContentNode( + content=text, + tokens=count_tokens(text, model=model), + title=doc_title, + index=1, + level=1, + )) + return nodes + + # Content before the first heading (abstract, preamble) would + # otherwise be silently dropped and become unretrievable. + preamble = "\n".join(lines[: headers[0]["line_num"] - 1]).strip() + if preamble: + nodes.append(ContentNode( + content=preamble, + tokens=count_tokens(preamble, model=model), + title=doc_title, + index=1, + level=headers[0]["level"], + )) + for i, header in enumerate(headers): start = header["line_num"] - 1 end = headers[i + 1]["line_num"] - 1 if i + 1 < len(headers) else len(lines) diff --git a/pageindex/storage/sqlite.py b/pageindex/storage/sqlite.py index eed1bc4..2ba4419 100644 --- a/pageindex/storage/sqlite.py +++ b/pageindex/storage/sqlite.py @@ -16,7 +16,12 @@ class SQLiteStorage: def _get_conn(self) -> sqlite3.Connection: """Return a thread-local SQLite connection.""" if not hasattr(self._local, "conn"): - conn = sqlite3.connect(str(self._db_path)) + # Each thread gets its own connection (threading.local), so + # statements never race. check_same_thread=False exists solely so + # close() can close every tracked connection from whichever thread + # calls it — with the default True those closes raise + # ProgrammingError and the connections leak. + conn = sqlite3.connect(str(self._db_path), check_same_thread=False) conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys=ON") self._local.conn = conn diff --git a/tests/test_agent.py b/tests/test_agent.py index 16c98f3..a22b1db 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -18,3 +18,30 @@ def test_scoped_prompt_omits_list_documents(): assert "list_documents" not in SCOPED_SYSTEM_PROMPT assert "get_document_structure" in SCOPED_SYSTEM_PROMPT assert "get_page_content" in SCOPED_SYSTEM_PROMPT + + +def test_run_works_inside_running_event_loop(monkeypatch): + """Regression: Runner.run_sync raises RuntimeError under a running loop + (Jupyter/FastAPI); AgentRunner.run must offload to a worker thread.""" + import asyncio + agents = __import__("agents") + + class FakeResult: + final_output = "ok" + + async def fake_run(agent, question): + return FakeResult() + + def fail_run_sync(agent, question): + raise AssertionError("run_sync must not be called inside a running loop") + + monkeypatch.setattr(agents.Runner, "run", fake_run) + monkeypatch.setattr(agents.Runner, "run_sync", fail_run_sync) + monkeypatch.setattr(agents, "Agent", lambda **kwargs: object()) + + runner = AgentRunner(tools=AgentTools(function_tools=[]), model="gpt-4o") + + async def main(): + return runner.run("question") # sync call from inside a running loop + + assert asyncio.run(main()) == "ok" diff --git a/tests/test_cloud_backend.py b/tests/test_cloud_backend.py index 9bcdfba..9bb1105 100644 --- a/tests/test_cloud_backend.py +++ b/tests/test_cloud_backend.py @@ -202,3 +202,20 @@ def test_query_stream_connect_failure_raises_instead_of_hanging(monkeypatch): with pytest.raises(CloudAPIError, match="request failed"): asyncio.run(_run()) + + +def test_query_uses_long_timeout_and_single_attempt(monkeypatch): + """Non-streaming chat completion is non-idempotent and slow: it must get + a long timeout and must NOT be retried (each retry re-bills the query).""" + backend = CloudBackend(api_key="pi-test") + calls = [] + + def fake_request(method, url, headers=None, **kwargs): + calls.append(kwargs) + raise cloud_mod.requests.ConnectionError("boom") + + monkeypatch.setattr(cloud_mod.requests, "request", fake_request) + with pytest.raises(CloudAPIError): + backend.query("col", "q", doc_ids=["d1"]) + assert len(calls) == 1 + assert calls[0]["timeout"] == 300 diff --git a/tests/test_legacy_sdk_contract.py b/tests/test_legacy_sdk_contract.py index 12b76e5..210769b 100644 --- a/tests/test_legacy_sdk_contract.py +++ b/tests/test_legacy_sdk_contract.py @@ -324,3 +324,14 @@ def test_empty_api_key_warns_and_falls_back_to_local(caplog, tmp_path, monkeypat assert any("empty api_key" in r.message for r in caplog.records) assert client._legacy_cloud_api is None + + +def test_is_retrieval_ready_propagates_api_errors(monkeypatch): + """Regression: API failures (revoked key etc.) were swallowed as False, + turning `while not is_retrieval_ready(...)` into an infinite poll.""" + def fake_request(method, url, **kwargs): + return FakeResponse(status_code=401, text="invalid api key") + + monkeypatch.setattr("pageindex.cloud_api.requests.request", fake_request) + with pytest.raises(PageIndexAPIError): + PageIndexClient("pi-test").is_retrieval_ready("doc-1") diff --git a/tests/test_markdown_parser.py b/tests/test_markdown_parser.py index cbd06af..6d337c1 100644 --- a/tests/test_markdown_parser.py +++ b/tests/test_markdown_parser.py @@ -53,3 +53,21 @@ def test_parse_nodes_have_index(sample_md): result = parser.parse(sample_md) for node in result.nodes: assert node.index is not None + + +def test_preamble_before_first_header_is_kept(tmp_path): + md = tmp_path / "pre.md" + md.write_text("Abstract: important preamble text.\n\n# Chapter 1\nBody.\n") + result = MarkdownParser().parse(str(md)) + assert result.nodes[0].title == "pre" + assert "important preamble text" in result.nodes[0].content + assert result.nodes[1].title == "Chapter 1" + + +def test_headerless_file_yields_single_node(tmp_path): + md = tmp_path / "plain.md" + md.write_text("Just some text.\nNo headings at all.\n") + result = MarkdownParser().parse(str(md)) + assert len(result.nodes) == 1 + assert result.nodes[0].title == "plain" + assert "No headings at all" in result.nodes[0].content diff --git a/tests/test_sqlite_storage.py b/tests/test_sqlite_storage.py index 3e89845..9921f92 100644 --- a/tests/test_sqlite_storage.py +++ b/tests/test_sqlite_storage.py @@ -59,3 +59,23 @@ def test_delete_collection_cascades_documents(storage): storage.save_document("papers", "doc-1", {"doc_name": "test.pdf", "doc_type": "pdf", "file_path": "/tmp/test.pdf", "structure": []}) storage.delete_collection("papers") assert "papers" not in storage.list_collections() + + +def test_close_closes_connections_created_in_other_threads(storage): + """Regression: with check_same_thread=True, close() from another thread + raised ProgrammingError (swallowed) and leaked every worker connection.""" + import sqlite3 + import threading + + conns = {} + + def worker(): + conns["worker"] = storage._get_conn() + + t = threading.Thread(target=worker) + t.start() + t.join() + + storage.close() # main thread closes the worker's connection too + with pytest.raises(sqlite3.ProgrammingError): + conns["worker"].execute("SELECT 1")