From d96c4678c803eea134dfd0ab8aa329ad673389c5 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 1 Aug 2026 13:43:20 +0800 Subject: [PATCH 1/4] fix: release retrieval sessions before workflow waits --- .../tests/contract/test_retrieval_contract.py | 76 ++++ ...est_retrieval_workflow_session_contract.py | 334 ++++++++++++++++++ .../shared-python/shared/core/config/ai.py | 4 + .../services/retrieval/execution/routes.py | 69 ++-- .../retrieval/workflow/orchestrator.py | 18 +- .../retrieval/workflow/plan_service.py | 2 + .../services/retrieval/workflow/planner.py | 57 ++- .../retrieval/workflow/runtime_config.py | 12 + 8 files changed, 521 insertions(+), 51 deletions(-) create mode 100644 apps/api/tests/contract/test_retrieval_workflow_session_contract.py diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 58de969c7..9b8eb1074 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -781,6 +781,82 @@ async def run_request( assert results == [] +@pytest.mark.asyncio +async def test_agentic_retrieval_should_fail_when_final_hydration_db_fails( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + class FakeWorkflowOrchestrator: + async def run_request( + self, + _db: AsyncSession, + *, + request: WorkflowRunRequest, + ) -> WorkflowResult: + return WorkflowResult( + namespace=request.namespace, + query=request.query, + router_used="workflow_single_step", + answer_text="", + referenced_chunks=[ + { + "chunk_id": visible_document["chunk_id"], + "document_id": visible_document["document_id"], + "chunk_type": "text", + "section_path": visible_document["section_path"], + "file_path": None, + "job_id": visible_document["job_id"], + } + ], + ) + + async def fail_final_hydration(**_kwargs: object) -> object: + raise RuntimeError("forced final hydration database failure") + + async with developer_api_client_factory() as api_client: + visible_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-final-hydration-failure", + source_file_name="visible.pdf", + section_path="visible/section", + content="visible scoped content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-final-hydration-failure", + source_file_name="filler.pdf", + section_path="filler/section", + content="filler content", + ) + from shared.services.retrieval.execution import routes as retrieval_routes + + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", + FakeWorkflowOrchestrator, + ) + monkeypatch.setattr( + retrieval_routes, + "resolve_workflow_references", + fail_final_hydration, + ) + + with pytest.raises( + RuntimeError, + match="forced final hydration database failure", + ): + await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-final-hydration-failure", + "query": "visible", + "top_k": 1, + "use_agentic": True, + }, + ) + + @pytest.mark.asyncio async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_id_across_documents( developer_api_client_factory: Callable[ diff --git a/apps/api/tests/contract/test_retrieval_workflow_session_contract.py b/apps/api/tests/contract/test_retrieval_workflow_session_contract.py new file mode 100644 index 000000000..c222d4457 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_workflow_session_contract.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import cast + +import pytest +from pytest import MonkeyPatch +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import LLMServiceException +from shared.services.retrieval.execution import routes as route_module +from shared.services.retrieval.execution.reference_resolver import ( + ResolvedWorkflowReferences, +) +from shared.services.retrieval.execution.route_types import RetrievalRouteContext +from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.workflow.orchestrator import DbSessionFactory, WorkflowOrchestrator +from shared.services.retrieval.workflow.plan_service import WorkflowPlanService +from shared.services.retrieval.workflow.planner import QueryPlanner +from shared.services.retrieval.workflow.run_request import WorkflowRunRequest +from shared.services.retrieval.workflow.step_runner import WorkflowStepRunner +from shared.services.retrieval.workflow.types import ( + PlannedStep, + QueryPlan, + StepResult, + WorkflowResult, +) + +RouteRow = dict[str, object] + + +def _build_planner( + llm_fn: LLMFn, + *, + timeout_seconds: float = 1.0, +) -> QueryPlanner: + return QueryPlanner( + llm_fn=llm_fn, + planner_ledger=None, + max_steps=3, + total_budget=100, + per_step_budget=10, + timeout_seconds=timeout_seconds, + ) + + +@pytest.mark.asyncio +async def test_workflow_planner_timeout_should_return_single_step_fallback() -> None: + async def slow_llm(_prompt: object) -> str: + await asyncio.sleep(0.05) + return "{}" + + planner = _build_planner(slow_llm, timeout_seconds=0.001) + + plan = await planner.plan(query="original query") + + assert plan.planner_status == "fallback" + assert plan.steps[0].sub_query == "original query" + assert plan.planner_error is not None + assert "timed out" in plan.planner_error + + +@pytest.mark.asyncio +async def test_workflow_planner_provider_error_should_return_single_step_fallback() -> None: + async def failing_llm(_prompt: object) -> str: + raise LLMServiceException(internal_message="provider unavailable") + + planner = _build_planner(failing_llm) + + plan = await planner.plan(query="provider failure query") + + assert plan.planner_status == "fallback" + assert plan.steps[0].sub_query == "provider failure query" + assert plan.planner_error is not None + + +@pytest.mark.asyncio +async def test_workflow_planner_invalid_json_should_return_single_step_fallback() -> None: + async def invalid_llm(_prompt: object) -> str: + return "not json" + + planner = _build_planner(invalid_llm) + + plan = await planner.plan(query="invalid planner output") + + assert plan.planner_status == "fallback" + assert plan.steps[0].sub_query == "invalid planner output" + assert plan.planner_error is not None + assert "JSON" in plan.planner_error + + +@pytest.mark.asyncio +async def test_workflow_planner_unexpected_code_error_should_propagate() -> None: + async def buggy_llm(_prompt: object) -> str: + raise RuntimeError("unexpected planner bug") + + planner = _build_planner(buggy_llm) + + with pytest.raises(RuntimeError, match="unexpected planner bug"): + await planner.plan(query="buggy planner query") + + +@pytest.mark.asyncio +async def test_workflow_inventory_session_should_close_before_planner_starts( + monkeypatch: MonkeyPatch, +) -> None: + events: list[str] = [] + inventory_db = object() + + @asynccontextmanager + async def fake_db_factory() -> AsyncGenerator[AsyncSession, None]: + events.append("inventory_open") + try: + yield cast(AsyncSession, inventory_db) + finally: + events.append("inventory_close") + + async def fake_load_budget_inventory( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], + ) -> tuple[int, int, dict[str, int]]: + assert db is inventory_db + assert user_id == "contract-user" + assert namespace == "contract-namespace" + assert exclude_document_ids == [] + events.append("inventory_loaded") + return 3, 2, {} + + class RecordingPlanService: + async def load_or_create(self, **kwargs: object) -> QueryPlan: + events.append("planner_start") + assert events == [ + "inventory_open", + "inventory_loaded", + "inventory_close", + "planner_start", + ] + return QueryPlan.single_step(str(kwargs["query"])) + + class RecordingStepRunner: + async def run_step(self, **kwargs: object) -> None: + step = cast(PlannedStep, kwargs["step"]) + results_by_id = cast(dict[str, StepResult], kwargs["results_by_id"]) + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="done", + answer_text="", + ) + + def create_step_runner( + db_factory: DbSessionFactory, + parent_run_id: str, + ) -> WorkflowStepRunner: + assert db_factory is fake_db_factory + assert parent_run_id + return cast(WorkflowStepRunner, RecordingStepRunner()) + + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator._load_budget_inventory", + fake_load_budget_inventory, + ) + orchestrator = WorkflowOrchestrator( + db_factory=fake_db_factory, + plan_service=cast(WorkflowPlanService, RecordingPlanService()), + step_runner_factory=create_step_runner, + ) + + await orchestrator.run_request( + cast(AsyncSession, object()), + request=WorkflowRunRequest( + user_id="contract-user", + namespace="contract-namespace", + query="session lifetime", + top_k=1, + exclude_document_ids=[], + exclude_sections=[], + ), + llm_fn=None, + ) + + assert events[:4] == [ + "inventory_open", + "inventory_loaded", + "inventory_close", + "planner_start", + ] + + +@pytest.mark.asyncio +async def test_agentic_route_should_release_route_session_before_fresh_final_hydration( + monkeypatch: MonkeyPatch, +) -> None: + events: list[str] = [] + route_db = _RecordingRouteSession(events) + hydration_db = object() + + @asynccontextmanager + async def fake_get_db_context() -> AsyncGenerator[AsyncSession, None]: + events.append("hydration_open") + try: + yield cast(AsyncSession, hydration_db) + finally: + events.append("hydration_close") + + class FakeWorkflowOrchestrator: + async def run_request( + self, + db: AsyncSession, + *, + request: WorkflowRunRequest, + ) -> WorkflowResult: + events.append("workflow_start") + assert db is route_db + assert events[:2] == ["route_rollback", "workflow_start"] + return WorkflowResult( + namespace=request.namespace, + query=request.query, + router_used="workflow_single_step", + answer_text="", + referenced_chunks=[ + { + "chunk_id": "chunk_contract", + "document_id": "doc_contract", + "chunk_type": "text", + "section_path": "contract/section", + "file_path": None, + "job_id": "job_contract", + } + ], + ) + + async def fake_resolve_workflow_references( + *, + db: AsyncSession, + user_id: str, + namespace: str, + refs: list[RouteRow], + score_by_chunk_id: dict[str, float] | None = None, + ) -> ResolvedWorkflowReferences: + assert db is hydration_db + assert user_id == "contract-user" + assert namespace == "contract-namespace" + assert score_by_chunk_id is None + events.append("resolve_references") + row = { + "document_id": "doc_contract", + "chunk_id": "chunk_contract", + "source_file_name": "contract.pdf", + "section_path": "contract/section", + "chunk_type": "text", + "content": "contract content", + } + return ResolvedWorkflowReferences(refs=refs, rows=[row]) + + async def fake_assemble_retrieval_results( + *, + db: AsyncSession, + rows: list[RouteRow], + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + allowed_chunk_types: set[str] | None, + ) -> list[RouteRow]: + assert db is hydration_db + assert exclude_document_ids == [] + assert exclude_sections == [] + assert allowed_chunk_types is None + events.append("assemble_results") + return rows + + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", + FakeWorkflowOrchestrator, + ) + monkeypatch.setattr("shared.core.database.get_db_context", fake_get_db_context) + monkeypatch.setattr( + route_module, + "resolve_workflow_references", + fake_resolve_workflow_references, + ) + monkeypatch.setattr( + route_module, + "assemble_retrieval_results", + fake_assemble_retrieval_results, + ) + + outcome = await route_module._run_agentic_route( + RetrievalRouteContext( + db=cast(AsyncSession, route_db), + user_id="contract-user", + namespace="contract-namespace", + query="session lifetime", + top_k=1, + exclude_document_ids=[], + exclude_sections=[], + allowed_chunk_types=None, + chunk_types=None, + signal_paths=None, + filter_mode="delete", + channels=None, + channel_weights=None, + rerank=False, + threshold=0.0, + internal_recall_k=None, + effective_recall_k=3, + use_agentic=True, + ) + ) + + assert events == [ + "route_rollback", + "workflow_start", + "hydration_open", + "resolve_references", + "assemble_results", + "hydration_close", + ] + assert outcome.response["results"][0]["citation"]["document_id"] == "doc_contract" + + +class _RecordingRouteSession: + def __init__(self, events: list[str]) -> None: + self._events = events + + async def rollback(self) -> None: + self._events.append("route_rollback") diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index b2f612d2c..90ca10bbf 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -68,6 +68,10 @@ class AIConfig(BaseModel): default=3, description="Maximum concurrent workflow steps in the same DAG batch.", ) + RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS: float = Field( + default=25.0, + description="Timeout for the optional retrieval workflow planner LLM call.", + ) # Runtime LLM controls. LLM_MOCK_ENABLED: bool = Field( diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 08b33a2a4..b93fbc806 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -36,19 +36,13 @@ async def run_retrieval_route( async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: - try: - total_chunk_count = await count_scoped_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - allowed_chunk_types=context.allowed_chunk_types, - ) - except Exception as exc: - logger.warning( - f"Failed to count scoped chunks, skipping small corpus optimization: {exc}" - ) - total_chunk_count = context.top_k + 1 + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + ) logger.info(f"\n Total chunks in scope: {total_chunk_count}") if total_chunk_count > context.top_k: @@ -161,6 +155,8 @@ async def _run_agentic_route( from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator from shared.services.retrieval.workflow.run_request import WorkflowRunRequest + await context.db.rollback() + workflow = WorkflowOrchestrator() workflow_result = await workflow.run_request( context.db, @@ -199,30 +195,33 @@ async def _run_agentic_route( except (TypeError, ValueError): pass - resolved_references = await resolve_workflow_references( - db=context.db, - user_id=context.user_id, - namespace=context.namespace, - refs=workflow_result.referenced_chunks, - score_by_chunk_id=score_by_chunk_id if score_by_chunk_id else None, - ) + from shared.core.database import get_db_context - # Backfill doc-level confidence for chunks that have no discovery score - if doc_confidence: - for row in resolved_references.rows: - cid = row.get('chunk_id', '') - if cid and row.get('score') is None: - doc_id = row.get('document_id', '') - if doc_id in doc_confidence: - row['score'] = doc_confidence[doc_id] + async with get_db_context() as hydration_db: + resolved_references = await resolve_workflow_references( + db=hydration_db, + user_id=context.user_id, + namespace=context.namespace, + refs=workflow_result.referenced_chunks, + score_by_chunk_id=score_by_chunk_id if score_by_chunk_id else None, + ) - assembled_workflow_rows = await assemble_retrieval_results( - db=context.db, - rows=resolved_references.rows, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - allowed_chunk_types=context.allowed_chunk_types, - ) + # Backfill doc-level confidence for chunks that have no discovery score + if doc_confidence: + for row in resolved_references.rows: + cid = row.get('chunk_id', '') + if cid and row.get('score') is None: + doc_id = row.get('document_id', '') + if doc_id in doc_confidence: + row['score'] = doc_confidence[doc_id] + + assembled_workflow_rows = await assemble_retrieval_results( + db=hydration_db, + rows=resolved_references.rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) response = workflow_result.to_api_response() response["answer_text"] = "" response["referenced_chunks"] = resolved_references.refs diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index d8a791878..10c04fff6 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -104,6 +104,7 @@ async def run_request( request: WorkflowRunRequest, llm_fn=None, ) -> WorkflowResult: + _ = db t0 = time.monotonic() config = WorkflowRuntimeConfig.from_env() llm_fn = llm_fn or create_retrieval_llm_fn() @@ -115,12 +116,14 @@ async def run_request( bootstrap=config.planner_budget, per_doc_min_share=0, ) - total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory( - db, - user_id=request.user_id, - namespace=request.namespace, - exclude_document_ids=request.exclude_document_ids, - ) + db_factory = self._get_db_factory() + async with db_factory() as inventory_db: + total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory( + inventory_db, + user_id=request.user_id, + namespace=request.namespace, + exclude_document_ids=request.exclude_document_ids, + ) planner_ledger.total_chunks = total_chunks planner_ledger.total_docs = total_docs plan = await self._plan_service.load_or_create( @@ -137,6 +140,7 @@ async def run_request( per_retrieve=config.per_retrieve_step_budget, corpus_total_docs=total_docs, corpus_total_chunks=total_chunks, + planner_timeout_seconds=config.planner_timeout_seconds, ) # TODO(retrieval-agentic-nav): redesign this outer workflow as a real # observe-act agent that can pass evidence between steps, decide whether @@ -152,7 +156,7 @@ async def run_request( results_by_id: dict[str, StepResult] = {} sem = asyncio.Semaphore(config.parallel_max) step_runner = self._step_runner_factory( - self._get_db_factory(), + db_factory, self.parent_run_id, ) diff --git a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py index a887ec5f4..188206546 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py +++ b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py @@ -30,6 +30,7 @@ async def load_or_create( per_retrieve: int, corpus_total_docs: int, corpus_total_chunks: int, + planner_timeout_seconds: float, ) -> QueryPlan: try: cached = await get_cached_workflow_plan( @@ -51,6 +52,7 @@ async def load_or_create( max_steps=max_steps, total_budget=wallet_total, per_step_budget=per_retrieve, + timeout_seconds=planner_timeout_seconds, ) plan = await planner.plan( query=query, diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py index 53364ac30..1f951445c 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/planner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -1,6 +1,7 @@ """Query planner for decomposed retrieval workflows.""" from __future__ import annotations +import asyncio import json import re import time @@ -8,6 +9,7 @@ from loguru import logger +from shared.core.exceptions.domain_exceptions import LLMServiceException, UnavailableException from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetLedger from shared.services.retrieval.agentic.core.runtime import _extract_actual_tokens from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage @@ -62,6 +64,18 @@ """ +class PlannerExpectedFailure(Exception): + """Expected planner failure that can fall back to single-step retrieval.""" + + +class PlannerProviderFailure(PlannerExpectedFailure): + """Planner LLM provider call failed or timed out.""" + + +class PlannerOutputFailure(PlannerExpectedFailure): + """Planner LLM returned invalid JSON or an invalid query plan.""" + + class QueryPlanner: """LLM-backed planner with strict fallback to a single retrieve step.""" @@ -73,12 +87,14 @@ def __init__( max_steps: int, total_budget: int, per_step_budget: int, + timeout_seconds: float, ) -> None: self._llm_fn = llm_fn self._ledger = planner_ledger self._max_steps = max(max_steps, 1) self._total_budget = max(total_budget, 1) self._per_step_budget = max(per_step_budget, 1) + self._timeout_seconds = max(timeout_seconds, 0.001) async def plan( self, @@ -103,14 +119,17 @@ async def plan( try: raw = await self._call_llm_with_budget(prompt) - plan = _parse_plan_response( - raw, - original_query=query, - max_steps=self._max_steps, - ) - plan.validate() + try: + plan = _parse_plan_response( + raw, + original_query=query, + max_steps=self._max_steps, + ) + plan.validate() + except (TypeError, ValueError) as exc: + raise PlannerOutputFailure(str(exc)) from exc return plan - except Exception as exc: + except (BudgetExceeded, PlannerExpectedFailure) as exc: logger.warning(f"workflow planner failed, falling back to single step: {exc}") plan = QueryPlan.single_step(query, reason="planner_fallback_single_step") plan.planner_status = "fallback" @@ -130,7 +149,7 @@ def _planner_used(self) -> int: async def _call_llm_with_budget(self, prompt: str) -> str: if self._ledger is None: - return await self._llm_fn(prompt) # type: ignore[misc] + return await self._call_planner_llm(prompt) est = estimate_tokens(prompt) reserved = await self._ledger.try_reserve("bootstrap", est) @@ -139,7 +158,16 @@ async def _call_llm_with_budget(self, prompt: str) -> str: t0 = time.monotonic() try: - response = await self._llm_fn(prompt) # type: ignore[misc] + response = await asyncio.wait_for( + self._llm_fn(prompt), # type: ignore[misc] + timeout=self._timeout_seconds, + ) + except asyncio.TimeoutError as exc: + await self._ledger.refund("bootstrap", est=est) + raise PlannerProviderFailure("planner LLM call timed out") from exc + except (LLMServiceException, UnavailableException) as exc: + await self._ledger.refund("bootstrap", est=est) + raise PlannerProviderFailure(str(exc)) from exc except Exception: await self._ledger.refund("bootstrap", est=est) raise @@ -154,6 +182,17 @@ async def _call_llm_with_budget(self, prompt: str) -> str: ) return response + async def _call_planner_llm(self, prompt: str) -> str: + try: + return await asyncio.wait_for( + self._llm_fn(prompt), # type: ignore[misc] + timeout=self._timeout_seconds, + ) + except asyncio.TimeoutError as exc: + raise PlannerProviderFailure("planner LLM call timed out") from exc + except (LLMServiceException, UnavailableException) as exc: + raise PlannerProviderFailure(str(exc)) from exc + def _parse_plan_response(text: str, *, original_query: str, max_steps: int) -> QueryPlan: data = _extract_json_object(text) diff --git a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py index 28630dff0..e73f82c2e 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py +++ b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py @@ -12,6 +12,7 @@ class WorkflowRuntimeConfig: per_retrieve_step_budget: int = 40000 max_steps: int = 5 parallel_max: int = 3 + planner_timeout_seconds: float = 25.0 @classmethod def from_env(cls) -> "WorkflowRuntimeConfig": @@ -21,6 +22,10 @@ def from_env(cls) -> "WorkflowRuntimeConfig": per_retrieve_step_budget=_env_int("RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET", 40000), max_steps=_env_int("RETRIEVAL_DECOMPOSITION_MAX_STEPS", 5), parallel_max=_env_int("RETRIEVAL_WORKFLOW_PARALLEL_MAX", 3), + planner_timeout_seconds=_env_float( + "RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS", + 25.0, + ), ) @@ -29,3 +34,10 @@ def _env_int(name: str, default: int) -> int: return int(os.environ.get(name, str(default))) except (TypeError, ValueError): return default + + +def _env_float(name: str, default: float) -> float: + try: + return float(os.environ.get(name, str(default))) + except (TypeError, ValueError): + return default From 8da8aa065e289502861525c4ad709ebb712954cc Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 2 Aug 2026 22:03:21 +0800 Subject: [PATCH 2/4] fix: stabilize retrieval workflow timeout handling --- apps/api/.env.example | 1 + ...est_retrieval_workflow_session_contract.py | 43 ++++++++++++++++++- .../services/retrieval/execution/routes.py | 13 ++++-- .../shared/services/retrieval/llm_adapter.py | 10 +++++ .../retrieval/workflow/orchestrator.py | 5 ++- 5 files changed, 66 insertions(+), 6 deletions(-) diff --git a/apps/api/.env.example b/apps/api/.env.example index 04f297b40..e362f3ce2 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -102,6 +102,7 @@ ARK_API_KEY= # evidence_text is the primary output and answer_text is always empty. Set # RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy # 3-channel RRF mode. +# RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS=25.0 # File handling defaults SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md,.html,.htm diff --git a/apps/api/tests/contract/test_retrieval_workflow_session_contract.py b/apps/api/tests/contract/test_retrieval_workflow_session_contract.py index c222d4457..892393afa 100644 --- a/apps/api/tests/contract/test_retrieval_workflow_session_contract.py +++ b/apps/api/tests/contract/test_retrieval_workflow_session_contract.py @@ -15,7 +15,7 @@ ResolvedWorkflowReferences, ) from shared.services.retrieval.execution.route_types import RetrievalRouteContext -from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.llm_adapter import LLMFn, create_retrieval_planner_fn from shared.services.retrieval.workflow.orchestrator import DbSessionFactory, WorkflowOrchestrator from shared.services.retrieval.workflow.plan_service import WorkflowPlanService from shared.services.retrieval.workflow.planner import QueryPlanner @@ -102,6 +102,45 @@ async def buggy_llm(_prompt: object) -> str: await planner.plan(query="buggy planner query") +@pytest.mark.asyncio +async def test_workflow_planner_llm_should_pass_timeout_to_provider_client( + monkeypatch: MonkeyPatch, +) -> None: + observed_timeouts: list[object] = [] + + class FakeClient: + def chat_completion_with_usage( + self, + _prompt: object, + **kwargs: object, + ) -> tuple[str, dict[str, int]]: + observed_timeouts.append(kwargs.get("timeout")) + return "{}", {"total_tokens": 1} + + def fake_build_client_for_channel( + *, + channel: str, + model: str, + ) -> tuple[FakeClient, str]: + assert channel == "text" + return FakeClient(), model + + monkeypatch.setattr( + "shared.services.retrieval.llm_adapter._has_llm_credentials", + lambda: True, + ) + monkeypatch.setattr( + "shared.services.retrieval.llm_adapter._build_client_for_channel", + fake_build_client_for_channel, + ) + + planner_llm = create_retrieval_planner_fn(timeout_seconds=2.1) + + assert planner_llm is not None + await planner_llm("timeout contract") + assert observed_timeouts == [3] + + @pytest.mark.asyncio async def test_workflow_inventory_session_should_close_before_planner_starts( monkeypatch: MonkeyPatch, @@ -280,7 +319,7 @@ async def fake_assemble_retrieval_results( "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", FakeWorkflowOrchestrator, ) - monkeypatch.setattr("shared.core.database.get_db_context", fake_get_db_context) + monkeypatch.setattr(route_module, "open_hydration_db_context", fake_get_db_context) monkeypatch.setattr( route_module, "resolve_workflow_references", diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index b93fbc806..7725f22cc 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -1,6 +1,9 @@ from __future__ import annotations +from contextlib import AbstractAsyncContextManager + from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.agentic.discovery.tools import bottom_discovery from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references @@ -20,6 +23,12 @@ ) +def open_hydration_db_context() -> AbstractAsyncContextManager[AsyncSession]: + from shared.core.database import get_db_context + + return get_db_context() + + async def run_retrieval_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: @@ -195,9 +204,7 @@ async def _run_agentic_route( except (TypeError, ValueError): pass - from shared.core.database import get_db_context - - async with get_db_context() as hydration_db: + async with open_hydration_db_context() as hydration_db: resolved_references = await resolve_workflow_references( db=hydration_db, user_id=context.user_id, diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index 8c8cb627d..2a72ccf11 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio +import math import os from contextvars import ContextVar from typing import Any, Callable, Coroutine, Union, Sequence, cast @@ -177,6 +178,7 @@ def create_retrieval_planner_fn( thinking: bool = True, model: str | None = None, max_tokens: int = 8192, + timeout_seconds: float | None = None, ) -> LLMFn | None: """Create a reasoning-capable LLM callable for query planning.""" if not _has_llm_credentials(): @@ -184,6 +186,7 @@ def create_retrieval_planner_fn( return None effective_model = model or _resolve_planner_model(thinking=thinking) + request_timeout = _coerce_provider_timeout_seconds(timeout_seconds) async def llm_fn(prompt: LLMFnInput) -> str: client, resolved_model = _build_client_for_channel( @@ -197,6 +200,7 @@ async def llm_fn(prompt: LLMFnInput) -> str: model=resolved_model, temperature=0.0, max_tokens=max_tokens, + timeout=request_timeout, ) current_llm_usage.set(usage) return result @@ -204,6 +208,12 @@ async def llm_fn(prompt: LLMFnInput) -> str: return llm_fn +def _coerce_provider_timeout_seconds(timeout_seconds: float | None) -> int | None: + if timeout_seconds is None or not math.isfinite(timeout_seconds): + return None + return max(1, math.ceil(timeout_seconds)) + + def create_retrieval_vlm_fn( *, model: str | None = None, diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index 10c04fff6..aa4db437a 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -108,7 +108,10 @@ async def run_request( t0 = time.monotonic() config = WorkflowRuntimeConfig.from_env() llm_fn = llm_fn or create_retrieval_llm_fn() - planner_llm = create_retrieval_planner_fn(thinking=True) + planner_llm = create_retrieval_planner_fn( + thinking=True, + timeout_seconds=config.planner_timeout_seconds, + ) planner_ledger = BudgetLedger( total=config.planner_budget, From d490b268f080069079ae0769b151788de94fcf17 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 3 Aug 2026 00:04:31 +0800 Subject: [PATCH 3/4] fix: clarify retrieval workflow database context --- apps/api/.env.example | 2 +- ...test_retrieval_workflow_session_contract.py | 18 +++++++++--------- .../shared-python/shared/core/config/ai.py | 2 +- .../services/retrieval/execution/routes.py | 9 +++++---- .../retrieval/workflow/runtime_config.py | 4 ++-- 5 files changed, 18 insertions(+), 17 deletions(-) diff --git a/apps/api/.env.example b/apps/api/.env.example index e362f3ce2..e33e8bc37 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -102,7 +102,7 @@ ARK_API_KEY= # evidence_text is the primary output and answer_text is always empty. Set # RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy # 3-channel RRF mode. -# RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS=25.0 +# RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS=10.0 # File handling defaults SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md,.html,.htm diff --git a/apps/api/tests/contract/test_retrieval_workflow_session_contract.py b/apps/api/tests/contract/test_retrieval_workflow_session_contract.py index 892393afa..cbe7e472c 100644 --- a/apps/api/tests/contract/test_retrieval_workflow_session_contract.py +++ b/apps/api/tests/contract/test_retrieval_workflow_session_contract.py @@ -240,15 +240,15 @@ async def test_agentic_route_should_release_route_session_before_fresh_final_hyd ) -> None: events: list[str] = [] route_db = _RecordingRouteSession(events) - hydration_db = object() + fresh_db = object() @asynccontextmanager async def fake_get_db_context() -> AsyncGenerator[AsyncSession, None]: - events.append("hydration_open") + events.append("fresh_db_open") try: - yield cast(AsyncSession, hydration_db) + yield cast(AsyncSession, fresh_db) finally: - events.append("hydration_close") + events.append("fresh_db_close") class FakeWorkflowOrchestrator: async def run_request( @@ -285,7 +285,7 @@ async def fake_resolve_workflow_references( refs: list[RouteRow], score_by_chunk_id: dict[str, float] | None = None, ) -> ResolvedWorkflowReferences: - assert db is hydration_db + assert db is fresh_db assert user_id == "contract-user" assert namespace == "contract-namespace" assert score_by_chunk_id is None @@ -308,7 +308,7 @@ async def fake_assemble_retrieval_results( exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None, ) -> list[RouteRow]: - assert db is hydration_db + assert db is fresh_db assert exclude_document_ids == [] assert exclude_sections == [] assert allowed_chunk_types is None @@ -319,7 +319,7 @@ async def fake_assemble_retrieval_results( "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", FakeWorkflowOrchestrator, ) - monkeypatch.setattr(route_module, "open_hydration_db_context", fake_get_db_context) + monkeypatch.setattr(route_module, "open_fresh_database_context", fake_get_db_context) monkeypatch.setattr( route_module, "resolve_workflow_references", @@ -357,10 +357,10 @@ async def fake_assemble_retrieval_results( assert events == [ "route_rollback", "workflow_start", - "hydration_open", + "fresh_db_open", "resolve_references", "assemble_results", - "hydration_close", + "fresh_db_close", ] assert outcome.response["results"][0]["citation"]["document_id"] == "doc_contract" diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index 90ca10bbf..487a35808 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -69,7 +69,7 @@ class AIConfig(BaseModel): description="Maximum concurrent workflow steps in the same DAG batch.", ) RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS: float = Field( - default=25.0, + default=10.0, description="Timeout for the optional retrieval workflow planner LLM call.", ) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 7725f22cc..7436dd4dd 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -23,7 +23,8 @@ ) -def open_hydration_db_context() -> AbstractAsyncContextManager[AsyncSession]: +def open_fresh_database_context() -> AbstractAsyncContextManager[AsyncSession]: + """Open a fresh session for final reference resolution after workflow waits.""" from shared.core.database import get_db_context return get_db_context() @@ -204,9 +205,9 @@ async def _run_agentic_route( except (TypeError, ValueError): pass - async with open_hydration_db_context() as hydration_db: + async with open_fresh_database_context() as final_db: resolved_references = await resolve_workflow_references( - db=hydration_db, + db=final_db, user_id=context.user_id, namespace=context.namespace, refs=workflow_result.referenced_chunks, @@ -223,7 +224,7 @@ async def _run_agentic_route( row['score'] = doc_confidence[doc_id] assembled_workflow_rows = await assemble_retrieval_results( - db=hydration_db, + db=final_db, rows=resolved_references.rows, exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, diff --git a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py index e73f82c2e..4e6a3f353 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py +++ b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py @@ -12,7 +12,7 @@ class WorkflowRuntimeConfig: per_retrieve_step_budget: int = 40000 max_steps: int = 5 parallel_max: int = 3 - planner_timeout_seconds: float = 25.0 + planner_timeout_seconds: float = 10.0 @classmethod def from_env(cls) -> "WorkflowRuntimeConfig": @@ -24,7 +24,7 @@ def from_env(cls) -> "WorkflowRuntimeConfig": parallel_max=_env_int("RETRIEVAL_WORKFLOW_PARALLEL_MAX", 3), planner_timeout_seconds=_env_float( "RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS", - 25.0, + 10.0, ), ) From d05f4394b2045015fd919a1802528c672ee3477f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 3 Aug 2026 09:52:58 +0800 Subject: [PATCH 4/4] docs: clarify retrieval route session rollback --- .../shared/services/retrieval/execution/routes.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 7436dd4dd..f980dbebc 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -165,6 +165,9 @@ async def _run_agentic_route( from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator from shared.services.retrieval.workflow.run_request import WorkflowRunRequest + # The small-corpus count above may leave a read transaction checked out on + # the request session. End it before planner/navigation LLM waits; workflow + # steps and final reference resolution open their own short-lived sessions. await context.db.rollback() workflow = WorkflowOrchestrator()