From 1fc66578fbebcebf1659183f6cb2cab23e973faa Mon Sep 17 00:00:00 2001 From: Yi Lu Date: Mon, 10 Aug 2026 00:26:43 -0700 Subject: [PATCH 1/2] fix: make durable playbook recovery deterministic Reconstruct review evidence from the extraction run's persisted window, require an unambiguous user owner, and discover resumable work from the authoritative provider before relying on a bootstrap org. Preserve exception metadata only when it belongs to the current LiteLLM failure. --- reflexio/server/api.py | 9 +- .../server/llm/_litellm_text_generation.py | 27 ++- reflexio/server/llm/_litellm_types.py | 2 + reflexio/server/services/extraction/README.md | 16 +- .../services/extraction/agent_run_records.py | 15 +- .../services/extraction/resumable_agent.py | 2 +- .../services/extraction/resume_scheduler.py | 67 ++++--- .../services/extraction/resume_worker.py | 38 +++- reflexio/server/services/playbook/README.md | 14 +- .../playbook/playbook_service_utils.py | 3 +- .../services/playbook/review_service.py | 82 ++------- .../server/services/playbook/review_window.py | 166 ++++++++++++++++++ reflexio/server/services/playbook/service.py | 107 +++++++---- tests/eval/extraction/providers.py | 1 + tests/server/llm/test_litellm_client_unit.py | 66 +++++++ .../extraction/test_resume_scheduler.py | 104 +++++++++++ .../services/extraction/test_resume_worker.py | 76 +++++++- .../test_extractor_polarity_integration.py | 17 +- .../playbook/test_playbook_extractor.py | 41 ++--- .../test_playbook_generation_service.py | 157 ++++++++++------- ...playbook_generation_service_integration.py | 38 +++- .../storage/test_agent_run_helpers.py | 36 ++++ 22 files changed, 845 insertions(+), 239 deletions(-) create mode 100644 reflexio/server/services/playbook/review_window.py diff --git a/reflexio/server/api.py b/reflexio/server/api.py index 9fc9bdb26..5117b2aac 100644 --- a/reflexio/server/api.py +++ b/reflexio/server/api.py @@ -414,10 +414,11 @@ def create_app( # noqa: C901 logic. It is only consulted when ``REFLEXIO_DURABLE_LEARNING_QUEUE`` is on — when the flag is off ``maybe_start_durable_learning`` returns None without ever calling the provider. - resume_org_ids_provider: Optional zero-arg callable used only when the - bootstrap org is not available yet. This lets a multi-tenant deployment - defer the resume scheduler on an empty fleet and adopt its first real - org without restarting. The OSS default remains unchanged. + resume_org_ids_provider: Optional zero-arg callable consulted before the + bootstrap context on every scheduler tick. Its actionable org list is + authoritative, allowing a multi-tenant deployment to recover when the + prior bootstrap org disappears and to discover work across data refs. + The OSS default remains unchanged when no provider is supplied. Returns: Configured FastAPI application. diff --git a/reflexio/server/llm/_litellm_text_generation.py b/reflexio/server/llm/_litellm_text_generation.py index 124aa5181..2f2ab54d5 100644 --- a/reflexio/server/llm/_litellm_text_generation.py +++ b/reflexio/server/llm/_litellm_text_generation.py @@ -151,10 +151,26 @@ class _StructuredAttempt: def _is_expected_transient_llm_error(exc: BaseException) -> bool: """True for expected transient upstream failures (timeout / connection / rate-limit / overload), including our own ``LLMHardTimeoutError`` (a - ``TimeoutError`` subclass raised when a provider hang is killed).""" - if isinstance(exc, TimeoutError): # incl. LLMHardTimeoutError - return True - return type(exc).__name__ in _TRANSIENT_LLM_ERROR_NAMES + ``TimeoutError`` subclass raised when a provider hang is killed). + + Subprocess-isolated calls cross a pickle boundary, so their concrete + provider exception cannot safely be re-raised in the parent. The wrapper + retains ``upstream_error_type`` explicitly; checking it here prevents an + expected provider outage from being promoted to an application ERROR just + because isolation made the outer type ``LiteLLMClientError``. + """ + current: BaseException | None = exc + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + if isinstance(current, TimeoutError): # incl. LLMHardTimeoutError + return True + if type(current).__name__ in _TRANSIENT_LLM_ERROR_NAMES: + return True + if getattr(current, "upstream_error_type", None) in _TRANSIENT_LLM_ERROR_NAMES: + return True + current = current.__cause__ or current.__context__ + return False def _rung_reason(error: Exception | None) -> str: @@ -860,7 +876,8 @@ def _completion_with_hard_timeout( raise LiteLLMClientError( "litellm.completion failed in isolated worker: " f"{payload.type_name}: {payload.message} " - f"({', '.join(context_parts)})" + f"({', '.join(context_parts)})", + upstream_error_type=payload.type_name, ) finally: result_queue.close() diff --git a/reflexio/server/llm/_litellm_types.py b/reflexio/server/llm/_litellm_types.py index c1d0a40d2..f7d0bb7be 100644 --- a/reflexio/server/llm/_litellm_types.py +++ b/reflexio/server/llm/_litellm_types.py @@ -131,9 +131,11 @@ def __init__( message: str, *, first_parsed_provenance: ModelProvenance | None = None, + upstream_error_type: str | None = None, ) -> None: super().__init__(message) self.first_parsed_provenance = first_parsed_provenance + self.upstream_error_type = upstream_error_type class StructuredOutputRepairError(LiteLLMClientError): diff --git a/reflexio/server/services/extraction/README.md b/reflexio/server/services/extraction/README.md index 3acb5a553..668455202 100644 --- a/reflexio/server/services/extraction/README.md +++ b/reflexio/server/services/extraction/README.md @@ -15,7 +15,7 @@ information, and resumes outside the request path. | `pending_tool_call_dispatch.py` | Implements the `ask_human` and pending-info tool dispatch flow. | | `prior_answer_search.py` | Finds and formats previous human answers for async extraction context. | | `agent_run_records.py` | Builds durable extraction-agent run records and source interaction identity. | -| `resume_scheduler.py` | Schedules due paused extraction runs in a background singleton. | +| `resume_scheduler.py` | Discovers and schedules due paused/finalization work in a background singleton. | | `resume_worker.py` | Resumes paused runs, rebuilds request context, and records retry state. | | `outcome.py` | Provides the generic extraction outcome wrapper used by callers. | @@ -29,3 +29,17 @@ information, and resumes outside the request path. flat package stops being easier to scan. - Do not recreate removed legacy modules such as `tools.py`, `plan.py`, or `invariants.py`; use the current focused files above. +- New durable extraction runs require a non-empty `user_id`. Nullable stored + bindings remain readable for backward compatibility; playbook resume derives + an in-memory owner only from complete, unanimous persisted source evidence. + +## Resume discovery + +When an `org_id_provider` is installed, the scheduler calls it on every tick and +treats its actionable org list as authoritative for cross-ref discovery. Without +a provider, local storage discovery remains the fallback. Each provider result +selects a current org for reading scheduler configuration, so an org retained +from an earlier tick cannot block later discovery if it becomes stale. If the +provider itself raises, the scheduler still attempts the last bootstrap org for +that tick. Before draining each discovered org context, it expires pending tool +calls on that context's storage ref; separate refs have separate queues. diff --git a/reflexio/server/services/extraction/agent_run_records.py b/reflexio/server/services/extraction/agent_run_records.py index 5945ba1c2..41ff61e5c 100644 --- a/reflexio/server/services/extraction/agent_run_records.py +++ b/reflexio/server/services/extraction/agent_run_records.py @@ -46,7 +46,7 @@ def build_extractor_agent_run_record( *, org_id: str, extractor_kind: str, - user_id: str | None, + user_id: str, agent_version: str | None, source: str | None, request_interaction_data_models: list[RequestInteractionDataModel], @@ -57,6 +57,19 @@ def build_extractor_agent_run_record( generation_request_id: str | None = None, request_id: str | None = None, ) -> AgentRunRecord: + user_id = user_id.strip() + if not user_id: + raise ValueError("Durable extraction runs require a non-empty user_id") + if any( + data_model.request.user_id != user_id + or any( + interaction.user_id != user_id for interaction in data_model.interactions + ) + for data_model in request_interaction_data_models + ): + raise ValueError( + "Durable extraction run source evidence must belong to its user_id" + ) if generation_request_id is not None: if request_id is not None and request_id != generation_request_id: raise TypeError( diff --git a/reflexio/server/services/extraction/resumable_agent.py b/reflexio/server/services/extraction/resumable_agent.py index 1f85ad0fa..48896389c 100644 --- a/reflexio/server/services/extraction/resumable_agent.py +++ b/reflexio/server/services/extraction/resumable_agent.py @@ -191,7 +191,7 @@ def run_resumable_extraction_agent( request_context: RequestContext, client: LiteLLMClient, extractor_kind: str, - user_id: str | None, + user_id: str, request_id: str, agent_version: str | None, source: str | None, diff --git a/reflexio/server/services/extraction/resume_scheduler.py b/reflexio/server/services/extraction/resume_scheduler.py index 8ff7a6d51..28fb2fcda 100644 --- a/reflexio/server/services/extraction/resume_scheduler.py +++ b/reflexio/server/services/extraction/resume_scheduler.py @@ -51,8 +51,8 @@ def _on_started(self) -> None: def _on_stopped(self) -> None: logger.info("event=extraction_resume_scheduler_stopped") - def _discover_org_ids(self, bootstrap_ctx: RequestContext) -> list[str]: - """Return every org with actionable work, always including the bootstrap org.""" + def _discover_local_org_ids(self, bootstrap_ctx: RequestContext) -> list[str]: + """Return local actionable orgs plus the bootstrap org.""" org_ids: list[str] = [] storage = getattr(bootstrap_ctx, "storage", None) if storage is not None: @@ -60,18 +60,37 @@ def _discover_org_ids(self, bootstrap_ctx: RequestContext) -> list[str]: org_ids = storage.list_resumable_work_org_ids(now=datetime.now(UTC)) except NotImplementedError: org_ids = [] - # Always sweep the bootstrap org so the maintenance loop runs even when - # the cross-org discovery query is empty or unsupported. if bootstrap_ctx.org_id not in org_ids: org_ids = [bootstrap_ctx.org_id, *org_ids] - return org_ids + return list(dict.fromkeys(org_ids)) - def _expire_pending_tool_calls(self, bootstrap_ctx: RequestContext) -> None: - storage = getattr(bootstrap_ctx, "storage", None) + def _discover_provider_org_ids(self) -> list[str] | None: + """Return the provider's authoritative list, or ``None`` on failure.""" + if self.org_id_provider is None: + return None + try: + return list( + dict.fromkeys( + org_id + for org_id in self.org_id_provider() + if org_id != DEFAULT_ORG_ID + ) + ) + except Exception as exc: + with error_tags( + subsystem="extraction", + op="scheduler_org_discovery", + error_type=type(exc).__name__, + ): + logger.exception("event=extraction_resume_scheduler_provider_failed") + return None + + def _expire_pending_tool_calls(self, ctx: RequestContext) -> None: + storage = getattr(ctx, "storage", None) if storage is None: return - # ``expire_pending_tool_calls`` is not org-scoped, so a single call - # sweeps every tenant's overdue pending rows for this tick. + # One storage ref can contain several tenants, but another ref is an + # independent queue. Sweep every discovered ref before draining it. try: expired = storage.expire_pending_tool_calls(now=datetime.now(UTC)) except NotImplementedError: @@ -84,6 +103,7 @@ def _drain_org(self, org_id: str) -> None: ctx = self.request_context_factory(org_id) if not pending_tool_calls_enabled(ctx): return + self._expire_pending_tool_calls(ctx) resumed = ExtractionResumeWorker(request_context=ctx).drain( max_runs=self.max_runs_per_tick ) @@ -108,23 +128,26 @@ def _drain_org(self, org_id: str) -> None: def _run_once(self) -> float: poll_interval = _DEFAULT_POLL_INTERVAL_SECONDS try: - if ( - self.bootstrap_org_id == DEFAULT_ORG_ID - and self.org_id_provider is not None - ): - org_ids = [ - org_id - for org_id in self.org_id_provider() - if org_id != DEFAULT_ORG_ID - ] - if not org_ids: + provider_org_ids = self._discover_provider_org_ids() + if self.org_id_provider is not None and provider_org_ids is not None: + if not provider_org_ids: return poll_interval - self.bootstrap_org_id = org_ids[0] + # Resolve config through an org that the provider proved is + # actionable on this tick. The previous bootstrap may have + # been deleted or moved and must not gate future discovery. + self.bootstrap_org_id = provider_org_ids[0] bootstrap_ctx = self.request_context_factory(self.bootstrap_org_id) config = bootstrap_ctx.configurator.get_config() poll_interval = config.pending_tool_call_config.resume_poll_interval_seconds - self._expire_pending_tool_calls(bootstrap_ctx) - for org_id in self._discover_org_ids(bootstrap_ctx): + if provider_org_ids is not None: + org_ids = provider_org_ids + elif self.org_id_provider is not None: + # A raised provider cannot authoritatively replace the list; + # preserve the last known bootstrap as a one-org fallback. + org_ids = [bootstrap_ctx.org_id] + else: + org_ids = self._discover_local_org_ids(bootstrap_ctx) + for org_id in org_ids: if self._stop_event.is_set(): break self._drain_org(org_id) diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index c66b9daa4..ff1732027 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -6,6 +6,7 @@ import uuid from collections import defaultdict from collections.abc import Callable, Sequence +from dataclasses import replace from datetime import UTC, datetime, timedelta from typing import Any @@ -46,6 +47,9 @@ has_expert_content, uses_evidence_grounded_extraction, ) +from reflexio.server.services.playbook.review_window import ( + infer_playbook_review_user_id, +) from reflexio.server.services.playbook.service import ( PlaybookGenerationService, PlaybookGenerationServiceConfig, @@ -256,6 +260,19 @@ def drain(self, *, max_runs: int = 10) -> int: resumed += 1 return resumed + def _with_resolved_playbook_user_id(self, run: AgentRunRecord) -> AgentRunRecord: + """Return a playbook run with a proven owner for legacy nullable rows.""" + if run.binding.extractor_kind != "playbook": + return run + if run.binding.user_id and run.binding.user_id.strip(): + return run + user_id = infer_playbook_review_user_id( + storage=self.storage, + source_interaction_ids=run.binding.source_interaction_ids, + subject=f"Playbook extraction run {run.id}", + ) + return replace(run, binding=replace(run.binding, user_id=user_id)) + def run_once(self) -> AgentRunRecord | None: config = self.request_context.configurator.get_config() pending_config = config.pending_tool_call_config @@ -283,6 +300,7 @@ def run_once(self) -> AgentRunRecord | None: ) try: + run = self._with_resolved_playbook_user_id(run) resolved_calls = self._load_resolved_tool_calls(run) if not resolved_calls: raise ResumeWorkerError( @@ -360,6 +378,7 @@ def _retry_finalization(self, run: AgentRunRecord) -> AgentRunRecord | None: config = self.request_context.configurator.get_config() pending_config = config.pending_tool_call_config try: + run = self._with_resolved_playbook_user_id(run) items, pending_tool_call_ids, model_provenance = ( self._items_from_committed_output(run) ) @@ -557,12 +576,15 @@ def _resume_playbook( ) -> tuple[list[Any], list[str], ModelProvenance | None]: if not isinstance(extractor_config, PlaybookConfig): raise ResumeWorkerError("Expected playbook extractor config") + user_id = run.binding.user_id + if not user_id: + raise ResumeWorkerError("Playbook resume requires user_id") agent_context = self.request_context.configurator.get_agent_context() service_config = PlaybookGenerationServiceConfig( request_id=run.binding.request_id, agent_version=run.binding.agent_version or "", - user_id=run.binding.user_id, + user_id=user_id, source=run.binding.source, auto_run=False, force_extraction=True, @@ -802,6 +824,9 @@ def _playbook_items_from_output( ) -> tuple[list[Any], list[str]]: if not isinstance(extractor_config, PlaybookConfig): raise ResumeWorkerError("Expected playbook extractor config") + user_id = run.binding.user_id + if not user_id: + raise ResumeWorkerError("Playbook finalization retry requires user_id") expert_mode = has_expert_content( extract_interactions_from_request_interaction_data_models( request_interaction_data_models @@ -840,7 +865,7 @@ def _playbook_items_from_output( service_config = PlaybookGenerationServiceConfig( request_id=run.binding.request_id, agent_version=run.binding.agent_version or "", - user_id=run.binding.user_id, + user_id=user_id, source=run.binding.source, auto_run=False, force_extraction=True, @@ -892,6 +917,9 @@ def _finalize_items( ) return if run.binding.extractor_kind == "playbook": + user_id = run.binding.user_id + if not user_id: + raise ResumeWorkerError("Playbook finalization requires user_id") service = PlaybookGenerationService( llm_client=self.client, request_context=self.request_context, @@ -899,13 +927,15 @@ def _finalize_items( service.service_config = PlaybookGenerationServiceConfig( request_id=run.binding.request_id, agent_version=run.binding.agent_version or "", - user_id=run.binding.user_id, + user_id=user_id, source=run.binding.source, auto_run=False, force_extraction=True, ) persisted_items = service._finalize_extracted_items( - items, model_provenance=model_provenance + items, + model_provenance=model_provenance, + extraction_run=run, ) self._record_finalized_learnings( run, persisted_items or [], entity_type="user_playbook" diff --git a/reflexio/server/services/playbook/README.md b/reflexio/server/services/playbook/README.md index cd41b6f80..b3896212a 100644 --- a/reflexio/server/services/playbook/README.md +++ b/reflexio/server/services/playbook/README.md @@ -20,6 +20,7 @@ Description: Evidence-grounded playbook extraction, candidate review, aggregatio | `playbook_service_constants.py` | Prompt IDs for all playbook operations | | `playbook_service_utils.py` | Request dataclasses, Pydantic output schemas, message construction utilities | | `playbook_evidence.py` | Strict evidence validation, call-local reference checks, and persisted provenance helpers | +| `review_window.py` | Shared fail-closed reconstruction of persisted review chronology for automatic and manual review | | `review_service.py` | Time-window selection, persisted evidence reconstruction, reporting, and newest-first per-playbook apply | | `aggregation_trigger.py` | Converts post-generation activity into an idempotent durable scheduling signal | | `aggregation_scheduler.py` | Polling, fleet claim/lease handling, retries, and structured aggregation progress telemetry | @@ -66,7 +67,10 @@ Every candidate must be accounted for exactly once as `accept`, `revise`, or `reject`. Revisions may narrow unsupported wording but cannot add evidence or create a lesson that extraction missed. Reviewer output receives one bounded repair attempt and otherwise fails closed. Expert and legacy extraction paths -do not use this reviewer. +do not use this reviewer. New user-playbook generation requires a non-empty +`user_id`. A legacy durable run whose stored binding is null may resume only +when every persisted source interaction and request proves the same owner; +missing or mixed-owner evidence fails closed before review or persistence. ### Persisted Review (`review_service.py`) @@ -77,9 +81,11 @@ generation window or cited evidence can no longer be reconstructed yields a Manual review reconstructs context from the full interaction window persisted on the row's finalized playbook-extraction run, plus any extra cited interactions -retained through consolidation. It never substitutes the current extractor -window or silently falls back to the smaller cited-evidence subset. Automatic -post-generation review continues to use the generation call's configured window. +retained through consolidation. Automatic review uses the exact +`source_interaction_ids` on the extraction agent run, including when finalization +is retried later. Neither path re-runs the sliding last-K query or silently falls +back to a smaller evidence subset, so interactions published after extraction +cannot change the review chronology. Only the playbook row's cited interaction IDs become candidate evidence units; the rest of the generation window is ancillary chronology. Evidence spans are rebuilt from those exact stored interactions instead of requiring every cited diff --git a/reflexio/server/services/playbook/playbook_service_utils.py b/reflexio/server/services/playbook/playbook_service_utils.py index 90b9b7337..3fe8c4e44 100644 --- a/reflexio/server/services/playbook/playbook_service_utils.py +++ b/reflexio/server/services/playbook/playbook_service_utils.py @@ -9,6 +9,7 @@ from reflexio.models.api_schema.domain.entities import Interaction, UserPlaybook from reflexio.models.api_schema.internal_schema import RequestInteractionDataModel +from reflexio.models.api_schema.validators import NonEmptyStr from reflexio.models.structured_output import ( StrictStructuredOutput, normalize_provider_keys, @@ -807,7 +808,7 @@ def ensure_playbook_content( class PlaybookGenerationRequest(BaseModel): request_id: str agent_version: str - user_id: str | None = None # for per-user playbook extraction + user_id: NonEmptyStr source: str | None = None rerun_start_time: int | None = None # Unix timestamp for rerun flows rerun_end_time: int | None = None # Unix timestamp for rerun flows diff --git a/reflexio/server/services/playbook/review_service.py b/reflexio/server/services/playbook/review_service.py index c20009d22..f7ccbc54e 100644 --- a/reflexio/server/services/playbook/review_service.py +++ b/reflexio/server/services/playbook/review_service.py @@ -10,7 +10,6 @@ from typing import Literal from reflexio.models.api_schema.domain.entities import ( - Interaction, LineageContext, ReviewUserPlaybookEdit, ReviewUserPlaybookResult, @@ -33,6 +32,10 @@ from reflexio.server.services.playbook.playbook_service_utils import ( is_evidence_validated, ) +from reflexio.server.services.playbook.review_window import ( + PlaybookReviewWindowError, + reconstruct_playbook_review_window, +) logger = logging.getLogger(__name__) @@ -53,10 +56,6 @@ class _PlaybookReviewRaceError(RuntimeError): """A reviewed incumbent stopped being current before its apply.""" -class _PlaybookNotReviewableError(ValueError): - """A selected row's generation window or cited evidence is unavailable.""" - - @dataclass class _ReviewedPlaybook: original: UserPlaybook @@ -84,11 +83,11 @@ def _review_window( playbook: UserPlaybook, ) -> list[RequestInteractionDataModel]: if not playbook.user_id: - raise _PlaybookNotReviewableError( + raise PlaybookReviewWindowError( f"User playbook {playbook.user_playbook_id} has no owning user_id" ) if not is_evidence_validated(playbook): - raise _PlaybookNotReviewableError( + raise PlaybookReviewWindowError( f"User playbook {playbook.user_playbook_id} lacks validated evidence" ) @@ -99,7 +98,7 @@ def _review_window( request_id=playbook.request_id, ) if generation_run is None or not generation_run.binding.source_interaction_ids: - raise _PlaybookNotReviewableError( + raise PlaybookReviewWindowError( f"User playbook {playbook.user_playbook_id} has no complete " "generation-window provenance" ) @@ -116,63 +115,12 @@ def _review_window( ] ) ) - interactions_by_id = { - interaction.interaction_id: interaction - for interaction in self.storage.get_interactions_by_ids(source_ids) - } - missing_interaction_ids = [ - interaction_id - for interaction_id in source_ids - if interaction_id not in interactions_by_id - ] - if missing_interaction_ids: - raise _PlaybookNotReviewableError( - f"User playbook {playbook.user_playbook_id} is missing persisted " - f"generation-window interactions: {missing_interaction_ids}" - ) - - interactions_by_request: dict[str, list[Interaction]] = {} - for interaction in interactions_by_id.values(): - if interaction.user_id != playbook.user_id: - raise _PlaybookNotReviewableError( - f"User playbook {playbook.user_playbook_id} has source interaction " - f"{interaction.interaction_id} owned by another user" - ) - interactions_by_request.setdefault(interaction.request_id, []).append( - interaction - ) - - review_window: list[RequestInteractionDataModel] = [] - for request_id, interactions in interactions_by_request.items(): - request = self.storage.get_request(request_id) - if request is None: - raise _PlaybookNotReviewableError( - f"User playbook {playbook.user_playbook_id} is missing persisted " - f"source request {request_id}" - ) - if request.user_id != playbook.user_id: - raise _PlaybookNotReviewableError( - f"User playbook {playbook.user_playbook_id} has source request " - f"{request_id} owned by another user" - ) - review_window.append( - RequestInteractionDataModel( - session_id=request.session_id, - request=request, - interactions=sorted( - interactions, - key=lambda item: (item.created_at, item.interaction_id), - ), - ) - ) - - review_window.sort( - key=lambda group: ( - group.interactions[0].created_at, - group.interactions[0].interaction_id, - ) + return reconstruct_playbook_review_window( + storage=self.storage, + source_interaction_ids=source_ids, + user_id=playbook.user_id, + subject=f"User playbook {playbook.user_playbook_id}", ) - return review_window def _existing_playbooks( self, @@ -219,7 +167,7 @@ def _public_result( @staticmethod def _skipped_result( playbook: UserPlaybook, - exc: _PlaybookNotReviewableError, + exc: PlaybookReviewWindowError, ) -> _ReviewedPlaybook: # The reason text is this service's own fail-closed message about the # row's provenance, never a raw storage/model error, so it is safe to @@ -268,7 +216,7 @@ def _review_each( candidates = [playbook] try: window = self._review_window(playbook) - except _PlaybookNotReviewableError as exc: + except PlaybookReviewWindowError as exc: # Selection is "newest current rows in a window", which cannot # know whether a row is still reviewable. One unreviewable row # must not cost the rest of the batch. @@ -295,7 +243,7 @@ def _review_each( tool_context=tool_context, ) except PlaybookCandidateEvidenceError as exc: - unavailable = _PlaybookNotReviewableError(str(exc)) + unavailable = PlaybookReviewWindowError(str(exc)) logger.info( "event=playbook_review_skipped user_playbook_id=%d reason=%s", playbook.user_playbook_id, diff --git a/reflexio/server/services/playbook/review_window.py b/reflexio/server/services/playbook/review_window.py new file mode 100644 index 000000000..909632a26 --- /dev/null +++ b/reflexio/server/services/playbook/review_window.py @@ -0,0 +1,166 @@ +"""Reconstruct playbook-review chronology from persisted interaction provenance.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from reflexio.models.api_schema.domain.entities import Interaction +from reflexio.models.api_schema.internal_schema import RequestInteractionDataModel +from reflexio.server.services.playbook.components.reviewer import ( + PlaybookCandidateEvidenceError, +) +from reflexio.server.services.storage.storage_base import BaseStorage + + +class PlaybookReviewWindowError(PlaybookCandidateEvidenceError): + """Persisted review provenance is incomplete or crosses a user boundary.""" + + +def infer_playbook_review_user_id( + *, + storage: BaseStorage, + source_interaction_ids: Sequence[int], + subject: str, +) -> str: + """Infer one legacy run owner from complete persisted source evidence. + + New playbook runs always persist ``user_id``. This compatibility path is + only for older nullable rows and remains fail-closed unless every cited + interaction and request exists and names the same non-empty owner. + """ + source_ids = list(dict.fromkeys(source_interaction_ids)) + if not source_ids: + raise PlaybookReviewWindowError( + f"{subject} has no complete generation-window provenance" + ) + + interactions_by_id = { + interaction.interaction_id: interaction + for interaction in storage.get_interactions_by_ids(source_ids) + } + missing_interaction_ids = [ + interaction_id + for interaction_id in source_ids + if interaction_id not in interactions_by_id + ] + if missing_interaction_ids: + raise PlaybookReviewWindowError( + f"{subject} is missing persisted generation-window interactions: " + f"{missing_interaction_ids}" + ) + + interaction_owners = { + interaction.user_id.strip() + for interaction in interactions_by_id.values() + if interaction.user_id and interaction.user_id.strip() + } + if len(interaction_owners) != 1 or any( + not interaction.user_id or not interaction.user_id.strip() + for interaction in interactions_by_id.values() + ): + raise PlaybookReviewWindowError( + f"{subject} does not have one unambiguous interaction owner" + ) + user_id = next(iter(interaction_owners)) + + for request_id in { + interaction.request_id for interaction in interactions_by_id.values() + }: + request = storage.get_request(request_id) + if request is None: + raise PlaybookReviewWindowError( + f"{subject} is missing persisted source request {request_id}" + ) + if not request.user_id or request.user_id.strip() != user_id: + raise PlaybookReviewWindowError( + f"{subject} has source request {request_id} owned by another user" + ) + + return user_id + + +def reconstruct_playbook_review_window( + *, + storage: BaseStorage, + source_interaction_ids: Sequence[int], + user_id: str, + subject: str, +) -> list[RequestInteractionDataModel]: + """Load and validate one exact persisted interaction window. + + Args: + storage: Storage containing the persisted requests and interactions. + source_interaction_ids: Exact interaction IDs recorded by extraction, + optionally extended with candidate-cited evidence IDs. + user_id: Owner every interaction and request must match. + subject: Safe identifier used in fail-closed error messages. + + Returns: + Request groups ordered by their earliest interaction, with interactions + inside each request ordered chronologically. + + Raises: + PlaybookReviewWindowError: If any interaction/request is missing or is + owned by another user. + """ + source_ids = list(dict.fromkeys(source_interaction_ids)) + if not source_ids: + raise PlaybookReviewWindowError( + f"{subject} has no complete generation-window provenance" + ) + + interactions_by_id = { + interaction.interaction_id: interaction + for interaction in storage.get_interactions_by_ids(source_ids) + } + missing_interaction_ids = [ + interaction_id + for interaction_id in source_ids + if interaction_id not in interactions_by_id + ] + if missing_interaction_ids: + raise PlaybookReviewWindowError( + f"{subject} is missing persisted generation-window interactions: " + f"{missing_interaction_ids}" + ) + + interactions_by_request: dict[str, list[Interaction]] = {} + for interaction in interactions_by_id.values(): + if interaction.user_id != user_id: + raise PlaybookReviewWindowError( + f"{subject} has source interaction {interaction.interaction_id} " + "owned by another user" + ) + interactions_by_request.setdefault(interaction.request_id, []).append( + interaction + ) + + review_window: list[RequestInteractionDataModel] = [] + for request_id, interactions in interactions_by_request.items(): + request = storage.get_request(request_id) + if request is None: + raise PlaybookReviewWindowError( + f"{subject} is missing persisted source request {request_id}" + ) + if request.user_id != user_id: + raise PlaybookReviewWindowError( + f"{subject} has source request {request_id} owned by another user" + ) + review_window.append( + RequestInteractionDataModel( + session_id=request.session_id, + request=request, + interactions=sorted( + interactions, + key=lambda item: (item.created_at, item.interaction_id), + ), + ) + ) + + review_window.sort( + key=lambda group: ( + group.interactions[0].created_at, + group.interactions[0].interaction_id, + ) + ) + return review_window diff --git a/reflexio/server/services/playbook/service.py b/reflexio/server/services/playbook/service.py index f3ba0cc67..ce12ec491 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -10,7 +10,10 @@ from reflexio.server.api_endpoints.request_context import RequestContext from reflexio.server.llm.litellm_client import LiteLLMClient from reflexio.server.services.deferred_learning_plan import GenerationComputePlan - from reflexio.server.services.storage.storage_base import BaseStorage + from reflexio.server.services.storage.storage_base import ( + AgentRunRecord, + BaseStorage, + ) from reflexio.models.api_schema.common import sanitise_for_log from reflexio.models.api_schema.domain.entities import LineageContext @@ -46,6 +49,10 @@ is_evidence_validated, uses_evidence_grounded_extraction, ) +from reflexio.server.services.playbook.review_window import ( + PlaybookReviewWindowError, + reconstruct_playbook_review_window, +) from reflexio.server.services.service_utils import ( extract_interactions_from_request_interaction_data_models, format_sessions_to_history_string, @@ -98,7 +105,7 @@ class PlaybookGenerationServiceConfig: request_id: str agent_version: str - user_id: str | None = None + user_id: str source: str | None = None allow_manual_trigger: bool = False rerun_start_time: int | None = None @@ -106,6 +113,13 @@ class PlaybookGenerationServiceConfig: auto_run: bool = True force_extraction: bool = False + def __post_init__(self) -> None: + if not isinstance(self.user_id, str): + raise ValueError("Playbook generation requires a non-empty user_id") + self.user_id = self.user_id.strip() + if not self.user_id: + raise ValueError("Playbook generation requires a non-empty user_id") + def _consolidation_search_keys(playbooks: list[UserPlaybook]) -> set[str]: """Return the strings consolidation would use as its hybrid-search queries. @@ -155,6 +169,7 @@ def __init__( self.output_pending_status = output_pending_status self.skip_aggregation = skip_aggregation self._review_window_cache: list[RequestInteractionDataModel] | None = None + self._review_run: AgentRunRecord | None = None def _load_generation_service_config( self, request: PlaybookGenerationRequest @@ -171,6 +186,7 @@ def _load_generation_service_config( # One service instance can process multiple users in a batch. Never let # a review window loaded for the previous request cross that boundary. self._review_window_cache = None + self._review_run = None generation_request_id = request.request_id return PlaybookGenerationServiceConfig( request_id=generation_request_id, @@ -189,37 +205,52 @@ def _configured_playbook_config(self) -> PlaybookConfig | None: return getattr(root_config, "user_playbook_extractor_config", None) def _review_interaction_window( - self, playbook_config: PlaybookConfig + self, _candidates: list[UserPlaybook] ) -> list[RequestInteractionDataModel]: - """Load the interaction window the reviewer reasons over, once per run. - - The reviewer needs the same chronology extraction saw. Extraction runs in - a separate step that does not hand its window back, so this reloads it — - memoized because ``_resolve_write_plan`` may be reached more than once for - one generation. - - Args: - playbook_config (PlaybookConfig): Config whose window/source filters - define the interaction window. - - Returns: - list[RequestInteractionDataModel]: The reloaded interaction window. - - Raises: - RuntimeError: If the configured source filter excludes the reviewer - reload after extraction already produced candidates. - """ + """Reconstruct the exact persisted extraction window once per run.""" if self._review_window_cache is None: if self.service_config is None: - raise RuntimeError("service_config must be set before review") - loaded_window = self._create_extractor( - playbook_config, self.service_config - )._get_interactions() - if loaded_window is None: - raise RuntimeError( - "Normal playbook review source filter excluded its interaction window" + raise PlaybookReviewWindowError( + "service_config must be set before playbook review" + ) + if not self.service_config.user_id: + raise PlaybookReviewWindowError( + "Normal playbook review has no owning user_id" ) - self._review_window_cache = loaded_window + + run = self._review_run + if run is None: + if len(self._last_extraction_run_ids) != 1: + raise PlaybookReviewWindowError( + "Normal playbook review has no unambiguous extraction run" + ) + run = self.storage.get_agent_run( # type: ignore[reportOptionalMemberAccess] + self._last_extraction_run_ids[0] + ) + if run is None: + raise PlaybookReviewWindowError( + "Normal playbook review extraction run is missing" + ) + + binding = run.binding + if ( + binding.extractor_kind != "playbook" + or binding.org_id != self.request_context.org_id + or binding.user_id != self.service_config.user_id + or binding.request_id != self.service_config.request_id + ): + raise PlaybookReviewWindowError( + "Normal playbook review extraction run does not match its request" + ) + + self._review_window_cache = reconstruct_playbook_review_window( + storage=self.storage, # type: ignore[arg-type] + source_interaction_ids=binding.source_interaction_ids, + user_id=self.service_config.user_id, + subject=( + f"Playbook generation request {self.service_config.request_id}" + ), + ) return self._review_window_cache def _load_extractor_config(self) -> PlaybookConfig | None: @@ -406,10 +437,10 @@ def _resolve_write_plan( and self.service_config is not None and os.getenv("MOCK_LLM_RESPONSE", "").lower() != "true" ): - review_interactions = self._review_interaction_window(playbook_config) + review_interactions = self._review_interaction_window(all_playbooks) if not review_interactions: - raise RuntimeError( - "Normal playbook review could not reload its interaction window" + raise PlaybookReviewWindowError( + "Normal playbook review could not reconstruct its interaction window" ) flat_review_interactions = ( extract_interactions_from_request_interaction_data_models( @@ -637,6 +668,7 @@ def _finalize_extracted_items( all_playbooks: list[UserPlaybook], *, model_provenance: ModelProvenance | None = None, + extraction_run: AgentRunRecord | None = None, ) -> list[UserPlaybook]: """Permanent V3 wrapper: compute→persist→schedulers together (no fence). @@ -649,7 +681,16 @@ def _finalize_extracted_items( """ if model_provenance is not None: self._last_model_provenance = model_provenance - plan = self._resolve_write_plan([all_playbooks]) + previous_review_run = self._review_run + previous_review_window = self._review_window_cache + if extraction_run is not None: + self._review_run = extraction_run + self._review_window_cache = None + try: + plan = self._resolve_write_plan([all_playbooks]) + finally: + self._review_run = previous_review_run + self._review_window_cache = previous_review_window if plan is None: return [] with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] diff --git a/tests/eval/extraction/providers.py b/tests/eval/extraction/providers.py index df1a84045..c096b187f 100644 --- a/tests/eval/extraction/providers.py +++ b/tests/eval/extraction/providers.py @@ -129,6 +129,7 @@ def _minimal_playbook_service_config() -> PlaybookGenerationServiceConfig: return PlaybookGenerationServiceConfig( agent_version="1.0.0", request_id=_EVAL_REQUEST_ID, + user_id=_EVAL_USER_ID, source="api", ) diff --git a/tests/server/llm/test_litellm_client_unit.py b/tests/server/llm/test_litellm_client_unit.py index cdeae0431..37dc68cee 100644 --- a/tests/server/llm/test_litellm_client_unit.py +++ b/tests/server/llm/test_litellm_client_unit.py @@ -3665,6 +3665,39 @@ def _raise_connection_error(**_params): assert payload.model == "gpt-5-mini" assert payload.llm_provider == "openai" + def test_isolated_worker_error_retains_upstream_type(self, monkeypatch): + """The parent wrapper must retain enough type data for log severity.""" + client = LiteLLMClient(LiteLLMConfig(model="gpt-5-mini")) + result_queue = MagicMock() + result_queue.get.return_value = ( + "error", + _CompletionErrorSnapshot( + type_name="APIConnectionError", + message="connection failed", + model="gpt-5-mini", + llm_provider="openai", + ), + ) + process = MagicMock() + process.is_alive.return_value = False + process_context = MagicMock() + process_context.Queue.return_value = result_queue + process_context.Process.return_value = process + monkeypatch.setattr(multiprocessing, "get_context", lambda: process_context) + monkeypatch.setattr( + client, "_should_process_isolate_completion", lambda *_args: True + ) + params = { + "model": "gpt-5-mini", + "messages": self._messages(), + "timeout": 0.5, + } + + with pytest.raises(LiteLLMClientError) as exc_info: + client._completion_with_hard_timeout(params, hard_timeout=5.0) + + assert exc_info.value.upstream_error_type == "APIConnectionError" + def test_hard_timeout_not_retried_at_client_level(self, monkeypatch): """A hard timeout is NOT retried at the client level. Same-model retry of a hang is exactly what produced the 490s in PYTHON-FASTAPI-62; the @@ -3903,6 +3936,39 @@ def _complete(params, _hard_timeout): for record in caplog.records ) + def test_isolated_transport_failure_is_warning_when_fallback_serves( + self, monkeypatch, caplog + ): + """A subprocess wrapper must not turn a recoverable outage into ERROR.""" + client = LiteLLMClient( + LiteLLMConfig( + model="minimax/MiniMax-M3", + fallback_models=["zai/glm-5.2"], + ) + ) + + def _complete(params, _hard_timeout): + if params["model"] == "minimax/MiniMax-M3": + raise LiteLLMClientError( + "isolated worker failed", + upstream_error_type="APIConnectionError", + ) + return _make_completion_response("from-glm") + + monkeypatch.setattr(client, "_completion_with_hard_timeout", _complete) + with caplog.at_level(logging.INFO): + result = client.generate_chat_response(self._messages()) + + assert result == "from-glm" + failed = [ + record + for record in caplog.records + if "event=llm_request_end" in record.message + and "success=False" in record.message + ] + assert len(failed) == 1 + assert failed[0].levelno == logging.WARNING + def test_all_rungs_fail_raises_last_error(self, monkeypatch): client = LiteLLMClient( LiteLLMConfig(model="minimax/MiniMax-M3", fallback_models=["zai/glm-5.2"]) diff --git a/tests/server/services/extraction/test_resume_scheduler.py b/tests/server/services/extraction/test_resume_scheduler.py index 44d5e2ce2..7a04057a6 100644 --- a/tests/server/services/extraction/test_resume_scheduler.py +++ b/tests/server/services/extraction/test_resume_scheduler.py @@ -134,3 +134,107 @@ def drain(self, *, max_runs: int) -> int: assert storage.expire_pending_tool_calls.called # The bootstrap org plus every discovered org are all drained. assert {"org_1", "org_2", "org_3"} <= set(drained_orgs) + + +def test_resume_scheduler_stale_bootstrap_cannot_block_future_discovery(monkeypatch): + provider = MagicMock(side_effect=[["org_2"], ["org_3"]]) + drained: list[str] = [] + live_org_ids = {"org_2"} + factory_calls: list[str] = [] + + class FakeWorker: + def __init__(self, *, request_context): + self.request_context = request_context + + def drain(self, *, max_runs: int) -> int: + drained.append(self.request_context.org_id) + return 0 + + monkeypatch.setattr(resume_scheduler, "pending_tool_calls_enabled", lambda _: True) + monkeypatch.setattr(resume_scheduler, "ExtractionResumeWorker", FakeWorker) + + def factory(org_id: str): + factory_calls.append(org_id) + if org_id not in live_org_ids: + raise RuntimeError(f"stale org: {org_id}") + return _request_context( + org_id, + storage=SimpleNamespace( + expire_pending_tool_calls=MagicMock(return_value=0), + list_resumable_work_org_ids=MagicMock(return_value=["wrong-org"]), + ), + ) + + scheduler = resume_scheduler.ExtractionResumeScheduler( + request_context_factory=cast(Callable[[str], RequestContext], factory), + bootstrap_org_id="org_1", + org_id_provider=provider, + ) + + scheduler._run_once() + live_org_ids.clear() + live_org_ids.add("org_3") + scheduler._run_once() + + assert provider.call_count == 2 + assert drained == ["org_2", "org_3"] + assert factory_calls == ["org_2", "org_2", "org_3", "org_3"] + assert scheduler.bootstrap_org_id == "org_3" + + +def test_resume_scheduler_provider_failure_still_drains_bootstrap(monkeypatch): + drained: list[str] = [] + storage = SimpleNamespace(expire_pending_tool_calls=MagicMock(return_value=0)) + + class FakeWorker: + def __init__(self, *, request_context): + self.request_context = request_context + + def drain(self, *, max_runs: int) -> int: + drained.append(self.request_context.org_id) + return 0 + + monkeypatch.setattr(resume_scheduler, "pending_tool_calls_enabled", lambda _: True) + monkeypatch.setattr(resume_scheduler, "ExtractionResumeWorker", FakeWorker) + scheduler = resume_scheduler.ExtractionResumeScheduler( + request_context_factory=cast( + Callable[[str], RequestContext], + lambda org_id: _request_context(org_id, storage=storage), + ), + bootstrap_org_id="org_1", + org_id_provider=MagicMock(side_effect=RuntimeError("registry unavailable")), + ) + + scheduler._run_once() + + assert drained == ["org_1"] + + +def test_resume_scheduler_expires_pending_calls_on_each_discovered_ref(monkeypatch): + storages = { + org_id: SimpleNamespace(expire_pending_tool_calls=MagicMock(return_value=0)) + for org_id in ("org_1", "org_2") + } + + class FakeWorker: + def __init__(self, *, request_context): + self.request_context = request_context + + def drain(self, *, max_runs: int) -> int: + return 0 + + monkeypatch.setattr(resume_scheduler, "pending_tool_calls_enabled", lambda _: True) + monkeypatch.setattr(resume_scheduler, "ExtractionResumeWorker", FakeWorker) + scheduler = resume_scheduler.ExtractionResumeScheduler( + request_context_factory=cast( + Callable[[str], RequestContext], + lambda org_id: _request_context(org_id, storage=storages[org_id]), + ), + bootstrap_org_id="org_1", + org_id_provider=lambda: ["org_1", "org_2"], + ) + + scheduler._run_once() + + storages["org_1"].expire_pending_tool_calls.assert_called_once() + storages["org_2"].expire_pending_tool_calls.assert_called_once() diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index be2880686..773aa97af 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -167,6 +167,76 @@ def _seed_interactions(storage: SQLiteStorage) -> None: ) +def _legacy_ownerless_playbook_run(source_interaction_ids: list[int]) -> AgentRunRecord: + return AgentRunRecord( + id="legacy_playbook_run", + binding=AgentBinding( + org_id="org_1", + extractor_kind="playbook", + user_id=None, + request_id="request_1", + agent_version="v1", + source="api", + source_interaction_ids=source_interaction_ids, + ), + status=AgentRunStatus.FINALIZATION_FAILED, + generation_request_snapshot={}, + ) + + +def test_legacy_playbook_run_infers_owner_from_complete_source_evidence( + request_context, storage +): + _seed_interactions(storage) + worker = ExtractionResumeWorker(request_context=request_context) + + resolved = worker._with_resolved_playbook_user_id( + _legacy_ownerless_playbook_run([1, 2]) + ) + + assert resolved.binding.user_id == "user_1" + + +def test_legacy_playbook_run_fails_closed_when_source_evidence_is_missing( + request_context, storage +): + _seed_interactions(storage) + worker = ExtractionResumeWorker(request_context=request_context) + + with pytest.raises(ValueError, match="missing persisted.*interactions"): + worker._with_resolved_playbook_user_id(_legacy_ownerless_playbook_run([1, 999])) + + +def test_legacy_playbook_run_fails_closed_for_multiple_source_owners( + request_context, storage +): + _seed_interactions(storage) + storage.add_request( + Request( + request_id="request_2", + user_id="user_2", + created_at=1_002, + source="api", + agent_version="v1", + session_id="request_2", + ) + ) + storage._insert_interaction( + Interaction( + interaction_id=3, + user_id="user_2", + request_id="request_2", + created_at=1_002, + role="user", + content="Use a different tenant's evidence.", + ) + ) + worker = ExtractionResumeWorker(request_context=request_context) + + with pytest.raises(ValueError, match="unambiguous interaction owner"): + worker._with_resolved_playbook_user_id(_legacy_ownerless_playbook_run([1, 3])) + + def _seed_ready_run(storage: SQLiteStorage) -> None: storage.create_agent_run( AgentRunRecord( @@ -375,12 +445,16 @@ def test_resume_bills_only_items_that_survive_finalization( worker = ExtractionResumeWorker(request_context=request_context) with ( - patch(finalize_path, return_value=[survivor]), + patch(finalize_path, return_value=[survivor]) as finalize, patch.object(worker, "_record_finalized_learnings") as record, ): worker._finalize_items(run, [dropped, survivor]) record.assert_called_once_with(run, [survivor], entity_type=entity_type) + if extractor_kind == "playbook": + assert finalize.call_args.kwargs["extraction_run"] is run + else: + assert "extraction_run" not in finalize.call_args.kwargs def test_resume_worker_tagging_schedule_failure_is_best_effort( diff --git a/tests/server/services/playbook/test_extractor_polarity_integration.py b/tests/server/services/playbook/test_extractor_polarity_integration.py index aed2c3b73..ff62b50a4 100644 --- a/tests/server/services/playbook/test_extractor_polarity_integration.py +++ b/tests/server/services/playbook/test_extractor_polarity_integration.py @@ -116,6 +116,7 @@ def service_config(): return PlaybookGenerationServiceConfig( agent_version="1.0.0", request_id="test_request", + user_id="test_user", source="api", ) @@ -125,7 +126,7 @@ def neutral_request_interaction_models(): """A neutral window — user is satisfied, no pushback / failure evidence.""" request = Request( request_id="req_neutral", - user_id="user_neutral", + user_id="test_user", session_id="test_session", created_at=1000, source="api", @@ -133,7 +134,7 @@ def neutral_request_interaction_models(): interactions = [ Interaction( interaction_id=101, - user_id="user_neutral", + user_id="test_user", content="Can you summarize my last invoice?", request_id="req_neutral", created_at=1000, @@ -141,7 +142,7 @@ def neutral_request_interaction_models(): ), Interaction( interaction_id=102, - user_id="user_neutral", + user_id="test_user", content="Here is the summary of your last invoice: ...", request_id="req_neutral", created_at=1001, @@ -149,7 +150,7 @@ def neutral_request_interaction_models(): ), Interaction( interaction_id=103, - user_id="user_neutral", + user_id="test_user", content="Thanks, that's exactly what I needed.", request_id="req_neutral", created_at=1002, @@ -170,7 +171,7 @@ def failure_request_interaction_models(): """A failure-evidence window — user pushes back on the agent's response.""" request = Request( request_id="req_failure", - user_id="user_failure", + user_id="test_user", session_id="test_session", created_at=2000, source="api", @@ -178,7 +179,7 @@ def failure_request_interaction_models(): interactions = [ Interaction( interaction_id=201, - user_id="user_failure", + user_id="test_user", content="Please cancel my subscription.", request_id="req_failure", created_at=2000, @@ -186,7 +187,7 @@ def failure_request_interaction_models(): ), Interaction( interaction_id=202, - user_id="user_failure", + user_id="test_user", content="Can you confirm you want to cancel? Are you sure?", request_id="req_failure", created_at=2001, @@ -194,7 +195,7 @@ def failure_request_interaction_models(): ), Interaction( interaction_id=203, - user_id="user_failure", + user_id="test_user", content="I already said yes, stop asking me to confirm.", request_id="req_failure", created_at=2002, diff --git a/tests/server/services/playbook/test_playbook_extractor.py b/tests/server/services/playbook/test_playbook_extractor.py index f82c87068..494f8bd83 100644 --- a/tests/server/services/playbook/test_playbook_extractor.py +++ b/tests/server/services/playbook/test_playbook_extractor.py @@ -124,13 +124,14 @@ def service_config(): return PlaybookGenerationServiceConfig( agent_version="1.0.0", request_id="test_request", + user_id="user1", source="api", ) @pytest.fixture def sample_interactions(): - """Create sample interactions from multiple users for testing.""" + """Create sample interactions for the configured playbook owner.""" return [ Interaction( interaction_id=1, @@ -150,7 +151,7 @@ def sample_interactions(): ), Interaction( interaction_id=3, - user_id="user2", + user_id="user1", content="Could be faster", request_id="req2", created_at=1002, @@ -171,7 +172,7 @@ def sample_request_interaction_models(sample_interactions): ) request2 = Request( request_id="req2", - user_id="user2", + user_id="user1", session_id="test_session", created_at=1002, source="api", @@ -257,18 +258,18 @@ def test_different_playbook_names_have_different_keys( # =============================== -# Test: Get Interactions (Not User-Scoped) +# Test: Get Interactions # =============================== class TestGetInteractions: - """Tests for interaction collection logic (not user-scoped). + """Tests for user-scoped interaction collection logic. Note: Stride checking is handled upstream by BaseGenerationService._filter_configs_by_stride() before the extractor is created, so stride_size tests are at the service level. """ - def test_passes_none_user_id_to_storage( + def test_passes_user_id_to_storage( self, request_context, mock_llm_client, @@ -300,7 +301,7 @@ def test_passes_none_user_id_to_storage( call_kwargs = request_context.storage.get_last_k_interactions_grouped.call_args[ 1 ] - assert call_kwargs["user_id"] is None # service_config.user_id is None + assert call_kwargs["user_id"] == "user1" def test_returns_interactions( self, @@ -333,14 +334,14 @@ def test_returns_interactions( assert result is not None assert len(result) == 2 # Two sessions - def test_uses_window_size_with_none_user_id( + def test_uses_window_size_with_user_id( self, request_context, mock_llm_client, service_config, sample_request_interaction_models, ): - """Test that window size is used with user_id=None for all users.""" + """Test that window size and required user scope are passed together.""" config = PlaybookConfig( extractor_name="quality_playbook", extraction_definition_prompt="Evaluate agent quality", @@ -362,12 +363,12 @@ def test_uses_window_size_with_none_user_id( extractor._get_interactions() - # Verify get_last_k_interactions_grouped was called with user_id=None + # Verify get_last_k_interactions_grouped received the required user scope. request_context.storage.get_last_k_interactions_grouped.assert_called_once() call_kwargs = request_context.storage.get_last_k_interactions_grouped.call_args[ 1 ] - assert call_kwargs["user_id"] is None + assert call_kwargs["user_id"] == "user1" assert call_kwargs["k"] == 50 def test_none_sources_enabled_gets_all_sources( @@ -414,14 +415,14 @@ def test_none_sources_enabled_gets_all_sources( class TestUpdateOperationState: """Tests for operation state update logic.""" - def test_run_bookmark_advance_carries_all_users_interactions( + def test_run_bookmark_advance_carries_returned_interactions( self, request_context, mock_llm_client, service_config, sample_request_interaction_models, ): - """run()'s outcome.bookmark_advance carries interactions from all users. + """run()'s outcome.bookmark_advance carries the returned interactions. The extractor no longer self-advances the bookmark (F1); it defers the advance onto the ExtractionOutcome for persist to apply atomically with @@ -468,14 +469,14 @@ def test_run_bookmark_advance_carries_all_users_interactions( class TestRun: """Integration tests for the run() method.""" - def test_run_collects_interactions_from_all_users( + def test_run_queries_interactions_for_required_user( self, request_context, mock_llm_client, service_config, sample_request_interaction_models, ): - """Test that run() collects interactions from all users.""" + """Test that run() keeps storage discovery scoped to its user.""" config = PlaybookConfig( extractor_name="quality_playbook", extraction_definition_prompt="Evaluate agent quality", @@ -497,11 +498,11 @@ def test_run_collects_interactions_from_all_users( with patch.dict(os.environ, {"MOCK_LLM_RESPONSE": "true"}): extractor.run() - # Verify storage was queried with user_id=None + # Verify storage was queried with the required user ID. call_kwargs = request_context.storage.get_last_k_interactions_grouped.call_args[ 1 ] - assert call_kwargs["user_id"] is None + assert call_kwargs["user_id"] == "user1" def test_run_returns_user_playbook( self, @@ -703,7 +704,7 @@ def test_generates_playbook_and_finalizes_agent_run( patch.dict(os.environ, {"MOCK_LLM_RESPONSE": "false"}), ): playbooks = extractor.extract_playbook_entries( - sample_request_interaction_models + sample_request_interaction_models[:1] ) assert len(playbooks) == 1 @@ -718,9 +719,9 @@ def test_generates_playbook_and_finalizes_agent_run( assert run is not None assert run.status == AgentRunStatus.AGENT_COMPLETED assert run.binding.org_id == "test_org" - assert run.binding.user_id is None + assert run.binding.user_id == "user1" assert run.binding.extractor_kind == "playbook" - assert run.binding.source_interaction_ids == [1, 2, 3] + assert run.binding.source_interaction_ids == [1, 2] assert run.generation_request_snapshot["output_schema_name"] == ( "StructuredReferencedExtractedPlaybookList" ) diff --git a/tests/server/services/playbook/test_playbook_generation_service.py b/tests/server/services/playbook/test_playbook_generation_service.py index 20833ea4f..def0fb41c 100644 --- a/tests/server/services/playbook/test_playbook_generation_service.py +++ b/tests/server/services/playbook/test_playbook_generation_service.py @@ -30,6 +30,11 @@ PlaybookGenerationService, PlaybookGenerationServiceConfig, ) +from reflexio.server.services.storage.storage_base import ( + AgentBinding, + AgentRunRecord, + AgentRunStatus, +) def create_request_interaction_data_model( @@ -55,6 +60,29 @@ def _storage(service: PlaybookGenerationService) -> Any: return cast(Any, service.storage) +def _playbook_run( + *, + org_id: str = "0", + user_id: str | None = "test_user", + request_id: str = "test_request_id", + source_interaction_ids: list[int], +) -> AgentRunRecord: + return AgentRunRecord( + id="run_for_review", + binding=AgentBinding( + org_id=org_id, + extractor_kind="playbook", + user_id=user_id, + request_id=request_id, + agent_version="1.0", + source="test_source", + source_interaction_ids=source_interaction_ids, + ), + status=AgentRunStatus.AGENT_COMPLETED, + generation_request_snapshot={}, + ) + + def _aggregation_enabled_config() -> Config: return Config( storage_config=StorageConfigSQLite(), @@ -695,6 +723,7 @@ def test_resolve_write_plan_reviews_grounded_normal_candidates_before_consolidat reviewed = candidate.model_copy( update={"content": "Treat the supplied answer as binding."} ) + service._review_run = _playbook_run(source_interaction_ids=[44]) with ( patch( @@ -726,76 +755,74 @@ def test_resolve_write_plan_reviews_grounded_normal_candidates_before_consolidat assert consolidator.deduplicate.call_args.args[0] == [[reviewed]] -def test_review_interaction_window_raises_when_source_filter_excludes_window(): - service = PlaybookGenerationService( - llm_client=MagicMock(), request_context=MagicMock() - ) - service.service_config = PlaybookGenerationServiceConfig( - request_id="request", - agent_version="v1", - ) - playbook_config = PlaybookConfig( - extractor_name="test_playbook", - extraction_definition_prompt="Review grounded lessons", - ) - extractor = MagicMock() - extractor._get_interactions.return_value = None - - with ( - patch.object(service, "_create_extractor", return_value=extractor), - pytest.raises(RuntimeError, match="source filter excluded"), - ): - service._review_interaction_window(playbook_config) - +@pytest.mark.parametrize("user_id", [None, "", " "]) +def test_playbook_generation_request_rejects_missing_user_id(user_id): + with pytest.raises(ValueError, match="user_id|String must not be empty"): + PlaybookGenerationRequest.model_validate( + { + "request_id": "test_request_id", + "agent_version": "1.0", + "user_id": user_id, + } + ) -def test_review_interaction_window_returns_honest_empty_window(): - service = PlaybookGenerationService( - llm_client=MagicMock(), request_context=MagicMock() - ) - service.service_config = PlaybookGenerationServiceConfig( - request_id="request", - agent_version="v1", - ) - playbook_config = PlaybookConfig( - extractor_name="test_playbook", - extraction_definition_prompt="Review grounded lessons", - ) - extractor = MagicMock() - extractor._get_interactions.return_value = [] - with patch.object(service, "_create_extractor", return_value=extractor): - assert service._review_interaction_window(playbook_config) == [] +@pytest.mark.parametrize("user_id", [None, "", " "]) +def test_playbook_generation_service_config_rejects_missing_user_id(user_id): + with pytest.raises(ValueError, match="non-empty user_id"): + PlaybookGenerationServiceConfig( + request_id="test_request_id", + agent_version="1.0", + user_id=cast(Any, user_id), + ) -def test_automatic_review_reloads_the_configured_extraction_window(): - request_context = MagicMock() - service = PlaybookGenerationService( - llm_client=MagicMock(), request_context=request_context - ) - service.service_config = PlaybookGenerationServiceConfig( - request_id="request", - agent_version="v1", - user_id="user-1", - source="chat", - ) - playbook_config = PlaybookConfig( - extractor_name="test_playbook", - extraction_definition_prompt="Review grounded lessons", - request_sources_enabled=["chat"], - window_size_override=3, - ) - expected_window = [MagicMock(spec=RequestInteractionDataModel)] - extractor = MagicMock() - extractor._get_interactions.return_value = expected_window +def test_automatic_review_reconstructs_exact_persisted_extraction_window(): + with tempfile.TemporaryDirectory() as temp_dir: + service = PlaybookGenerationService( + llm_client=MagicMock(), + request_context=RequestContext(org_id="0", storage_base_dir=temp_dir), + ) + service.service_config = PlaybookGenerationServiceConfig( + request_id="test_request_id", + agent_version="1.0", + user_id="test_user", + source="test_source", + ) + request = Request( + request_id="source_request", + user_id="test_user", + source="test_source", + agent_version="1.0", + session_id="session_1", + ) + _storage(service).add_request(request) + for interaction_id in (44, 45): + _storage(service).add_user_interaction( + "test_user", + Interaction( + interaction_id=interaction_id, + user_id="test_user", + request_id="source_request", + content=f"interaction {interaction_id}", + role="user", + created_at=interaction_id, + ), + ) + service._review_run = _playbook_run(source_interaction_ids=[44]) + candidate = UserPlaybook( + agent_version="1.0", + request_id="test_request_id", + user_id="test_user", + content="Use interaction 44.", + source_interaction_ids=[44], + ) - with patch.object( - service, "_create_extractor", return_value=extractor - ) as create_extractor: - assert service._review_interaction_window(playbook_config) == expected_window + window = service._review_interaction_window([candidate]) - create_extractor.assert_called_once_with(playbook_config, service.service_config) - extractor._get_interactions.assert_called_once_with() - request_context.storage.get_interactions_by_ids.assert_not_called() + assert [ + item.interaction_id for group in window for item in group.interactions + ] == [44] def test_loading_a_new_generation_request_clears_the_review_window_cache(): @@ -814,6 +841,7 @@ def test_loading_a_new_generation_request_clears_the_review_window_cache(): assert config.request_id == "request-2" assert service._review_window_cache is None + assert service._review_run is None def test_resolve_write_plan_fails_closed_for_strict_candidate_without_evidence(): @@ -864,6 +892,7 @@ def test_resolve_write_plan_fails_closed_for_strict_candidate_without_evidence() source_span=None, reader_angle="correction", ) + service._review_run = _playbook_run(source_interaction_ids=[45]) with ( patch( diff --git a/tests/server/services/playbook/test_playbook_generation_service_integration.py b/tests/server/services/playbook/test_playbook_generation_service_integration.py index 31030d76b..ffd8e77f9 100644 --- a/tests/server/services/playbook/test_playbook_generation_service_integration.py +++ b/tests/server/services/playbook/test_playbook_generation_service_integration.py @@ -1,6 +1,7 @@ """Integration tests for PlaybookGenerationService.""" import contextlib +from dataclasses import replace from datetime import UTC, datetime from unittest.mock import MagicMock @@ -79,11 +80,32 @@ def mock_request_context(): context = MagicMock(spec=RequestContext) context.org_id = "test_org_123" context.storage = MagicMock() + agent_runs = {} + + def create_agent_run(record): + agent_runs[record.id] = record + return record + + def update_agent_run_status(run_id, status, **updates): + record = agent_runs[run_id] + expected_statuses = updates.pop("expected_statuses", None) + if expected_statuses and record.status not in expected_statuses: + return record + if updates.pop("increment_resume_attempts", False): + updates["resume_attempts"] = record.resume_attempts + 1 + if updates.pop("increment_finalization_attempts", False): + updates["finalization_attempts"] = record.finalization_attempts + 1 + updates = {key: value for key, value in updates.items() if value is not None} + updated = replace(record, status=status, **updates) + agent_runs[run_id] = updated + return updated + + context.storage.create_agent_run.side_effect = create_agent_run + context.storage.get_agent_run.side_effect = agent_runs.get # Mock get_operation_state to return None by default (no in-progress state) context.storage.get_operation_state.return_value = None - context.storage.update_agent_run_status.side_effect = ( - lambda _run_id, status, **_kwargs: MagicMock(status=status) - ) + context.storage.claim_due_playbook_aggregation.return_value = None + context.storage.update_agent_run_status.side_effect = update_agent_run_status # Mock try_acquire_in_progress_lock to return success context.storage.try_acquire_in_progress_lock.return_value = {"acquired": True} # Mock get_user_playbooks to return empty list (for existing playbooks check) @@ -250,6 +272,16 @@ def test_playbook_generation_with_storage( [request_interaction_data_model], test_interactions, ) + mock_request_context.storage.get_interactions_by_ids.side_effect = ( + lambda interaction_ids: [ + interaction + for interaction in test_interactions + if interaction.interaction_id in interaction_ids + ] + ) + mock_request_context.storage.get_request.return_value = ( + request_interaction_data_model.request + ) # Create playbook generation request with new API request = PlaybookGenerationRequest( diff --git a/tests/server/services/storage/test_agent_run_helpers.py b/tests/server/services/storage/test_agent_run_helpers.py index 7f6cbc7b7..74c909bdd 100644 --- a/tests/server/services/storage/test_agent_run_helpers.py +++ b/tests/server/services/storage/test_agent_run_helpers.py @@ -126,6 +126,42 @@ def test_build_extractor_agent_run_record_accepts_legacy_request_id_keyword(): assert run.generation_request_snapshot["request_id"] == "legacy_req" +@pytest.mark.parametrize("user_id", ["", " "]) +def test_build_extractor_agent_run_record_requires_non_empty_user_id(user_id): + with pytest.raises(ValueError, match="non-empty user_id"): + build_extractor_agent_run_record( + org_id="org_1", + extractor_kind="playbook", + user_id=user_id, + generation_request_id="request_1", + agent_version="v1", + source="api", + request_interaction_data_models=_request_interaction_data_models(), + extractor_config=_ExtractorConfig(), + service_config={"request_id": "request_1"}, + agent_context="context", + ) + + +def test_build_extractor_agent_run_record_rejects_cross_user_source_evidence(): + request_models = _request_interaction_data_models() + request_models[0].interactions[0].user_id = "user_2" + + with pytest.raises(ValueError, match="source evidence must belong"): + build_extractor_agent_run_record( + org_id="org_1", + extractor_kind="playbook", + user_id="user_1", + generation_request_id="request_1", + agent_version="v1", + source="api", + request_interaction_data_models=request_models, + extractor_config=_ExtractorConfig(), + service_config={"request_id": "request_1"}, + agent_context="context", + ) + + def test_build_extractor_agent_run_record_rejects_mismatched_request_id_alias(): request_interaction_data_models = _request_interaction_data_models() From 7a5c438f31ce05449f969a6221fc7cd93a6a97c4 Mon Sep 17 00:00:00 2001 From: Yi Lu Date: Mon, 10 Aug 2026 10:26:36 -0700 Subject: [PATCH 2/2] test: keep review-window regex lint-clean --- tests/server/services/extraction/test_resume_worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 773aa97af..f26a88d00 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -203,7 +203,7 @@ def test_legacy_playbook_run_fails_closed_when_source_evidence_is_missing( _seed_interactions(storage) worker = ExtractionResumeWorker(request_context=request_context) - with pytest.raises(ValueError, match="missing persisted.*interactions"): + with pytest.raises(ValueError, match=r"missing persisted.*interactions"): worker._with_resolved_playbook_user_id(_legacy_ownerless_playbook_run([1, 999]))