From e0b0fbc67090d1ad317c7f6ad8bf2f656befe6b3 Mon Sep 17 00:00:00 2001 From: guangyu-reflexio Date: Sun, 2 Aug 2026 22:23:07 -0800 Subject: [PATCH 01/63] feat: add open-world evidence foundation (#407) - Add the shared evidence foundation for a future offline tuner without enabling a production tuner or scheduler. - Make session outcomes immutable, conflict-aware, reconstructable records with server-owned contract and trajectory identities. - Record synchronous search-exposure evidence and retain it under governance-safe lifecycle rules. - Keep billing limited to durable learnings created by online/resumable extraction; derived aggregation and tuning remain unmetered. - Add durable session-outcome identity fields, canonical digest helpers, exact-retry behavior, and conflicting-finalization rejection. - Add synchronous search-exposure recording with deterministic identities and incomplete-evidence classification. - Extend retention contracts for session outcomes and exposure evidence. - Bind erasure to authoritative tenant-scoped users. - Add purge execution claims, lease renewal, fencing, stale-claim recovery, and deterministic retry behavior. - Preserve missing-schema tolerance and stable deletion receipts. - Emit learning usage only for durable extraction-created rows. - Remove derived aggregation billing and make resumable retry keys deterministic. - Document immutable outcomes, extraction-only billing, and the intentionally unavailable tuner capability. - Add focused SQLite, API, governance, concurrency, migration, and compatibility coverage. ```mermaid flowchart LR A[Online extraction] --> B[Durable profiles and playbooks] A --> C[Session outcomes] D[Search] --> E[Exposure evidence] C --> F[Governance and retention] E --> F B --> G[Future offline analysis] C --> G E --> G G -. disabled in Phase 1 .-> H[Candidate publication] ``` - Phase matrix: `146 passed, 9 skipped` across SQLite, Supabase/PostgREST, native PostgreSQL, governance, exposure, and migration cases. - Migration/schema/placement gate: `64 passed`. - Shared governance and storage focused suite: `58 passed`. - Tuner availability/config gate: `7 passed`. - Changed Python scope: Ruff clean, format clean, Pyright `0 errors, 0 warnings`. - Changed SQL migration lint: `7 files, 0 findings`. - Import and range diff checks passed. * **New Features** * Session outcomes support `unknown` results with stable identity and integrity details. * Exact retries remain idempotent; conflicting finalizations are rejected. * Search results can be recorded as user-playbook exposures. * User-data erasure supports safer retries and recovery during concurrent processing. * Finalized learning records support reliable retry handling and accurate billing. * **Bug Fixes** * Session outcomes are removed using the authoritative user identity. * Billing records avoid duplicate or unattributable finalized-learning entries. * **Migration** * Existing session-outcome data is preserved and upgraded automatically. * **Documentation** * Local SQLite storage now requires version 3.35.0 or newer. --- README.md | 7 + docs/lib/methods/requests-sessions.ts | 10 +- reflexio/client/client.py | 15 +- reflexio/lib/_session_outcome.py | 81 +- reflexio/models/api_schema/domain/entities.py | 62 +- reflexio/models/api_schema/domain/enums.py | 2 + reflexio/server/billing_meter.py | 47 +- reflexio/server/routes/search.py | 14 + .../base_generation/_usage_billing.py | 23 +- .../services/base_generation_service.py | 33 +- .../server/services/deferred_learning_plan.py | 17 +- reflexio/server/services/extraction/README.md | 5 +- .../services/extraction/resume_worker.py | 65 +- .../server/services/governance/service.py | 332 +++- .../playbook/components/aggregator.py | 68 +- reflexio/server/services/playbook/service.py | 72 +- reflexio/server/services/profile/service.py | 45 +- reflexio/server/services/search_exposure.py | 155 ++ .../services/storage/governance_claims.py | 30 + .../services/storage/governance_validation.py | 8 + reflexio/server/services/storage/retention.py | 12 + .../storage/session_outcome_identity.py | 188 +++ .../services/storage/sqlite_storage/_base.py | 342 ++-- .../storage/sqlite_storage/_governance.py | 27 + .../sqlite_storage/_session_outcomes.py | 179 ++- .../agent_run/_agent_run_store.py | 87 + .../governance/_erase_execution.py | 35 +- .../sqlite_storage/governance/_purge.py | 361 ++++- .../governance/_subject_barrier.py | 106 +- .../services/storage/storage_base/__init__.py | 2 +- .../storage/storage_base/_session_outcomes.py | 18 + .../agent_run/_agent_run_store.py | 19 + .../governance/_erase_execution.py | 14 +- .../storage/storage_base/governance/_purge.py | 42 +- .../governance/_subject_barrier.py | 16 +- tests/client/test_session_outcomes_client.py | 23 + tests/models/test_session_outcome_identity.py | 336 ++++ .../test_session_outcomes_integration.py | 172 +- .../routes/test_search_exposure_boundary.py | 153 ++ .../services/extraction/test_resume_worker.py | 419 ++++- .../governance/test_governance_local_e2e.py | 703 ++++++++- .../test_subject_write_barrier_sqlite.py | 358 ++++- .../sqlite_storage/test_agent_run_storage.py | 206 +++ .../test_governance_retrieved_learning.py | 21 +- .../sqlite_storage/test_governance_storage.py | 1398 ++++++++++++++++- .../test_session_id_migration.py | 56 +- .../test_session_outcome_migration.py | 310 ++++ .../services/storage/test_sqlite_storage.py | 37 + .../test_storage_contract_clear_user_data.py | 65 + .../test_storage_contract_session_outcomes.py | 338 ++++ .../test_generation_billing_emission.py | 149 +- .../test_non_extraction_learning_metering.py | 205 ++- tests/server/services/test_search_exposure.py | 244 +++ tests/server/test_billing_meter_events.py | 11 +- 54 files changed, 7061 insertions(+), 682 deletions(-) create mode 100644 reflexio/server/services/search_exposure.py create mode 100644 reflexio/server/services/storage/governance_claims.py create mode 100644 reflexio/server/services/storage/session_outcome_identity.py create mode 100644 tests/models/test_session_outcome_identity.py create mode 100644 tests/server/routes/test_search_exposure_boundary.py create mode 100644 tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py create mode 100644 tests/server/services/test_search_exposure.py diff --git a/README.md b/README.md index 035ba2b1d..5289a8459 100644 --- a/README.md +++ b/README.md @@ -105,8 +105,15 @@ Publish conversations from your agent, and Reflexio closes the self-improvement | --- | --- | | [Python](https://www.python.org/) >= 3.12 | Required for PyPI and source installs | | [uv](https://docs.astral.sh/uv/getting-started/installation/) | Required when running from source | +| Python's linked SQLite runtime | >= 3.35.0 (required for local SQLite storage; not the standalone `sqlite3` CLI) | | [Node.js](https://nodejs.org/) >= 18 | Required only for the local docs site in a source checkout | +Check the SQLite runtime linked to Python with: + +```shell +uv run python -c "import sqlite3; print(sqlite3.sqlite_version_info)" +``` +

Reflexio Demo

diff --git a/docs/lib/methods/requests-sessions.ts b/docs/lib/methods/requests-sessions.ts index 4efbe928e..9b7be37e7 100644 --- a/docs/lib/methods/requests-sessions.ts +++ b/docs/lib/methods/requests-sessions.ts @@ -7,7 +7,7 @@ export const requestSessionMethods: MethodDef[] = [ displayName: "Mark Session Outcome", group: "requests-sessions", description: - "Record the lifetime-durable first success or failure marker for an existing session. User and source are derived from the first request.", + "Record the lifetime-durable first terminal outcome for an existing session. Idempotency requires unchanged payload, governance context, and finalized trajectory/session context; changed values conflict. User and source are derived from the first request.", httpMethod: "POST", endpoint: "/api/session_outcome", requestStyle: "json_body", @@ -17,8 +17,8 @@ export const requestSessionMethods: MethodDef[] = [ name: "outcome", type: "enum", required: true, - description: "Terminal outcome", - enumValues: ["success", "failure"], + description: "Terminal outcome: success, failure, or unknown", + enumValues: ["success", "failure", "unknown"], }, { name: "occurred_at", type: "number", required: true, description: "Unix epoch seconds when the outcome occurred" }, { name: "label", type: "string", required: false, description: "Optional label, at most 128 characters" }, @@ -43,8 +43,8 @@ export const requestSessionMethods: MethodDef[] = [ name: "outcome", type: "enum", required: false, - description: "Exact outcome filter", - enumValues: ["success", "failure"], + description: "Exact outcome filter: success, failure, or unknown", + enumValues: ["success", "failure", "unknown"], }, { name: "label", type: "string", required: false, description: "Exact label filter" }, { name: "start_time", type: "number", required: false, description: "Inclusive minimum event time" }, diff --git a/reflexio/client/client.py b/reflexio/client/client.py index 71a2d8697..684c8e728 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -1210,13 +1210,20 @@ def mark_session_outcome( value: float | None = None, metadata: dict[str, Any] | None = None, ) -> SetSessionOutcomeResponse: - """Record the first terminal outcome for a published session. + """Record the immutable first outcome for a published session. The session must already contain at least one published request. Reflexio derives both ``user_id`` and ``source`` from the earliest request ordered - by ``(created_at, request_id)``. Only the first outcome is recorded; - retries return ``success=True`` and ``recorded=False``. Sessions are not - required to report an outcome. + by ``(created_at, request_id)``. New canonical rows bind the outcome to + the server-owned outcome contract and canonical finalized trajectory. An + exact canonical retry must match the payload, contract, and trajectory; + otherwise it is rejected with ``reason="conflicting_finalization"``. + Rolling-upgrade rows with all four identity fields null compare the + caller payload and any available server-derived session context, but + cannot compare absent contract or trajectory digests. An accepted retry + preserves all four null identity fields and returns ``success=True`` and + ``recorded=False``. Sessions may report ``success``, ``failure``, or + ``unknown`` and are not required to report an outcome. """ request = SetSessionOutcomeRequest( session_id=session_id, diff --git a/reflexio/lib/_session_outcome.py b/reflexio/lib/_session_outcome.py index d49cbcead..7a39ba518 100644 --- a/reflexio/lib/_session_outcome.py +++ b/reflexio/lib/_session_outcome.py @@ -53,51 +53,44 @@ def mark_session_outcome( try: for _attempt in range(3): context = storage.get_session_outcome_context(request.session_id) - if context.existing: - return SetSessionOutcomeResponse( - success=True, - recorded=False, - user_id=context.user_id, - source=context.source, - message="Outcome already exists", - ) - if context.user_id is None or context.first_request_at is None: - return SetSessionOutcomeResponse( - success=False, - reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, - message="Session has no published requests", - ) - if context.user_contract_violation: - logger.warning( - "Session outcome contract violation: multiple users for session %s", - sanitise_for_log(request.session_id), - ) - if context.source_contract_violation: - logger.warning( - "Session outcome contract violation: multiple sources for session %s", - sanitise_for_log(request.session_id), - ) - if request.occurred_at < context.first_request_at: - return SetSessionOutcomeResponse( - success=False, - reason=SessionOutcomeFailureReason.OCCURRED_BEFORE_SESSION, - message="Outcome occurred before the session began", - user_id=context.user_id, - source=context.source, - ) - provider = get_service(SESSION_OUTCOME_ACCEPTANCE) - if provider is not None: - reason = provider( - self.org_id, request, received_at, context.user_id - ) - if reason is not None: + if not context.existing: + if context.user_id is None or context.first_request_at is None: + return SetSessionOutcomeResponse( + success=False, + reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, + message="Session has no published requests", + ) + if context.user_contract_violation: + logger.warning( + "Session outcome contract violation: multiple users for session %s", + sanitise_for_log(request.session_id), + ) + if context.source_contract_violation: + logger.warning( + "Session outcome contract violation: multiple sources for session %s", + sanitise_for_log(request.session_id), + ) + if request.occurred_at < context.first_request_at: return SetSessionOutcomeResponse( success=False, - reason=reason, - message="Outcome was not accepted", + reason=SessionOutcomeFailureReason.OCCURRED_BEFORE_SESSION, + message="Outcome occurred before the session began", user_id=context.user_id, source=context.source, ) + provider = get_service(SESSION_OUTCOME_ACCEPTANCE) + if provider is not None: + reason = provider( + self.org_id, request, received_at, context.user_id + ) + if reason is not None: + return SetSessionOutcomeResponse( + success=False, + reason=reason, + message="Outcome was not accepted", + user_id=context.user_id, + source=context.source, + ) result = storage.record_session_outcome( request, created_at=received_at, @@ -112,6 +105,10 @@ def mark_session_outcome( message="Outcome was not recorded", user_id=result.user_id, source=result.source, + outcome_id=result.outcome_id, + outcome_revision=result.outcome_revision, + outcome_contract_digest=result.outcome_contract_digest, + finalized_trajectory_digest=result.finalized_trajectory_digest, ) return SetSessionOutcomeResponse( success=True, @@ -123,6 +120,10 @@ def mark_session_outcome( if result.recorded else "Outcome already exists" ), + outcome_id=result.outcome_id, + outcome_revision=result.outcome_revision, + outcome_contract_digest=result.outcome_contract_digest, + finalized_trajectory_digest=result.finalized_trajectory_digest, ) except Exception: logger.exception("Failed to record session outcome") diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 295f0ac2c..f4f280945 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -825,6 +825,8 @@ class DeleteSessionResponse(BaseModel): class SessionOutcomeRecord(BaseModel): + outcome_id: NonEmptyStr | None = None + outcome_revision: int | None = Field(default=None, ge=1) user_id: str session_id: NonEmptyStr outcome: SessionOutcomeKind @@ -833,8 +835,35 @@ class SessionOutcomeRecord(BaseModel): label: str | None = Field(default=None, max_length=128) value: float | None = Field(default=None, allow_inf_nan=False) metadata: dict[str, Any] | None = None + outcome_contract_digest: Sha256Digest | None = None + finalized_trajectory_digest: Sha256Digest | None = None created_at: int = Field(ge=0) + @field_validator("outcome_contract_digest", "finalized_trajectory_digest") + @classmethod + def validate_sha256_digest(cls, value: str | None) -> str | None: + if value is not None and ( + len(value) != 64 or any(char not in "0123456789abcdef" for char in value) + ): + raise ValueError("outcome identity digests must be lowercase SHA-256 hex") + return value + + @model_validator(mode="after") + def validate_identity_shape(self) -> Self: + identity = ( + self.outcome_id, + self.outcome_revision, + self.outcome_contract_digest, + self.finalized_trajectory_digest, + ) + if any(value is None for value in identity) and not all( + value is None for value in identity + ): + raise ValueError( + "outcome identity fields must be all populated or all null" + ) + return self + class SetSessionOutcomeRequest(CapturesUnknownFields): session_id: NonEmptyStr @@ -880,6 +909,35 @@ class SetSessionOutcomeResponse(BaseModel): message: str = "" user_id: str | None = None source: str | None = None + outcome_id: NonEmptyStr | None = None + outcome_revision: int | None = Field(default=None, ge=1) + outcome_contract_digest: Sha256Digest | None = None + finalized_trajectory_digest: Sha256Digest | None = None + + @field_validator("outcome_contract_digest", "finalized_trajectory_digest") + @classmethod + def validate_sha256_digest(cls, value: str | None) -> str | None: + if value is not None and ( + len(value) != 64 or any(char not in "0123456789abcdef" for char in value) + ): + raise ValueError("outcome identity digests must be lowercase SHA-256 hex") + return value + + @model_validator(mode="after") + def validate_identity_shape(self) -> Self: + identity = ( + self.outcome_id, + self.outcome_revision, + self.outcome_contract_digest, + self.finalized_trajectory_digest, + ) + if any(value is None for value in identity) and not all( + value is None for value in identity + ): + raise ValueError( + "outcome identity fields must be all populated or all null" + ) + return self class GetSessionOutcomesRequest(CapturesUnknownFields): @@ -956,8 +1014,8 @@ class DeleteUserPlaybooksByIdsRequest(BaseModel): user_playbook_ids: list[int] = Field(min_length=1, max_length=10_000) -# Clear all data scoped to a single user_id (interactions, requests, session -# outcomes, user playbooks, profiles). Used by paired-protocol harnesses (e.g. SWE-bench) to +# Clear all data scoped to a single user_id (session outcomes, interactions, +# requests, user playbooks, profiles). Used by paired-protocol harnesses (e.g. SWE-bench) to # isolate per-task data on a shared storage backend without nuking sibling # tasks' rows. Intentionally does NOT touch agent_playbooks — they are the # cross-project rollup of skills and have no user_id column. diff --git a/reflexio/models/api_schema/domain/enums.py b/reflexio/models/api_schema/domain/enums.py index 41c99be37..6ea0d9fa8 100644 --- a/reflexio/models/api_schema/domain/enums.py +++ b/reflexio/models/api_schema/domain/enums.py @@ -52,9 +52,11 @@ class SessionOutcomeKind(StrEnum): SUCCESS = "success" FAILURE = "failure" + UNKNOWN = "unknown" class SessionOutcomeFailureReason(StrEnum): + CONFLICTING_FINALIZATION = "conflicting_finalization" UNKNOWN_SESSION = "unknown_session" OCCURRED_BEFORE_SESSION = "occurred_before_session" OCCURRED_IN_FUTURE = "occurred_in_future" diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py index 4e0380d5e..9a3633328 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -91,11 +91,10 @@ def record_learnings_generated( ) -> None: """Emit the Learning value facet — number of profiles/playbooks generated. - Documented FALLBACK for callers that genuinely lack a per-record id list - (e.g. dedup/consolidation can reduce the persisted count below the raw - extracted count, so there is no safe 1:1 id per unit of ``count``). Prefer - :func:`record_learnings_generated_records` whenever the caller has the - durable learning ids in scope. No-op when ``count <= 0``. + Intended for online extraction paths that have a known billable count but + do not retain a complete per-record id list. Resumable finalization must + use :func:`record_learnings_generated_records` and skip items without + durable ids. No-op when ``count <= 0``. Emits a single event carrying ``event_key`` when the caller has a durable retry identity, otherwise synthesizes ``f"learn-batch:{uuid4()}"``. This @@ -181,7 +180,7 @@ def record_learnings_generated_records( Args: org_id: Organisation identifier. learning_ids: Ids of the learnings durably generated in this run - (e.g. ``profile_id`` / ``user_playbook_id`` / ``agent_playbook_id``). + (e.g. ``profile_id`` / ``user_playbook_id``). platform_llm: True iff the platform supplies the LLM for this org. platform_storage: True iff the platform supplies storage; None defers to rollup. pipeline: Optional pipeline tag (e.g. ``"playbook"``). @@ -235,20 +234,20 @@ def emit_learnings_generated( ) -> None: """Resolve ``platform_llm`` from config and emit the Learning value facet. - Convenience wrapper for non-extraction learning-mutation paths such as - resumable-extraction finalization, aggregation, and offline-tuner auto-apply. It - owns the ``configurator.get_config()`` + ``platform_llm_from_config`` lookup so - each call site stays a thin one-liner, and — critically — is **guarded**: the - product path must never fail because metering failed, so config resolution and - emission are wrapped and any exception is logged and swallowed (mirroring the - extraction path's ``_record_billing_learning_events``). No-op when - ``count <= 0``. + Convenience wrapper for count-based online extraction callers. It owns the + ``configurator.get_config()`` + ``platform_llm_from_config`` lookup so the + call site stays a thin one-liner, and — critically — is **guarded**: the + product path must never fail because metering failed, so config resolution + and emission are wrapped and any exception is logged and swallowed + (mirroring the extraction path's ``_record_billing_learning_events``). + Resumable finalization must use :func:`emit_learnings_generated_records`. + No-op when ``count <= 0``. Args: org_id: Organisation identifier. configurator: Object exposing ``get_config()`` for platform-LLM resolution. count: Number of learnings durably produced by this path. - source: Metering source/path label (e.g. ``"offline_optimizer"``). + source: Metering source/path label (e.g. ``"online_extraction"``). pipeline: Optional pipeline tag (e.g. ``"playbook"``). user_id: Optional user ID tied to the generated learning. request_id: Optional request correlation ID. @@ -305,22 +304,20 @@ def emit_learnings_generated_records( ) -> None: """Resolve ``platform_llm`` from config and emit one event per learning id. - Entity-backed counterpart to :func:`emit_learnings_generated`, currently - adopted by two of the non-extraction learning-mutation paths — - resumable-extraction finalization and aggregation — the callers with - durable per-record ids in scope. Extraction and offline-tuner auto-apply do - not have a safe 1:1 id per unit of count (see - :func:`record_learnings_generated_records`) and use the count-based - :func:`emit_learnings_generated` fallback instead. Same guard semantics: - config resolution and emission are wrapped and any exception is logged - and swallowed — the product path must never fail because metering + Entity-backed counterpart to :func:`emit_learnings_generated`, used by + resumable-extraction finalization for every created user learning with a + durable id. Items without ids are not billable on that path. Online + extraction uses the count-based :func:`record_learnings_generated` helper + because it does not retain a safe 1:1 id per generated unit. Same guard + semantics: config resolution and emission are wrapped and any exception is + logged and swallowed — the product path must never fail because metering failed. No-op when ``learning_ids`` is empty. Args: org_id: Organisation identifier. configurator: Object exposing ``get_config()`` for platform-LLM resolution. learning_ids: Ids of the learnings durably produced by this path. - source: Metering source/path label (e.g. ``"aggregation"``). + source: Metering source/path label (e.g. ``"resumable_extraction"``). pipeline: Optional pipeline tag (e.g. ``"playbook"``). user_id: Optional user ID tied to the generated learning. request_id: Optional request correlation ID. diff --git a/reflexio/server/routes/search.py b/reflexio/server/routes/search.py index cb83bf6b1..ad098cebe 100644 --- a/reflexio/server/routes/search.py +++ b/reflexio/server/routes/search.py @@ -49,6 +49,10 @@ active_retrieval_experiment_assignment, ) from reflexio.server.services.search_metering_worker import enqueue_search_metering +from reflexio.server.services.search_exposure import ( + SearchExposureBatch, + record_search_exposures, +) from reflexio.server.tracing import profile_step logger = logging.getLogger(__name__) @@ -412,6 +416,16 @@ def run_search() -> Any: rehydrated_text=response.rehydrated_text, experiment=assignment, ) + record_search_exposures( + SearchExposureBatch( + org_id=org_id, + request_id=payload.request_id, + session_id=payload.session_id, + interaction_id=payload.interaction_id, + user_id=payload.user_id, + user_playbooks=tuple(response.user_playbooks), + ) + ) enqueue_search_metering( org_id=org_id, caller_type=caller_type, diff --git a/reflexio/server/services/base_generation/_usage_billing.py b/reflexio/server/services/base_generation/_usage_billing.py index 1b1304f90..5aee8a7d9 100644 --- a/reflexio/server/services/base_generation/_usage_billing.py +++ b/reflexio/server/services/base_generation/_usage_billing.py @@ -188,14 +188,14 @@ def _record_billing_learning_events( (cost facet) via the OSS emission helpers. ``platform_storage`` is left ``None`` here and resolved enterprise-side at rollup (Phase 1). - Gated by ``EMITS_LEARNING_BILLING`` — only profile/playbook generation - services opt in here. Non-extraction learning mutation paths emit their - own ``learnings_generated`` value-facet events when they durably apply - revisions/successors. + Gated by ``EMITS_LEARNING_BILLING`` — only online profile/playbook + extraction services opt in here. Resumable-extraction finalization emits + the same value facet separately. Derived mutation paths emit no additional + ``learnings_generated`` events. Args: prepared: The prepared generation run (used for input-text computation). - generated_count: Number of learnings produced by this extraction run. + generated_count: Number of retained write-plan learnings eligible for billing. """ if not self.EMITS_LEARNING_BILLING: return @@ -260,3 +260,16 @@ def _count_generated_results(result: Any) -> int: if isinstance(result, list): return len(result) return 1 if result else 0 + + @staticmethod + def _count_retained_online_learnings(write_plan: Any) -> int: + from reflexio.server.services.deferred_learning_plan import ( + PlaybookWritePlan, + ProfileWritePlan, + ) + + if isinstance(write_plan, ProfileWritePlan): + return len(write_plan.new_profiles) + if isinstance(write_plan, PlaybookWritePlan): + return len(write_plan.new_playbooks) + return 0 diff --git a/reflexio/server/services/base_generation_service.py b/reflexio/server/services/base_generation_service.py index 2334d3741..9fd3f9fe6 100644 --- a/reflexio/server/services/base_generation_service.py +++ b/reflexio/server/services/base_generation_service.py @@ -210,9 +210,9 @@ class BaseGenerationService( ABC, Generic[TExtractorConfig, TExtractor, TGenerationServiceConfig, TRequest], # noqa: UP046 ): - # Only profile/playbook GENERATION services emit extraction-run billing here. - # Non-extraction learning mutation paths emit their value facet at their own - # durable-success point. + # Only online profile/playbook extraction services emit extraction-run billing + # here. Resumable-extraction finalization emits separately; derived mutation + # paths emit no additional learnings_generated events. # Default is False so any future subclass is safe by default (opt-IN). EMITS_LEARNING_BILLING: bool = False """ @@ -342,11 +342,20 @@ def _process_results(self, results: list) -> None: results: List of all results from extractors (one per successful extractor) """ - def _finalize_extracted_items(self, items: list) -> list: + def _finalize_extracted_items( + self, + items: list, + *, + finalization_run_id: str | None = None, + ) -> list[str] | None: """Persist already-flattened extracted items through the service path.""" + if finalization_run_id is not None: + raise NotImplementedError( + "Receipt-aware finalization must be implemented by resumable services" + ) if items: self._process_results([items]) - return items + return None @abstractmethod def _should_track_in_progress(self) -> bool: @@ -634,7 +643,6 @@ def compute_generation(self, request: TRequest) -> GenerationComputePlan | None: self._last_bookmark_advance = None self._last_model_provenance = None result = self._execute_extractor(prepared.extractor_config, prepared.identifier) - generated_count = self._count_generated_results(result) try: write_plan = self._resolve_write_plan([result]) if result else None @@ -643,9 +651,17 @@ def compute_generation(self, request: TRequest) -> GenerationComputePlan | None: self._mark_extraction_runs_finalization_failed(exc) raise + generated_count = self._count_generated_results(result) + billable_count = ( + self._count_retained_online_learnings(write_plan) + if self.EMITS_LEARNING_BILLING + else 0 + ) + return GenerationComputePlan( prepared=prepared, generated_count=generated_count, + billable_count=billable_count, write_plan=write_plan, bookmark_advance=self._last_bookmark_advance, generation_start=generation_start, @@ -676,7 +692,8 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: Runs only for a fence-winning job (the durable worker calls it after the scope commits; ``.run()`` calls it inline). Reads the plan's compute-time - snapshot (``generated_count`` / ``prepared`` / ``generation_start``) so a + snapshot (``generated_count`` / ``billable_count`` / ``prepared`` / + ``generation_start``) so a fence-lost job never emits. Billing purity note (round-2 finding): ``_record_billing_learning_events`` @@ -708,7 +725,7 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: }, ) self._record_billing_learning_events( - prepared=plan.prepared, generated_count=plan.generated_count + prepared=plan.prepared, generated_count=plan.billable_count ) @abstractmethod diff --git a/reflexio/server/services/deferred_learning_plan.py b/reflexio/server/services/deferred_learning_plan.py index eb0a87cd8..a5db43ad1 100644 --- a/reflexio/server/services/deferred_learning_plan.py +++ b/reflexio/server/services/deferred_learning_plan.py @@ -134,16 +134,20 @@ class GenerationComputePlan: ``emit_generation_side_effects`` fires the post-commit telemetry + billing. The billing inputs (``extraction_run_ids`` / ``token_totals`` / - ``generated_count`` / ``prepared``) are **snapshotted at compute time** so - the fence-crossing emit reads this plan rather than the reused service - instance's mutable ``_last_*`` accumulators (purity contract, plan §File - Structure). See ``emit_generation_side_effects`` for the single-use-instance - invariant that also keeps the money helper's ``self._last_*`` reads safe. + ``billable_count`` / ``prepared``) and telemetry's ``generated_count`` are + **snapshotted at compute time** so the fence-crossing emit reads this plan + rather than the reused service instance's mutable ``_last_*`` accumulators + (purity contract, plan §File Structure). See + ``emit_generation_side_effects`` for the single-use-instance invariant that + also keeps the money helper's ``self._last_*`` reads safe. Attributes: prepared: The prepared generation run (identifier / extractor_name / extractor_config), reused by emit for telemetry + billing input. - generated_count: Learnings produced by this extraction run. + generated_count: Raw learnings produced by the extractor, used for + generation-success telemetry. + billable_count: Retained write-plan learnings eligible for online billing; + zero for services that do not emit learning billing. write_plan: The resolved write-plan (``ProfileWritePlan`` / ``PlaybookWritePlan`` in Tasks 6-7, a ``_LegacyItems`` shim marker until then) or ``None`` when the extractor produced nothing. @@ -159,6 +163,7 @@ class GenerationComputePlan: prepared: PreparedGenerationRun[Any] generated_count: int + billable_count: int write_plan: Any bookmark_advance: ExtractorBookmarkAdvance | None generation_start: float diff --git a/reflexio/server/services/extraction/README.md b/reflexio/server/services/extraction/README.md index 668455202..a0dfe2b5f 100644 --- a/reflexio/server/services/extraction/README.md +++ b/reflexio/server/services/extraction/README.md @@ -16,13 +16,16 @@ information, and resumes outside the request path. | `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` | Discovers and schedules due paused/finalization work in a background singleton. | -| `resume_worker.py` | Resumes paused runs, rebuilds request context, and records retry state. | +| `resume_worker.py` | Resumes paused runs, rebuilds request context, and records retry state. Finalization uses an immutable run-keyed receipt so learning writes and retry billing reuse the same persisted IDs. | | `outcome.py` | Provides the generic extraction outcome wrapper used by callers. | ## Boundary Rules - Keep profile-specific and playbook-specific extraction behavior in their own modules; call this package only for shared async runtime concerns. +- Commit resumable learning writes, lineage changes, and the agent-run + finalization receipt in one storage transaction. Retries must return the + receipt's persisted IDs instead of repeating finalization. - Add a new file here when the behavior is shared by more than one extraction caller or is part of the resumable runtime itself. - Split into subpackages only when a responsibility grows large enough that a diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index ff1732027..3d750f36d 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -909,12 +909,12 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - persisted_items = service._finalize_extracted_items( - items, model_provenance=model_provenance - ) - self._record_finalized_learnings( - run, persisted_items or [], entity_type="profile" + learning_ids = service._finalize_extracted_items( + items, + model_provenance=model_provenance, + finalization_run_id=run.id, ) + self._record_finalized_learnings(run, learning_ids, entity_type="profile") return if run.binding.extractor_kind == "playbook": user_id = run.binding.user_id @@ -932,13 +932,14 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - persisted_items = service._finalize_extracted_items( + learning_ids = service._finalize_extracted_items( items, model_provenance=model_provenance, extraction_run=run, + finalization_run_id=run.id, ) self._record_finalized_learnings( - run, persisted_items or [], entity_type="user_playbook" + run, learning_ids, entity_type="user_playbook" ) return raise ResumeWorkerError( @@ -946,67 +947,31 @@ def _finalize_items( ) def _record_finalized_learnings( - self, run: AgentRunRecord, items: list[Any], *, entity_type: str + self, run: AgentRunRecord, learning_ids: list[str], *, entity_type: str ) -> None: """Emit ``learnings_generated`` for a finalized resumable-extraction batch. - Prefers one event per learning id (``entity_id``/``profile_id`` for - profiles, ``user_playbook_id`` for playbooks) when every item in - ``items`` carries a durable id — the common case, since these ids are - assigned by the extractor (profile) or by ``save_user_playbooks`` - in-place during ``_finalize_extracted_items`` (playbook), which has - already run by the time this is called. Falls back to the - count-based aggregate event when any item lacks one (e.g. dropped by - within-batch/consolidator dedup before persist, leaving a default - ``user_playbook_id=0``) — this avoids both fabricating an id for a row - that never persisted and colliding on the shared default-id key. - Totals are preserved either way: ``len(items)`` learnings are counted - whether via ``len(learning_ids)`` per-record events or one aggregate - ``count=len(items)`` event. + Emits one event per durable learning id returned by finalization. + Per-record keys make finalization retries idempotent downstream. """ - from reflexio.server.billing_meter import ( - emit_learnings_generated, - emit_learnings_generated_records, - ) + from reflexio.server.billing_meter import emit_learnings_generated_records - if not items: + if not learning_ids: return - - id_attr = "profile_id" if entity_type == "profile" else "user_playbook_id" - learning_ids = [ - str(getattr(item, id_attr)) - for item in items - if getattr(item, id_attr, None) - ] metadata = { "run_id": run.id, "extractor_kind": run.binding.extractor_kind, } - if len(learning_ids) == len(items): - emit_learnings_generated_records( - org_id=self.request_context.org_id, - configurator=self.request_context.configurator, - learning_ids=learning_ids, - source="resumable_extraction", - pipeline=run.binding.extractor_kind, - user_id=run.binding.user_id, - request_id=run.binding.request_id, - agent_version=run.binding.agent_version, - entity_type=entity_type, - metadata=metadata, - ) - return - emit_learnings_generated( + emit_learnings_generated_records( org_id=self.request_context.org_id, configurator=self.request_context.configurator, - count=len(items), + learning_ids=learning_ids, source="resumable_extraction", pipeline=run.binding.extractor_kind, user_id=run.binding.user_id, request_id=run.binding.request_id, agent_version=run.binding.agent_version, entity_type=entity_type, - event_key=f"learn-batch:resumable:{run.id}:{entity_type}", metadata=metadata, ) diff --git a/reflexio/server/services/governance/service.py b/reflexio/server/services/governance/service.py index fca64cdbf..de3004e7a 100644 --- a/reflexio/server/services/governance/service.py +++ b/reflexio/server/services/governance/service.py @@ -1,5 +1,8 @@ from __future__ import annotations +import threading +import time +import uuid from contextlib import suppress from typing import Any, Literal, Protocol, TypedDict @@ -16,6 +19,7 @@ governance_subject_ref, ) from reflexio.server.services.governance.subject_refs import stable_id +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim _DELETE_TARGET_NAME_TO_RESULT_KEY = { "interaction": "interactions", @@ -29,12 +33,18 @@ "offline_tuner_reward_label_target_by_target_owner": ( "offline_tuner_reward_label_targets_by_target_owner" ), + "session_outcome": "session_outcomes", "profile_purge": "purged_profiles", "user_playbook_purge": "purged_user_playbooks", } _REQUIRED_DELETE_TARGET_NAMES = tuple(_DELETE_TARGET_NAME_TO_RESULT_KEY) _USER_PLAYBOOK_PAGE_SIZE = 1000 _LIFECYCLE_COMPLETION_STATUS = "complete" +_DUPLICATE_ERASE_POLL_SECONDS = 0.05 +_DUPLICATE_ERASE_MAX_POLL_SECONDS = 1.0 +_DUPLICATE_ERASE_WAIT_SECONDS = 5.0 +_PURGE_EXECUTION_LEASE_SECONDS = 300 +_PURGE_EXECUTION_HEARTBEAT_SECONDS = 30 class GovernanceActorContext(TypedDict): @@ -43,7 +53,7 @@ class GovernanceActorContext(TypedDict): class SubjectErasureLifecycle(Protocol): - """Deployment-specific erasure work that must precede barrier completion.""" + """External erasure work invoked only after a synchronous live-claim check.""" def erase_subject( self, @@ -51,9 +61,78 @@ def erase_subject( storage: Any, subject_ref: str, purge_id: str, + execution_claim: PurgeExecutionClaim, ) -> None: ... +class _PurgeExecutionHeartbeatLostError(ValueError): + pass + + +class GovernanceEraseRetryLaterError(RuntimeError): + pass + + +class _PurgeExecutionHeartbeat: + def __init__( + self, + *, + storage: Any, + purge_id: str, + execution_claim: PurgeExecutionClaim, + ) -> None: + self._storage = storage + self._purge_id = purge_id + self._claim = execution_claim + self._lock = threading.Lock() + self._renewal_error: Exception | None = None + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def __enter__(self) -> _PurgeExecutionHeartbeat: + self.renew_now() + self._thread.start() + return self + + def __exit__(self, *_exc: object) -> None: + self._stop.set() + self._thread.join(timeout=1) + + def claim(self) -> PurgeExecutionClaim: + with self._lock: + if self._renewal_error is not None: + raise _PurgeExecutionHeartbeatLostError( + "purge execution heartbeat renewal was lost" + ) from self._renewal_error + return self._claim + + def renew_now(self) -> PurgeExecutionClaim: + try: + renewed = self._storage.renew_purge_operation_execution_claim( + self._purge_id, + self.claim(), + lease_ttl_seconds=_PURGE_EXECUTION_LEASE_SECONDS, + ) + except _PurgeExecutionHeartbeatLostError: + raise + except Exception as exc: + with self._lock: + self._renewal_error = exc + raise _PurgeExecutionHeartbeatLostError( + "purge execution heartbeat renewal was lost" + ) from exc + with self._lock: + self._claim = renewed + return renewed + + def _run(self) -> None: + while not self._stop.wait(_PURGE_EXECUTION_HEARTBEAT_SECONDS): + try: + self.renew_now() + except _PurgeExecutionHeartbeatLostError: + return + + class GovernanceService: def __init__( self, @@ -129,84 +208,145 @@ def erase_user( f"{self.org_id}:user_erasure:{subref}:{reqref}", ) purge_id = stable_id("purge", idempotency_key) - purge = self.storage.begin_purge_operation( - purge_id=purge_id, - idempotency_key=idempotency_key, - operation_type="user_erasure", - scope_type="user", - subject_ref=subref, - request_ref=reqref, - ) + try: + purge = self.storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key=idempotency_key, + operation_type="user_erasure", + scope_type="user", + subject_ref=subref, + request_ref=reqref, + authoritative_user_id=user_id, + ) + except Exception as begin_exc: + try: + purge = self._matching_user_erasure_purge_for_retry( + purge_id=purge_id, + operation_type="user_erasure", + scope_type="user", + subject_ref=subref, + request_ref=reqref, + authoritative_user_id=user_id, + ) + except Exception: + raise begin_exc from None + if purge.status == "complete": + raise begin_exc from None if purge.status == "complete": - barrier = self._completed_barrier_for_retry( + return self._completed_erase_result_for_retry( subject_ref=subref, purge_id=purge_id ) - if barrier.status != "erased": - raise ValueError( - "Completed purge retry requires an erased subject barrier" + lease_owner = f"governance-erase-{uuid.uuid4().hex}" + execution_claim: PurgeExecutionClaim | None = None + claim_deadline = self._monotonic() + _DUPLICATE_ERASE_WAIT_SECONDS + poll_seconds = _DUPLICATE_ERASE_POLL_SECONDS + while execution_claim is None: + execution_claim = self.storage.claim_purge_operation_execution( + purge_id, + lease_owner=lease_owner, + lease_ttl_seconds=_PURGE_EXECUTION_LEASE_SECONDS, + ) + if execution_claim is not None: + break + purge = self.storage.get_purge_operation(purge_id) + if purge.status == "complete": + return self._completed_erase_result_for_retry( + subject_ref=subref, + purge_id=purge_id, ) - return UserEraseResult( - subject_ref=subref, - purge_id=purge_id, - status="complete", - deleted_counts=self._deleted_counts_from_targets(purge_id), - rebuilt_agent_playbook_ids=( - self._rebuilt_agent_playbook_ids_from_targets(purge_id) - ), + if purge.status not in {"pending", "running", "failed"}: + raise ValueError(f"Unsupported purge operation status: {purge.status}") + remaining_seconds = claim_deadline - self._monotonic() + if remaining_seconds <= 0: + raise GovernanceEraseRetryLaterError( + "Another erase request still owns the execution claim; retry later" + ) + self._sleep(min(poll_seconds, remaining_seconds)) + poll_seconds = min( + poll_seconds * 2, + _DUPLICATE_ERASE_MAX_POLL_SECONDS, ) try: - self.storage.begin_subject_erasure_barrier(subref, purge_id) - if not self.storage.purge_targets_prepared(purge_id): - self.storage.prepare_governance_erase_targets( + with _PurgeExecutionHeartbeat( + storage=self.storage, + purge_id=purge_id, + execution_claim=execution_claim, + ) as heartbeat: + self.storage.begin_subject_erasure_barrier( + subref, purge_id, - user_id, + execution_claim=heartbeat.claim(), ) + if not self.storage.purge_targets_prepared(purge_id): + self.storage.prepare_governance_erase_targets( + purge_id, + user_id, + execution_claim=heartbeat.claim(), + ) - if not self._delete_targets_complete(purge_id): - self.storage.apply_governance_user_data_delete(purge_id, user_id) - if ( - self.subject_erasure_lifecycle is not None - and not self._subject_erasure_lifecycle_complete(purge_id) - ): - self.subject_erasure_lifecycle.erase_subject( - storage=self.storage, - subject_ref=subref, - purge_id=purge_id, - ) - self._record_subject_erasure_lifecycle_complete(purge_id) - deleted_counts = self._deleted_counts_from_targets(purge_id) + if not self._delete_targets_complete(purge_id): + self.storage.apply_governance_user_data_delete( + purge_id, + user_id, + execution_claim=heartbeat.claim(), + ) + if ( + self.subject_erasure_lifecycle is not None + and not self._subject_erasure_lifecycle_complete(purge_id) + ): + heartbeat.renew_now() + self._assert_execution_claim(purge_id, heartbeat.claim()) + self.subject_erasure_lifecycle.erase_subject( + storage=self.storage, + subject_ref=subref, + purge_id=purge_id, + execution_claim=heartbeat.claim(), + ) + self._record_subject_erasure_lifecycle_complete( + purge_id, + execution_claim=heartbeat.claim(), + ) + deleted_counts = self._deleted_counts_from_targets(purge_id) - rebuilt_agent_playbook_ids: list[int] = [] - completed = self.storage.complete_subject_erasure_barrier_after_empty_check( - purge_id, - AuditEvent( - org_id=self.org_id, - actor_type=actor_type, - actor_ref=actor_ref, - operation="ERASE", - entity_type="request", - subject_ref=subref, - request_ref=reqref, - idempotency_key=purge_id, - detail={ - "deleted_counts": deleted_counts, - "rebuilt_agent_playbook_ids": rebuilt_agent_playbook_ids, - }, - ), - ) + rebuilt_agent_playbook_ids: list[int] = [] + completed = self.storage.complete_subject_erasure_barrier_after_empty_check( + purge_id, + AuditEvent( + org_id=self.org_id, + actor_type=actor_type, + actor_ref=actor_ref, + operation="ERASE", + entity_type="request", + subject_ref=subref, + request_ref=reqref, + idempotency_key=purge_id, + detail={ + "deleted_counts": deleted_counts, + "rebuilt_agent_playbook_ids": rebuilt_agent_playbook_ids, + }, + ), + authoritative_user_id=user_id, + execution_claim=heartbeat.claim(), + ) except Exception as exc: + if isinstance(exc, _PurgeExecutionHeartbeatLostError): + raise + if not self._execution_claim_is_current(purge_id, execution_claim): + raise with suppress(Exception): self.storage.fail_subject_erasure_barrier( subref, purge_id, error_code="governance_erase_failed", error_detail=type(exc).__name__, + execution_claim=execution_claim, ) with suppress(Exception): self.storage.fail_purge_operation( purge_id, error_code="governance_erase_failed", error_detail=type(exc).__name__, + execution_claim=execution_claim, ) raise return UserEraseResult( @@ -217,6 +357,28 @@ def erase_user( rebuilt_agent_playbook_ids=rebuilt_agent_playbook_ids, ) + @staticmethod + def _monotonic() -> float: + return time.monotonic() + + @staticmethod + def _sleep(seconds: float) -> None: + time.sleep(seconds) + + def _assert_execution_claim( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> None: + self.storage.assert_purge_operation_execution_claim(purge_id, execution_claim) + + def _execution_claim_is_current( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> bool: + try: + self._assert_execution_claim(purge_id, execution_claim) + except Exception: + return False + return True + def _assert_storage_ref_secret_matches(self) -> None: storage_secret = get_governance_ref_secret() if storage_secret != self.ref_secret: @@ -235,6 +397,57 @@ def _completed_barrier_for_retry( ) return barrier + def _completed_erase_result_for_retry( + self, *, subject_ref: str, purge_id: str + ) -> UserEraseResult: + barrier = self._completed_barrier_for_retry( + subject_ref=subject_ref, purge_id=purge_id + ) + if barrier.status != "erased": + raise ValueError("Completed purge retry requires an erased subject barrier") + return UserEraseResult( + subject_ref=subject_ref, + purge_id=purge_id, + status="complete", + deleted_counts=self._deleted_counts_from_targets(purge_id), + rebuilt_agent_playbook_ids=( + self._rebuilt_agent_playbook_ids_from_targets(purge_id) + ), + ) + + def _matching_user_erasure_purge_for_retry( + self, + *, + purge_id: str, + operation_type: str, + scope_type: str, + subject_ref: str, + request_ref: str, + authoritative_user_id: str, + ) -> Any: + purge = self.storage.get_purge_operation(purge_id) + expected_identity = { + "purge_id": purge_id, + "operation_type": operation_type, + "scope_type": scope_type, + "subject_ref": subject_ref, + "request_ref": request_ref, + } + for field_name, expected_value in expected_identity.items(): + if getattr(purge, field_name) != expected_value: + raise ValueError( + "Existing purge operation for idempotency_key has " + f"mismatched {field_name}" + ) + if ( + governance_subject_ref(self.org_id, authoritative_user_id, self.ref_secret) + != purge.subject_ref + ): + raise ValueError( + "Existing purge operation has mismatched authoritative user" + ) + return purge + def _load_user_requests_and_sessions( self, user_id: str ) -> tuple[list[Any], list[dict[str, Any]]]: @@ -310,7 +523,11 @@ def _subject_erasure_lifecycle_complete(self, purge_id: str) -> bool: and (snapshot.detail or {}).get("status") == _LIFECYCLE_COMPLETION_STATUS ) - def _record_subject_erasure_lifecycle_complete(self, purge_id: str) -> None: + def _record_subject_erasure_lifecycle_complete( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + ) -> None: snapshot = self._prepared_target_snapshot(purge_id) if snapshot is None or snapshot.status != "complete": raise ValueError( @@ -327,6 +544,7 @@ def _record_subject_erasure_lifecycle_complete(self, purge_id: str) -> None: detail=detail, deleted_count=snapshot.deleted_count, error_detail=snapshot.error_detail, + execution_claim=execution_claim, ) def _prepared_target_snapshot(self, purge_id: str) -> PurgeOperationTarget | None: diff --git a/reflexio/server/services/playbook/components/aggregator.py b/reflexio/server/services/playbook/components/aggregator.py index b0887d0e7..11bd366ce 100644 --- a/reflexio/server/services/playbook/components/aggregator.py +++ b/reflexio/server/services/playbook/components/aggregator.py @@ -5,7 +5,7 @@ import os import time import uuid -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Sequence from contextlib import AbstractContextManager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Protocol, cast @@ -1989,17 +1989,6 @@ def run(self, playbook_aggregator_request: PlaybookAggregatorRequest) -> dict: duration_ms=int((time.perf_counter() - aggregation_start) * 1000), metadata=stats, ) - self._record_learnings_generated( - learning_ids=[ - str(saved.agent_playbook_id) - for saved in saved_playbook_list - if getattr(saved, "agent_playbook_id", None) - ], - playbook_name=playbook_name, - request_id=_run_id, - metadata=stats, - total_count=len(saved_playbook_list), - ) return stats except Exception as e: @@ -2042,61 +2031,6 @@ def run(self, playbook_aggregator_request: PlaybookAggregatorRequest) -> dict: # Re-raise the exception after restoring raise - def _record_learnings_generated( - self, - *, - learning_ids: list[str], - playbook_name: str, - request_id: str, - metadata: Mapping[str, Any], - total_count: int | None = None, - ) -> None: - """Emit ``learnings_generated`` for a completed aggregation run. - - Prefers one event per learning id (entity-backed) when every saved - playbook in this run carries a durable ``agent_playbook_id`` — the - common case, since ``save_agent_playbooks`` - raises rather than returning a partial row. Falls back to the - count-based aggregate event when ``learning_ids`` is short of - ``total_count`` (a falsy/unset id slipped through), mirroring - ``ExtractionResumeWorker._record_finalized_learnings`` — this avoids - emitting a colliding ``learn:agent_playbook:0`` key. ``total_count`` - defaults to ``len(learning_ids)`` so callers that already guarantee a - complete id list (e.g. existing tests) are unaffected. - """ - from reflexio.server.billing_meter import ( - emit_learnings_generated, - emit_learnings_generated_records, - ) - - total = len(learning_ids) if total_count is None else total_count - if len(learning_ids) == total: - emit_learnings_generated_records( - org_id=self.request_context.org_id, - configurator=self.configurator, - learning_ids=learning_ids, - source="aggregation", - pipeline="playbook", - request_id=request_id, - agent_version=self.agent_version, - playbook_name=playbook_name, - entity_type="agent_playbook", - metadata=metadata, - ) - return - emit_learnings_generated( - org_id=self.request_context.org_id, - configurator=self.configurator, - count=total, - source="aggregation", - pipeline="playbook", - request_id=request_id, - agent_version=self.agent_version, - playbook_name=playbook_name, - entity_type="agent_playbook", - metadata=metadata, - ) - def get_clusters( self, user_playbooks: list[UserPlaybook], diff --git a/reflexio/server/services/playbook/service.py b/reflexio/server/services/playbook/service.py index ce12ec491..1954278c4 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -628,10 +628,10 @@ def _dispatch_playbook_schedulers(self, plan: PlaybookWritePlan) -> None: Phantom-billing gate: on the durable / ``.run()`` path this is invoked from ``emit_generation_side_effects`` (post-commit), so a fence-lost (superseded) job never enqueues optimization or triggers aggregation. On - the synchronous resume/manual path the permanent - ``_finalize_extracted_items`` wrapper invokes it right after persist, - keeping that path identical to the pre-split monolith. The two callers - are mutually exclusive, so the schedulers fire exactly once per run. + the synchronous resume/manual path ``_finalize_extracted_items`` invokes + it after persistence. Dispatch is best-effort and at most once per + committed finalization attempt; derived scheduler work has no durable + replay idempotency. """ try: self._enqueue_user_playbook_optimization(plan.new_playbooks) @@ -669,16 +669,28 @@ def _finalize_extracted_items( *, model_provenance: ModelProvenance | None = None, extraction_run: AgentRunRecord | None = None, - ) -> list[UserPlaybook]: - """Permanent V3 wrapper: compute→persist→schedulers together (no fence). + finalization_run_id: str | None = None, + ) -> list[str]: + """Finalize extracted playbooks for synchronous resume/manual callers. Kept for the synchronous resume/manual callers (``ExtractionResumeWorker`` calls this directly). Routes them through the same ``_resolve_write_plan`` (compute) + ``_persist_write_plan`` - (persist) split the durable worker uses — with no external - ``commit_scope`` — then dispatches the same off-thread schedulers, so the - result is identical to the pre-split monolith. + (persist) split the durable worker uses. Derived schedulers dispatch + best-effort after the finalization transaction commits. When an existing + finalization receipt is found, the method intentionally returns its + learning ids without replaying those schedulers. """ + entity_type = "user_playbook" + if finalization_run_id is not None: + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + # Receipts make persistence/billing idempotent. Derived schedulers + # are best-effort at-most-once and lack durable replay idempotency. + return receipt if model_provenance is not None: self._last_model_provenance = model_provenance previous_review_run = self._review_run @@ -691,12 +703,44 @@ def _finalize_extracted_items( finally: self._review_run = previous_review_run self._review_window_cache = previous_review_window - if plan is None: - return [] + learning_ids: list[str] = [] + if finalization_run_id is None: + if plan is not None: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + self._persist_write_plan(plan) + self._dispatch_playbook_schedulers(plan) + return ( + [ + str(playbook.user_playbook_id) + for playbook in plan.new_playbooks + if playbook.user_playbook_id + ] + if plan is not None + else [] + ) + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] - self._persist_write_plan(plan) - self._dispatch_playbook_schedulers(plan) - return plan.new_playbooks + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return receipt + if plan is not None: + self._persist_write_plan(plan) + learning_ids = [ + str(playbook.user_playbook_id) + for playbook in plan.new_playbooks + if playbook.user_playbook_id + ] + self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + learning_ids=learning_ids, + ) + if plan is not None: + self._dispatch_playbook_schedulers(plan) + return learning_ids def _apply_consolidation_lineage( self, diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index 4b2bd1053..691f8d34c 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -337,7 +337,8 @@ def _finalize_extracted_items( all_new_profiles: list[UserProfile], *, model_provenance: ModelProvenance | None = None, - ) -> list[UserProfile]: + finalization_run_id: str | None = None, + ) -> list[str]: """Permanent V3 wrapper: compute-then-persist together (no external fence). Kept for the synchronous resume/manual callers @@ -346,14 +347,48 @@ def _finalize_extracted_items( (persist) split the durable worker uses — with no external ``commit_scope`` — so the result is identical to the pre-split monolith. """ + entity_type = "profile" + if finalization_run_id is not None: + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return receipt if model_provenance is not None: self._last_model_provenance = model_provenance plan = self._resolve_write_plan([all_new_profiles]) - if plan is None: - return [] + # Profile IDs are assigned before persistence, unlike database-assigned playbook IDs. + learning_ids = ( + [ + str(profile.profile_id) + for profile in plan.new_profiles + if profile.profile_id + ] + if plan is not None + else [] + ) + if finalization_run_id is None: + if plan is not None: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + self._persist_write_plan(plan) + return learning_ids + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] - self._persist_write_plan(plan) - return plan.new_profiles + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return receipt + if plan is not None: + self._persist_write_plan(plan) + self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + learning_ids=learning_ids, + ) + return learning_ids def check_and_update_profiles(self, profiles: list[UserProfile]) -> None: """check if the profiles are expired and update them if they are""" diff --git a/reflexio/server/services/search_exposure.py b/reflexio/server/services/search_exposure.py new file mode 100644 index 000000000..aae7fc4ff --- /dev/null +++ b/reflexio/server/services/search_exposure.py @@ -0,0 +1,155 @@ +"""Optional synchronous recording boundary for served user playbooks.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from hashlib import sha256 +from secrets import token_hex +from typing import Protocol + +from reflexio.models.api_schema.domain import UserPlaybook +from reflexio.server.extensions import ServiceKey, get_service +from reflexio.server.services.playbook.publication import ( + canonical_json_bytes, + incumbent_user_playbook_semantic_digest, +) + + +@dataclass(frozen=True) +class SearchExposureBatch: + """The final user-playbook set returned by one unified search.""" + + org_id: str + request_id: str | None + session_id: str | None + interaction_id: int | None + user_id: str | None + user_playbooks: tuple[UserPlaybook, ...] + invocation_id: str = field(default_factory=lambda: token_hex(16)) + + def __post_init__(self) -> None: + object.__setattr__( + self, "request_id", _normalize_correlation_id(self.request_id) + ) + object.__setattr__( + self, "session_id", _normalize_correlation_id(self.session_id) + ) + if self.interaction_id is not None and self.interaction_id <= 0: + object.__setattr__(self, "interaction_id", None) + + +@dataclass(frozen=True) +class UserPlaybookExposureEvent: + """Immutable durable envelope for one served user playbook.""" + + exposure_event_id: str + request_id: str | None + session_id: str | None + user_id: str | None + user_playbook_id: int | None + served_semantic_digest: str | None + served_full_version_fingerprint: str | None + exposed_at: int | None + ingested_at: int + governance_subject_ref: str | None + + +@dataclass(frozen=True) +class ExposureEventWriteResult: + """Durable write result returned by the enterprise ledger store.""" + + recorded: bool + integrity_state: str + integrity_reasons: tuple[str, ...] + + +class SearchExposureRecorder(Protocol): + """Durably record a final search result set before response release.""" + + def record(self, batch: SearchExposureBatch) -> None: ... + + +SEARCH_EXPOSURE_RECORDER = ServiceKey[SearchExposureRecorder]( + "search_exposure_recorder" +) + + +def _normalize_correlation_id(value: str | None) -> str | None: + normalized = value.strip() if value is not None else "" + return normalized or None + + +def record_search_exposures(batch: SearchExposureBatch) -> None: + """Synchronously invoke the optional enterprise exposure recorder.""" + recorder = get_service(SEARCH_EXPOSURE_RECORDER) + if recorder is not None: + recorder.record(batch) + + +def user_playbook_full_version_fingerprint(playbook: UserPlaybook) -> str: + """Bind every persisted playbook field except its derived embedding vector. + + Adding or changing persisted ``UserPlaybook`` fields requires bumping + ``user-playbook-full-version-v1``; cross-version fingerprint comparisons are + undefined. + """ + payload = { + "schema_version": "user-playbook-full-version-v1", + "user_playbook": playbook.model_dump(mode="json", exclude={"embedding"}), + } + return sha256(canonical_json_bytes(payload)).hexdigest() + + +def build_user_playbook_exposure_event( + batch: SearchExposureBatch, + playbook: UserPlaybook, + *, + exposed_at: int, + ingested_at: int, + governance_subject_ref: str | None, +) -> UserPlaybookExposureEvent: + """Build one deterministic event identity from retrieval-owned correlation.""" + identity: dict[str, object] = { + "schema_version": "user-playbook-exposure-event-v1", + "org_id": batch.org_id, + "request_id": batch.request_id, + "session_id": batch.session_id, + "interaction_id": batch.interaction_id, + "user_playbook_id": playbook.user_playbook_id, + } + if ( + batch.request_id is None + and batch.session_id is None + and batch.interaction_id is None + ): + identity["invocation_id"] = batch.invocation_id + content_digest = sha256(playbook.content.encode("utf-8")).hexdigest() + return UserPlaybookExposureEvent( + exposure_event_id=sha256(canonical_json_bytes(identity)).hexdigest(), + request_id=batch.request_id, + session_id=batch.session_id, + user_id=playbook.user_id, + user_playbook_id=playbook.user_playbook_id, + served_semantic_digest=incumbent_user_playbook_semantic_digest( + content_digest=content_digest, + trigger=playbook.trigger, + ), + served_full_version_fingerprint=user_playbook_full_version_fingerprint( + playbook + ), + exposed_at=exposed_at, + ingested_at=ingested_at, + governance_subject_ref=governance_subject_ref, + ) + + +__all__ = [ + "SEARCH_EXPOSURE_RECORDER", + "ExposureEventWriteResult", + "SearchExposureBatch", + "SearchExposureRecorder", + "UserPlaybookExposureEvent", + "build_user_playbook_exposure_event", + "record_search_exposures", + "user_playbook_full_version_fingerprint", +] diff --git a/reflexio/server/services/storage/governance_claims.py b/reflexio/server/services/storage/governance_claims.py new file mode 100644 index 000000000..147d3386b --- /dev/null +++ b/reflexio/server/services/storage/governance_claims.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class PurgeExecutionClaim: + purge_id: str + owner: str + fence: int + expires_at: int + + +def validate_purge_execution_claim( + purge_id: str, + execution_claim: PurgeExecutionClaim | None, +) -> PurgeExecutionClaim: + if execution_claim is None: + raise ValueError("purge execution claim is required") + if type(execution_claim) is not PurgeExecutionClaim: + raise ValueError("purge execution claim must be typed") + if execution_claim.purge_id != purge_id: + raise ValueError("purge execution claim purge_id mismatch") + if not execution_claim.owner.strip(): + raise ValueError("purge execution claim owner is required") + if execution_claim.fence <= 0: + raise ValueError("purge execution claim fence is invalid") + if execution_claim.expires_at <= 0: + raise ValueError("purge execution claim expiry is invalid") + return execution_claim diff --git a/reflexio/server/services/storage/governance_validation.py b/reflexio/server/services/storage/governance_validation.py index a89fa1def..3358604ec 100644 --- a/reflexio/server/services/storage/governance_validation.py +++ b/reflexio/server/services/storage/governance_validation.py @@ -90,6 +90,7 @@ { "affected_agent_playbook_ids", "agent_playbook_id", + "authoritative_user_digest", "count", "deleted_counts", "deleted_count", @@ -166,6 +167,7 @@ "user_playbooks", "profiles", "requests", + "session_outcomes", "agent_success_evaluation_results", "retrieved_learning_evaluation_results", "evaluation_operation_states", @@ -476,6 +478,12 @@ def _validate_governance_detail_entry( if key in {"agent_playbook_id", "user_playbook_id"}: _validate_governance_int(field_name, value) return cast(int, value) + if key == "authoritative_user_digest": + if not isinstance(value, str) or re.fullmatch(r"[0-9a-f]{64}", value) is None: + _raise_governance_validation_error( + field_name, "expected 64 lowercase hex chars" + ) + return value if key == "deleted_counts": return _validate_governance_deleted_counts(field_name, value) if key in { diff --git a/reflexio/server/services/storage/retention.py b/reflexio/server/services/storage/retention.py index d48035689..3a58a9251 100644 --- a/reflexio/server/services/storage/retention.py +++ b/reflexio/server/services/storage/retention.py @@ -48,6 +48,12 @@ class OptimizationRetentionClass: ), RetentionTarget("interactions", "interactions", "created_at", ("interaction_id",)), RetentionTarget("requests", "requests", "created_at", ("request_id",)), + RetentionTarget( + "session_outcomes", + "session_outcomes", + "created_at", + ("user_id", "session_id"), + ), RetentionTarget( "user_playbooks", "user_playbooks", @@ -130,6 +136,12 @@ class OptimizationRetentionClass: "created_at", ("retrieval_log_id",), ), + RetentionTarget( + "user_playbook_exposure_events", + "user_playbook_exposure_events", + "ingested_at", + ("exposure_event_id",), + ), RetentionTarget("skills", "skills", "created_at", ("skill_id",)), ) diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py new file mode 100644 index 000000000..273b0195a --- /dev/null +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -0,0 +1,188 @@ +"""Canonical identities for immutable session outcomes.""" + +import json +from collections.abc import Collection, Mapping, Sequence +from datetime import UTC, datetime +from hashlib import sha256 +from typing import TypedDict + +from reflexio.server.services.playbook.publication import canonical_json_bytes + +__all__ = [ + "CanonicalSessionTrajectory", + "canonical_session_trajectory", + "canonical_json_bytes", + "outcome_contract_digest", + "trajectory_digest", +] + + +class CanonicalRequest(TypedDict): + request_id: str + user_id: str + created_at: str + source: str + agent_version: str + session_id: str + evaluation_only: bool + retrieval_experiment_id: str | None + retrieval_experiment_arm: str | None + + +class CanonicalInteraction(TypedDict): + interaction_id: int + user_id: str + request_id: str + created_at: str + content: str + role: str + token_count: int | None + user_action: str + user_action_description: str + interacted_image_url: str + image_encoding: str + shadow_content: str + expert_content: str + tools_used: object + citations: object + retrieved_learnings: object + + +class CanonicalTrajectoryRequest(TypedDict): + request: CanonicalRequest + interactions: list[CanonicalInteraction] + + +class CanonicalSessionTrajectory(TypedDict): + session_id: str + requests: list[CanonicalTrajectoryRequest] + + +def _canonical_timestamp(value: object) -> str: + if isinstance(value, datetime): + if value.tzinfo is None: + value = value.replace(tzinfo=UTC) + return value.astimezone(UTC).isoformat() + return str(value) + + +def _canonical_json_column(value: object, *, default: object) -> object: + if value is None or value == "": + return default + if isinstance(value, str): + return json.loads(value) + return value + + +def _canonical_int(value: object) -> int: + if not isinstance(value, int | str): + raise TypeError("canonical integer field must be an integer or integer text") + return int(value) + + +def _canonical_interaction( + interaction: Mapping[str, object], +) -> CanonicalInteraction: + return { + "interaction_id": _canonical_int(interaction["interaction_id"]), + "user_id": str(interaction["user_id"]), + "request_id": str(interaction["request_id"]), + "created_at": _canonical_timestamp(interaction["created_at"]), + "content": str(interaction["content"]), + "role": str(interaction["role"]), + "token_count": ( + _canonical_int(interaction["token_count"]) + if interaction["token_count"] is not None + else None + ), + "user_action": str(interaction["user_action"]), + "user_action_description": str(interaction["user_action_description"] or ""), + "interacted_image_url": str(interaction["interacted_image_url"] or ""), + "image_encoding": str(interaction["image_encoding"] or ""), + "shadow_content": str(interaction["shadow_content"] or ""), + "expert_content": str(interaction["expert_content"] or ""), + "tools_used": _canonical_json_column(interaction["tools_used"], default=[]), + "citations": _canonical_json_column(interaction["citations"], default=[]), + "retrieved_learnings": _canonical_json_column( + interaction["retrieved_learnings"], default=[] + ), + } + + +def canonical_session_trajectory( + session_id: str, + request_rows: Sequence[Mapping[str, object]], + interactions_by_request: Mapping[str, Sequence[Mapping[str, object]]], +) -> CanonicalSessionTrajectory: + """Project adapter-specific durable rows into one trajectory identity shape.""" + requests: list[CanonicalTrajectoryRequest] = [] + for row in request_rows: + request_id = str(row["request_id"]) + request: CanonicalRequest = { + "request_id": request_id, + "user_id": str(row["user_id"]), + "created_at": _canonical_timestamp(row["created_at"]), + "source": str(row["source"] or ""), + "agent_version": str(row["agent_version"] or ""), + "session_id": str(row["session_id"]), + "evaluation_only": bool(row["evaluation_only"]), + "retrieval_experiment_id": ( + str(row["retrieval_experiment_id"]) + if row["retrieval_experiment_id"] is not None + else None + ), + "retrieval_experiment_arm": ( + str(row["retrieval_experiment_arm"]) + if row["retrieval_experiment_arm"] is not None + else None + ), + } + interactions = [ + _canonical_interaction(interaction) + for interaction in interactions_by_request.get(request_id, ()) + ] + requests.append({"request": request, "interactions": interactions}) + return {"session_id": session_id, "requests": requests} + + +def outcome_contract_digest( + *, + source: str, + schema_version: int | str, + allowed_values: Collection[str], + finalization_rule: str, +) -> str: + """Hash a server-owned structured outcome contract.""" + payload = { + "allowed_values": sorted(set(allowed_values)), + "finalization_rule": finalization_rule, + "schema_version": schema_version, + "source": source, + } + return sha256(canonical_json_bytes(payload)).hexdigest() + + +def _canonical_trajectory_json(value: object) -> str: + """Encode trajectory JSON, including finite floats, deterministically.""" + if isinstance(value, float): + return json.dumps(value, allow_nan=False, separators=(",", ":")) + if isinstance(value, tuple | list): + return "[" + ",".join(_canonical_trajectory_json(item) for item in value) + "]" + if isinstance(value, Mapping): + if not all(isinstance(key, str) for key in value): + raise TypeError("canonical trajectory object keys must be strings") + keys = sorted(value, key=lambda key: key.encode("utf-16be")) + return ( + "{" + + ",".join( + f"{canonical_json_bytes(key).decode()}:{_canonical_trajectory_json(value[key])}" + for key in keys + ) + + "}" + ) + return canonical_json_bytes(value).decode() + + +def trajectory_digest(trajectory: object) -> str: + """Hash the canonical finalized session trajectory.""" + return sha256(_canonical_trajectory_json(trajectory).encode()).hexdigest() diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index faba66279..e6d4ab9f0 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -18,10 +18,12 @@ import unicodedata from collections.abc import Callable, Generator, Sequence from datetime import UTC, datetime +from hashlib import sha256 from pathlib import Path from typing import Any, ClassVar, Literal from reflexio.models.api_schema.common import BlockingIssue +from reflexio.models.api_schema.domain import SessionOutcomeKind from reflexio.models.api_schema.service_schemas import ( AgentPlaybook, AgentSuccessEvaluationResult, @@ -59,6 +61,13 @@ require_non_empty_session_id, ) from reflexio.server.services.storage.retention_mixin import RetentionMixin +from reflexio.server.services.storage.session_outcome_identity import ( + CanonicalSessionTrajectory, + canonical_json_bytes, + canonical_session_trajectory, + outcome_contract_digest, + trajectory_digest, +) from reflexio.server.services.storage.storage_base import BaseStorage from reflexio.server.site_var.site_var_manager import SiteVarManager @@ -67,6 +76,7 @@ logger = logging.getLogger(__name__) +_OUTCOME_ALLOWED_VALUES = tuple(kind.value for kind in SessionOutcomeKind) _MINIMUM_SQLITE_VERSION = (3, 35, 0) _SQLITE_INITIALIZATION_LOCK_STRIPES = 64 _sqlite_initialization_locks = tuple( @@ -101,6 +111,47 @@ def _json_loads(text: str | None) -> Any: return json.loads(text) +def _canonical_session_snapshot( + conn: sqlite3.Connection, session_id: str +) -> CanonicalSessionTrajectory: + """Return the durable session state used to bind an outcome finalization.""" + request_rows = conn.execute( + """SELECT request_id, user_id, created_at, source, agent_version, session_id, + evaluation_only, retrieval_experiment_id, retrieval_experiment_arm + FROM requests WHERE session_id = ? + ORDER BY created_at ASC, request_id ASC""", + (session_id,), + ).fetchall() + request_payloads = [dict(request) for request in request_rows] + interactions_by_request: dict[str, list[dict[str, object]]] = { + str(request["request_id"]): [] for request in request_rows + } + interaction_rows = conn.execute( + """SELECT interaction_id, user_id, request_id, created_at, content, role, + token_count, user_action, user_action_description, + interacted_image_url, image_encoding, shadow_content, + expert_content, tools_used, citations, retrieved_learnings + FROM interactions + WHERE request_id IN ( + SELECT request_id FROM requests WHERE session_id = ? + ) + ORDER BY request_id ASC, created_at ASC, interaction_id ASC""", + (session_id,), + ).fetchall() + for interaction in interaction_rows: + interactions_by_request[str(interaction["request_id"])].append( + dict(interaction) + ) + return canonical_session_trajectory( + session_id, request_payloads, interactions_by_request + ) + + +def _legacy_session_outcome_id(user_id: str, session_id: str) -> str: + """Return a delimiter-safe immutable identity for a migrated legacy outcome.""" + return sha256(canonical_json_bytes([user_id, session_id])).hexdigest() + + _FTS5_OPERATORS = frozenset({"OR", "AND", "NOT"}) _FTS5_RESERVED = _FTS5_OPERATORS | {"NEAR"} _TOKEN_RE = re.compile(r"\w+", re.UNICODE) @@ -915,13 +966,27 @@ def migrate(self) -> bool: # _DDL creates an index over these columns. Upgrade legacy request tables # before executescript so index creation cannot fail on missing columns. self._migrate_request_retrieval_experiment() - self._migrate_session_outcomes_schema() with self._lock: + session_outcome_columns = { + row["name"] + for row in self.conn.execute( + "PRAGMA table_info(session_outcomes)" + ).fetchall() + } + if ( + session_outcome_columns + and "governance_subject_ref" not in session_outcome_columns + ): + self.conn.execute( + "ALTER TABLE session_outcomes " + "ADD COLUMN governance_subject_ref TEXT" + ) cur = self.conn.cursor() cur.executescript(_DDL) init_governance_tables(self.conn) init_playbook_aggregation_tables(self.conn) self.conn.commit() + self._migrate_session_outcomes_schema() if self._has_sqlite_vec: self._create_vec_tables() self._migrate_vec_tables() @@ -1020,120 +1085,6 @@ def _migrate_unicode_lexical_indexes(self) -> None: self.conn.rollback() raise - def _migrate_session_outcomes_schema(self) -> None: - """Restore the pre-identity outcome schema after downgrading #407.""" - with self._lock: - table_info = self.conn.execute( - "PRAGMA table_info(session_outcomes)" - ).fetchall() - if not table_info: - return - expected_columns = { - "user_id", - "session_id", - "outcome", - "occurred_at", - "source", - "label", - "value", - "metadata", - "governance_subject_ref", - "created_at", - } - columns = {str(row["name"]): row for row in table_info} - table = self.conn.execute( - "SELECT sql FROM sqlite_master WHERE type = 'table' " - "AND name = 'session_outcomes'" - ).fetchone() - table_sql = str(table["sql"] or "") if table is not None else "" - governance_column = columns.get("governance_subject_ref") - has_empty_subject_default = ( - governance_column is not None - and governance_column["dflt_value"] in ("''", '""') - ) - identity_columns = { - "outcome_id", - "outcome_revision", - "outcome_contract_digest", - "finalized_trajectory_digest", - } - if ( - expected_columns.issubset(columns) - and not identity_columns.intersection(columns) - and governance_column is not None - and int(governance_column["notnull"]) == 1 - and not has_empty_subject_default - and "'unknown'" not in table_sql - ): - return - - self.conn.execute("BEGIN IMMEDIATE") - try: - legacy_rows = self.conn.execute( - "SELECT * FROM session_outcomes" - ).fetchall() - self.conn.execute( - "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" - ) - self.conn.execute( - """CREATE TABLE session_outcomes ( - user_id TEXT NOT NULL, - session_id TEXT NOT NULL, - outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), - occurred_at INTEGER NOT NULL, - source TEXT NOT NULL, - label TEXT, - value REAL, - metadata TEXT, - governance_subject_ref TEXT NOT NULL, - created_at INTEGER NOT NULL, - PRIMARY KEY (user_id, session_id) - )""" - ) - compatible_rows = [ - row - for row in legacy_rows - if row["outcome"] in ("success", "failure") - ] - for row in compatible_rows: - subject_ref = ( - row["governance_subject_ref"] - if "governance_subject_ref" in columns - else None - ) - if subject_ref is None or not str(subject_ref).strip(): - subject_ref = self._subject_ref_for_user_id(str(row["user_id"])) - self.conn.execute( - """INSERT INTO session_outcomes ( - user_id, session_id, outcome, occurred_at, source, - label, value, metadata, governance_subject_ref, created_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - row["user_id"], - row["session_id"], - row["outcome"], - row["occurred_at"], - row["source"], - row["label"] if "label" in columns else None, - row["value"] if "value" in columns else None, - row["metadata"] if "metadata" in columns else None, - subject_ref, - row["created_at"], - ), - ) - dropped_unknown = len(legacy_rows) - len(compatible_rows) - if dropped_unknown: - logger.warning( - "Dropped %d unrepresentable 'unknown' session outcomes while " - "restoring the pre-#407 schema", - dropped_unknown, - ) - self.conn.execute("DROP TABLE session_outcomes_legacy") - self.conn.commit() - except Exception: - self.conn.rollback() - raise - def _try_load_sqlite_vec(self) -> bool: """Attempt to load the sqlite-vec extension for native KNN search. @@ -1237,6 +1188,139 @@ def _migrate_vec_tables(self) -> None: if emb: self._vec_upsert(vec_table, r["rid"], emb) + def _migrate_session_outcomes_schema(self) -> None: + """Rebuild legacy outcome rows with immutable v1 finalization identities.""" + with self._lock: + self.conn.execute("BEGIN IMMEDIATE") + try: + columns = { + row["name"] + for row in self.conn.execute( + "PRAGMA table_info(session_outcomes)" + ).fetchall() + } + table = self.conn.execute( + "SELECT sql FROM sqlite_master WHERE type = 'table' " + "AND name = 'session_outcomes'" + ).fetchone() + table_sql = (table["sql"] if table is not None else "") or "" + required = { + "outcome_id", + "outcome_revision", + "outcome_contract_digest", + "finalized_trajectory_digest", + } + if required.issubset(columns) and "'unknown'" in table_sql: + rows_missing_subject_ref = self.conn.execute( + """SELECT user_id, session_id FROM session_outcomes + WHERE governance_subject_ref IS NULL + OR trim(governance_subject_ref) = ''""" + ).fetchall() + self.conn.executemany( + """UPDATE session_outcomes SET governance_subject_ref = ? + WHERE user_id = ? AND session_id = ?""", + [ + ( + self._subject_ref_for_user_id(str(row["user_id"])), + row["user_id"], + row["session_id"], + ) + for row in rows_missing_subject_ref + ], + ) + self.conn.commit() + return + + legacy_rows = self.conn.execute( + "SELECT * FROM session_outcomes" + ).fetchall() + self.conn.execute( + "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" + ) + self.conn.execute( + """CREATE TABLE session_outcomes ( + outcome_id TEXT NOT NULL UNIQUE, + outcome_revision INTEGER NOT NULL CHECK (outcome_revision >= 1), + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure', 'unknown')), + occurred_at INTEGER NOT NULL, + source TEXT NOT NULL, + label TEXT, + value REAL, + metadata TEXT, + outcome_contract_digest TEXT NOT NULL, + finalized_trajectory_digest TEXT NOT NULL, + governance_subject_ref TEXT NOT NULL, + created_at INTEGER NOT NULL, + PRIMARY KEY (user_id, session_id) + )""" + ) + for row in legacy_rows: + source = str(row["source"]) + subject_ref = row["governance_subject_ref"] + if subject_ref is None or not str(subject_ref).strip(): + subject_ref = self._subject_ref_for_user_id(str(row["user_id"])) + self.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + _legacy_session_outcome_id( + str(row["user_id"]), str(row["session_id"]) + ), + 1, + row["user_id"], + row["session_id"], + row["outcome"], + row["occurred_at"], + source, + row["label"], + row["value"], + row["metadata"], + outcome_contract_digest( + source=source, + schema_version=1, + allowed_values=_OUTCOME_ALLOWED_VALUES, + finalization_rule="first_write", + ), + trajectory_digest( + _canonical_session_snapshot( + self.conn, str(row["session_id"]) + ) + ), + subject_ref, + row["created_at"], + ), + ) + self.conn.execute("DROP TABLE session_outcomes_legacy") + self.conn.execute( + "CREATE INDEX idx_session_outcomes_occurred_at " + "ON session_outcomes(occurred_at)" + ) + self.conn.execute( + "CREATE INDEX idx_session_outcomes_session_id " + "ON session_outcomes(session_id)" + ) + self.conn.execute( + "CREATE INDEX idx_session_outcomes_source_outcome " + "ON session_outcomes(source, outcome)" + ) + self.conn.execute( + "CREATE INDEX idx_session_outcomes_label ON session_outcomes(label)" + ) + self.conn.execute( + "CREATE INDEX idx_session_outcomes_subject_ref " + "ON session_outcomes(governance_subject_ref)" + ) + self.conn.commit() + except Exception: + self.conn.rollback() + raise + def _migrate_interactions_schema(self) -> None: """Add new columns to existing interactions table if missing.""" with self._lock: @@ -2274,6 +2358,9 @@ def _migrate_request_session_id_required(self) -> None: retrieval_experiment_arm_expr = ( "retrieval_experiment_arm" if "retrieval_experiment_arm" in cols else "NULL" ) + governance_subject_ref_expr = ( + "governance_subject_ref" if "governance_subject_ref" in cols else "NULL" + ) # NOTE: this rebuild hardcodes the full `requests` column set. If a # future migration adds a column to `requests`, it MUST be added here # too (and to the SELECT below) or the rebuild will silently drop it. @@ -2288,7 +2375,8 @@ def _migrate_request_session_id_required(self) -> None: session_id TEXT NOT NULL CHECK (trim(session_id) != ''), evaluation_only INTEGER NOT NULL DEFAULT 0, retrieval_experiment_id TEXT, - retrieval_experiment_arm TEXT + retrieval_experiment_arm TEXT, + governance_subject_ref TEXT ); INSERT INTO requests_new ( @@ -2300,7 +2388,8 @@ def _migrate_request_session_id_required(self) -> None: session_id, evaluation_only, retrieval_experiment_id, - retrieval_experiment_arm + retrieval_experiment_arm, + governance_subject_ref ) SELECT request_id, @@ -2315,7 +2404,8 @@ def _migrate_request_session_id_required(self) -> None: END, {evaluation_only_expr}, {retrieval_experiment_id_expr}, - {retrieval_experiment_arm_expr} + {retrieval_experiment_arm_expr}, + {governance_subject_ref_expr} FROM requests; DROP TABLE requests; ALTER TABLE requests_new RENAME TO requests; @@ -2324,6 +2414,8 @@ def _migrate_request_session_id_required(self) -> None: CREATE INDEX IF NOT EXISTS idx_requests_created_at ON requests(created_at); CREATE INDEX IF NOT EXISTS idx_requests_retrieval_experiment ON requests(retrieval_experiment_id, user_id, session_id, created_at, request_id); + CREATE INDEX IF NOT EXISTS idx_requests_governance_subject_ref + ON requests(governance_subject_ref); """ ) self.conn.commit() @@ -2800,14 +2892,18 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: ON requests(retrieval_experiment_id, user_id, session_id, created_at, request_id); CREATE TABLE IF NOT EXISTS session_outcomes ( + outcome_id TEXT NOT NULL UNIQUE, + outcome_revision INTEGER NOT NULL CHECK (outcome_revision >= 1), user_id TEXT NOT NULL, session_id TEXT NOT NULL, - outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure', 'unknown')), occurred_at INTEGER NOT NULL, source TEXT NOT NULL, label TEXT, value REAL, metadata TEXT, + outcome_contract_digest TEXT NOT NULL, + finalized_trajectory_digest TEXT NOT NULL, governance_subject_ref TEXT NOT NULL, created_at INTEGER NOT NULL, PRIMARY KEY (user_id, session_id) @@ -3187,6 +3283,14 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: CREATE INDEX IF NOT EXISTS idx_agent_runs_ready ON _agent_runs(status, next_resume_at, updated_at); CREATE INDEX IF NOT EXISTS idx_agent_runs_binding ON _agent_runs(org_id, extractor_kind, user_id); +CREATE TABLE IF NOT EXISTS _agent_run_finalization_receipts ( + run_id TEXT PRIMARY KEY, + entity_type TEXT NOT NULL CHECK (entity_type IN ('profile', 'user_playbook')), + learning_ids TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + FOREIGN KEY (run_id) REFERENCES _agent_runs(id) ON DELETE CASCADE +); + CREATE TABLE IF NOT EXISTS _pending_tool_calls ( id TEXT PRIMARY KEY, org_id TEXT NOT NULL, diff --git a/reflexio/server/services/storage/sqlite_storage/_governance.py b/reflexio/server/services/storage/sqlite_storage/_governance.py index 5e4abad09..77f8fec91 100644 --- a/reflexio/server/services/storage/sqlite_storage/_governance.py +++ b/reflexio/server/services/storage/sqlite_storage/_governance.py @@ -72,12 +72,16 @@ subject_ref TEXT, request_ref TEXT NOT NULL, idempotency_key TEXT NOT NULL, + authoritative_user_digest TEXT, status TEXT NOT NULL DEFAULT 'pending', error_code TEXT, error_detail TEXT, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, completed_at INTEGER, + execution_claim_owner TEXT, + execution_claim_fence INTEGER NOT NULL DEFAULT 0, + execution_claim_expires_at INTEGER, PRIMARY KEY (org_id, purge_id) ); CREATE UNIQUE INDEX IF NOT EXISTS idx_purge_operations_org_idem @@ -109,6 +113,29 @@ def init_governance_tables(conn: sqlite3.Connection) -> None: conn.executescript(GOVERNANCE_DDL) _enforce_audit_request_ref_not_null(conn) _ensure_governance_subject_ref_columns(conn) + _ensure_purge_operation_execution_claim_columns(conn) + + +def _ensure_purge_operation_execution_claim_columns(conn: sqlite3.Connection) -> None: + columns = [row[1] for row in conn.execute("PRAGMA table_info(purge_operations)")] + if not columns: + return + if "execution_claim_owner" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN execution_claim_owner TEXT" + ) + if "execution_claim_fence" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN execution_claim_fence INTEGER NOT NULL DEFAULT 0" + ) + if "execution_claim_expires_at" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN execution_claim_expires_at INTEGER" + ) + if "authoritative_user_digest" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN authoritative_user_digest TEXT" + ) def _ensure_governance_subject_ref_columns(conn: sqlite3.Connection) -> None: diff --git a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py index aaec5a1b2..e41fa2fa0 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -3,6 +3,7 @@ import json import sqlite3 from typing import Any +from uuid import uuid4 from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, @@ -11,12 +12,51 @@ SetSessionOutcomeRequest, ) from reflexio.server.services.storage.error import SubjectWriteBarrierError +from reflexio.server.services.storage.session_outcome_identity import ( + outcome_contract_digest, + trajectory_digest, +) from reflexio.server.services.storage.storage_base._session_outcomes import ( SessionOutcomeContext, SessionOutcomeWriteResult, ) -from ._base import SQLiteStorageBase, _iso_to_epoch +from ._base import ( + _OUTCOME_ALLOWED_VALUES, + SQLiteStorageBase, + _canonical_session_snapshot, + _iso_to_epoch, +) + +_OUTCOME_SCHEMA_VERSION = 1 +_OUTCOME_FINALIZATION_RULE = "first_write" + + +def _canonical_metadata_json(metadata: object) -> str | None: + if metadata is None: + return None + return json.dumps( + metadata, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + + +def _metadata_matches(*, stored_metadata: str | None, request_metadata: object) -> bool: + if stored_metadata is None: + stored_value = None + else: + try: + stored_value = json.loads(stored_metadata) + except (RecursionError, TypeError, ValueError): + return False + try: + return _canonical_metadata_json(stored_value) == _canonical_metadata_json( + request_metadata + ) + except (RecursionError, TypeError, ValueError): + return False class SessionOutcomeStoreMixin: @@ -67,15 +107,94 @@ def record_session_outcome( try: self.conn.execute("BEGIN IMMEDIATE") existing = self.conn.execute( - "SELECT user_id, source FROM session_outcomes WHERE session_id = ?", + "SELECT * FROM session_outcomes WHERE session_id = ?", (request.session_id,), ).fetchone() if existing is not None: + first = self.conn.execute( + """SELECT user_id, source, governance_subject_ref + FROM requests WHERE session_id = ? + ORDER BY created_at ASC, request_id ASC LIMIT 1""", + (request.session_id,), + ).fetchone() + source = ( + str(first["source"]) + if first is not None + else str(existing["source"]) + ) + subject_ref = ( + str( + first["governance_subject_ref"] + or self._subject_ref_for_user_id(str(first["user_id"])) + ) + if first is not None + else str(existing["governance_subject_ref"]) + ) + contract_digest = self._outcome_contract_digest(source) + current_snapshot_digest = ( + trajectory_digest( + _canonical_session_snapshot(self.conn, request.session_id) + ) + if first is not None + else None + ) + stored_contract_digest = existing["outcome_contract_digest"] + stored_snapshot_digest = existing["finalized_trajectory_digest"] + server_context_matches = first is None or ( + str(existing["user_id"]) == str(first["user_id"]) + and str(existing["source"]) == str(first["source"]) + and str(existing["governance_subject_ref"]) == subject_ref + ) + exact_retry = ( + existing["outcome"] == str(request.outcome) + and int(existing["occurred_at"]) == request.occurred_at + and existing["label"] == request.label + and existing["value"] == request.value + and _metadata_matches( + stored_metadata=existing["metadata"], + request_metadata=request.metadata, + ) + and server_context_matches + and ( + stored_contract_digest is None + or stored_contract_digest == contract_digest + ) + and ( + stored_snapshot_digest is None + or current_snapshot_digest is None + or stored_snapshot_digest == current_snapshot_digest + ) + ) self.conn.rollback() return SessionOutcomeWriteResult( recorded=False, user_id=str(existing["user_id"]), source=str(existing["source"]), + reason=( + None + if exact_retry + else SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + ), + outcome_id=( + str(existing["outcome_id"]) + if existing["outcome_id"] is not None + else None + ), + outcome_revision=( + int(existing["outcome_revision"]) + if existing["outcome_revision"] is not None + else None + ), + outcome_contract_digest=( + str(stored_contract_digest) + if stored_contract_digest is not None + else None + ), + finalized_trajectory_digest=( + str(stored_snapshot_digest) + if stored_snapshot_digest is not None + else None + ), ) first = self.conn.execute( """SELECT user_id, source, created_at, request_id @@ -122,31 +241,44 @@ def record_session_outcome( user_id=user_id, reason=SessionOutcomeFailureReason.SUBJECT_NOT_WRITABLE, ) + contract_digest = self._outcome_contract_digest(source) + snapshot_digest = trajectory_digest( + _canonical_session_snapshot(self.conn, request.session_id) + ) + outcome_id = uuid4().hex self.conn.execute( """INSERT INTO session_outcomes - (user_id, session_id, outcome, occurred_at, source, label, value, - metadata, governance_subject_ref, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, + governance_subject_ref, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( + outcome_id, + 1, user_id, request.session_id, - request.outcome.value, + str(request.outcome), request.occurred_at, source, request.label, request.value, - json.dumps( - request.metadata, sort_keys=True, separators=(",", ":") - ) - if request.metadata is not None - else None, + self._metadata_json(request), + contract_digest, + snapshot_digest, subject_ref, created_at, ), ) self.conn.commit() return SessionOutcomeWriteResult( - recorded=True, user_id=user_id, source=source + recorded=True, + user_id=user_id, + source=source, + outcome_id=outcome_id, + outcome_revision=1, + outcome_contract_digest=contract_digest, + finalized_trajectory_digest=snapshot_digest, ) except Exception: self.conn.rollback() @@ -179,13 +311,17 @@ def get_session_outcomes( params.append(request.end_time) where = f" WHERE {' AND '.join(clauses)}" if clauses else "" rows = self.conn.execute( - f"""SELECT user_id, session_id, outcome, occurred_at, source, label, value, - metadata, created_at FROM session_outcomes{where} + f"""SELECT outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, created_at + FROM session_outcomes{where} ORDER BY occurred_at DESC, user_id ASC, session_id ASC LIMIT ? OFFSET ?""", [*params, request.top_k, request.offset], ).fetchall() return [ SessionOutcomeRecord( + outcome_id=row["outcome_id"], + outcome_revision=row["outcome_revision"], user_id=row["user_id"], session_id=row["session_id"], outcome=row["outcome"], @@ -194,6 +330,8 @@ def get_session_outcomes( label=row["label"], value=row["value"], metadata=json.loads(row["metadata"]) if row["metadata"] else None, + outcome_contract_digest=row["outcome_contract_digest"], + finalized_trajectory_digest=row["finalized_trajectory_digest"], created_at=row["created_at"], ) for row in rows @@ -212,3 +350,16 @@ def clear_session_outcomes_for_user(self, user_id: str) -> dict[str, int]: return { "session_outcomes": int(outcome_cursor.rowcount or 0), } + + @staticmethod + def _metadata_json(request: SetSessionOutcomeRequest) -> str | None: + return _canonical_metadata_json(request.metadata) + + @staticmethod + def _outcome_contract_digest(source: str) -> str: + return outcome_contract_digest( + source=source, + schema_version=_OUTCOME_SCHEMA_VERSION, + allowed_values=_OUTCOME_ALLOWED_VALUES, + finalization_rule=_OUTCOME_FINALIZATION_RULE, + ) diff --git a/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py b/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py index 3858b29b3..299b20871 100644 --- a/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py +++ b/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py @@ -15,6 +15,7 @@ from __future__ import annotations +import json import sqlite3 from datetime import UTC, datetime, timedelta from typing import Any @@ -29,6 +30,13 @@ from .._base import SQLiteStorageBase, _json_dumps +def _valid_finalized_learning_ids(value: object) -> bool: + return isinstance(value, list) and all( + isinstance(learning_id, str) and bool(learning_id.strip()) + for learning_id in value + ) + + class SQLiteAgentRunStoreMixin: """SQLite-backed resumable extraction run store primitives.""" @@ -37,6 +45,7 @@ class SQLiteAgentRunStoreMixin: _fetchone: Any _fetchall: Any _current_timestamp: Any + _own_transaction: Any org_id: str def _finalize_runs_without_pending_dependencies_unlocked(self, now_s: str) -> None: @@ -213,6 +222,84 @@ def get_agent_run(self, run_id: str) -> AgentRunRecord | None: row = self._fetchone("SELECT * FROM _agent_runs WHERE id = ?", (run_id,)) return _row_to_agent_run(row) if row else None + @SQLiteStorageBase.handle_exceptions + def get_agent_run_finalization_receipt( + self, + *, + run_id: str, + entity_type: str, + ) -> list[str] | None: + row = self._fetchone( + """ + SELECT receipt.entity_type, receipt.learning_ids + FROM _agent_run_finalization_receipts AS receipt + JOIN _agent_runs AS run ON run.id = receipt.run_id + WHERE receipt.run_id = ? AND run.org_id = ? + """, + (run_id, self.org_id), + ) + if row is None: + return None + if row["entity_type"] != entity_type: + raise ValueError("agent-run finalization receipt entity type changed") + learning_ids = json.loads(row["learning_ids"]) + if not _valid_finalized_learning_ids(learning_ids): + raise ValueError("agent-run finalization receipt is corrupt") + return learning_ids + + @SQLiteStorageBase.handle_exceptions + def save_agent_run_finalization_receipt( + self, + *, + run_id: str, + entity_type: str, + learning_ids: list[str], + ) -> None: + expected_by_extractor = { + "profile": "profile", + "playbook": "user_playbook", + } + if not _valid_finalized_learning_ids(learning_ids): + raise ValueError( + "agent-run finalization receipt learning ids must be non-empty strings" + ) + encoded_ids = _json_dumps(learning_ids) + with self._lock: + run = self.conn.execute( + "SELECT org_id, extractor_kind FROM _agent_runs WHERE id = ?", + (run_id,), + ).fetchone() + if run is None or run["org_id"] != self.org_id: + raise ValueError("agent-run finalization receipt owner is invalid") + if expected_by_extractor.get(run["extractor_kind"]) != entity_type: + raise ValueError( + "agent-run finalization receipt entity type is invalid" + ) + self.conn.execute( + """ + INSERT OR IGNORE INTO _agent_run_finalization_receipts + (run_id, entity_type, learning_ids) + VALUES (?, ?, ?) + """, + (run_id, entity_type, encoded_ids), + ) + stored = self.conn.execute( + """ + SELECT entity_type, learning_ids + FROM _agent_run_finalization_receipts + WHERE run_id = ? + """, + (run_id,), + ).fetchone() + if ( + stored is None + or stored["entity_type"] != entity_type + or json.loads(stored["learning_ids"]) != learning_ids + ): + raise ValueError("agent-run finalization receipt is immutable") + if self._own_transaction(): + self.conn.commit() + @SQLiteStorageBase.handle_exceptions def get_latest_finalized_agent_run_for_request( self, diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py index 954eae6c1..6d1cdf726 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py @@ -39,6 +39,7 @@ AuditEvent, PurgeOperation, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, _PREPARE_PHASE, @@ -88,6 +89,11 @@ class GovernanceEraseExecutionMixin: ] get_purge_operation: Callable[[str], PurgeOperation] _record_purge_target_locked: Callable[..., None] + _assert_purge_operation_execution_claim_locked: Callable[ + [str, PurgeExecutionClaim | None], None + ] + _assert_authoritative_user_identity_locked: Callable[[str, str], str] + _assert_bound_authoritative_user_identity_locked: Callable[[str, str, str], None] def _purge_governance_entity_content_locked( self, @@ -346,7 +352,11 @@ def _delete_evaluation_operation_states_locked( return deleted def apply_governance_user_data_delete( - self, purge_id: str, user_id: str + self, + purge_id: str, + user_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> dict[str, int]: purge_id = _validate_governance_purge_id("purge_id", purge_id) name_map = { @@ -370,6 +380,10 @@ def apply_governance_user_data_delete( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) + self._assert_authoritative_user_identity_locked(purge_id, user_id) self._validate_prepared_delete_target_matrix_locked(purge_id) self._validate_hide_for_rebuild_targets_locked(purge_id) expected_user_playbook_ids = ( @@ -404,7 +418,12 @@ def apply_governance_user_data_delete( return counts def complete_purge_operation_with_audit( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) if audit_event.org_id != self.org_id: @@ -420,6 +439,9 @@ def complete_purge_operation_with_audit( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) row = self.conn.execute( "SELECT * FROM purge_operations WHERE purge_id = ? AND org_id = ?", (purge_id, self.org_id), @@ -452,6 +474,11 @@ def complete_purge_operation_with_audit( raise ValueError( "Cannot complete purge without target snapshot marker" ) + self._assert_bound_authoritative_user_identity_locked( + purge_id, + audit_event.subject_ref or "", + authoritative_user_id, + ) delete_rows = self.conn.execute( """SELECT target_name, status FROM purge_operation_targets WHERE org_id = ? AND purge_id = ? AND phase = 'delete' @@ -531,7 +558,9 @@ def complete_purge_operation_with_audit( error_code = NULL, error_detail = NULL, updated_at = ?, - completed_at = ? + completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ?""", (now, now, purge_id, self.org_id), ) diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py index 1f2a7f0e9..2087ee2c6 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py @@ -18,6 +18,8 @@ from __future__ import annotations +import hashlib +import hmac import sqlite3 import threading from collections.abc import Callable @@ -27,6 +29,11 @@ PurgeOperation, PurgeOperationTarget, ) +from reflexio.server.services.governance.config import get_governance_ref_secret +from reflexio.server.services.storage.governance_claims import ( + PurgeExecutionClaim, + validate_purge_execution_claim, +) from reflexio.server.services.storage.governance_validation import ( _ALLOWED_PURGE_OPERATION_TYPES, _ALLOWED_PURGE_SCOPE_TYPES, @@ -48,7 +55,12 @@ _validate_governance_target_ref, ) -from .._governance import _json_dumps, _row_to_purge_operation, _row_to_purge_target +from .._governance import ( + _json_dumps, + _json_loads, + _row_to_purge_operation, + _row_to_purge_target, +) if TYPE_CHECKING: from .._governance import _SQLiteGovernanceDeps @@ -65,6 +77,114 @@ class PurgeOperationStoreMixin: _deps: Callable[[], _SQLiteGovernanceDeps] _owned_user_playbook_ids_locked: Callable[[str], set[int]] _planned_governance_delete_counts: Callable[[str, set[int]], dict[str, int]] + _subject_ref_for_user_id: Callable[[str], str] + + def _authoritative_user_digest(self, purge_id: str, user_id: str) -> str: + material = f"authoritative-user-v1\0{self.org_id}\0{purge_id}\0{user_id}" + return hmac.new( + get_governance_ref_secret().encode(), + material.encode(), + hashlib.sha256, + ).hexdigest() + + @staticmethod + def _legacy_authoritative_user_digest(purge_id: str, user_id: str) -> str: + return hashlib.sha256(f"{purge_id}\0{user_id}".encode()).hexdigest() + + def _assert_authoritative_user_identity_locked( + self, purge_id: str, user_id: str + ) -> str: + row = self.conn.execute( + """SELECT operation_type, scope_type, subject_ref, + authoritative_user_digest + FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (self.org_id, purge_id), + ).fetchone() + expected_digest = self._authoritative_user_digest(purge_id, user_id) + if ( + row is None + or row["operation_type"] != "user_erasure" + or row["scope_type"] != "user" + or row["subject_ref"] != self._subject_ref_for_user_id(user_id) + or row["authoritative_user_digest"] != expected_digest + ): + raise ValueError("Purge authoritative user identity does not match") + return expected_digest + + def _adopt_authoritative_user_digest_bindings_locked( + self, + *, + purge_id: str, + user_id: str, + existing_digest: object, + authoritative_user_digest: str, + now: int, + ) -> None: + legacy_digest = self._legacy_authoritative_user_digest(purge_id, user_id) + + def is_recognized(binding: object) -> bool: + return binding is None or ( + isinstance(binding, str) + and ( + hmac.compare_digest(binding, authoritative_user_digest) + or hmac.compare_digest(binding, legacy_digest) + ) + ) + + if not is_recognized(existing_digest): + raise ValueError( + "Existing purge operation has mismatched authoritative user identity" + ) + + snapshot_row = self.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = ? + AND target_ref = 'all' AND phase = ?""", + (self.org_id, purge_id, _SNAPSHOT_TARGET_NAME, _PREPARE_PHASE), + ).fetchone() + snapshot_detail = None + if snapshot_row is not None: + snapshot_detail = _json_loads(snapshot_row["detail"]) + if not isinstance(snapshot_detail, dict) or not is_recognized( + snapshot_detail.get("authoritative_user_digest") + ): + raise ValueError( + "Existing purge snapshot has mismatched authoritative user identity" + ) + + if existing_digest != authoritative_user_digest: + self.conn.execute( + """UPDATE purge_operations + SET authoritative_user_digest = ?, updated_at = ? + WHERE org_id = ? AND purge_id = ? + AND authoritative_user_digest IS ?""", + ( + authoritative_user_digest, + now, + self.org_id, + purge_id, + existing_digest, + ), + ) + if ( + snapshot_detail is not None + and snapshot_detail.get("authoritative_user_digest") + != authoritative_user_digest + ): + snapshot_detail["authoritative_user_digest"] = authoritative_user_digest + self.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = ? + AND target_ref = 'all' AND phase = ?""", + ( + _json_dumps(snapshot_detail), + self.org_id, + purge_id, + _SNAPSHOT_TARGET_NAME, + _PREPARE_PHASE, + ), + ) def _record_purge_target_locked( self, @@ -165,6 +285,7 @@ def begin_purge_operation( scope_type: Literal["user", "org"], subject_ref: str | None, request_ref: str, + authoritative_user_id: str | None = None, ) -> PurgeOperation: _validate_governance_enum( "operation_type", @@ -187,6 +308,20 @@ def begin_purge_operation( str, _validate_governance_idempotency_key("idempotency_key", idempotency_key), ) + if operation_type == "user_erasure" and scope_type == "user": + if not authoritative_user_id: + raise ValueError("authoritative user identity is required") + if subject_ref != self._subject_ref_for_user_id(authoritative_user_id): + raise ValueError("authoritative user identity must match subject_ref") + elif authoritative_user_id: + raise ValueError( + "authoritative user identity is only valid for user erasure" + ) + authoritative_user_digest = ( + self._authoritative_user_digest(validated_purge_id, authoritative_user_id) + if authoritative_user_id + else None + ) now = _epoch_now() with self._lock: try: @@ -211,13 +346,26 @@ def begin_purge_operation( "Existing purge operation for idempotency_key has " f"mismatched {field_name}" ) - self.conn.rollback() - return existing_operation + if authoritative_user_id and authoritative_user_digest: + self._adopt_authoritative_user_digest_bindings_locked( + purge_id=validated_purge_id, + user_id=authoritative_user_id, + existing_digest=existing["authoritative_user_digest"], + authoritative_user_digest=authoritative_user_digest, + now=now, + ) + elif existing["authoritative_user_digest"] is not None: + raise ValueError( + "Existing purge operation has mismatched authoritative user identity" + ) + self.conn.commit() + return _row_to_purge_operation(existing) self.conn.execute( """INSERT INTO purge_operations ( purge_id, org_id, operation_type, scope_type, subject_ref, - request_ref, idempotency_key, status, created_at, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)""", + request_ref, idempotency_key, authoritative_user_digest, + status, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)""", ( validated_purge_id, self.org_id, @@ -226,6 +374,7 @@ def begin_purge_operation( subject_ref, request_ref, validated_idempotency_key, + authoritative_user_digest, now, now, ), @@ -236,12 +385,177 @@ def begin_purge_operation( raise return self.get_purge_operation(validated_purge_id) + def claim_purge_operation_execution( + self, + purge_id: str, + *, + lease_owner: str, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim | None: + validated_purge_id = _validate_governance_purge_id("purge_id", purge_id) + if not lease_owner.strip(): + raise ValueError("lease_owner is required") + if lease_ttl_seconds <= 0: + raise ValueError("lease_ttl_seconds must be positive") + now = _epoch_now() + expires_at = now + lease_ttl_seconds + with self._lock: + try: + self.conn.execute("BEGIN IMMEDIATE") + cursor = self.conn.execute( + """UPDATE purge_operations + SET status = 'running', error_code = NULL, error_detail = NULL, + completed_at = NULL, updated_at = ?, + execution_claim_owner = ?, + execution_claim_fence = execution_claim_fence + 1, + execution_claim_expires_at = ? + WHERE purge_id = ? AND org_id = ? + AND ( + status IN ('pending', 'failed') + OR ( + status = 'running' + AND ( + execution_claim_expires_at IS NULL + OR execution_claim_expires_at <= ? + ) + ) + ) + RETURNING execution_claim_owner, + execution_claim_fence, + execution_claim_expires_at""", + ( + now, + lease_owner, + expires_at, + validated_purge_id, + self.org_id, + now, + ), + ) + row = cursor.fetchone() + self.conn.commit() + if row is None: + return None + return PurgeExecutionClaim( + purge_id=validated_purge_id, + owner=str(row["execution_claim_owner"]), + fence=int(row["execution_claim_fence"]), + expires_at=int(row["execution_claim_expires_at"]), + ) + except Exception: + self.conn.rollback() + raise + + def assert_purge_operation_execution_claim( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> None: + purge_id = _validate_governance_purge_id("purge_id", purge_id) + claim = validate_purge_execution_claim(purge_id, execution_claim) + now = _epoch_now() + row = self._deps()._fetchone( + """SELECT status, execution_claim_owner, execution_claim_fence, + execution_claim_expires_at + FROM purge_operations + WHERE purge_id = ? AND org_id = ?""", + (purge_id, self.org_id), + ) + if row is None: + raise ValueError(f"Purge operation {purge_id!r} not found") + if ( + row["status"] != "running" + or row["execution_claim_owner"] != claim.owner + or int(row["execution_claim_fence"]) != claim.fence + or row["execution_claim_expires_at"] is None + or int(row["execution_claim_expires_at"]) <= now + ): + raise ValueError("purge execution claim is no longer active") + + def _assert_purge_operation_execution_claim_locked( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + ) -> None: + purge_id = _validate_governance_purge_id("purge_id", purge_id) + claim = validate_purge_execution_claim(purge_id, execution_claim) + now = _epoch_now() + row = self.conn.execute( + """SELECT status, execution_claim_owner, execution_claim_fence, + execution_claim_expires_at + FROM purge_operations + WHERE purge_id = ? AND org_id = ?""", + (purge_id, self.org_id), + ).fetchone() + if row is None: + raise ValueError(f"Purge operation {purge_id!r} not found") + if ( + row["status"] != "running" + or row["execution_claim_owner"] != claim.owner + or int(row["execution_claim_fence"]) != claim.fence + or row["execution_claim_expires_at"] is None + or int(row["execution_claim_expires_at"]) <= now + ): + raise ValueError("purge execution claim is no longer active") + + def renew_purge_operation_execution_claim( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + *, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim: + purge_id = _validate_governance_purge_id("purge_id", purge_id) + claim = validate_purge_execution_claim(purge_id, execution_claim) + if lease_ttl_seconds <= 0: + raise ValueError("lease_ttl_seconds must be positive") + now = _epoch_now() + expires_at = now + lease_ttl_seconds + with self._lock: + try: + self.conn.execute("BEGIN IMMEDIATE") + cursor = self.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = ?, updated_at = ? + WHERE purge_id = ? AND org_id = ? + AND status = 'running' + AND execution_claim_owner = ? + AND execution_claim_fence = ? + AND execution_claim_expires_at IS NOT NULL + AND execution_claim_expires_at > ? + RETURNING execution_claim_owner, + execution_claim_fence, + execution_claim_expires_at""", + ( + expires_at, + now, + purge_id, + self.org_id, + claim.owner, + claim.fence, + now, + ), + ) + row = cursor.fetchone() + self.conn.commit() + except Exception: + self.conn.rollback() + raise + if row is None: + raise ValueError("purge execution claim is no longer active") + return PurgeExecutionClaim( + purge_id=purge_id, + owner=str(row["execution_claim_owner"]), + fence=int(row["execution_claim_fence"]), + expires_at=int(row["execution_claim_expires_at"]), + ) + def record_purge_target( self, purge_id: str, target_name: str, phase: str, status: Literal["pending", "running", "failed", "complete"], + *, + execution_claim: PurgeExecutionClaim, target_ref: str = "", detail: dict[str, object] | None = None, deleted_count: int = 0, @@ -265,6 +579,10 @@ def record_purge_target( ) with self._lock: try: + self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) self._record_purge_target_locked( purge_id=purge_id, target_name=target_name, @@ -308,14 +626,28 @@ def prepare_governance_erase_targets( self, purge_id: str, user_id: str, + *, + execution_claim: PurgeExecutionClaim, owned_user_playbook_ids: set[int] | None = None, ) -> None: purge_id = _validate_governance_purge_id("purge_id", purge_id) with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") - if self.purge_targets_prepared(purge_id): - self.conn.rollback() + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) + authoritative_user_digest = ( + self._assert_authoritative_user_identity_locked(purge_id, user_id) + ) + prepared = self.conn.execute( + """SELECT 1 FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = ? AND target_ref = 'all' + AND phase = ? AND status = 'complete'""", + (self.org_id, purge_id, _SNAPSHOT_TARGET_NAME, _PREPARE_PHASE), + ).fetchone() + if prepared is not None: + self.conn.commit() return owned_user_playbook_ids = ( set(owned_user_playbook_ids) @@ -344,6 +676,7 @@ def prepare_governance_erase_targets( phase=_PREPARE_PHASE, status="complete", detail={ + "authoritative_user_digest": authoritative_user_digest, "owned_user_playbook_ids": sorted(owned_user_playbook_ids), }, deleted_count=0, @@ -355,7 +688,12 @@ def prepare_governance_erase_targets( raise def fail_purge_operation( - self, purge_id: str, error_code: str, error_detail: str + self, + purge_id: str, + error_code: str, + error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) validated_error_code = _validate_governance_error_code(error_code) @@ -364,10 +702,15 @@ def fail_purge_operation( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) cur = self.conn.execute( """UPDATE purge_operations SET status = 'failed', error_code = ?, error_detail = ?, - updated_at = ?, completed_at = ? + updated_at = ?, completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ? AND status != 'complete'""", ( validated_error_code, diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py b/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py index 3390533f0..da3506a99 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py @@ -38,6 +38,7 @@ governance_subject_ref, ) from reflexio.server.services.storage.error import SubjectWriteBarrierError +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, _PREPARE_PHASE, @@ -53,6 +54,7 @@ ) from .._governance import ( + _json_loads, _row_to_audit_event, _row_to_purge_operation, _row_to_subject_write_barrier, @@ -75,6 +77,10 @@ class SubjectBarrierMixin: [sqlite3.Connection | sqlite3.Cursor, AuditEvent], bool ] get_purge_operation: Callable[[str], PurgeOperation] + _assert_purge_operation_execution_claim_locked: Callable[ + [str, PurgeExecutionClaim | None], None + ] + _authoritative_user_digest: Callable[[str, str], str] def _barrier_from_purge( self, @@ -158,7 +164,59 @@ def _legacy_user_id_rows_remain_locked( return True return False - def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: + def _authoritative_user_session_outcome_remains_locked(self, user_id: str) -> bool: + return ( + self.conn.execute( + "SELECT 1 FROM session_outcomes WHERE user_id = ? LIMIT 1", + (user_id,), + ).fetchone() + is not None + ) + + def _assert_bound_authoritative_user_identity_locked( + self, purge_id: str, subject_ref: str, authoritative_user_id: str + ) -> None: + purge_row = self.conn.execute( + """SELECT operation_type, scope_type, subject_ref, + authoritative_user_digest + FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (self.org_id, purge_id), + ).fetchone() + if purge_row is None: + raise ValueError(f"Purge operation {purge_id!r} not found") + if ( + purge_row["operation_type"] != "user_erasure" + or purge_row["scope_type"] != "user" + ): + raise ValueError("Completion requires a user erasure purge") + authoritative_user_digest = purge_row["authoritative_user_digest"] + snapshot_row = self.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = ? + AND target_ref = 'all' AND phase = ? AND status = 'complete'""", + (self.org_id, purge_id, _SNAPSHOT_TARGET_NAME, _PREPARE_PHASE), + ).fetchone() + snapshot_detail = ( + _json_loads(snapshot_row["detail"]) if snapshot_row is not None else None + ) + expected_digest = self._authoritative_user_digest( + purge_id, authoritative_user_id + ) + if ( + purge_row["scope_type"] != "user" + or purge_row["subject_ref"] != subject_ref + or self._subject_ref_for_user_id(authoritative_user_id) != subject_ref + or not isinstance(authoritative_user_digest, str) + or authoritative_user_digest != expected_digest + or not isinstance(snapshot_detail, dict) + or snapshot_detail.get("authoritative_user_digest") != expected_digest + ): + raise ValueError("Purge authoritative user identity does not match") + + def _same_subject_rows_remain_locked( + self, subject_ref: str, authoritative_user_id: str + ) -> bool: legacy_request_ids = self._legacy_request_ids_for_subject_locked(subject_ref) for table in ( "requests", @@ -167,7 +225,6 @@ def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: "user_playbooks", "agent_success_evaluation_result", "retrieved_learning_evaluation", - "session_outcomes", ): row = self.conn.execute( f"""SELECT 1 FROM {table} @@ -179,6 +236,10 @@ def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: return True if legacy_request_ids: return True + if self._authoritative_user_session_outcome_remains_locked( + authoritative_user_id + ): + return True if self._legacy_user_id_rows_remain_locked( table="interactions", subject_ref=subject_ref, @@ -211,7 +272,11 @@ def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: ) def begin_subject_erasure_barrier( - self, subject_ref: str, purge_id: str + self, + subject_ref: str, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: _validate_governance_prefixed_ref( "subject_ref", subject_ref, prefix="subref_v1_" @@ -221,6 +286,9 @@ def begin_subject_erasure_barrier( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + validated_purge_id, execution_claim + ) purge_row = self.conn.execute( """SELECT * FROM purge_operations WHERE purge_id = ? AND org_id = ?""", @@ -293,7 +361,12 @@ def assert_subject_writable(self, subject_ref: str) -> None: raise def complete_subject_erasure_barrier_after_empty_check( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) if audit_event.org_id != self.org_id: @@ -309,6 +382,9 @@ def complete_subject_erasure_barrier_after_empty_check( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) row = self.conn.execute( "SELECT * FROM purge_operations WHERE purge_id = ? AND org_id = ?", (purge_id, self.org_id), @@ -334,7 +410,14 @@ def complete_subject_erasure_barrier_after_empty_check( raise ValueError( "Cannot complete purge without target snapshot marker" ) - if self._same_subject_rows_remain_locked(audit_event.subject_ref or ""): + self._assert_bound_authoritative_user_identity_locked( + purge_id, + audit_event.subject_ref or "", + authoritative_user_id, + ) + if self._same_subject_rows_remain_locked( + audit_event.subject_ref or "", authoritative_user_id + ): raise ValueError("same-subject rows remain") delete_rows = self.conn.execute( """SELECT target_name, status FROM purge_operation_targets @@ -424,7 +507,9 @@ def complete_subject_erasure_barrier_after_empty_check( error_code = NULL, error_detail = NULL, updated_at = ?, - completed_at = ? + completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ?""", (now, now, purge_id, self.org_id), ) @@ -440,6 +525,8 @@ def fail_subject_erasure_barrier( purge_id: str, error_code: str, error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: _validate_governance_prefixed_ref( "subject_ref", subject_ref, prefix="subref_v1_" @@ -451,6 +538,9 @@ def fail_subject_erasure_barrier( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + validated_purge_id, execution_claim + ) update_cursor = self.conn.execute( """UPDATE subject_write_barriers SET status = 'failed', @@ -480,7 +570,9 @@ def fail_subject_erasure_barrier( self.conn.execute( """UPDATE purge_operations SET status = 'failed', error_code = ?, error_detail = ?, - updated_at = ?, completed_at = ? + updated_at = ?, completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ?""", ( validated_error_code, diff --git a/reflexio/server/services/storage/storage_base/__init__.py b/reflexio/server/services/storage/storage_base/__init__.py index 4a790ee8d..2ac8a0f6b 100644 --- a/reflexio/server/services/storage/storage_base/__init__.py +++ b/reflexio/server/services/storage/storage_base/__init__.py @@ -147,7 +147,7 @@ def _partition_purge_vs_delete( def clear_user_data(self, user_id: str) -> dict[str, int]: """Delete all rows scoped to a single ``user_id``. - Removes the user's interactions, session outcomes, user playbooks, + Removes the user's session outcomes, interactions, user playbooks, profiles, and requests. Intentionally does NOT touch ``agent_playbooks`` — those are the cross-project rollup of skills and have no ``user_id`` column. This is the data-isolation primitive used diff --git a/reflexio/server/services/storage/storage_base/_session_outcomes.py b/reflexio/server/services/storage/storage_base/_session_outcomes.py index d417e6065..14695d9a3 100644 --- a/reflexio/server/services/storage/storage_base/_session_outcomes.py +++ b/reflexio/server/services/storage/storage_base/_session_outcomes.py @@ -18,6 +18,24 @@ class SessionOutcomeWriteResult: source: str | None = None reason: SessionOutcomeFailureReason | None = None context_changed: bool = False + outcome_id: str | None = None + outcome_revision: int | None = None + outcome_contract_digest: str | None = None + finalized_trajectory_digest: str | None = None + + def __post_init__(self) -> None: + identity = ( + self.outcome_id, + self.outcome_revision, + self.outcome_contract_digest, + self.finalized_trajectory_digest, + ) + if any(value is None for value in identity) and not all( + value is None for value in identity + ): + raise ValueError( + "outcome identity fields must be all populated or all null" + ) @dataclass(frozen=True) diff --git a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py index b8bba0371..39edaf1fd 100644 --- a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py +++ b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py @@ -24,6 +24,25 @@ def create_agent_run(self, record: AgentRunRecord) -> AgentRunRecord: def get_agent_run(self, run_id: str) -> AgentRunRecord | None: raise NotImplementedError(f"{type(self).__name__} does not support agent runs") + def get_agent_run_finalization_receipt( + self, + *, + run_id: str, + entity_type: str, + ) -> list[str] | None: + """Return persisted learning ids for a completed run finalization.""" + raise NotImplementedError(f"{type(self).__name__} does not support agent runs") + + def save_agent_run_finalization_receipt( + self, + *, + run_id: str, + entity_type: str, + learning_ids: list[str], + ) -> None: + """Persist an immutable run-to-learning binding with the learning writes.""" + raise NotImplementedError(f"{type(self).__name__} does not support agent runs") + def get_latest_finalized_agent_run_for_request( self, *, diff --git a/reflexio/server/services/storage/storage_base/governance/_erase_execution.py b/reflexio/server/services/storage/storage_base/governance/_erase_execution.py index cc4016f72..4650f726a 100644 --- a/reflexio/server/services/storage/storage_base/governance/_erase_execution.py +++ b/reflexio/server/services/storage/storage_base/governance/_erase_execution.py @@ -6,6 +6,7 @@ AuditEvent, PurgeOperation, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class GovernanceEraseExecutionMixin(ABC): @@ -19,12 +20,21 @@ class GovernanceEraseExecutionMixin(ABC): @abstractmethod def apply_governance_user_data_delete( - self, purge_id: str, user_id: str + self, + purge_id: str, + user_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> dict[str, int]: raise NotImplementedError @abstractmethod def complete_purge_operation_with_audit( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: raise NotImplementedError diff --git a/reflexio/server/services/storage/storage_base/governance/_purge.py b/reflexio/server/services/storage/storage_base/governance/_purge.py index ef63fdc2c..f3f1b4620 100644 --- a/reflexio/server/services/storage/storage_base/governance/_purge.py +++ b/reflexio/server/services/storage/storage_base/governance/_purge.py @@ -7,6 +7,7 @@ PurgeOperation, PurgeOperationTarget, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class PurgeOperationStoreMixin(ABC): @@ -27,9 +28,39 @@ def begin_purge_operation( scope_type: Literal["user", "org"], subject_ref: str | None, request_ref: str, + authoritative_user_id: str | None = None, ) -> PurgeOperation: raise NotImplementedError + @abstractmethod + def claim_purge_operation_execution( + self, + purge_id: str, + *, + lease_owner: str, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim | None: + """Atomically claim or take over a stale purge execution.""" + raise NotImplementedError + + @abstractmethod + def assert_purge_operation_execution_claim( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> None: + """Raise when the purge execution claim no longer owns the live fence.""" + raise NotImplementedError + + @abstractmethod + def renew_purge_operation_execution_claim( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + *, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim: + """Atomically renew an active purge execution claim.""" + raise NotImplementedError + @abstractmethod def record_purge_target( self, @@ -37,6 +68,8 @@ def record_purge_target( target_name: str, phase: str, status: Literal["pending", "running", "failed", "complete"], + *, + execution_claim: PurgeExecutionClaim, target_ref: str = "", detail: dict[str, object] | None = None, deleted_count: int = 0, @@ -59,13 +92,20 @@ def prepare_governance_erase_targets( self, purge_id: str, user_id: str, + *, + execution_claim: PurgeExecutionClaim, owned_user_playbook_ids: set[int] | None = None, ) -> None: raise NotImplementedError @abstractmethod def fail_purge_operation( - self, purge_id: str, error_code: str, error_detail: str + self, + purge_id: str, + error_code: str, + error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: raise NotImplementedError diff --git a/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py b/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py index 7752ec5c9..2de2f5a26 100644 --- a/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py +++ b/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py @@ -7,6 +7,7 @@ PurgeOperation, SubjectWriteBarrier, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class SubjectBarrierMixin(ABC): @@ -20,7 +21,11 @@ class SubjectBarrierMixin(ABC): @abstractmethod def begin_subject_erasure_barrier( - self, subject_ref: str, purge_id: str + self, + subject_ref: str, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: raise NotImplementedError @@ -30,7 +35,12 @@ def assert_subject_writable(self, subject_ref: str) -> None: @abstractmethod def complete_subject_erasure_barrier_after_empty_check( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: raise NotImplementedError @@ -41,6 +51,8 @@ def fail_subject_erasure_barrier( purge_id: str, error_code: str, error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: raise NotImplementedError diff --git a/tests/client/test_session_outcomes_client.py b/tests/client/test_session_outcomes_client.py index 662a7436b..9b72405e1 100644 --- a/tests/client/test_session_outcomes_client.py +++ b/tests/client/test_session_outcomes_client.py @@ -3,6 +3,9 @@ import inspect from typing import Any +import pytest +from pydantic import ValidationError + from reflexio import ReflexioClient @@ -53,3 +56,23 @@ def test_get_session_outcomes_has_no_untyped_filter_kwargs() -> None: parameter.kind is not inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() ) + + +def test_mark_session_outcome_rejects_partial_server_identity(monkeypatch) -> None: + client = ReflexioClient(api_key="test-key", url_endpoint="http://localhost:8000") + monkeypatch.setattr( + client, + "_make_request", + lambda *_args, **_kwargs: { + "success": True, + "recorded": False, + "outcome_id": "partial-outcome", + }, + ) + + with pytest.raises(ValidationError, match="all populated or all null"): + client.mark_session_outcome( + session_id="session-1", + outcome="success", + occurred_at=1, + ) diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py new file mode 100644 index 000000000..41e55df8f --- /dev/null +++ b/tests/models/test_session_outcome_identity.py @@ -0,0 +1,336 @@ +from datetime import UTC, datetime +from hashlib import sha256 +from itertools import combinations +from typing import Any + +import pytest +from pydantic import ValidationError + +from reflexio.models.api_schema.domain.entities import ( + GetSessionOutcomesResponse, + SessionOutcomeRecord, + SetSessionOutcomeResponse, +) +from reflexio.models.api_schema.domain.enums import SessionOutcomeKind +from reflexio.server.services.storage.session_outcome_identity import ( + canonical_json_bytes, + canonical_session_trajectory, + outcome_contract_digest, + trajectory_digest, +) +from reflexio.server.services.storage.storage_base import SessionOutcomeWriteResult + + +def _outcome_contract_digest(**changes: object) -> str: + payload: dict[str, object] = { + "source": "customer_webhook", + "schema_version": 1, + "allowed_values": ("success", "failure", "unknown"), + "finalization_rule": "first_write", + } + payload.update(changes) + return outcome_contract_digest(**payload) # type: ignore[arg-type] + + +def test_canonical_json_bytes_ignores_object_key_order() -> None: + assert canonical_json_bytes({"b": [2, {"d": 4, "c": 3}], "a": 1}) == ( + b'{"a":1,"b":[2,{"c":3,"d":4}]}' + ) + + +def test_outcome_contract_digest_changes_when_source_changes() -> None: + assert _outcome_contract_digest( + source="customer_webhook" + ) != _outcome_contract_digest(source="customer_batch") + + +def test_outcome_contract_digest_changes_when_schema_version_changes() -> None: + assert _outcome_contract_digest(schema_version=1) != _outcome_contract_digest( + schema_version=2 + ) + + +def test_outcome_contract_digest_normalizes_allowed_value_order() -> None: + assert _outcome_contract_digest( + allowed_values=("success", "failure", "unknown") + ) == _outcome_contract_digest(allowed_values=("unknown", "success", "failure")) + + +def test_outcome_contract_digest_changes_when_allowed_values_change() -> None: + assert _outcome_contract_digest( + allowed_values=("success", "failure", "unknown") + ) != _outcome_contract_digest(allowed_values=("success", "failure")) + + +def test_outcome_contract_digest_changes_when_finalization_rule_changes() -> None: + assert _outcome_contract_digest( + finalization_rule="first_write" + ) != _outcome_contract_digest(finalization_rule="replaceable") + + +def test_trajectory_digest_changes_when_trajectory_data_changes() -> None: + assert trajectory_digest({"messages": [{"role": "user", "content": "one"}]}) != ( + trajectory_digest({"messages": [{"role": "user", "content": "two"}]}) + ) + + +def test_canonical_session_trajectory_normalizes_sqlite_and_postgres_rows() -> None: + sqlite_request = { + "request_id": "parity-request", + "user_id": "parity-user", + "created_at": "2023-11-14T22:13:20+00:00", + "source": "parity-source", + "agent_version": "parity-agent", + "session_id": "parity-session", + "evaluation_only": 0, + "retrieval_experiment_id": None, + "retrieval_experiment_arm": None, + } + postgres_request = { + **sqlite_request, + "created_at": datetime(2023, 11, 14, 22, 13, 20, tzinfo=UTC), + "evaluation_only": False, + } + sqlite_interaction = { + "interaction_id": 4242, + "user_id": "parity-user", + "request_id": "parity-request", + "created_at": "2023-11-14T22:13:21+00:00", + "content": "Parity trajectory", + "role": "User", + "token_count": 3, + "user_action": "none", + "user_action_description": "", + "interacted_image_url": "", + "image_encoding": "", + "shadow_content": "", + "expert_content": "", + "tools_used": '[{"tool_data":{"confidence":0.75},"tool_name":"rank"}]', + "citations": "[]", + "retrieved_learnings": "[]", + } + postgres_interaction = { + **sqlite_interaction, + "created_at": datetime(2023, 11, 14, 22, 13, 21, tzinfo=UTC), + "tools_used": [{"tool_name": "rank", "tool_data": {"confidence": 0.75}}], + "citations": [], + "retrieved_learnings": [], + } + + sqlite_projection = canonical_session_trajectory( + "parity-session", + [sqlite_request], + {"parity-request": [sqlite_interaction]}, + ) + postgres_projection = canonical_session_trajectory( + "parity-session", + [postgres_request], + {"parity-request": [postgres_interaction]}, + ) + + assert sqlite_projection == postgres_projection + assert sqlite_projection["requests"][0]["request"]["evaluation_only"] is False + assert trajectory_digest(sqlite_projection) == ( + "73d0f738bb5a3c7668787c678230c4758b68923e12a16e3851b401db780e4272" + ) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_trajectory_digest_rejects_non_finite_nested_floats(value: float) -> None: + with pytest.raises(ValueError, match="Out of range float values"): + trajectory_digest({"nested": [{"value": value}]}) + + +def test_session_outcome_record_accepts_unknown_and_serializes_identities() -> None: + record = SessionOutcomeRecord( + outcome_id="outcome-1", + outcome_revision=1, + user_id="user-1", + session_id="session-1", + outcome=SessionOutcomeKind.UNKNOWN, + occurred_at=1, + source="customer_webhook", + outcome_contract_digest="a" * 64, + finalized_trajectory_digest="b" * 64, + created_at=2, + ) + + response = GetSessionOutcomesResponse(success=True, session_outcomes=[record]) + + assert response.model_dump(mode="json")["session_outcomes"] == [ + { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "user_id": "user-1", + "session_id": "session-1", + "outcome": "unknown", + "occurred_at": 1, + "source": "customer_webhook", + "label": None, + "value": None, + "metadata": None, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + "created_at": 2, + } + ] + + +def test_session_outcome_record_accepts_legacy_all_null_identity() -> None: + record = SessionOutcomeRecord( + outcome_id=None, + outcome_revision=None, + user_id="legacy-user", + session_id="legacy-session", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source="customer_webhook", + outcome_contract_digest=None, + finalized_trajectory_digest=None, + created_at=2, + ) + + assert record.outcome_id is None + assert record.outcome_revision is None + assert record.outcome_contract_digest is None + assert record.finalized_trajectory_digest is None + + +def test_session_outcome_record_rejects_partial_legacy_identity() -> None: + with pytest.raises(ValidationError, match="all populated or all null"): + SessionOutcomeRecord( + outcome_id="outcome-1", + outcome_revision=None, + user_id="legacy-user", + session_id="legacy-session", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source="customer_webhook", + outcome_contract_digest=None, + finalized_trajectory_digest=None, + created_at=2, + ) + + +@pytest.mark.parametrize( + "populated_fields", + [ + fields + for populated_count in range(1, 4) + for fields in combinations( + ( + "outcome_id", + "outcome_revision", + "outcome_contract_digest", + "finalized_trajectory_digest", + ), + populated_count, + ) + ], +) +def test_session_outcome_write_result_rejects_partial_identity( + populated_fields: tuple[str, ...], +) -> None: + identity = { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + } + + with pytest.raises(ValueError, match="all populated or all null"): + SessionOutcomeWriteResult( + recorded=False, + **{ + field_name: value + for field_name, value in identity.items() + if field_name in populated_fields + }, + ) + + +@pytest.mark.parametrize( + "populated_fields", + [ + fields + for populated_count in range(1, 4) + for fields in combinations( + ( + "outcome_id", + "outcome_revision", + "outcome_contract_digest", + "finalized_trajectory_digest", + ), + populated_count, + ) + ], +) +def test_set_session_outcome_response_rejects_partial_identity( + populated_fields: tuple[str, ...], +) -> None: + identity = { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + } + + with pytest.raises(ValidationError, match="all populated or all null"): + SetSessionOutcomeResponse( + success=True, + **{ + field_name: value + for field_name, value in identity.items() + if field_name in populated_fields + }, + ) + + +@pytest.mark.parametrize( + "identity", + [ + {}, + { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + }, + ], +) +def test_write_result_and_response_accept_complete_identity_shapes( + identity: dict[str, Any], +) -> None: + write_result = SessionOutcomeWriteResult(recorded=False, **identity) + response = SetSessionOutcomeResponse(success=True, **identity) + + assert write_result.outcome_id == response.outcome_id + + +@pytest.mark.parametrize( + "field_name", ["outcome_contract_digest", "finalized_trajectory_digest"] +) +@pytest.mark.parametrize("digest", ["A" * 64, "a" * 63, "not-a-digest"]) +def test_session_outcome_record_rejects_non_sha256_identity_digests( + field_name: str, digest: str +) -> None: + payload: dict[str, object] = { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "user_id": "user-1", + "session_id": "session-1", + "outcome": "success", + "occurred_at": 1, + "source": "customer_webhook", + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + "created_at": 2, + } + payload[field_name] = digest + + with pytest.raises(ValidationError, match="lowercase SHA-256 hex"): + SessionOutcomeRecord(**payload) # type: ignore[arg-type] + + +def test_trajectory_digest_matches_sha256_of_canonical_json() -> None: + assert trajectory_digest({"b": 2, "a": 1}) == sha256(b'{"a":1,"b":2}').hexdigest() diff --git a/tests/server/api_endpoints/test_session_outcomes_integration.py b/tests/server/api_endpoints/test_session_outcomes_integration.py index 5471dd743..1b3d408a9 100644 --- a/tests/server/api_endpoints/test_session_outcomes_integration.py +++ b/tests/server/api_endpoints/test_session_outcomes_integration.py @@ -54,13 +54,18 @@ def test_source_is_derived_from_tiebroken_first_request( ) assert response.status_code == 200 - assert response.json() == { + body = response.json() + assert { "success": True, "recorded": True, "message": "Outcome recorded", "user_id": "u1", "source": "canonical-source", - } + }.items() <= body.items() + assert body["outcome_id"] + assert body["outcome_revision"] == 1 + assert len(body["outcome_contract_digest"]) == 64 + assert len(body["finalized_trajectory_digest"]) == 64 assert "stripped unknown fields: source" in caplog.text assert "multiple sources for session source-session" in caplog.text @@ -105,6 +110,169 @@ def test_retry_survives_ordinary_session_deletion( assert retry.json()["source"] == "published" +def test_conflicting_retry_is_not_accepted_only_by_session_identity( + client_with_org: tuple[TestClient, str], +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + storage.add_request( + Request( + request_id="conflict-r1", + user_id="u1", + session_id="conflict-session", + source="published", + created_at=100, + ) + ) + first = client.post( + "/api/session_outcome", + json={ + "session_id": "conflict-session", + "outcome": "success", + "occurred_at": 101, + }, + ) + conflict = client.post( + "/api/session_outcome", + json={ + "session_id": "conflict-session", + "outcome": "failure", + "occurred_at": 101, + }, + ) + + assert first.status_code == 200 + assert first.json()["recorded"] is True + assert conflict.status_code == 200 + assert conflict.json()["success"] is False + assert conflict.json()["recorded"] is False + assert conflict.json()["reason"] == "conflicting_finalization" + + +def test_retry_compares_metadata_as_json_values( + client_with_org: tuple[TestClient, str], +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + storage.add_request( + Request( + request_id="semantic-metadata-r1", + user_id="u1", + session_id="semantic-metadata-session", + source="published", + created_at=100, + ) + ) + payload = { + "session_id": "semantic-metadata-session", + "outcome": "success", + "occurred_at": 101, + "metadata": {"label": "same", "nested": {"one": 1, "two": 2}}, + } + first = client.post("/api/session_outcome", json=payload) + assert first.json()["recorded"] is True + storage.conn.execute( # type: ignore[attr-defined] + "UPDATE session_outcomes SET metadata = ? WHERE session_id = ?", + ( + '{ "nested": { "two": 2, "one": 1 }, "label": "same" }', + "semantic-metadata-session", + ), + ) + storage.conn.commit() # type: ignore[attr-defined] + + retry = client.post("/api/session_outcome", json=payload) + + assert retry.status_code == 200 + assert retry.json()["success"] is True + assert retry.json()["recorded"] is False + assert "reason" not in retry.json() + + +def test_retry_rejects_metadata_with_different_json_value_types( + client_with_org: tuple[TestClient, str], +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + storage.add_request( + Request( + request_id="typed-metadata-r1", + user_id="u1", + session_id="typed-metadata-session", + source="published", + created_at=100, + ) + ) + first = client.post( + "/api/session_outcome", + json={ + "session_id": "typed-metadata-session", + "outcome": "success", + "occurred_at": 101, + "metadata": {"nested": {"value": True}}, + }, + ) + assert first.json()["recorded"] is True + + retry = client.post( + "/api/session_outcome", + json={ + "session_id": "typed-metadata-session", + "outcome": "success", + "occurred_at": 101, + "metadata": {"nested": {"value": 1}}, + }, + ) + + assert retry.status_code == 200 + assert retry.json()["success"] is False + assert retry.json()["recorded"] is False + assert retry.json()["reason"] == "conflicting_finalization" + + +@pytest.mark.parametrize( + "stored_metadata", + [ + pytest.param("", id="empty"), + pytest.param("{malformed", id="malformed"), + pytest.param("[" * 10_000 + "]" * 10_000, id="pathological-nesting"), + ], +) +def test_retry_rejects_invalid_stored_metadata( + client_with_org: tuple[TestClient, str], stored_metadata: str +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + session_id = f"invalid-stored-metadata-{len(stored_metadata)}" + storage.add_request( + Request( + request_id=f"{session_id}-r1", + user_id="u1", + session_id=session_id, + source="published", + created_at=100, + ) + ) + payload = { + "session_id": session_id, + "outcome": "success", + "occurred_at": 101, + } + first = client.post("/api/session_outcome", json=payload) + assert first.json()["recorded"] is True + storage.conn.execute( # type: ignore[attr-defined] + "UPDATE session_outcomes SET metadata = ? WHERE session_id = ?", + (stored_metadata, session_id), + ) + storage.conn.commit() # type: ignore[attr-defined] + + retry = client.post("/api/session_outcome", json=payload) + + assert retry.status_code == 200 + assert retry.json()["success"] is False + assert retry.json()["recorded"] is False + assert retry.json()["reason"] == "conflicting_finalization" + + def test_outcome_validation_boundaries( client_with_org: tuple[TestClient, str], ) -> None: diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py new file mode 100644 index 000000000..94cef8c68 --- /dev/null +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -0,0 +1,153 @@ +"""Search-route contract for synchronous user-playbook exposure recording.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import MagicMock, patch + +from fastapi.testclient import TestClient + +from reflexio.models.api_schema.domain import UserPlaybook +from reflexio.models.config_schema import Config, StorageConfigSQLite +from reflexio.server.api import create_app +from reflexio.server.extensions import register_service +from reflexio.server.services.search_exposure import SEARCH_EXPOSURE_RECORDER + + +def _playbook(playbook_id: int, content: str) -> UserPlaybook: + return UserPlaybook( + user_playbook_id=playbook_id, + user_id="user-1", + agent_version="agent-v1", + request_id=f"source-{playbook_id}", + playbook_name=f"Playbook {playbook_id}", + created_at=1_700_000_000 + playbook_id, + content=content, + trigger=f"Trigger {playbook_id}", + tags=["support"], + ) + + +@contextmanager +def _search_results(playbooks: list[UserPlaybook]) -> Iterator[None]: + reflexio = MagicMock() + reflexio.request_context.configurator.get_config.return_value = Config( + storage_config=StorageConfigSQLite() + ) + result = MagicMock( + success=True, + profiles=[], + agent_playbooks=[], + user_playbooks=playbooks, + reformulated_query=None, + msg="OK", + agent_trace=None, + rehydrated_text=None, + ) + reflexio.unified_search.return_value = result + with patch( + "reflexio.server.routes.search.reflexio_cache.get_reflexio", + return_value=reflexio, + ): + yield + + +def _client() -> TestClient: + return TestClient( + create_app( + get_org_id=lambda: "org-1", + get_caller_type=lambda: "production_agent", + ), + raise_server_exceptions=False, + ) + + +@dataclass +class _Recorder: + batches: list[Any] = field(default_factory=list) + completed: bool = False + + def record(self, batch: Any) -> None: + self.batches.append(batch) + self.completed = True + + +def test_unified_search_records_the_final_user_playbook_set_before_return() -> None: + playbooks = [_playbook(11, "First"), _playbook(12, "Second")] + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results(playbooks): + response = _client().post( + "/api/search", + json={ + "query": "answer", + "user_id": "user-1", + "request_id": "request-1", + "session_id": "session-1", + "interaction_id": 41, + }, + ) + + assert response.status_code == 200, response.text + assert recorder.completed is True + assert len(recorder.batches) == 1 + batch = recorder.batches[0] + assert batch.org_id == "org-1" + assert batch.request_id == "request-1" + assert batch.session_id == "session-1" + assert batch.interaction_id == 41 + assert batch.user_id == "user-1" + assert batch.user_playbooks == tuple(playbooks) + + +def test_recorder_failure_prevents_a_successful_search_response() -> None: + class _FailingRecorder: + def record(self, _batch: Any) -> None: + raise RuntimeError("ledger unavailable") + + register_service(SEARCH_EXPOSURE_RECORDER, _FailingRecorder()) + + with _search_results([_playbook(11, "First")]): + response = _client().post( + "/api/search", + json={ + "query": "answer", + "user_id": "user-1", + "request_id": "request-1", + }, + ) + + assert response.status_code == 500 + + +def test_no_user_playbook_results_record_one_empty_synchronous_batch() -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([]): + response = _client().post( + "/api/search", + json={"query": "answer", "user_id": "user-1"}, + ) + + assert response.status_code == 200, response.text + assert recorder.completed is True + assert len(recorder.batches) == 1 + assert recorder.batches[0].user_playbooks == () + + +def test_oss_search_succeeds_when_no_recorder_is_registered() -> None: + with _search_results([_playbook(11, "First")]): + response = _client().post( + "/api/search", + json={"query": "answer", "user_id": "user-1"}, + ) + + assert response.status_code == 200, response.text + assert [item["user_playbook_id"] for item in response.json()["user_playbooks"]] == [ + 11 + ] diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index f26a88d00..348919ed7 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -2,15 +2,23 @@ import json import tempfile +import threading +from collections.abc import Callable from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock, patch import pytest -from reflexio.models.api_schema.service_schemas import Interaction, Request +from reflexio.models.api_schema.service_schemas import ( + Interaction, + Request, + UserPlaybook, + UserProfile, +) from reflexio.models.config_schema import ( Config, PendingToolCallConfig, + PlaybookConfig, ProfileExtractorConfig, StorageConfigSQLite, ) @@ -20,6 +28,15 @@ _run_playbook_contract_selection, _run_uses_strict_playbook_evidence, ) +from reflexio.server.services.playbook.components.consolidator import ( + PlaybookConsolidationOutput, + UnifyDecision, +) +from reflexio.server.services.playbook.service import PlaybookGenerationService +from reflexio.server.services.profile.service import ( + ProfileGenerationService, + ProfileGenerationServiceConfig, +) from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.storage_base import ( AgentBinding, @@ -32,6 +49,10 @@ build_scope_hash, human_feedback_scope, ) +from reflexio.server.usage_metrics import ( + UsageEvent, + configure_usage_event_recorder, +) @pytest.fixture @@ -482,6 +503,402 @@ def test_resume_worker_tagging_schedule_failure_is_best_effort( worker._schedule_finalized_tagging(run) +def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( + request_context, +): + """A mixed batch charges its persisted profile once across finalization retries.""" + run = AgentRunRecord( + id="run_mixed_billing", + binding=AgentBinding( + org_id="org_1", + extractor_kind="profile", + user_id="user_1", + request_id="request_1", + agent_version="v1", + source="api", + ), + status=AgentRunStatus.FINALIZATION_FAILED, + generation_request_snapshot={"request_id": "request_1"}, + ) + learning_ids = ["profile_1"] + worker = ExtractionResumeWorker(request_context=request_context) + + with patch("reflexio.server.billing_meter.record_usage_event") as record_event: + worker._record_finalized_learnings(run, learning_ids, entity_type="profile") + worker._record_finalized_learnings(run, learning_ids, entity_type="profile") + + assert [call.kwargs["event_key"] for call in record_event.call_args_list] == [ + "learn:profile:profile_1", + "learn:profile:profile_1", + ] + assert [call.kwargs["count_value"] for call in record_event.call_args_list] == [ + 1, + 1, + ] + assert [call.kwargs["entity_id"] for call in record_event.call_args_list] == [ + "profile_1", + "profile_1", + ] + + +def test_retry_after_billing_reuses_ids_without_replaying_playbook_schedulers( + request_context, + storage, +): + """A post-billing retry reuses durable IDs and skips derived schedulers.""" + _seed_interactions(storage) + request_context.configurator.get_config.return_value = Config( + storage_config=StorageConfigSQLite(), + profile_extractor_config=ProfileExtractorConfig( + extraction_definition_prompt="Extract durable user facts.", + ), + user_playbook_extractor_config=PlaybookConfig( + extraction_definition_prompt="Extract durable operating rules.", + ), + pending_tool_call_config=PendingToolCallConfig(enabled=True), + ) + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + + def _run_failed_then_retried( + run_id: str, + *, + assert_scheduler_calls: Callable[[], None] | None = None, + ) -> None: + worker = ExtractionResumeWorker( + request_context=request_context, + llm_client=MagicMock(), + ) + with ( + patch.object(worker, "_schedule_finalized_tagging"), + patch.object( + storage, + "consume_run_tool_dependencies", + side_effect=[RuntimeError("failed after billing"), 0], + ), + ): + first_attempt = worker.run_once() + if assert_scheduler_calls is not None: + assert_scheduler_calls() + assert first_attempt is not None + assert first_attempt.status == AgentRunStatus.FINALIZATION_FAILED + storage.update_agent_run_status( + run_id, + AgentRunStatus.FINALIZATION_FAILED, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + retry = worker.run_once() + if assert_scheduler_calls is not None: + assert_scheduler_calls() + assert retry is not None + assert retry.status == AgentRunStatus.FINALIZED + + try: + profile_request_id = "request_profile_retry" + storage.create_agent_run( + AgentRunRecord( + id="run_profile_retry", + binding=AgentBinding( + org_id="org_1", + extractor_kind="profile", + user_id="user_1", + request_id=profile_request_id, + agent_version="v1", + source="api", + source_interaction_ids=[1, 2], + ), + status=AgentRunStatus.FINALIZATION_FAILED, + generation_request_snapshot={"request_id": profile_request_id}, + committed_output={ + "profiles": [ + { + "content": "User deployment target is AWS ECS.", + "time_to_live": "infinity", + } + ] + }, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + ) + with patch( + "reflexio.server.services.profile.components.consolidator." + "ProfileConsolidator.deduplicate", + side_effect=lambda profiles, _user_id, _request_id: ( + profiles, + [], + [], + ), + ): + _run_failed_then_retried("run_profile_retry") + + seed = UserPlaybook( + user_id="user_1", + agent_version="v1", + request_id="seed_request", + content="Prefer the current deployment default.", + trigger="when selecting a deployment target", + rationale="Existing operating rule.", + source="api", + ) + storage.save_user_playbooks([seed]) + playbook_request_id = "request_playbook_retry" + storage.create_agent_run( + AgentRunRecord( + id="run_playbook_retry", + binding=AgentBinding( + org_id="org_1", + extractor_kind="playbook", + user_id="user_1", + request_id=playbook_request_id, + agent_version="v1", + source="api", + source_interaction_ids=[1, 2], + ), + status=AgentRunStatus.FINALIZATION_FAILED, + generation_request_snapshot={ + "request_id": playbook_request_id, + "output_schema_name": "StructuredPlaybookList", + }, + committed_output={ + "playbooks": [ + { + "content": "Prefer AWS ECS as the deployment target.", + "trigger": "when selecting a deployment target", + "rationale": "The team standardizes on AWS.", + } + ] + }, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + ) + consolidation = PlaybookConsolidationOutput( + decisions=[ + UnifyDecision( + new_id="NEW-0", + archive_existing_ids=[0], + content="Prefer AWS ECS as the deployment target.", + trigger="when selecting a deployment target", + rationale="The team standardizes on AWS.", + ) + ] + ) + with ( + patch.object( + PlaybookGenerationService, + "_configured_playbook_config", + return_value=None, + ), + patch( + "reflexio.server.services.playbook.components.consolidator." + "PlaybookConsolidator.retrieve_existing_playbooks", + side_effect=lambda _new, **_kwargs: storage.get_user_playbooks( + user_id="user_1", + agent_version="v1", + ), + ), + patch( + "reflexio.server.services.playbook.components.consolidator." + "PlaybookConsolidator._consolidation_decisions", + return_value=consolidation, + ), + patch.object( + PlaybookGenerationService, + "_enqueue_user_playbook_optimization", + ) as enqueue_optimization, + patch.object( + PlaybookGenerationService, + "_trigger_playbook_aggregation", + ) as trigger_aggregation, + ): + + def assert_playbook_scheduler_calls() -> None: + enqueue_optimization.assert_called_once() + trigger_aggregation.assert_called_once() + + _run_failed_then_retried( + "run_playbook_retry", + assert_scheduler_calls=assert_playbook_scheduler_calls, + ) + finally: + configure_usage_event_recorder(None) + + profile_keys = [ + event.event_key for event in events if event.entity_type == "profile" + ] + profile_event_ids = [ + event.entity_id for event in events if event.entity_type == "profile" + ] + playbook_keys = [ + event.event_key for event in events if event.entity_type == "user_playbook" + ] + playbook_event_ids = [ + event.entity_id for event in events if event.entity_type == "user_playbook" + ] + profile_survivor_ids = [ + row[0] + for row in storage.conn.execute( + "SELECT profile_id FROM profiles WHERE generated_from_request_id = ?", + (profile_request_id,), + ).fetchall() + ] + playbook_survivor_ids = [ + str(row[0]) + for row in storage.conn.execute( + "SELECT user_playbook_id FROM user_playbooks WHERE request_id = ?", + (playbook_request_id,), + ).fetchall() + ] + profile_receipt_ids = storage.get_agent_run_finalization_receipt( + run_id="run_profile_retry", entity_type="profile" + ) + playbook_receipt_ids = storage.get_agent_run_finalization_receipt( + run_id="run_playbook_retry", entity_type="user_playbook" + ) + playbook_lineage_ids = [ + event.event_id + for event in storage.get_lineage_events(request_id=playbook_request_id) + ] + observed = { + "profile_persisted": len(profile_survivor_ids), + "profile_events": len(profile_keys), + "profile_distinct_keys": len(set(profile_keys)), + "playbook_persisted": len(playbook_survivor_ids), + "playbook_lineage_events": len(playbook_lineage_ids), + "playbook_events": len(playbook_keys), + "playbook_distinct_keys": len(set(playbook_keys)), + } + assert observed == { + "profile_persisted": 1, + "profile_events": 2, + "profile_distinct_keys": 1, + "playbook_persisted": 1, + "playbook_lineage_events": 1, + "playbook_events": 2, + "playbook_distinct_keys": 1, + } + assert profile_receipt_ids == profile_survivor_ids + assert profile_event_ids == profile_survivor_ids * 2 + assert profile_keys == [f"learn:profile:{profile_survivor_ids[0]}"] * 2 + assert playbook_receipt_ids == playbook_survivor_ids + assert playbook_event_ids == playbook_survivor_ids * 2 + assert playbook_keys == [f"learn:user_playbook:{playbook_survivor_ids[0]}"] * 2 + + +def test_two_stale_sqlite_workers_share_one_finalization_receipt(tmp_path): + db_path = str(tmp_path / "stale-finalizers.db") + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + storage_a = SQLiteStorage(org_id="org_1", db_path=db_path) + storage_b = SQLiteStorage(org_id="org_1", db_path=db_path) + storage_a.create_agent_run( + AgentRunRecord( + id="run_race", + binding=AgentBinding( + org_id="org_1", + extractor_kind="profile", + user_id="user_1", + request_id="request_race", + agent_version="v1", + source="api", + ), + status=AgentRunStatus.FINALIZING, + generation_request_snapshot={"request_id": "request_race"}, + ) + ) + barrier = threading.Barrier(2) + + def _service(storage_instance: SQLiteStorage) -> ProfileGenerationService: + context = RequestContext.__new__(RequestContext) + context.org_id = "org_1" + context.storage = storage_instance + context.configurator = MagicMock() + context.configurator.get_config.return_value = Config( + storage_config=StorageConfigSQLite(), + profile_extractor_config=ProfileExtractorConfig( + extraction_definition_prompt="Extract durable user facts.", + ), + ) + service = ProfileGenerationService( + llm_client=MagicMock(), + request_context=context, + ) + service.service_config = ProfileGenerationServiceConfig( + user_id="user_1", + request_id="request_race", + source="api", + auto_run=False, + force_extraction=True, + ) + return service + + def _gate_first_receipt_read(storage_instance: SQLiteStorage) -> None: + original = storage_instance.get_agent_run_finalization_receipt + first_read = True + + def gated_read(*, run_id: str, entity_type: str) -> list[str] | None: + nonlocal first_read + receipt = original(run_id=run_id, entity_type=entity_type) + if first_read: + first_read = False + barrier.wait(timeout=5) + return receipt + + storage_instance.get_agent_run_finalization_receipt = MagicMock( # type: ignore[method-assign] + side_effect=gated_read + ) + + _gate_first_receipt_read(storage_a) + _gate_first_receipt_read(storage_b) + services = [_service(storage_a), _service(storage_b)] + candidates = [ + UserProfile( + profile_id=f"profile-race-{index}", + user_id="user_1", + content="User deployment target is AWS ECS.", + last_modified_timestamp=1_000, + generated_from_request_id="request_race", + ) + for index in range(2) + ] + results: list[list[str]] = [] + errors: list[BaseException] = [] + + def finalize(index: int) -> None: + try: + results.append( + services[index]._finalize_extracted_items( + [candidates[index]], + finalization_run_id="run_race", + ) + ) + except BaseException as exc: # noqa: BLE001 - intentional thread error capture + errors.append(exc) + + with patch( + "reflexio.server.services.profile.components.consolidator." + "ProfileConsolidator.deduplicate", + side_effect=lambda profiles, _user_id, _request_id: (profiles, [], []), + ): + workers = [ + threading.Thread(target=finalize, args=(index,)) for index in range(2) + ] + for worker in workers: + worker.start() + for worker in workers: + worker.join(timeout=10) + + assert all(not worker.is_alive() for worker in workers) + assert errors == [] + persisted_ids = [ + profile.profile_id for profile in storage_a.get_user_profile("user_1") + ] + receipt_ids = storage_a.get_agent_run_finalization_receipt( + run_id="run_race", entity_type="profile" + ) + assert len(persisted_ids) == 1 + assert receipt_ids == persisted_ids + assert results == [persisted_ids, persisted_ids] + + def test_resume_worker_fails_run_when_step_budget_exhausted( monkeypatch, request_context, diff --git a/tests/server/services/governance/test_governance_local_e2e.py b/tests/server/services/governance/test_governance_local_e2e.py index d8ee06b06..576419903 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -1,5 +1,7 @@ from __future__ import annotations +import json +import threading from collections.abc import Generator from datetime import UTC, datetime from pathlib import Path @@ -18,12 +20,14 @@ UserProfile, ) from reflexio.models.api_schema.domain.enums import PlaybookStatus +from reflexio.models.api_schema.domain.governance import UserEraseResult from reflexio.models.api_schema.retriever_schema import SearchAgentPlaybookRequest from reflexio.models.config_schema import SearchMode from reflexio.server.services.governance import service as governance_service_module from reflexio.server.services.governance.config import governance_subject_ref from reflexio.server.services.governance.service import GovernanceService from reflexio.server.services.storage.error import SubjectWriteBarrierError +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.sqlite_storage import SQLiteStorage pytestmark = pytest.mark.integration @@ -117,6 +121,27 @@ def _eval_result( ) +def _insert_session_outcome( + storage: SQLiteStorage, + *, + outcome_id: str, + user_id: str, + session_id: str, + subject_ref: str, +) -> None: + storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, + governance_subject_ref, created_at + ) VALUES (?, 1, ?, ?, 'success', 100, 'test', NULL, NULL, NULL, + ?, ?, ?, 101)""", + (outcome_id, user_id, session_id, "a" * 64, "b" * 64, subject_ref), + ) + storage.conn.commit() + + @pytest.fixture def storage( tmp_path: Path, @@ -306,6 +331,7 @@ def test_local_governance_e2e_erases_exports_audits_and_preserves_org_agent_play assert erased.deleted_counts["requests"] == 1 assert erased.deleted_counts["user_playbooks"] == 2 assert erased.deleted_counts["agent_success_evaluation_results"] == 1 + assert erased.deleted_counts["session_outcomes"] == 0 assert erased.rebuilt_agent_playbook_ids == [] assert storage.get_user_interaction("alice") == [] @@ -437,6 +463,55 @@ def test_local_governance_e2e_erases_exports_audits_and_preserves_org_agent_play assert len(erase_events_after_retry) == 1 +def test_governance_erasure_uses_authoritative_user_for_session_outcomes_and_receipts( + storage: SQLiteStorage, +) -> None: + alice_ref = governance_subject_ref( + storage.org_id, "alice", "test-governance-secret" + ) + bob_ref = governance_subject_ref(storage.org_id, "bob", "test-governance-secret") + _insert_session_outcome( + storage, + outcome_id="alice-stale-ref", + user_id="alice", + session_id="alice-session", + subject_ref=bob_ref, + ) + _insert_session_outcome( + storage, + outcome_id="bob-conflicting-ref", + user_id="bob", + session_id="bob-session", + subject_ref=alice_ref, + ) + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + + erased = service.erase_user(user_id="alice", request_id="erase-outcomes") + retried = service.erase_user(user_id="alice", request_id="erase-outcomes") + + remaining = storage.conn.execute( + "SELECT outcome_id, user_id FROM session_outcomes ORDER BY outcome_id" + ).fetchall() + assert [(row["outcome_id"], row["user_id"]) for row in remaining] == [ + ("bob-conflicting-ref", "bob") + ] + assert erased.deleted_counts["session_outcomes"] == 1 + assert retried.deleted_counts == erased.deleted_counts + audit = next( + event + for event in storage.list_audit_events(subject_ref=alice_ref) + if event.operation == "ERASE" + ) + assert audit.detail is not None + deleted_counts = audit.detail["deleted_counts"] + assert isinstance(deleted_counts, dict) + assert deleted_counts["session_outcomes"] == 1 + + def test_governance_service_persists_actor_context_in_audit( storage: SQLiteStorage, ) -> None: @@ -517,6 +592,54 @@ def test_completed_erase_retry_reconstructs_response( assert second.rebuilt_agent_playbook_ids == first.rebuilt_agent_playbook_ids +@pytest.mark.parametrize("corrupt_binding", ["purge", "snapshot"]) +def test_completed_erase_retry_rejects_corrupt_authoritative_binding( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, + corrupt_binding: str, +) -> None: + monkeypatch.setenv("REFLEXIO_GOVERNANCE_REF_SECRET", "test-governance-secret") + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + completed = service.erase_user( + user_id="alice", + request_id=f"erase-retry-corrupt-{corrupt_binding}", + ) + + if corrupt_binding == "purge": + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + ("a" * 64, storage.org_id, completed.purge_id), + ) + else: + snapshot = next( + target + for target in storage.list_purge_targets( + completed.purge_id, phase="prepare_targets" + ) + if target.target_name == "target_snapshot" + ) + detail = dict(snapshot.detail or {}) + detail["authoritative_user_digest"] = "a" * 64 + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(detail), storage.org_id, completed.purge_id), + ) + storage.conn.commit() + + with pytest.raises(ValueError, match="authoritative user identity"): + service.erase_user( + user_id="alice", + request_id=f"erase-retry-corrupt-{corrupt_binding}", + ) + + def test_erase_fails_fast_when_service_and_storage_ref_secrets_differ( storage: SQLiteStorage, ) -> None: @@ -629,10 +752,21 @@ def test_second_erase_conflict_preserves_original_barrier_and_write_block( idempotency_key="idem_conflict_first", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000061", ) - storage.begin_subject_erasure_barrier(subject_ref, first_purge.purge_id) + first_claim = storage.claim_purge_operation_execution( + first_purge.purge_id, + lease_owner="test-conflict-first", + lease_ttl_seconds=30, + ) + assert first_claim is not None + storage.begin_subject_erasure_barrier( + subject_ref, + first_purge.purge_id, + execution_claim=first_claim, + ) service = GovernanceService( storage=storage, org_id=storage.org_id, @@ -721,6 +855,7 @@ def erase_subject( storage: SQLiteStorage, subject_ref: str, purge_id: str, + execution_claim: PurgeExecutionClaim, ) -> None: del subject_ref self.calls += 1 @@ -740,6 +875,7 @@ def erase_subject( status="complete", detail={"count": 2}, deleted_count=2, + execution_claim=execution_claim, ) lifecycle = RetrySafeLifecycle() @@ -785,6 +921,571 @@ def fail_after_first_lifecycle(*args, **kwargs): assert snapshot.detail["status"] == "complete" +def test_duplicate_erase_waits_for_lifecycle_winner_beyond_old_deadline( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + class SlowSingleUseLifecycle: + def __init__(self) -> None: + self.first_call_started = threading.Event() + self.release_first_call = threading.Event() + self._lock = threading.Lock() + self.calls = 0 + self.duplicate_rejections = 0 + + def erase_subject( + self, + *, + storage: SQLiteStorage, + subject_ref: str, + purge_id: str, + execution_claim: object, + ) -> None: + del storage, subject_ref, purge_id, execution_claim + with self._lock: + self.calls += 1 + call_number = self.calls + if call_number == 1: + self.first_call_started.set() + assert self.release_first_call.wait(timeout=5) + return + self.duplicate_rejections += 1 + raise RuntimeError("provider lifecycle already running") + + lifecycle = SlowSingleUseLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + + def release_winner_on_duplicate_poll(_seconds: float) -> None: + lifecycle.release_first_call.set() + threading.Event().wait(0.001) + + monkeypatch.setattr(service, "_sleep", release_winner_on_duplicate_poll) + winner_results: list[UserEraseResult] = [] + winner_errors: list[BaseException] = [] + + def run_winner() -> None: + try: + winner_results.append( + service.erase_user(user_id="alice", request_id="erase-slow-duplicate") + ) + except BaseException as exc: + winner_errors.append(exc) + + winner = threading.Thread(target=run_winner) + winner.start() + assert lifecycle.first_call_started.wait(timeout=5) + + try: + duplicate = service.erase_user( + user_id="alice", request_id="erase-slow-duplicate" + ) + finally: + lifecycle.release_first_call.set() + winner.join(timeout=5) + + assert not winner.is_alive() + assert winner_errors == [] + assert len(winner_results) == 1 + winner_result = winner_results[0] + assert duplicate.status == "complete" + assert duplicate.purge_id == winner_result.purge_id + assert storage.get_purge_operation(duplicate.purge_id).status == "complete" + barrier = storage.get_subject_write_barrier(duplicate.subject_ref) + assert barrier is not None + assert barrier.status == "erased" + assert lifecycle.calls == 1 + assert lifecycle.duplicate_rejections == 0 + + +def test_duplicate_erase_wait_is_bounded_with_exponential_backoff( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + class ControlledWaitGovernanceService(GovernanceService): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.now = 0.0 + self.sleep_delays: list[float] = [] + + def _monotonic(self) -> float: + return self.now + + def _sleep(self, seconds: float) -> None: + self.sleep_delays.append(seconds) + self.now += seconds + + service = ControlledWaitGovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + claim_attempts = 0 + + def reject_duplicate_claim(*args, **kwargs): + nonlocal claim_attempts + del args, kwargs + claim_attempts += 1 + if claim_attempts > 100: + raise AssertionError("duplicate claim wait exceeded its attempt bound") + + monkeypatch.setattr( + storage, + "claim_purge_operation_execution", + reject_duplicate_claim, + ) + with pytest.raises(RuntimeError, match="retry later") as exc_info: + service.erase_user(user_id="alice", request_id="erase-bounded-duplicate") + + assert ( + type(exc_info.value) is governance_service_module.GovernanceEraseRetryLaterError + ) + assert service.sleep_delays[:5] == pytest.approx([0.05, 0.1, 0.2, 0.4, 0.8]) + assert max(service.sleep_delays) == 1.0 + assert sum(service.sleep_delays) == pytest.approx(5.0) + + +def test_healthy_slow_lifecycle_renews_lease_and_duplicate_converges( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_now = {"value": 100} + monkeypatch.setattr( + governance_service_module, + "_PURGE_EXECUTION_HEARTBEAT_SECONDS", + 0.01, + ) + + from reflexio.server.services.storage.sqlite_storage.governance import ( + _purge as sqlite_purge_module, + ) + + monkeypatch.setattr(sqlite_purge_module, "_epoch_now", lambda: fake_now["value"]) + original_renew = storage.renew_purge_operation_execution_claim + renewal_confirmed = threading.Event() + renewals_by_owner: dict[str, int] = {} + + def observed_renew(*args, **kwargs): + claim = args[1] + renewals_by_owner[claim.owner] = renewals_by_owner.get(claim.owner, 0) + 1 + renewed = original_renew(*args, **kwargs) + if renewals_by_owner[claim.owner] >= 3: + renewal_confirmed.set() + return renewed + + monkeypatch.setattr( + storage, + "renew_purge_operation_execution_claim", + observed_renew, + ) + + class SlowSingleUseLifecycle: + def __init__(self) -> None: + self.first_call_started = threading.Event() + self.release_first_call = threading.Event() + self._lock = threading.Lock() + self.calls = 0 + + def erase_subject( + self, + *, + storage: SQLiteStorage, + subject_ref: str, + purge_id: str, + execution_claim: object, + ) -> None: + del storage, subject_ref, purge_id, execution_claim + with self._lock: + self.calls += 1 + call_number = self.calls + if call_number != 1: + raise RuntimeError("provider lifecycle already running") + self.first_call_started.set() + fake_now["value"] = 350 + assert renewal_confirmed.wait(timeout=5) + fake_now["value"] = 401 + assert self.release_first_call.wait(timeout=5) + + lifecycle = SlowSingleUseLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + winner_results: list[UserEraseResult] = [] + winner_errors: list[BaseException] = [] + duplicate_results: list[UserEraseResult] = [] + duplicate_errors: list[BaseException] = [] + + def run_winner() -> None: + try: + winner_results.append( + service.erase_user(user_id="alice", request_id="erase-renewed-slow") + ) + except BaseException as exc: + winner_errors.append(exc) + + def run_duplicate() -> None: + try: + duplicate_results.append( + service.erase_user(user_id="alice", request_id="erase-renewed-slow") + ) + except BaseException as exc: + duplicate_errors.append(exc) + + winner = threading.Thread(target=run_winner) + winner.start() + assert lifecycle.first_call_started.wait(timeout=5) + assert renewal_confirmed.wait(timeout=5) + + duplicate = threading.Thread(target=run_duplicate) + duplicate.start() + lifecycle.release_first_call.set() + winner.join(timeout=5) + duplicate.join(timeout=5) + + assert not winner.is_alive() + assert not duplicate.is_alive() + assert winner_errors == [] + assert duplicate_errors == [] + assert len(winner_results) == 1 + assert len(duplicate_results) == 1 + assert duplicate_results[0].status == "complete" + assert duplicate_results[0].purge_id == winner_results[0].purge_id + assert lifecycle.calls == 1 + + +def test_heartbeat_renewal_loss_fences_external_lifecycle_and_retry_converges( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_now = {"value": 100} + monkeypatch.setattr( + governance_service_module, + "_PURGE_EXECUTION_HEARTBEAT_SECONDS", + 0.01, + ) + from reflexio.server.services.storage.sqlite_storage.governance import ( + _purge as sqlite_purge_module, + ) + + monkeypatch.setattr(sqlite_purge_module, "_epoch_now", lambda: fake_now["value"]) + original_renew = storage.renew_purge_operation_execution_claim + first_owner: list[str] = [] + renewals_by_owner: dict[str, int] = {} + renewal_lost = threading.Event() + + def fail_first_owner_heartbeat(*args, **kwargs): + claim = args[1] + if not first_owner: + first_owner.append(claim.owner) + renewals_by_owner[claim.owner] = renewals_by_owner.get(claim.owner, 0) + 1 + if claim.owner == first_owner[0] and renewals_by_owner[claim.owner] == 2: + renewal_lost.set() + raise RuntimeError("simulated heartbeat renewal loss") + return original_renew(*args, **kwargs) + + monkeypatch.setattr( + storage, + "renew_purge_operation_execution_claim", + fail_first_owner_heartbeat, + ) + original_apply = storage.apply_governance_user_data_delete + + def apply_then_wait_for_renewal_loss(*args, **kwargs): + result = original_apply(*args, **kwargs) + assert renewal_lost.wait(timeout=5) + return result + + monkeypatch.setattr( + storage, + "apply_governance_user_data_delete", + apply_then_wait_for_renewal_loss, + ) + + class CountingLifecycle: + def __init__(self) -> None: + self.calls = 0 + + def erase_subject(self, **_kwargs) -> None: + self.calls += 1 + + lifecycle = CountingLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + + with pytest.raises(ValueError, match="heartbeat renewal was lost"): + service.erase_user(user_id="alice", request_id="erase-renewal-loss") + + assert lifecycle.calls == 0 + purge_id = str( + storage.conn.execute("SELECT purge_id FROM purge_operations").fetchone()[ + "purge_id" + ] + ) + assert storage.get_purge_operation(purge_id).status == "running" + assert storage.list_audit_events() == [] + + fake_now["value"] = 401 + recovered = service.erase_user( + user_id="alice", + request_id="erase-renewal-loss", + ) + + assert recovered.status == "complete" + assert lifecycle.calls == 1 + assert storage.get_purge_operation(purge_id).status == "complete" + + +def test_synchronous_renewal_loss_skips_lifecycle_and_retry_converges( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_renew = storage.renew_purge_operation_execution_claim + renewal_attempts = 0 + renewal_failed = threading.Event() + lifecycle_called = threading.Event() + + def fail_mandatory_lifecycle_renewal(*args, **kwargs): + nonlocal renewal_attempts + renewal_attempts += 1 + if renewal_attempts == 2: + renewal_failed.set() + raise RuntimeError("simulated synchronous renewal loss") + return original_renew(*args, **kwargs) + + monkeypatch.setattr( + storage, + "renew_purge_operation_execution_claim", + fail_mandatory_lifecycle_renewal, + ) + + class CountingLifecycle: + def __init__(self) -> None: + self.calls = 0 + + def erase_subject(self, **_kwargs) -> None: + self.calls += 1 + lifecycle_called.set() + + lifecycle = CountingLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + + with pytest.raises(ValueError, match="heartbeat renewal was lost"): + service.erase_user(user_id="alice", request_id="erase-sync-renewal-loss") + + purge_id = str( + storage.conn.execute("SELECT purge_id FROM purge_operations").fetchone()[ + "purge_id" + ] + ) + purge = storage.get_purge_operation(purge_id) + assert purge.subject_ref is not None + barrier = storage.get_subject_write_barrier(purge.subject_ref) + assert renewal_attempts == 2 + assert renewal_failed.is_set() + assert not lifecycle_called.is_set() + assert lifecycle.calls == 0 + assert purge.status == "running" + assert barrier is not None + assert barrier.status == "erasing" + assert storage.list_audit_events() == [] + + storage.conn.execute( + "UPDATE purge_operations SET execution_claim_expires_at = 0 WHERE purge_id = ?", + (purge_id,), + ) + storage.conn.commit() + + def fail_if_recovery_polls(_seconds: float) -> None: + raise AssertionError("expired synchronous-renewal claim did not recover") + + monkeypatch.setattr(service, "_sleep", fail_if_recovery_polls) + + recovered = service.erase_user( + user_id="alice", + request_id="erase-sync-renewal-loss", + ) + + assert recovered.status == "complete" + assert lifecycle_called.is_set() + assert lifecycle.calls == 1 + assert storage.get_purge_operation(purge_id).status == "complete" + + +def test_stale_running_erase_claim_recovers_after_crash( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + original_begin_barrier = storage.begin_subject_erasure_barrier + crash_once = True + + def crash_after_claim(subject_ref: str, purge_id: str, **kwargs): + nonlocal crash_once + if crash_once: + crash_once = False + raise SystemExit("simulated crash after claim") + return original_begin_barrier(subject_ref, purge_id, **kwargs) + + monkeypatch.setattr( + storage, + "begin_subject_erasure_barrier", + crash_after_claim, + ) + + with pytest.raises(SystemExit, match="simulated crash after claim"): + service.erase_user(user_id="alice", request_id="erase-crash-after-claim") + + storage.conn.execute("UPDATE purge_operations SET execution_claim_expires_at = 0") + storage.conn.commit() + + def fail_if_duplicate_polls(_seconds: float) -> None: + raise AssertionError("stale running claim did not recover") + + monkeypatch.setattr(service, "_sleep", fail_if_duplicate_polls) + + recovered = service.erase_user( + user_id="alice", + request_id="erase-crash-after-claim", + ) + + assert recovered.status == "complete" + assert storage.get_purge_operation(recovered.purge_id).status == "complete" + barrier = storage.get_subject_write_barrier(recovered.subject_ref) + assert barrier is not None + assert barrier.status == "erased" + + +def test_delete_committed_before_completion_retry_converges_idempotently( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + original_complete = storage.complete_subject_erasure_barrier_after_empty_check + completion_attempts = 0 + + def crash_after_delete_targets(*args, **kwargs): + nonlocal completion_attempts + completion_attempts += 1 + if completion_attempts == 1: + raise SystemExit("simulated crash after delete") + return original_complete(*args, **kwargs) + + monkeypatch.setattr( + storage, + "complete_subject_erasure_barrier_after_empty_check", + crash_after_delete_targets, + ) + + with pytest.raises(SystemExit, match="simulated crash after delete"): + service.erase_user(user_id="alice", request_id="erase-crash-after-delete") + + storage.conn.execute("UPDATE purge_operations SET execution_claim_expires_at = 0") + storage.conn.commit() + + delete_targets_after_crash = storage.list_purge_targets( + storage.conn.execute("SELECT purge_id FROM purge_operations").fetchone()[ + "purge_id" + ], + phase="delete", + ) + assert delete_targets_after_crash + assert all(target.status == "complete" for target in delete_targets_after_crash) + + def fail_if_duplicate_polls(_seconds: float) -> None: + raise AssertionError("stale post-delete claim did not recover") + + monkeypatch.setattr(service, "_sleep", fail_if_duplicate_polls) + + recovered = service.erase_user( + user_id="alice", + request_id="erase-crash-after-delete", + ) + + assert recovered.status == "complete" + assert storage.get_purge_operation(recovered.purge_id).status == "complete" + assert completion_attempts == 2 + + +def test_pending_duplicate_erase_has_one_durable_execution_owner( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_begin = storage.begin_purge_operation + both_pending = threading.Barrier(2) + + def synchronized_begin(*args, **kwargs): + purge = original_begin(*args, **kwargs) + both_pending.wait(timeout=5) + return purge + + monkeypatch.setattr(storage, "begin_purge_operation", synchronized_begin) + + class CountingLifecycle: + def __init__(self) -> None: + self.calls = 0 + self._lock = threading.Lock() + + def erase_subject(self, **_kwargs) -> None: + with self._lock: + self.calls += 1 + + lifecycle = CountingLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + results: list[UserEraseResult] = [] + errors: list[BaseException] = [] + + def erase() -> None: + try: + results.append( + service.erase_user(user_id="alice", request_id="erase-pending-race") + ) + except BaseException as exc: + errors.append(exc) + + callers = [threading.Thread(target=erase) for _ in range(2)] + for caller in callers: + caller.start() + for caller in callers: + caller.join(timeout=5) + + assert all(not caller.is_alive() for caller in callers) + assert errors == [] + assert len(results) == 2 + assert results[0].purge_id == results[1].purge_id + assert all(result.status == "complete" for result in results) + assert lifecycle.calls == 1 + assert storage.get_purge_operation(results[0].purge_id).status == "complete" + + def test_session_export_paginates_by_returned_rows_when_requests_are_missing() -> None: class _Storage: def __init__(self) -> None: diff --git a/tests/server/services/governance/test_subject_write_barrier_sqlite.py b/tests/server/services/governance/test_subject_write_barrier_sqlite.py index c649bcaaf..a0598f144 100644 --- a/tests/server/services/governance/test_subject_write_barrier_sqlite.py +++ b/tests/server/services/governance/test_subject_write_barrier_sqlite.py @@ -1,5 +1,7 @@ from __future__ import annotations +import hashlib +import hmac from datetime import UTC, datetime from pathlib import Path @@ -14,12 +16,13 @@ UserPlaybook, UserProfile, ) -from reflexio.models.api_schema.domain.governance import AuditEvent +from reflexio.models.api_schema.domain.governance import AuditEvent, SubjectWriteBarrier from reflexio.server.services.governance.config import governance_subject_ref from reflexio.server.services.storage.error import ( StorageError, SubjectWriteBarrierError, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, ) @@ -36,14 +39,78 @@ def _storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> SQLiteStorage: return SQLiteStorage(org_id="org-barrier", db_path=str(tmp_path / "barrier.db")) +def _claim_purge(storage: SQLiteStorage, purge_id: str) -> PurgeExecutionClaim: + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"test-{purge_id}", + lease_ttl_seconds=30, + ) + if claim is None: + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"test-{purge_id}", + lease_ttl_seconds=30, + ) + assert claim is not None + return claim + + +def _typed_test_claim_for_unvalidated_purge_id(purge_id: str) -> PurgeExecutionClaim: + return PurgeExecutionClaim( + purge_id=purge_id, + owner="test-unvalidated", + fence=1, + expires_at=1, + ) + + +def _begin_claimed_subject_erasure_barrier( + storage: SQLiteStorage, + subject_ref: str, + purge_id: str, +) -> SubjectWriteBarrier: + return storage.begin_subject_erasure_barrier( + subject_ref, + purge_id, + execution_claim=_claim_purge(storage, purge_id), + ) + + +def _authoritative_user_digest(storage: SQLiteStorage, purge_id: str) -> str: + return storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + + +def _expected_authoritative_user_digest( + *, secret: str, org_id: str, purge_id: str, user_id: str +) -> str: + material = f"authoritative-user-v1\0{org_id}\0{purge_id}\0{user_id}" + return hmac.new(secret.encode(), material.encode(), hashlib.sha256).hexdigest() + + def _mark_all_completion_targets(storage: SQLiteStorage, purge_id: str) -> None: + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - detail={"prepared": True}, + execution_claim=claim, + detail={ + "prepared": True, + "authoritative_user_digest": _authoritative_user_digest(storage, purge_id), + }, ) # Single source of truth — a stale local copy of the canonical tuple is # exactly how this suite went red when new delete targets landed. @@ -54,6 +121,7 @@ def _mark_all_completion_targets(storage: SQLiteStorage, purge_id: str) -> None: phase="delete", status="complete", target_ref="all", + execution_claim=claim, detail={"count": 0}, ) @@ -77,8 +145,190 @@ def _complete_empty_purge( idempotency_key=purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) + + +def test_begin_purge_operation_keys_authoritative_user_digest( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_keyed_authoritative_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_keyed_authoritative_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000056", ) + stored_digest = _authoritative_user_digest(storage, purge_id) + assert stored_digest == _expected_authoritative_user_digest( + secret="barrier-secret", + org_id="org-barrier", + purge_id=purge_id, + user_id="alice", + ) + assert stored_digest != hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + + +@pytest.mark.parametrize("legacy_digest", [None, "unkeyed"]) +def test_begin_purge_operation_upgrades_validated_legacy_authoritative_digest( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + legacy_digest: str | None, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_legacy_authoritative_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_authoritative_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000057", + ) + persisted_digest = ( + hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + if legacy_digest == "unkeyed" + else None + ) + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = NULLIF(?, '') + WHERE org_id = ? AND purge_id = ?""", + (persisted_digest or "", storage.org_id, purge_id), + ) + storage.conn.commit() + + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_authoritative_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000057", + ) + + assert _authoritative_user_digest( + storage, purge_id + ) == _expected_authoritative_user_digest( + secret="barrier-secret", + org_id="org-barrier", + purge_id=purge_id, + user_id="alice", + ) + + +def test_begin_purge_operation_does_not_upgrade_legacy_digest_for_wrong_user( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_legacy_wrong_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_wrong_identity", + operation_type="user_erasure", + scope_type="user", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000058", + authoritative_user_id="alice", + ) + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = NULL + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + + with pytest.raises(ValueError, match="must match subject_ref"): + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_wrong_identity", + operation_type="user_erasure", + scope_type="user", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000058", + authoritative_user_id="bob", + ) + + assert _authoritative_user_digest(storage, purge_id) is None + + +def test_completion_checks_session_outcomes_by_authoritative_user_id( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_exact_outcome_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + request_ref = "reqref_v1_00000000000000000000000000000059" + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_exact_outcome_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref=request_ref, + ) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge_id) + _mark_all_completion_targets(storage, purge_id) + storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES (?, 1, ?, ?, 'success', 1, 'test', ?, ?, ?, 1)""", + ( + "outcome-bob", + "bob", + "session-bob", + "contract-digest", + "trajectory-digest", + governance_subject_ref("org-barrier", "bob", "barrier-secret"), + ), + ) + storage.conn.commit() + original_subject_ref = storage._subject_ref_for_user_id + + def subject_ref_for_authoritative_user_only(user_id: str) -> str: + assert user_id == "alice", "completion enumerated an unrelated outcome user" + return original_subject_ref(user_id) + + monkeypatch.setattr( + storage, + "_subject_ref_for_user_id", + subject_ref_for_authoritative_user_only, + ) + + completed = storage.complete_subject_erasure_barrier_after_empty_check( + purge_id, + AuditEvent( + org_id="org-barrier", + operation="ERASE", + entity_type="request", + subject_ref=subject_ref, + request_ref=request_ref, + idempotency_key=purge_id, + detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, + ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) + + assert completed.status == "complete" + def test_barrier_blocks_request_interaction_and_profile_writes( tmp_path: Path, @@ -91,11 +341,14 @@ def test_barrier_blocks_request_interaction_and_profile_writes( idempotency_key="idem_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_11111111111111111111111111111111", ) - barrier = storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + barrier = _begin_claimed_subject_erasure_barrier( + storage, subject_ref, purge.purge_id + ) assert barrier.status == "erasing" with pytest.raises(SubjectWriteBarrierError): @@ -145,6 +398,7 @@ def test_barrier_blocks_playbook_eval_and_source_window_writes( idempotency_key="idem_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_11111111111111111111111111111111", ) @@ -174,7 +428,7 @@ def test_barrier_blocks_playbook_eval_and_source_window_writes( ] )[0] - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.save_user_playbooks( @@ -263,16 +517,21 @@ def test_begin_subject_erasure_barrier_requires_matching_purge( idempotency_key="idem_barrier_match", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=alice_subject_ref, request_ref="reqref_v1_00000000000000000000000000000021", ) with pytest.raises(ValueError, match="subject_ref must match"): - storage.begin_subject_erasure_barrier(bob_subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, bob_subject_ref, purge.purge_id) with pytest.raises(ValueError, match="not found"): storage.begin_subject_erasure_barrier( - alice_subject_ref, "purge_barrier_missing" + alice_subject_ref, + "purge_barrier_missing", + execution_claim=_typed_test_claim_for_unvalidated_purge_id( + "purge_barrier_missing" + ), ) @@ -287,6 +546,7 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( idempotency_key="idem_barrier_first", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000051", ) @@ -295,11 +555,12 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( idempotency_key="idem_barrier_second", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000052", ) - storage.begin_subject_erasure_barrier(subject_ref, first_purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, first_purge.purge_id) with pytest.raises(ValueError, match="matching barrier"): storage.fail_subject_erasure_barrier( @@ -307,6 +568,7 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( second_purge.purge_id, error_code="governance_erase_failed", error_detail="ValueError", + execution_claim=_claim_purge(storage, second_purge.purge_id), ) barrier = storage.get_subject_write_barrier(subject_ref) @@ -326,7 +588,7 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( ) -def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( +def test_begin_subject_erasure_barrier_rejects_inactive_claim_after_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -338,10 +600,11 @@ def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( idempotency_key="idem_barrier_terminal_begin", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref=request_ref, ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -349,7 +612,14 @@ def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( request_ref=request_ref, ) - barrier = storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + with pytest.raises(ValueError, match="purge execution claim"): + storage.begin_subject_erasure_barrier( + subject_ref, + purge.purge_id, + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge.purge_id), + ) + barrier = storage.get_subject_write_barrier(subject_ref) + assert barrier is not None stored_barrier = storage.get_subject_write_barrier(subject_ref) stored_purge = storage.get_purge_operation(purge.purge_id) @@ -360,7 +630,7 @@ def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( assert stored_purge.status == "complete" -def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( +def test_fail_subject_erasure_barrier_rejects_inactive_claim_after_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -372,10 +642,11 @@ def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( idempotency_key="idem_barrier_terminal_fail", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref=request_ref, ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -383,12 +654,13 @@ def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( request_ref=request_ref, ) - with pytest.raises(ValueError, match="matching barrier"): + with pytest.raises(ValueError, match="purge execution claim"): storage.fail_subject_erasure_barrier( subject_ref, purge.purge_id, error_code="governance_erase_failed", error_detail="late_failure", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge.purge_id), ) barrier = storage.get_subject_write_barrier(subject_ref) @@ -401,7 +673,7 @@ def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( assert purge_after_failure.error_code is None -def test_fail_purge_operation_rejects_terminal_complete_state( +def test_fail_purge_operation_rejects_inactive_claim_after_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -413,10 +685,11 @@ def test_fail_purge_operation_rejects_terminal_complete_state( idempotency_key="idem_barrier_terminal_purge_fail", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref=request_ref, ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -424,11 +697,12 @@ def test_fail_purge_operation_rejects_terminal_complete_state( request_ref=request_ref, ) - with pytest.raises(ValueError, match="already complete"): + with pytest.raises(ValueError, match="purge execution claim"): storage.fail_purge_operation( purge.purge_id, error_code="governance_erase_failed", error_detail="late_failure", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge.purge_id), ) barrier = storage.get_subject_write_barrier(subject_ref) @@ -470,10 +744,11 @@ def test_guarded_completion_allows_purged_retained_skeletons( idempotency_key="idem_purged_skeletons", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000061", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) assert storage.purge_content(entity_type="profile", entity_id=profile.profile_id) assert storage.purge_content( @@ -535,6 +810,7 @@ def test_guarded_completion_requires_empty_subject_rows( idempotency_key="idem_guarded_complete", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_0123456789abcdef0123456789abcdef", ) @@ -548,14 +824,20 @@ def test_guarded_completion_requires_empty_subject_rows( created_at=_now(), ) ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) storage.record_purge_target( purge.purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - detail={"prepared": True}, + execution_claim=_claim_purge(storage, purge.purge_id), + detail={ + "prepared": True, + "authoritative_user_digest": _authoritative_user_digest( + storage, purge.purge_id + ), + }, ) with pytest.raises(ValueError, match="same-subject rows remain"): @@ -570,6 +852,8 @@ def test_guarded_completion_requires_empty_subject_rows( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -584,6 +868,7 @@ def test_guarded_completion_requires_empty_legacy_null_subject_rows( idempotency_key="idem_guarded_legacy", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000031", ) @@ -605,14 +890,20 @@ def test_guarded_completion_requires_empty_legacy_null_subject_rows( ) storage.conn.commit() - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) storage.record_purge_target( purge.purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - detail={"prepared": True}, + execution_claim=_claim_purge(storage, purge.purge_id), + detail={ + "prepared": True, + "authoritative_user_digest": _authoritative_user_digest( + storage, purge.purge_id + ), + }, ) with pytest.raises(ValueError, match="same-subject rows remain"): @@ -627,6 +918,8 @@ def test_guarded_completion_requires_empty_legacy_null_subject_rows( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -641,6 +934,7 @@ def test_guarded_completion_requires_existing_erasing_subject_barrier( idempotency_key="idem_missing_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000032", ) @@ -658,6 +952,8 @@ def test_guarded_completion_requires_existing_erasing_subject_barrier( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -672,15 +968,17 @@ def test_guarded_completion_rejects_failed_subject_barrier( idempotency_key="idem_failed_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000035", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) storage.fail_subject_erasure_barrier( subject_ref, purge.purge_id, error_code="test_failed_barrier", error_detail="RuntimeError", + execution_claim=_claim_purge(storage, purge.purge_id), ) _mark_all_completion_targets(storage, purge.purge_id) @@ -696,6 +994,8 @@ def test_guarded_completion_rejects_failed_subject_barrier( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -718,10 +1018,11 @@ def test_barrier_blocks_profile_update_paths( idempotency_key="idem_profile_update", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000033", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.update_user_profile_tags("alice", profile.profile_id, ["blocked"]) @@ -758,10 +1059,11 @@ def test_barrier_blocks_user_playbook_update_paths( idempotency_key="idem_playbook_update", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000034", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.archive_user_playbook_by_id("alice", playbook.user_playbook_id) @@ -789,6 +1091,7 @@ def test_assert_subject_writable_blocks_only_barriered_subject( idempotency_key="idem_assert_writable", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=barriered_subject_ref, request_ref=request_ref, ) @@ -797,7 +1100,9 @@ def test_assert_subject_writable_blocks_only_barriered_subject( storage.assert_subject_writable(barriered_subject_ref) storage.assert_subject_writable(other_subject_ref) - storage.begin_subject_erasure_barrier(barriered_subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier( + storage, barriered_subject_ref, purge.purge_id + ) # The 'erasing' barrier blocks only its own subject. with pytest.raises(SubjectWriteBarrierError, match="blocked by erasure barrier"): @@ -831,6 +1136,7 @@ def test_source_window_write_blocks_legacy_null_subject_ref_user_playbook( idempotency_key="idem_source_window_legacy", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000041", ) @@ -867,7 +1173,7 @@ def test_source_window_write_blocks_legacy_null_subject_ref_user_playbook( ] )[0] - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.set_source_windows_for_agent_playbook( diff --git a/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py b/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py index 78300ee39..ce972c2de 100644 --- a/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py @@ -5,6 +5,8 @@ import pytest +from reflexio.models.api_schema.service_schemas import UserProfile +from reflexio.server.services.storage.error import StorageError from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.sqlite_storage._agent_run import _dt from reflexio.server.services.storage.storage_base import ( @@ -97,6 +99,210 @@ def test_sqlite_agent_run_crud_round_trip(storage): assert loaded.generation_request_snapshot == {"request_id": "request_1"} +def test_finalization_receipt_accepts_empty_ids_idempotently(storage): + storage.create_agent_run(_agent_run("run_empty", AgentRunStatus.FINALIZING)) + + storage.save_agent_run_finalization_receipt( + run_id="run_empty", entity_type="profile", learning_ids=[] + ) + storage.save_agent_run_finalization_receipt( + run_id="run_empty", entity_type="profile", learning_ids=[] + ) + + assert ( + storage.get_agent_run_finalization_receipt( + run_id="run_empty", entity_type="profile" + ) + == [] + ) + assert ( + storage.conn.execute( + "SELECT COUNT(*) FROM _agent_run_finalization_receipts " + "WHERE run_id = 'run_empty'" + ).fetchone()[0] + == 1 + ) + + +@pytest.mark.parametrize("malformed_id", [1, None, "", " "]) +def test_finalization_receipt_rejects_malformed_id_before_insert( + storage, + malformed_id, +): + storage.create_agent_run(_agent_run("run_invalid", AgentRunStatus.FINALIZING)) + + with pytest.raises(StorageError, match="non-empty strings"): + storage.save_agent_run_finalization_receipt( + run_id="run_invalid", + entity_type="profile", + learning_ids=[malformed_id], # type: ignore[list-item] + ) + + assert ( + storage.get_agent_run_finalization_receipt( + run_id="run_invalid", entity_type="profile" + ) + is None + ) + + +@pytest.mark.parametrize("encoded_ids", ["[1]", '[" "]', "null"]) +def test_finalization_receipt_rejects_corrupt_persisted_ids(storage, encoded_ids): + storage.create_agent_run(_agent_run("run_corrupt", AgentRunStatus.FINALIZING)) + storage.conn.execute( + """ + INSERT INTO _agent_run_finalization_receipts + (run_id, entity_type, learning_ids) + VALUES (?, ?, ?) + """, + ("run_corrupt", "profile", encoded_ids), + ) + storage.conn.commit() + + with pytest.raises(StorageError, match="corrupt"): + storage.get_agent_run_finalization_receipt( + run_id="run_corrupt", entity_type="profile" + ) + + +def test_finalization_receipt_rolls_back_with_learning(storage): + storage.create_agent_run(_agent_run("run_rollback", AgentRunStatus.FINALIZING)) + profile = UserProfile( + profile_id="profile-rollback", + user_id="user_1", + content="This row must roll back with its receipt.", + last_modified_timestamp=1_000, + generated_from_request_id="request_1", + ) + + with ( + pytest.raises(RuntimeError, match="force rollback"), + storage.commit_scope(), + ): + storage.add_user_profile("user_1", [profile]) + storage.save_agent_run_finalization_receipt( + run_id="run_rollback", + entity_type="profile", + learning_ids=[profile.profile_id], + ) + raise RuntimeError("force rollback") + + assert storage.get_user_profile("user_1") == [] + assert ( + storage.get_agent_run_finalization_receipt( + run_id="run_rollback", entity_type="profile" + ) + is None + ) + + +def test_finalization_receipt_rejects_conflicting_immutable_value(storage): + storage.create_agent_run(_agent_run("run_immutable", AgentRunStatus.FINALIZING)) + storage.save_agent_run_finalization_receipt( + run_id="run_immutable", + entity_type="profile", + learning_ids=["profile-1"], + ) + + with pytest.raises(StorageError, match="immutable"): + storage.save_agent_run_finalization_receipt( + run_id="run_immutable", + entity_type="profile", + learning_ids=["profile-2"], + ) + + assert storage.get_agent_run_finalization_receipt( + run_id="run_immutable", entity_type="profile" + ) == ["profile-1"] + + +def test_finalization_receipt_rejects_get_for_changed_entity_type(storage): + storage.create_agent_run(_agent_run("run_get_type", AgentRunStatus.FINALIZING)) + storage.save_agent_run_finalization_receipt( + run_id="run_get_type", + entity_type="profile", + learning_ids=["profile-1"], + ) + + with pytest.raises(StorageError, match="entity type changed"): + storage.get_agent_run_finalization_receipt( + run_id="run_get_type", entity_type="user_playbook" + ) + + assert storage.get_agent_run_finalization_receipt( + run_id="run_get_type", entity_type="profile" + ) == ["profile-1"] + + +def test_finalization_receipt_rejects_save_for_conflicting_extractor_type(storage): + storage.create_agent_run(_agent_run("run_save_type", AgentRunStatus.FINALIZING)) + storage.save_agent_run_finalization_receipt( + run_id="run_save_type", + entity_type="profile", + learning_ids=["profile-1"], + ) + + with pytest.raises(StorageError, match="entity type is invalid"): + storage.save_agent_run_finalization_receipt( + run_id="run_save_type", + entity_type="user_playbook", + learning_ids=["playbook-1"], + ) + + assert storage.get_agent_run_finalization_receipt( + run_id="run_save_type", entity_type="profile" + ) == ["profile-1"] + + +def test_finalization_receipt_isolated_by_org(tmp_path): + db_path = str(tmp_path / "receipt-orgs.db") + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + owner = SQLiteStorage(org_id="org_1", db_path=db_path) + peer = SQLiteStorage(org_id="org_2", db_path=db_path) + owner.create_agent_run(_agent_run("run_owned", AgentRunStatus.FINALIZING)) + owner.save_agent_run_finalization_receipt( + run_id="run_owned", + entity_type="profile", + learning_ids=["profile-owned"], + ) + + assert ( + peer.get_agent_run_finalization_receipt( + run_id="run_owned", entity_type="profile" + ) + is None + ) + with pytest.raises(StorageError, match="owner"): + peer.save_agent_run_finalization_receipt( + run_id="run_owned", + entity_type="profile", + learning_ids=["profile-peer"], + ) + assert owner.get_agent_run_finalization_receipt( + run_id="run_owned", entity_type="profile" + ) == ["profile-owned"] + + +def test_finalization_receipt_survives_reopen_and_repeated_migration(tmp_path): + db_path = str(tmp_path / "receipt-reopen.db") + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + first = SQLiteStorage(org_id="org_1", db_path=db_path) + first.create_agent_run(_agent_run("run_reopen", AgentRunStatus.FINALIZING)) + first.save_agent_run_finalization_receipt( + run_id="run_reopen", + entity_type="profile", + learning_ids=["profile-reopen"], + ) + first.conn.close() + + reopened = SQLiteStorage(org_id="org_1", db_path=db_path) + reopened.migrate() + + assert reopened.get_agent_run_finalization_receipt( + run_id="run_reopen", entity_type="profile" + ) == ["profile-reopen"] + + def test_sqlite_get_latest_finalized_agent_run_for_request_filters_binding(storage): matching = replace( _agent_run("run_matching", AgentRunStatus.FINALIZED), diff --git a/tests/server/services/storage/sqlite_storage/test_governance_retrieved_learning.py b/tests/server/services/storage/sqlite_storage/test_governance_retrieved_learning.py index 4ff0df379..17d1c385d 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_retrieved_learning.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_retrieved_learning.py @@ -117,11 +117,26 @@ def _erase(storage: SQLiteStorage, user_id: str, purge_id: str) -> dict[str, int idempotency_key=f"idem_{purge_id}", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + authoritative_user_id=user_id, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, ) - storage.prepare_governance_erase_targets(purge_id, user_id) - return storage.apply_governance_user_data_delete(purge_id, user_id) + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"test-{purge_id}", + lease_ttl_seconds=30, + ) + assert claim is not None + storage.prepare_governance_erase_targets( + purge_id, + user_id, + execution_claim=claim, + ) + return storage.apply_governance_user_data_delete( + purge_id, + user_id, + execution_claim=claim, + ) def test_erase_scrubs_rle_rows_and_all_state_namespaces(storage) -> None: diff --git a/tests/server/services/storage/sqlite_storage/test_governance_storage.py b/tests/server/services/storage/sqlite_storage/test_governance_storage.py index 880cbad09..2d2e45f77 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -1,10 +1,15 @@ from __future__ import annotations +import ast +import hashlib +import inspect import json import sqlite3 import threading import time from concurrent.futures import ThreadPoolExecutor +from collections.abc import Callable +from pathlib import Path from typing import Any, Literal, cast from unittest.mock import patch @@ -24,6 +29,8 @@ ) from reflexio.models.api_schema.retriever_schema import SearchAgentPlaybookRequest from reflexio.models.config_schema import GovernanceRetentionConfig +from reflexio.server.services.governance.config import governance_subject_ref +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, ) @@ -37,17 +44,50 @@ from reflexio.server.services.storage.sqlite_storage.governance import ( _purge as purge_module, ) +from reflexio.server.services.storage.storage_base.governance._erase_execution import ( + GovernanceEraseExecutionMixin, +) +from reflexio.server.services.storage.storage_base.governance._purge import ( + PurgeOperationStoreMixin, +) +from reflexio.server.services.storage.storage_base.governance._subject_barrier import ( + SubjectBarrierMixin, +) pytestmark = pytest.mark.integration -SUBJECT_REF = "subref_v1_" + "a" * 32 -OTHER_SUBJECT_REF = "subref_v1_" + "c" * 32 +SUBJECT_REF = governance_subject_ref("org1", "alice", "test-governance-secret") +OTHER_SUBJECT_REF = governance_subject_ref("org1", "bob", "test-governance-secret") REQUEST_REF = "reqref_v1_" + "b" * 32 OTHER_REQUEST_REF = "reqref_v1_" + "d" * 32 ACTOR_REF = "actref_v1_" + "e" * 32 # Single source of truth — a stale local copy of this tuple is exactly how # this suite went red when new canonical targets landed without test updates. CANONICAL_DELETE_TARGET_NAMES = _CANONICAL_DELETE_TARGET_NAMES +CLAIMED_ERASURE_MUTATIONS = { + "record_purge_target", + "prepare_governance_erase_targets", + "fail_purge_operation", + "begin_subject_erasure_barrier", + "complete_subject_erasure_barrier_after_empty_check", + "fail_subject_erasure_barrier", + "apply_governance_user_data_delete", + "complete_purge_operation_with_audit", +} + + +def _begin_test_purge_operation(storage: SQLiteStorage, **kwargs: Any): + if ( + kwargs.get("operation_type") == "user_erasure" + and kwargs.get("scope_type") == "user" + and "authoritative_user_id" not in kwargs + ): + subject_ref = kwargs.get("subject_ref") + for user_id in ("alice", "bob"): + if storage._subject_ref_for_user_id(user_id) == subject_ref: + kwargs["authoritative_user_id"] = user_id + break + return SQLiteStorage.begin_purge_operation(storage, **kwargs) @pytest.fixture @@ -68,29 +108,155 @@ def _make_storage(org_id: str) -> SQLiteStorage: yield _make_storage -def _begin_purge(storage: SQLiteStorage, purge_id: str) -> str: - purge = storage.begin_purge_operation( +def _begin_purge( + storage: SQLiteStorage, + purge_id: str, + *, + subject_ref: str | None = None, + authoritative_user_id: str = "alice", +) -> str: + subject_ref = subject_ref or storage._subject_ref_for_user_id(authoritative_user_id) + purge = _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key=f"idem_{purge_id}", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=subject_ref, request_ref=REQUEST_REF, + authoritative_user_id=authoritative_user_id, ) + claim = _claim_purge(storage, purge.purge_id) storage.record_purge_target( purge_id=purge.purge_id, target_name="target_snapshot", target_ref="all", phase="prepare_targets", status="complete", + execution_claim=claim, detail={ + "authoritative_user_digest": storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge.purge_id), + ).fetchone()["authoritative_user_digest"], "owned_user_playbook_ids": [11], }, ) return purge.purge_id -def _add_complete_delete_target_matrix(storage: SQLiteStorage, purge_id: str) -> None: +def _begin_raw_user_erasure_purge(storage: SQLiteStorage, purge_id: str) -> str: + return _begin_raw_user_erasure_purge_for_subject( + storage, purge_id, subject_ref=SUBJECT_REF + ) + + +def _begin_raw_user_erasure_purge_for_subject( + storage: SQLiteStorage, + purge_id: str, + *, + subject_ref: str, + authoritative_user_id: str = "alice", +) -> str: + purge = _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=subject_ref, + request_ref=REQUEST_REF, + authoritative_user_id=authoritative_user_id, + ) + return purge.purge_id + + +def _claim_then_take_over(storage: SQLiteStorage, purge_id: str): + first_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-a", + lease_ttl_seconds=30, + ) + if first_claim is None: + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + first_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-a", + lease_ttl_seconds=30, + ) + assert first_claim is not None + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + takeover_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-b", + lease_ttl_seconds=30, + ) + assert takeover_claim is not None + assert takeover_claim.fence == first_claim.fence + 1 + return first_claim, takeover_claim + + +def _claim_purge(storage: SQLiteStorage, purge_id: str) -> PurgeExecutionClaim: + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"owner-{purge_id}", + lease_ttl_seconds=30, + ) + if claim is None: + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"owner-{purge_id}", + lease_ttl_seconds=30, + ) + assert claim is not None + return claim + + +def _typed_test_claim_for_unvalidated_purge_id(purge_id: str) -> PurgeExecutionClaim: + return PurgeExecutionClaim( + purge_id=purge_id, + owner="test-invalid-purge-id", + fence=1, + expires_at=1, + ) + + +def _assert_rejects_missing_claim( + omitted_call: Callable[[], object], + none_call: Callable[[], object], +) -> None: + with pytest.raises((TypeError, ValueError)): + omitted_call() + with pytest.raises((TypeError, ValueError)): + none_call() + + +def _add_complete_delete_target_matrix( + storage: SQLiteStorage, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, +) -> None: for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -98,13 +264,34 @@ def _add_complete_delete_target_matrix(storage: SQLiteStorage, purge_id: str) -> target_ref="all", phase="delete", status="complete", + execution_claim=execution_claim, ) -def _begin_completeable_purge(storage: SQLiteStorage, purge_id: str) -> str: - purge_id = _begin_purge(storage, purge_id) - _add_complete_delete_target_matrix(storage, purge_id) - storage.begin_subject_erasure_barrier(SUBJECT_REF, purge_id) +def _begin_completeable_purge( + storage: SQLiteStorage, + purge_id: str, + *, + subject_ref: str = SUBJECT_REF, + authoritative_user_id: str = "alice", +) -> str: + purge_id = _begin_purge( + storage, + purge_id, + subject_ref=subject_ref, + authoritative_user_id=authoritative_user_id, + ) + claim = _claim_purge(storage, purge_id) + _add_complete_delete_target_matrix( + storage, + purge_id, + execution_claim=claim, + ) + storage.begin_subject_erasure_barrier( + subject_ref, + purge_id, + execution_claim=claim, + ) return purge_id @@ -113,12 +300,13 @@ def _erase_event( purge_id: str, status: AuditStatus = "ok", operation: AuditOperation = "ERASE", + subject_ref: str = SUBJECT_REF, ): return AuditEvent( org_id="org1", operation=operation, entity_type="request", - subject_ref=SUBJECT_REF, + subject_ref=subject_ref, request_ref=REQUEST_REF, idempotency_key=purge_id, status=status, @@ -346,6 +534,7 @@ def _record_agent_playbook_rebuild_target( target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status=status, + execution_claim=_claim_purge(storage, purge_id), detail={ "original_source_windows": original_windows or [ @@ -487,7 +676,8 @@ def test_list_audit_events_rejects_cross_org_override(storage_factory): def test_purge_targets_require_snapshot_marker(storage): - purge = storage.begin_purge_operation( + purge = _begin_test_purge_operation( + storage, purge_id="purge_snapshot_marker", idempotency_key="idem_snapshot_marker", operation_type="user_erasure", @@ -502,10 +692,15 @@ def test_purge_targets_require_snapshot_marker(storage): phase="delete", status="complete", deleted_count=1, + execution_claim=_claim_purge(storage, purge.purge_id), ) assert storage.purge_targets_prepared(purge.purge_id) is False - storage.begin_subject_erasure_barrier(SUBJECT_REF, purge.purge_id) + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge.purge_id, + execution_claim=_claim_purge(storage, purge.purge_id), + ) with pytest.raises(ValueError, match="target snapshot"): storage.complete_purge_operation_with_audit( purge.purge_id, @@ -517,25 +712,293 @@ def test_purge_targets_require_snapshot_marker(storage): request_ref=REQUEST_REF, idempotency_key=purge.purge_id, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) def test_complete_purge_operation_with_audit_is_atomic_success_path(storage): purge_id = _begin_completeable_purge(storage, "purge_atomic_success") + complete_claim = _claim_purge(storage, purge_id) complete = storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=complete_claim, ) assert complete.status == "complete" rows = storage.list_audit_events(subject_ref=SUBJECT_REF) assert [row.operation for row in rows] == ["ERASE"] - same = storage.complete_purge_operation_with_audit( + with pytest.raises(ValueError, match="purge execution claim"): + storage.complete_purge_operation_with_audit( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=complete_claim, + ) + assert len(storage.list_audit_events(subject_ref=SUBJECT_REF)) == 1 + + +def test_complete_purge_operation_with_audit_rejects_wrong_authoritative_user(storage): + purge_id = _begin_completeable_purge(storage, "purge_wrong_complete_user") + + with pytest.raises(ValueError, match="authoritative user identity"): + storage.complete_purge_operation_with_audit( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="bob", + execution_claim=_claim_purge(storage, purge_id), + ) + + assert storage.get_purge_operation(purge_id).status == "running" + assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] + + +@pytest.mark.parametrize( + "complete_method_name", + [ + "complete_purge_operation_with_audit", + "complete_subject_erasure_barrier_after_empty_check", + ], +) +def test_user_erasure_completion_contracts_reject_org_purge( + storage, + complete_method_name, +): + purge = _begin_test_purge_operation( + storage, + purge_id=f"purge_org_scope_{complete_method_name}", + idempotency_key=f"idem_org_scope_{complete_method_name}", + operation_type="org_purge", + scope_type="org", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + ) + claim = _claim_purge(storage, purge.purge_id) + storage.record_purge_target( + purge.purge_id, + target_name="target_snapshot", + target_ref="all", + phase="prepare_targets", + status="complete", + detail={"prepared": True}, + execution_claim=claim, + ) + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge.purge_id, + execution_claim=claim, + ) + + complete = getattr(storage, complete_method_name) + with pytest.raises(ValueError, match="user erasure"): + complete( + purge.purge_id, + _erase_event(purge_id=purge.purge_id), + authoritative_user_id="arbitrary-user", + execution_claim=claim, + ) + + assert storage.get_purge_operation(purge.purge_id).status == "running" + assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] + + +@pytest.mark.parametrize( + ("purge_binding", "snapshot_binding"), + [ + pytest.param("legacy", "legacy", id="both-unkeyed"), + pytest.param(None, None, id="both-null"), + pytest.param("current", "legacy", id="interrupted-row-only-upgrade"), + ], +) +def test_idempotent_retry_upgrades_legacy_identity_bindings_and_resumes_completion( + storage, + purge_binding, + snapshot_binding, +): + purge_id = _begin_completeable_purge(storage, "purge_interrupted_legacy_resume") + current_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + legacy_digest = hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + + def resolve_binding(binding): + return ( + legacy_digest + if binding == "legacy" + else current_digest + if binding == "current" + else None + ) + + snapshot_row = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + snapshot_detail = json.loads(snapshot_row["detail"]) + snapshot_detail["authoritative_user_digest"] = resolve_binding(snapshot_binding) + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + (resolve_binding(purge_binding), storage.org_id, purge_id), + ) + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(snapshot_detail), storage.org_id, purge_id), + ) + storage.conn.commit() + + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + + adopted_purge_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + adopted_snapshot = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + assert adopted_purge_digest == current_digest + assert ( + json.loads(adopted_snapshot["detail"])["authoritative_user_digest"] + == current_digest + ) + + completed = storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) + assert completed.status == "complete" + + +def test_idempotent_retry_rolls_back_purge_digest_when_snapshot_binding_mismatches( + storage, +): + purge_id = _begin_completeable_purge(storage, "purge_mismatched_legacy_snapshot") + legacy_digest = hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + snapshot_row = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + snapshot_detail = json.loads(snapshot_row["detail"]) + snapshot_detail["authoritative_user_digest"] = "mismatched-digest" + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + (legacy_digest, storage.org_id, purge_id), + ) + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(snapshot_detail), storage.org_id, purge_id), + ) + storage.conn.commit() + + with pytest.raises(ValueError, match="authoritative user identity"): + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + + persisted_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + assert persisted_digest == legacy_digest + + +def test_idempotent_retry_rolls_back_row_upgrade_when_snapshot_upgrade_fails(storage): + purge_id = _begin_completeable_purge(storage, "purge_snapshot_upgrade_failure") + legacy_digest = hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + snapshot_row = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + snapshot_detail = json.loads(snapshot_row["detail"]) + snapshot_detail["authoritative_user_digest"] = legacy_digest + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + (legacy_digest, storage.org_id, purge_id), + ) + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(snapshot_detail), storage.org_id, purge_id), + ) + storage.conn.execute( + f"""CREATE TRIGGER fail_snapshot_digest_upgrade + BEFORE UPDATE OF detail ON purge_operation_targets + WHEN OLD.org_id = '{storage.org_id}' AND OLD.purge_id = '{purge_id}' + AND OLD.target_name = 'target_snapshot' + BEGIN + SELECT RAISE(ABORT, 'snapshot upgrade failed'); + END""" + ) + storage.conn.commit() + + with pytest.raises(sqlite3.IntegrityError, match="snapshot upgrade failed"): + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + + persisted_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + persisted_snapshot = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + assert persisted_digest == legacy_digest + assert ( + json.loads(persisted_snapshot["detail"])["authoritative_user_digest"] + == legacy_digest ) - assert same.status == "complete" - assert len(storage.list_audit_events(subject_ref=SUBJECT_REF)) == 1 def test_complete_purge_operation_with_audit_begins_immediate_transaction_before_reads( @@ -548,6 +1011,8 @@ def test_complete_purge_operation_with_audit_begins_immediate_transaction_before storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) finally: storage.conn.set_trace_callback(None) @@ -619,6 +1084,8 @@ def test_complete_purge_operation_with_audit_accepts_planned_success_detail(stor "rebuilt_agent_playbook_ids": rebuilt_ids, }, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert complete.status == "complete" @@ -648,7 +1115,12 @@ def test_complete_purge_operation_rejects_audit_refs_that_mismatch_persisted_pur event = _erase_event(purge_id=purge_id).model_copy(update=event_kwargs) with pytest.raises(ValueError, match=match): - storage.complete_purge_operation_with_audit(purge_id, event) + storage.complete_purge_operation_with_audit( + purge_id, + event, + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) assert storage.get_purge_operation(purge_id).status == "running" assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] @@ -671,7 +1143,8 @@ def test_complete_purge_operation_rejects_audit_refs_that_mismatch_persisted_pur def test_begin_purge_operation_rejects_mismatched_idempotent_retry( storage, retry_kwargs, match ): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_begin_retry", idempotency_key="idem_begin_retry", operation_type="user_erasure", @@ -681,7 +1154,8 @@ def test_begin_purge_operation_rejects_mismatched_idempotent_retry( ) with pytest.raises(ValueError, match=match): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=retry_kwargs.get("purge_id", "purge_begin_retry"), idempotency_key="idem_begin_retry", operation_type=retry_kwargs.get("operation_type", "user_erasure"), @@ -698,7 +1172,8 @@ def test_begin_purge_operation_rejects_mismatched_idempotent_retry( def test_begin_purge_operation_rejects_numeric_idempotency_key(storage): with pytest.raises(ValueError, match="idempotency_key"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_numeric_idem", idempotency_key="12345", operation_type="user_erasure", @@ -714,7 +1189,8 @@ def test_begin_purge_operation_rejects_numeric_idempotency_key(storage): def test_begin_purge_operation_accepts_code_shaped_idempotency_key_with_content( storage, ): - purge = storage.begin_purge_operation( + purge = _begin_test_purge_operation( + storage, purge_id="purge_content_retry", idempotency_key="content_purge_retry_1", operation_type="user_erasure", @@ -729,7 +1205,8 @@ def test_begin_purge_operation_accepts_code_shaped_idempotency_key_with_content( @pytest.mark.parametrize("purge_id", ["purge_1", "purge_123"]) def test_begin_purge_operation_rejects_raw_numeric_purge_suffix(storage, purge_id): with pytest.raises(ValueError, match="purge_id"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key=f"idem_{purge_id}", operation_type="user_erasure", @@ -787,7 +1264,12 @@ def test_complete_purge_operation_rejects_invalid_audit_event(storage, event, ma purge_id = _begin_completeable_purge(storage, "purge_invalid") with pytest.raises(ValueError, match=match): - storage.complete_purge_operation_with_audit(purge_id, event) + storage.complete_purge_operation_with_audit( + purge_id, + event, + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) assert storage.get_purge_operation(purge_id).status == "running" assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] @@ -816,7 +1298,10 @@ def test_complete_purge_operation_requires_matching_existing_erase_row( with pytest.raises(ValueError, match=match): storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).status == "running" @@ -873,7 +1358,10 @@ def test_complete_purge_operation_rejects_mismatched_existing_erase_row( with pytest.raises(ValueError, match="matching successful ERASE"): storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).status == "running" @@ -904,12 +1392,18 @@ def test_append_audit_event_rejects_successful_erase_without_idempotency_key(sto def test_complete_purge_operation_requires_full_delete_target_matrix(storage): purge_id = _begin_purge(storage, "purge_snapshot_only") - storage.begin_subject_erasure_barrier(SUBJECT_REF, purge_id) + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge_id, + execution_claim=_claim_purge(storage, purge_id), + ) with pytest.raises(ValueError, match="delete target matrix"): storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).status == "running" @@ -923,6 +1417,7 @@ def test_complete_retry_replaces_failed_completed_at(storage): purge_id, error_code="governance_erase_failed", error_detail="RuntimeError", + execution_claim=_claim_purge(storage, purge_id), ) assert failed.completed_at == 111 @@ -930,25 +1425,641 @@ def test_complete_retry_replaces_failed_completed_at(storage): completed = storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert completed.status == "complete" assert completed.completed_at == 222 +def test_stale_execution_claim_takeover_fences_previous_owner(storage): + purge_id = _begin_purge(storage, "purge_stale_claim") + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + first_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-a", + lease_ttl_seconds=30, + ) + assert first_claim is not None + + live_duplicate = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-b", + lease_ttl_seconds=30, + ) + assert live_duplicate is None + + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + + takeover_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-b", + lease_ttl_seconds=30, + ) + assert takeover_claim is not None + assert takeover_claim.owner == "owner-b" + assert takeover_claim.fence == first_claim.fence + 1 + + with pytest.raises(ValueError, match="purge execution claim"): + storage.record_purge_target( + purge_id=purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=first_claim, + ) + + storage.record_purge_target( + purge_id=purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=takeover_claim, + ) + targets = storage.list_purge_targets(purge_id, phase="delete") + assert [(target.target_name, target.status) for target in targets] == [ + ("interaction", "complete") + ] + + +def test_shared_file_claim_takeover_fences_independent_storage_instance( + storage_factory, +) -> None: + storage_a = storage_factory("org1") + storage_b = storage_factory("org1") + subject_ref = governance_subject_ref("org1", "alice", "test-governance-secret") + purge = storage_a.begin_purge_operation( + purge_id="purge_cross_connection_claim", + idempotency_key="idem_cross_connection_claim", + operation_type="user_erasure", + scope_type="user", + subject_ref=subject_ref, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + ready = threading.Barrier(2) + claims: list[PurgeExecutionClaim | None] = [] + + def claim(storage_instance: SQLiteStorage, owner: str) -> None: + ready.wait(timeout=5) + claims.append( + storage_instance.claim_purge_operation_execution( + purge.purge_id, + lease_owner=owner, + lease_ttl_seconds=30, + ) + ) + + callers = [ + threading.Thread(target=claim, args=(storage_a, "owner-a")), + threading.Thread(target=claim, args=(storage_b, "owner-b")), + ] + for caller in callers: + caller.start() + for caller in callers: + caller.join(timeout=5) + + assert all(not caller.is_alive() for caller in callers) + live_claims = [claim for claim in claims if claim is not None] + assert len(live_claims) == 1 + first_claim = live_claims[0] + storage_b.conn.execute( + """UPDATE purge_operations SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage_b.org_id, purge.purge_id), + ) + storage_b.conn.commit() + takeover = storage_b.claim_purge_operation_execution( + purge.purge_id, + lease_owner="takeover", + lease_ttl_seconds=30, + ) + assert takeover is not None + assert takeover.fence == first_claim.fence + 1 + + with pytest.raises(ValueError, match="purge execution claim"): + storage_a.record_purge_target( + purge_id=purge.purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=first_claim, + ) + storage_b.record_purge_target( + purge_id=purge.purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=takeover, + ) + + +def test_user_erasure_rejects_mismatched_authoritative_identity_at_each_stage( + storage, +) -> None: + alice_ref = governance_subject_ref( + storage.org_id, "alice", "test-governance-secret" + ) + with pytest.raises(ValueError, match="authoritative user"): + _begin_test_purge_operation( + storage, + purge_id="purge_identity_begin_mismatch", + idempotency_key="idem_identity_begin_mismatch", + operation_type="user_erasure", + scope_type="user", + subject_ref=alice_ref, + request_ref=REQUEST_REF, + authoritative_user_id="bob", + ) + + purge = _begin_test_purge_operation( + storage, + purge_id="purge_identity_stage_mismatch", + idempotency_key="idem_identity_stage_mismatch", + operation_type="user_erasure", + scope_type="user", + subject_ref=alice_ref, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + claim = _claim_purge(storage, purge.purge_id) + with pytest.raises(ValueError, match="authoritative user"): + storage.prepare_governance_erase_targets( + purge.purge_id, + "bob", + execution_claim=claim, + ) + storage.prepare_governance_erase_targets( + purge.purge_id, + "alice", + execution_claim=claim, + ) + with pytest.raises(ValueError, match="authoritative user"): + storage.apply_governance_user_data_delete( + purge.purge_id, + "bob", + execution_claim=claim, + ) + + with pytest.raises(ValueError, match="authoritative user"): + _begin_test_purge_operation( + storage, + purge_id=purge.purge_id, + idempotency_key="idem_identity_stage_mismatch", + operation_type="user_erasure", + scope_type="user", + subject_ref=alice_ref, + request_ref=REQUEST_REF, + authoritative_user_id="bob", + ) + + +def test_claimed_erasure_mutation_signatures_and_callers_require_claim() -> None: + method_owners = { + PurgeOperationStoreMixin: { + "record_purge_target", + "prepare_governance_erase_targets", + "fail_purge_operation", + }, + SubjectBarrierMixin: { + "begin_subject_erasure_barrier", + "complete_subject_erasure_barrier_after_empty_check", + "fail_subject_erasure_barrier", + }, + GovernanceEraseExecutionMixin: { + "apply_governance_user_data_delete", + "complete_purge_operation_with_audit", + }, + SQLiteStorage: CLAIMED_ERASURE_MUTATIONS, + } + for owner, method_names in method_owners.items(): + for method_name in method_names: + parameter = inspect.signature(getattr(owner, method_name)).parameters[ + "execution_claim" + ] + assert parameter.kind is inspect.Parameter.KEYWORD_ONLY, method_name + assert parameter.default is inspect.Parameter.empty, method_name + assert parameter.annotation in {"PurgeExecutionClaim", PurgeExecutionClaim} + + production_root = Path(__file__).resolve().parents[5] / "reflexio" + violations: list[str] = [] + for path in production_root.rglob("*.py"): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + function = node.func + if not isinstance(function, ast.Attribute): + continue + if function.attr not in CLAIMED_ERASURE_MUTATIONS: + continue + claim_keywords = [ + keyword for keyword in node.keywords if keyword.arg == "execution_claim" + ] + if not claim_keywords: + violations.append(f"{path.relative_to(production_root)}:{node.lineno}") + continue + if any(keyword.arg is None for keyword in node.keywords): + violations.append( + f"{path.relative_to(production_root)}:{node.lineno}: **kwargs" + ) + claim_value = claim_keywords[0].value + if isinstance(claim_value, ast.Constant) and claim_value.value is None: + violations.append( + f"{path.relative_to(production_root)}:{node.lineno}: None" + ) + assert violations == [] + + +def test_sqlite_claimed_erasure_mutations_reject_omitted_and_none_claim(storage): + target_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_target") + _, target_claim = _claim_then_take_over(storage, target_purge_id) + _assert_rejects_missing_claim( + lambda: storage.record_purge_target( + purge_id=target_purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + ), + lambda: storage.record_purge_target( + purge_id=target_purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.list_purge_targets(target_purge_id, phase="delete") == [] + storage.record_purge_target( + purge_id=target_purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=target_claim, + ) + + barrier_user_id = "no-claim-barrier-user" + barrier_subject_ref = storage._subject_ref_for_user_id(barrier_user_id) + barrier_purge_id = _begin_raw_user_erasure_purge_for_subject( + storage, + "purge_no_claim_barrier", + subject_ref=barrier_subject_ref, + authoritative_user_id=barrier_user_id, + ) + _, barrier_claim = _claim_then_take_over(storage, barrier_purge_id) + _assert_rejects_missing_claim( + lambda: storage.begin_subject_erasure_barrier( + barrier_subject_ref, barrier_purge_id + ), + lambda: storage.begin_subject_erasure_barrier( + barrier_subject_ref, + barrier_purge_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_subject_write_barrier(barrier_subject_ref) is None + storage.begin_subject_erasure_barrier( + barrier_subject_ref, + barrier_purge_id, + execution_claim=barrier_claim, + ) + + prepare_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_prepare") + _, prepare_claim = _claim_then_take_over(storage, prepare_purge_id) + _assert_rejects_missing_claim( + lambda: storage.prepare_governance_erase_targets( + purge_id=prepare_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + ), + lambda: storage.prepare_governance_erase_targets( + purge_id=prepare_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.list_purge_targets(prepare_purge_id) == [] + storage.prepare_governance_erase_targets( + purge_id=prepare_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=prepare_claim, + ) + + delete_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_delete") + delete_claim = _claim_purge(storage, delete_purge_id) + storage.prepare_governance_erase_targets( + purge_id=delete_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=delete_claim, + ) + _assert_rejects_missing_claim( + lambda: storage.apply_governance_user_data_delete(delete_purge_id, "alice"), + lambda: storage.apply_governance_user_data_delete( + delete_purge_id, + "alice", + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert all( + target.status == "pending" + for target in storage.list_purge_targets(delete_purge_id, phase="delete") + ) + + complete_user_id = "no-claim-complete-user" + complete_subject_ref = storage._subject_ref_for_user_id(complete_user_id) + complete_purge_id = _begin_completeable_purge( + storage, + "purge_no_claim_complete", + subject_ref=complete_subject_ref, + authoritative_user_id=complete_user_id, + ) + _claim_purge(storage, complete_purge_id) + _assert_rejects_missing_claim( + lambda: storage.complete_subject_erasure_barrier_after_empty_check( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + ), + lambda: storage.complete_subject_erasure_barrier_after_empty_check( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_purge_operation(complete_purge_id).status == "running" + _assert_rejects_missing_claim( + lambda: storage.complete_purge_operation_with_audit( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + ), + lambda: storage.complete_purge_operation_with_audit( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_purge_operation(complete_purge_id).status == "running" + + fail_barrier_user_id = "no-claim-fail-barrier-user" + fail_barrier_subject_ref = storage._subject_ref_for_user_id(fail_barrier_user_id) + fail_barrier_purge_id = _begin_raw_user_erasure_purge_for_subject( + storage, + "purge_no_claim_fail_barrier", + subject_ref=fail_barrier_subject_ref, + authoritative_user_id=fail_barrier_user_id, + ) + fail_barrier_claim = _claim_purge(storage, fail_barrier_purge_id) + storage.begin_subject_erasure_barrier( + fail_barrier_subject_ref, + fail_barrier_purge_id, + execution_claim=fail_barrier_claim, + ) + _assert_rejects_missing_claim( + lambda: storage.fail_subject_erasure_barrier( + fail_barrier_subject_ref, + fail_barrier_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + ), + lambda: storage.fail_subject_erasure_barrier( + fail_barrier_subject_ref, + fail_barrier_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=None, # type: ignore[arg-type] + ), + ) + barrier = storage.get_subject_write_barrier(fail_barrier_subject_ref) + assert barrier is not None + assert barrier.status == "erasing" + + fail_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_fail") + _claim_purge(storage, fail_purge_id) + _assert_rejects_missing_claim( + lambda: storage.fail_purge_operation( + fail_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + ), + lambda: storage.fail_purge_operation( + fail_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_purge_operation(fail_purge_id).status == "running" + + +def test_stale_execution_claim_cannot_start_subject_barrier(storage): + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_barrier_start") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge_id, + execution_claim=first_claim, + ) + assert storage.get_subject_write_barrier(SUBJECT_REF) is None + + barrier = storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge_id, + execution_claim=takeover_claim, + ) + assert barrier.status == "erasing" + + +def test_stale_execution_claim_cannot_prepare_delete_targets(storage): + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_prepare") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.prepare_governance_erase_targets( + purge_id=purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=first_claim, + ) + assert storage.list_purge_targets(purge_id) == [] + + storage.prepare_governance_erase_targets( + purge_id=purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=takeover_claim, + ) + assert storage.purge_targets_prepared(purge_id) + + +def test_stale_execution_claim_cannot_apply_protected_delete(storage): + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_delete") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + storage.prepare_governance_erase_targets( + purge_id=purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=takeover_claim, + ) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.apply_governance_user_data_delete( + purge_id, + "alice", + execution_claim=first_claim, + ) + assert all( + target.status == "pending" + for target in storage.list_purge_targets(purge_id, phase="delete") + ) + + counts = storage.apply_governance_user_data_delete( + purge_id, + "alice", + execution_claim=takeover_claim, + ) + assert counts["requests"] == 0 + assert all( + target.status == "complete" + for target in storage.list_purge_targets(purge_id, phase="delete") + ) + + +def test_stale_execution_claim_cannot_complete_purge_or_barrier(storage): + purge_id = _begin_completeable_purge(storage, "purge_stale_complete") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.complete_subject_erasure_barrier_after_empty_check( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=first_claim, + ) + assert storage.get_purge_operation(purge_id).status == "running" + barrier = storage.get_subject_write_barrier(SUBJECT_REF) + assert barrier is not None + assert barrier.status == "erasing" + assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] + + completed = storage.complete_subject_erasure_barrier_after_empty_check( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=takeover_claim, + ) + assert completed.status == "complete" + + +def test_stale_execution_claim_cannot_fail_purge_or_barrier(storage): + barrier_purge_id = _begin_raw_user_erasure_purge( + storage, "purge_stale_barrier_failure" + ) + barrier_claim = _claim_purge(storage, barrier_purge_id) + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + barrier_purge_id, + execution_claim=barrier_claim, + ) + first_claim, takeover_claim = _claim_then_take_over(storage, barrier_purge_id) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.fail_subject_erasure_barrier( + SUBJECT_REF, + barrier_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=first_claim, + ) + barrier = storage.get_subject_write_barrier(SUBJECT_REF) + assert barrier is not None + assert barrier.status == "erasing" + + failed_barrier = storage.fail_subject_erasure_barrier( + SUBJECT_REF, + barrier_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=takeover_claim, + ) + assert failed_barrier.status == "failed" + + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_purge_failure") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + with pytest.raises(ValueError, match="purge execution claim"): + storage.fail_purge_operation( + purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=first_claim, + ) + assert storage.get_purge_operation(purge_id).status == "running" + + failed_purge = storage.fail_purge_operation( + purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=takeover_claim, + ) + assert failed_purge.status == "failed" + + def test_prepare_governance_erase_targets_sanitizes_snapshot_detail(storage): - storage.begin_purge_operation( + user_id = "user_123@example.com" + _begin_test_purge_operation( + storage, purge_id="purge_detail", idempotency_key="idem_purge_detail", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) storage.prepare_governance_erase_targets( purge_id="purge_detail", - user_id="user_123@example.com", + user_id=user_id, owned_user_playbook_ids={7}, + execution_claim=_claim_purge(storage, "purge_detail"), ) snapshot = next( @@ -958,23 +2069,33 @@ def test_prepare_governance_erase_targets_sanitizes_snapshot_detail(storage): ) if target.target_name == "target_snapshot" ) - assert snapshot.detail == {"owned_user_playbook_ids": [7]} + assert snapshot.detail == { + "authoritative_user_digest": storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = 'purge_detail'""", + (storage.org_id,), + ).fetchone()["authoritative_user_digest"], + "owned_user_playbook_ids": [7], + } def test_apply_governance_user_data_delete_rejects_playbook_snapshot_drift(storage): user_id = "user-snapshot-drift" owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_snapshot_drift", idempotency_key="idem_purge_snapshot_drift", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) storage.prepare_governance_erase_targets( purge_id="purge_snapshot_drift", user_id=user_id, + execution_claim=_claim_purge(storage, "purge_snapshot_drift"), ) storage.conn.execute( """INSERT INTO user_playbooks ( @@ -991,7 +2112,11 @@ def test_apply_governance_user_data_delete_rejects_playbook_snapshot_drift(stora storage.conn.commit() with pytest.raises(ValueError, match="prepared purge snapshot"): - storage.apply_governance_user_data_delete("purge_snapshot_drift", user_id) + storage.apply_governance_user_data_delete( + "purge_snapshot_drift", + user_id, + execution_claim=_claim_purge(storage, "purge_snapshot_drift"), + ) remaining_ids = { int(row["user_playbook_id"]) @@ -1006,13 +2131,16 @@ def test_apply_governance_user_data_delete_rejects_playbook_snapshot_drift(stora def test_prepare_governance_erase_targets_does_not_plan_org_agent_playbook_rebuilds( storage, ): - storage.begin_purge_operation( + user_id = "user-rebuild-windows" + _begin_test_purge_operation( + storage, purge_id="purge_rebuild_windows", idempotency_key="idem_purge_rebuild_windows", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) agent_playbook_id = _seed_agent_playbook( storage, @@ -1026,8 +2154,9 @@ def test_prepare_governance_erase_targets_does_not_plan_org_agent_playbook_rebui storage.prepare_governance_erase_targets( purge_id="purge_rebuild_windows", - user_id="user-rebuild-windows", + user_id=user_id, owned_user_playbook_ids={7}, + execution_claim=_claim_purge(storage, "purge_rebuild_windows"), ) assert ( @@ -1049,19 +2178,22 @@ def test_prepare_governance_erase_targets_records_full_delete_matrix_counts(stor session_id="session_seed", evaluation_name="governance_prepare_counts", ) - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_prepare_counts", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) storage.prepare_governance_erase_targets( purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) delete_targets = { @@ -1099,13 +2231,15 @@ def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_an storage, ): purge_id = "purge_hide_rebuild" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_hide_rebuild", operation_type="user_erasure", scope_type="user", subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, + authoritative_user_id="alice", ) agent_playbook_id = _seed_agent_playbook( storage, @@ -1161,7 +2295,8 @@ def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_an def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage): purge_id = "purge_rebuild_complete" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_complete", operation_type="user_erasure", @@ -1177,12 +2312,14 @@ def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1200,6 +2337,7 @@ def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) expected_detail = { "original_source_windows": [ @@ -1261,7 +2399,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_ad_hoc_rebuild_without_ storage, ): purge_id = "purge_rebuild_requires_target" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_requires_target", operation_type="user_erasure", @@ -1324,7 +2463,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_rebuild_before_hide_pha storage, ): purge_id = "purge_rebuild_requires_hide" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_requires_hide", operation_type="user_erasure", @@ -1340,12 +2480,14 @@ def test_apply_governance_agent_playbook_rebuild_rejects_rebuild_before_hide_pha AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1419,7 +2561,8 @@ def test_apply_governance_agent_playbook_rebuild_succeeds_after_prepare_and_hide storage, ): purge_id = "purge_rebuild_prepare_hide" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_prepare_hide", operation_type="user_erasure", @@ -1485,7 +2628,8 @@ def test_apply_governance_agent_playbook_rebuild_does_not_complete_target_when_s storage, monkeypatch ): purge_id = "purge_rebuild_search_refresh_failure" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_search_refresh_failure", operation_type="user_erasure", @@ -1501,12 +2645,14 @@ def test_apply_governance_agent_playbook_rebuild_does_not_complete_target_when_s AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1524,6 +2670,7 @@ def test_apply_governance_agent_playbook_rebuild_does_not_complete_target_when_s target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) original_row = storage.conn.execute( """SELECT content, trigger, rationale, blocking_issue, expanded_terms, tags, status @@ -1598,7 +2745,8 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when storage, ): purge_id = "purge_rebuild_remove_orphan" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_remove_orphan", operation_type="user_erasure", @@ -1613,12 +2761,14 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1633,6 +2783,7 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) storage.apply_governance_agent_playbook_rebuild( @@ -1699,7 +2850,8 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta storage, ): purge_id = "purge_rebuild_restore_archived" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_restore_archived", operation_type="user_erasure", @@ -1715,12 +2867,14 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1738,6 +2892,7 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) storage.apply_governance_agent_playbook_rebuild( @@ -1765,7 +2920,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl storage, ): purge_id = "purge_rebuild_second_call" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_second_call", operation_type="user_erasure", @@ -1877,7 +3033,8 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet storage, ): purge_id = "purge_hide_after_complete" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_hide_after_complete", operation_type="user_erasure", @@ -1969,7 +3126,8 @@ def test_hide_governance_agent_playbooks_for_rebuild_does_not_reopen_complete_ta storage, monkeypatch ): purge_id = "purge_hide_stale_prelock" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_hide_stale_prelock", operation_type="user_erasure", @@ -2080,19 +3238,22 @@ def test_prepare_governance_erase_targets_is_idempotent_after_completed_snapshot ): purge_id = "purge_prepare_idempotent_after_snapshot" user_id = "user-prepare-idempotent-after-snapshot" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_prepare_idempotent_after_snapshot", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) storage.prepare_governance_erase_targets( purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) before_targets = [ @@ -2112,6 +3273,7 @@ def test_prepare_governance_erase_targets_is_idempotent_after_completed_snapshot purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) after_targets = [ @@ -2144,8 +3306,9 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): idempotency_key=f"idem_{storage_instance.org_id}_{purge_id}", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage_instance._subject_ref_for_user_id("alice"), request_ref=request_ref, + authoritative_user_id="alice", ) storage_org1.record_purge_target( @@ -2154,6 +3317,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="prepare_targets", status="complete", + execution_claim=_claim_purge(storage_org1, purge_id), detail={"prepared": True}, ) storage_org1.record_purge_target( @@ -2162,6 +3326,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="delete", status="pending", + execution_claim=_claim_purge(storage_org1, purge_id), detail={"count": 1}, ) storage_org2.record_purge_target( @@ -2170,6 +3335,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage_org2, purge_id), detail={"count": 2}, deleted_count=2, ) @@ -2197,6 +3363,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="delete", status="running", + execution_claim=_claim_purge(storage_org2, purge_id), detail={"count": 3}, ) @@ -2293,6 +3460,7 @@ def test_record_purge_target_validates_governance_fields(storage, kwargs, match) "phase": "delete", "status": "running", "target_ref": "all", + "execution_claim": _claim_purge(storage, purge_id), } params.update(kwargs) @@ -2331,6 +3499,7 @@ def test_record_purge_target_rejects_invalid_deleted_count( phase="delete", status="complete", deleted_count=deleted_count, + execution_claim=_claim_purge(storage, purge_id), ) @@ -2348,6 +3517,7 @@ def test_record_purge_target_accepts_nonnegative_detail_deleted_count( target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail={"deleted_count": detail_deleted_count}, ) @@ -2370,6 +3540,7 @@ def test_record_purge_target_rejects_negative_detail_deleted_count(storage): phase="delete", status="complete", detail={"deleted_count": -1}, + execution_claim=_claim_purge(storage, purge_id), ) @@ -2428,6 +3599,16 @@ def test_record_purge_target_rejects_negative_detail_deleted_count(storage): None, id="audit-detail-allowed-deleted-counts", ), + pytest.param( + {"deleted_counts": {"session_outcomes": 1}}, + None, + id="audit-detail-allowed-session-outcome-counts", + ), + pytest.param( + {"deleted_counts": {"session_outcome": 1}}, + "session_outcome", + id="audit-detail-rejects-unknown-deleted-count-key", + ), pytest.param( {"agent_playbook_id": 7}, None, id="audit-detail-allowed-agent-playbook-id" ), @@ -2464,6 +3645,7 @@ def test_record_purge_target_accepts_target_detail_shapes(storage): purge_id = _begin_purge(storage, "purge_target_detail_shapes") detail = { + "authoritative_user_digest": "a" * 64, "owned_user_playbook_ids": [7], "source_interaction_ids": [11, 12], "original_source_windows": [ @@ -2481,6 +3663,7 @@ def test_record_purge_target_accepts_target_detail_shapes(storage): target_ref="7", phase="rebuild_without_erased_sources", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -2532,6 +3715,7 @@ def test_fail_purge_operation_rejects_raw_error_detail(storage): purge_id, error_code="boom", error_detail="RuntimeError: request reqref_123 for alice@example.com", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_detail is None @@ -2545,6 +3729,7 @@ def test_fail_purge_operation_rejects_freeform_error_detail(storage): purge_id, error_code="PURGE_TARGET_FAILED", error_detail="stable failure detail", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_detail is None @@ -2557,6 +3742,7 @@ def test_fail_purge_operation_persists_code_shaped_error_detail(storage): purge_id, error_code="PURGE_TARGET_FAILED", error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert failed.status == "failed" @@ -2719,6 +3905,7 @@ def test_fail_purge_operation_accepts_code_shaped_error_code_with_prompt_or_cont purge_id, error_code=error_code, error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert failed.status == "failed" @@ -2733,6 +3920,7 @@ def test_fail_purge_operation_rejects_prompt_content_prose_error_detail(storage) purge_id, error_code="PURGE_TARGET_FAILED", error_detail="prompt content leaked from upstream", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_detail is None @@ -2755,6 +3943,7 @@ def test_fail_purge_operation_validates_error_code(storage, error_code, match): purge_id, error_code=error_code, error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert failed.status == "failed" assert failed.error_code == error_code @@ -2765,6 +3954,7 @@ def test_fail_purge_operation_validates_error_code(storage, error_code, match): purge_id, error_code=error_code, error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_code is None @@ -2903,7 +4093,8 @@ def test_begin_purge_operation_validates_top_level_refs( storage, subject_ref, request_ref, match ): with pytest.raises(ValueError, match=match): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_top_level_refs", idempotency_key="idem_purge_top_level_refs", operation_type="user_erasure", @@ -2928,7 +4119,8 @@ def test_begin_purge_operation_rejects_invalid_enum_values( storage, operation_type, scope_type, match ): with pytest.raises(ValueError, match=match): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_invalid_enum", idempotency_key="idem_purge_invalid_enum", operation_type=operation_type, @@ -2949,7 +4141,8 @@ def test_begin_purge_operation_rejects_invalid_enum_values( ) def test_begin_purge_operation_rejects_unsafe_purge_id(storage, purge_id): with pytest.raises(ValueError, match="purge_id"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_invalid_id", operation_type="user_erasure", @@ -2990,6 +4183,7 @@ def test_record_purge_target_rejects_mixed_case_window_keys(storage, detail_key) target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={detail_key: [{"User_Playbook_Id": "alice@example.com"}]}, ) @@ -3025,6 +4219,7 @@ def test_record_purge_target_requires_window_user_playbook_id(storage, detail_ke target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={detail_key: [{"source_interaction_ids": [1, 2]}]}, ) @@ -3046,6 +4241,7 @@ def test_record_purge_target_accepts_previous_lifecycle_status_for_rebuild_targe target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [11, 12]} @@ -3124,6 +4320,7 @@ def test_record_purge_target_rejects_invalid_previous_lifecycle_status_detail( target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -3151,6 +4348,7 @@ def test_record_purge_target_validates_target_ref_contract(storage, target_ref, target_ref=target_ref, phase="delete", status="running", + execution_claim=_claim_purge(storage, purge_id), ) return @@ -3161,6 +4359,7 @@ def test_record_purge_target_validates_target_ref_contract(storage, target_ref, target_ref=target_ref, phase="delete", status="running", + execution_claim=_claim_purge(storage, purge_id), ) @@ -3195,6 +4394,7 @@ def test_persistence_paths_reject_unsafe_purge_id(storage, purge_id): target_ref="all", phase="delete", status="running", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge_id), ) with pytest.raises(ValueError, match="purge_id"): @@ -3208,6 +4408,8 @@ def test_persistence_paths_reject_unsafe_purge_id(storage, purge_id): request_ref=REQUEST_REF, idempotency_key=purge_id, ), + authoritative_user_id="alice", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge_id), ) with pytest.raises(ValueError, match="purge_id"): @@ -3225,6 +4427,9 @@ def test_apply_governance_user_data_delete_rejects_unsafe_purge_id_before_side_e storage.apply_governance_user_data_delete( purge_id="alice@example.com", user_id=user_id, + execution_claim=_typed_test_claim_for_unvalidated_purge_id( + "alice@example.com" + ), ) remaining = _user_scoped_row_counts(storage, user_id=user_id) @@ -3239,7 +4444,11 @@ def test_apply_governance_user_data_delete_rejects_unsafe_purge_id_before_side_e def test_apply_governance_user_data_delete_rejects_unexpected_target_name_from_internal_counts( storage, monkeypatch ): - purge_id = _begin_purge(storage, "purge_internal_target_name") + purge_id = _begin_purge( + storage, + "purge_internal_target_name", + authoritative_user_id="user-delete-seed", + ) for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -3247,6 +4456,7 @@ def test_apply_governance_user_data_delete_rejects_unexpected_target_name_from_i target_ref="all", phase="delete", status="pending", + execution_claim=_claim_purge(storage, purge_id), detail={"count": 0}, ) @@ -3269,6 +4479,7 @@ def _stub_clear_user_data_for_governance_locked( storage.apply_governance_user_data_delete( purge_id=purge_id, user_id="user-delete-seed", + execution_claim=_claim_purge(storage, purge_id), ) delete_targets = storage.list_purge_targets(purge_id, phase="delete") @@ -3278,8 +4489,12 @@ def _stub_clear_user_data_for_governance_locked( def test_apply_governance_user_data_delete_requires_complete_prepared_delete_matrix( storage, monkeypatch ): - purge_id = _begin_purge(storage, "purge_delete_requires_prepared_matrix") user_id = "user-delete-seed" + purge_id = _begin_purge( + storage, + "purge_delete_requires_prepared_matrix", + authoritative_user_id=user_id, + ) expected_user_id = user_id _seed_user_scoped_rows(storage, user_id=user_id) baseline_counts = _user_scoped_row_counts(storage, user_id=user_id) @@ -3289,6 +4504,7 @@ def test_apply_governance_user_data_delete_requires_complete_prepared_delete_mat target_ref="all", phase="delete", status="pending", + execution_claim=_claim_purge(storage, purge_id), detail={"count": 1}, ) storage.record_purge_target( @@ -3297,6 +4513,7 @@ def test_apply_governance_user_data_delete_requires_complete_prepared_delete_mat target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail={"count": 0}, deleted_count=0, ) @@ -3322,6 +4539,7 @@ def _stub_clear_user_data_for_governance_locked( storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert clear_locked_called is False @@ -3338,13 +4556,15 @@ def test_apply_governance_user_data_delete_preserves_org_agent_playbooks_without ): purge_id = "purge_delete_requires_hide" user_id = "user-delete-hide-required" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_delete_requires_hide", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) _seed_eval_result( @@ -3368,11 +4588,13 @@ def test_apply_governance_user_data_delete_preserves_org_agent_playbooks_without purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) counts = storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert counts["user_playbooks"] == 1 @@ -3400,13 +4622,15 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( """ purge_id = "purge_delete_after_hide" user_id = "user-delete-hide-complete" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_delete_after_hide", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) _seed_eval_result( @@ -3440,12 +4664,14 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) storage.hide_governance_agent_playbooks_for_rebuild(purge_id) counts = storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert counts == { @@ -3526,13 +4752,15 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( def test_apply_governance_user_data_delete_is_failure_atomic(storage, monkeypatch): purge_id = "purge_delete_atomic" user_id = "user-delete-atomic" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_delete_atomic", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) affected_user_playbook_id = min(owned_user_playbook_ids) @@ -3550,6 +4778,7 @@ def test_apply_governance_user_data_delete_is_failure_atomic(storage, monkeypatc purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) storage.hide_governance_agent_playbooks_for_rebuild(purge_id) @@ -3605,6 +4834,7 @@ def _raising_record_purge_target_locked( storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert _user_scoped_row_counts(storage, user_id=user_id) == before_counts @@ -3676,6 +4906,7 @@ def test_fail_purge_operation_rejects_unsafe_purge_id_before_side_effects(storag SUBJECT_REF, "governance.error", "detail.code", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(SUBJECT_REF), ) failed = storage.get_purge_operation(purge_id) @@ -3688,7 +4919,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so storage, ): purge_id = "purge_rebuild_windows_mismatch" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_windows_mismatch", operation_type="user_erasure", @@ -3710,6 +4942,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -3727,6 +4960,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=_claim_purge(storage, purge_id), ) original_row = storage.conn.execute( """SELECT content, trigger, rationale, blocking_issue, expanded_terms, tags, status @@ -3867,6 +5101,7 @@ def test_record_purge_target_rejects_invalid_enum_values( target_ref="all", phase=phase, status=status, + execution_claim=_claim_purge(storage, purge_id), ) @@ -3919,7 +5154,8 @@ def test_governance_persistence_rejects_unsafe_idempotency_keys( storage.append_audit_event(event) with pytest.raises(ValueError, match="idempotency_key"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_unsafe_idem", idempotency_key=idempotency_key, operation_type="user_erasure", @@ -4024,6 +5260,7 @@ def test_governance_detail_rejects_noncanonical_status_and_route( target_ref="all", phase="delete", status="running", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -4034,7 +5271,8 @@ def test_governance_detail_rejects_noncanonical_status_and_route( ) def test_begin_purge_operation_rejects_identifier_like_purge_suffix(storage, purge_id): with pytest.raises(ValueError, match="purge_id"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_identifier_suffix", operation_type="user_erasure", @@ -4070,6 +5308,7 @@ def test_record_purge_target_canonicalizes_detail_keys_before_persistence(storag target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail={" Deleted_Counts ": {"requests": 2}}, deleted_count=2, ) @@ -4105,6 +5344,7 @@ def test_governance_detail_rejects_duplicate_normalized_keys(storage, persistenc target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -4232,6 +5472,7 @@ def test_record_purge_target_validates_target_ref_by_phase_and_name( target_ref=target_ref, phase=phase, status="running", + execution_claim=_claim_purge(storage, purge_id), ) return @@ -4242,6 +5483,7 @@ def test_record_purge_target_validates_target_ref_by_phase_and_name( target_ref=target_ref, phase=phase, status="running", + execution_claim=_claim_purge(storage, purge_id), ) @@ -4302,13 +5544,15 @@ def test_init_governance_tables_upgrades_legacy_purge_target_table(tmp_path): with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): storage = SQLiteStorage(org_id="org1", db_path=str(db_path)) - purge_id = storage.begin_purge_operation( + purge_id = _begin_test_purge_operation( + storage, purge_id="purge_legacy_upgrade", idempotency_key="idem_legacy_upgrade", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id("alice"), request_ref=REQUEST_REF, + authoritative_user_id="alice", ).purge_id storage.record_purge_target( purge_id=purge_id, @@ -4316,6 +5560,7 @@ def test_init_governance_tables_upgrades_legacy_purge_target_table(tmp_path): target_ref="all", phase="prepare_targets", status="complete", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.purge_targets_prepared(purge_id) is True @@ -4522,18 +5767,23 @@ def test_successful_erase_audit_row_exists_only_after_complete_purge(storage): # (2) The one legitimate writer produces exactly one successful-ERASE row, # and only after the purge_operation transitions to 'complete'. completed = storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert completed.status == "complete" erase_rows = _successful_erase_audit_rows(storage) assert [event.idempotency_key for event in erase_rows] == [purge_id] _assert_successful_erase_rows_only_for_complete_purges(storage) - # Idempotent re-completion neither duplicates the row nor breaks the - # invariant. - storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) - ) + with pytest.raises(ValueError, match="purge execution claim"): + storage.complete_purge_operation_with_audit( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge_id), + ) assert [ event.idempotency_key for event in _successful_erase_audit_rows(storage) ] == [purge_id] diff --git a/tests/server/services/storage/sqlite_storage/test_session_id_migration.py b/tests/server/services/storage/sqlite_storage/test_session_id_migration.py index fb40107bd..2a6512833 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_id_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_id_migration.py @@ -20,6 +20,7 @@ source TEXT NOT NULL DEFAULT '', agent_version TEXT NOT NULL DEFAULT '', session_id TEXT, + governance_subject_ref TEXT, metadata TEXT NOT NULL DEFAULT '{}' ); """ @@ -29,12 +30,35 @@ def _seed_legacy_db(db_path: str) -> None: conn = sqlite3.connect(db_path) conn.executescript(_LEGACY_REQUESTS_DDL) conn.executemany( - "INSERT INTO requests (request_id, user_id, created_at, source, session_id) " - "VALUES (?, ?, ?, ?, ?)", + """INSERT INTO requests ( + request_id, user_id, created_at, source, session_id, + governance_subject_ref + ) VALUES (?, ?, ?, ?, ?, ?)""", [ - ("r-null", "u1", "2026-01-01T00:00:00+00:00", "web", None), - ("r-blank", "u1", "2026-01-01T00:00:01+00:00", "web", " "), - ("r-valid", "u1", "2026-01-01T00:00:02+00:00", "web", "s-valid"), + ( + "r-null", + "u1", + "2026-01-01T00:00:00+00:00", + "web", + None, + "subject-null", + ), + ( + "r-blank", + "u1", + "2026-01-01T00:00:01+00:00", + "web", + " ", + "subject-blank", + ), + ( + "r-valid", + "u1", + "2026-01-01T00:00:02+00:00", + "web", + "s-valid", + "subject-valid", + ), ], ) conn.commit() @@ -89,6 +113,28 @@ def test_migration_enforces_not_null_and_non_empty(tmp_path): conn.close() +def test_migration_preserves_governance_subject_ref(tmp_path): + db_path = str(tmp_path / "legacy.db") + _seed_legacy_db(db_path) + + SQLiteStorage(org_id="0", db_path=db_path) + + conn = sqlite3.connect(db_path) + try: + subject_refs = dict( + conn.execute( + "SELECT request_id, governance_subject_ref FROM requests" + ).fetchall() + ) + finally: + conn.close() + assert subject_refs == { + "r-null": "subject-null", + "r-blank": "subject-blank", + "r-valid": "subject-valid", + } + + def test_migration_is_idempotent(tmp_path): db_path = str(tmp_path / "legacy.db") _seed_legacy_db(db_path) diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py new file mode 100644 index 000000000..73439ffc7 --- /dev/null +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -0,0 +1,310 @@ +"""Regression coverage for the SQLite session-outcome identity migration.""" + +import json +from hashlib import sha256 + +import pytest + +from reflexio.models.api_schema.domain import ( + GetSessionOutcomesRequest, + Request, + SessionOutcomeKind, +) +from reflexio.server.services.storage.session_outcome_identity import ( + outcome_contract_digest, + trajectory_digest, +) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._base import ( + _canonical_session_snapshot, +) + +pytestmark = pytest.mark.integration + +_LEGACY_SESSION_OUTCOMES_DDL = """ +CREATE TABLE session_outcomes ( + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), + occurred_at INTEGER NOT NULL, + source TEXT NOT NULL, + label TEXT, + value REAL, + metadata TEXT, + governance_subject_ref TEXT NOT NULL, + created_at INTEGER NOT NULL, + PRIMARY KEY (user_id, session_id) +); +""" + + +def _legacy_session_outcomes_ddl(*, with_subject_column: bool) -> str: + subject_column = "governance_subject_ref TEXT," if with_subject_column else "" + return f""" +CREATE TABLE session_outcomes ( + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), + occurred_at INTEGER NOT NULL, + source TEXT NOT NULL, + label TEXT, + value REAL, + metadata TEXT, + {subject_column} + created_at INTEGER NOT NULL, + PRIMARY KEY (user_id, session_id) +); +""" + + +def _identity_complete_session_outcomes_ddl(*, with_subject_column: bool) -> str: + subject_column = "governance_subject_ref TEXT," if with_subject_column else "" + return f""" +CREATE TABLE session_outcomes ( + outcome_id TEXT NOT NULL UNIQUE, + outcome_revision INTEGER NOT NULL CHECK (outcome_revision >= 1), + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure', 'unknown')), + occurred_at INTEGER NOT NULL, + source TEXT NOT NULL, + label TEXT, + value REAL, + metadata TEXT, + outcome_contract_digest TEXT NOT NULL, + finalized_trajectory_digest TEXT NOT NULL, + {subject_column} + created_at INTEGER NOT NULL, + PRIMARY KEY (user_id, session_id) +); +""" + + +def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( + tmp_path, +) -> None: + db_path = str(tmp_path / "legacy-session-outcomes.db") + storage = SQLiteStorage(org_id="legacy-session-outcomes", db_path=db_path) + legacy_rows = [ + { + "user_id": "a:b", + "session_id": "c", + "outcome": "success", + "occurred_at": 101, + "source": "legacy-source-one", + "label": "resolved:one", + "value": 2.5, + "metadata": {"nested": {"a": 1}, "legacy": True}, + "governance_subject_ref": "subject:a:b:c", + "created_at": 102, + }, + { + "user_id": "a", + "session_id": "b:c", + "outcome": "failure", + "occurred_at": 201, + "source": "legacy-source-two", + "label": None, + "value": None, + "metadata": None, + "governance_subject_ref": "subject:a:b:c:two", + "created_at": 202, + }, + ] + for row in legacy_rows: + storage.add_request( + Request( + request_id=f"request-{row['session_id']}", + user_id=row["user_id"], + session_id=row["session_id"], + source=row["source"], + created_at=row["occurred_at"] - 1, + ) + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + for row in legacy_rows: + storage.conn.execute( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, label, value, + metadata, governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + row["user_id"], + row["session_id"], + row["outcome"], + row["occurred_at"], + row["source"], + row["label"], + row["value"], + json.dumps(row["metadata"], sort_keys=True) + if row["metadata"] is not None + else None, + row["governance_subject_ref"], + row["created_at"], + ), + ) + storage.conn.commit() + storage.conn.close() + + migrated = SQLiteStorage(org_id="legacy-session-outcomes", db_path=db_path) + records = migrated.get_session_outcomes( + GetSessionOutcomesRequest(session_ids=["c", "b:c"]) + ) + records_by_session = {record.session_id: record for record in records} + rows_by_session = { + row["session_id"]: row + for row in migrated.conn.execute( + "SELECT * FROM session_outcomes WHERE session_id IN (?, ?)", ("c", "b:c") + ).fetchall() + } + + assert set(records_by_session) == {"c", "b:c"} + assert records_by_session["c"].outcome_id == sha256(b'["a:b","c"]').hexdigest() + assert records_by_session["b:c"].outcome_id == sha256(b'["a","b:c"]').hexdigest() + assert records_by_session["c"].outcome_id != records_by_session["b:c"].outcome_id + + for row in legacy_rows: + record = records_by_session[row["session_id"]] + assert record.outcome_revision == 1 + assert record.user_id == row["user_id"] + assert record.session_id == row["session_id"] + assert record.outcome == SessionOutcomeKind(row["outcome"]) + assert record.occurred_at == row["occurred_at"] + assert record.source == row["source"] + assert record.label == row["label"] + assert record.value == row["value"] + assert record.metadata == row["metadata"] + assert record.created_at == row["created_at"] + assert ( + rows_by_session[row["session_id"]]["governance_subject_ref"] + == row["governance_subject_ref"] + ) + assert record.outcome_contract_digest == outcome_contract_digest( + source=row["source"], + schema_version=1, + allowed_values={"success", "failure", "unknown"}, + finalization_rule="first_write", + ) + assert record.finalized_trajectory_digest == trajectory_digest( + _canonical_session_snapshot(migrated.conn, row["session_id"]) + ) + + +@pytest.mark.parametrize( + ("with_subject_column", "stored_subject_ref"), + [(False, None), (True, None), (True, " ")], + ids=["absent", "null", "whitespace"], +) +def test_migration_derives_missing_legacy_governance_subject_ref( + tmp_path, with_subject_column: bool, stored_subject_ref: str | None +) -> None: + db_path = str(tmp_path / f"legacy-subject-{with_subject_column}.db") + storage = SQLiteStorage(org_id="legacy-subject", db_path=db_path) + storage.add_request( + Request( + request_id="legacy-subject-request", + user_id="legacy-user", + session_id="legacy-session", + source="legacy-source", + created_at=100, + ) + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript( + _legacy_session_outcomes_ddl(with_subject_column=with_subject_column) + ) + values = ( + "legacy-user", + "legacy-session", + "success", + 101, + "legacy-source", + 102, + ) + if with_subject_column: + storage.conn.execute( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, created_at, + governance_subject_ref + ) VALUES (?, ?, ?, ?, ?, ?, ?)""", + (*values, stored_subject_ref), + ) + else: + storage.conn.execute( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, created_at + ) VALUES (?, ?, ?, ?, ?, ?)""", + values, + ) + storage.conn.commit() + storage.conn.close() + + migrated = SQLiteStorage(org_id="legacy-subject", db_path=db_path) + + row = migrated.conn.execute( + "SELECT governance_subject_ref FROM session_outcomes WHERE session_id = ?", + ("legacy-session",), + ).fetchone() + assert row is not None + assert row["governance_subject_ref"] == migrated._subject_ref_for_user_id( + "legacy-user" + ) + + +@pytest.mark.parametrize("with_subject_column", [False, True]) +def test_identity_complete_migration_backfills_missing_governance_subject_ref( + tmp_path, with_subject_column: bool +) -> None: + db_path = str(tmp_path / f"complete-subject-{with_subject_column}.db") + storage = SQLiteStorage(org_id="complete-subject", db_path=db_path) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript( + _identity_complete_session_outcomes_ddl(with_subject_column=with_subject_column) + ) + values = ( + "stable-outcome-id", + 1, + "complete-user", + "complete-session", + "unknown", + 101, + "complete-source", + "a" * 64, + "b" * 64, + 102, + ) + if with_subject_column: + storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, created_at, + governance_subject_ref + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (*values, None), + ) + else: + storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + values, + ) + storage.conn.commit() + storage.conn.close() + + migrated = SQLiteStorage(org_id="complete-subject", db_path=db_path) + + row = migrated.conn.execute( + """SELECT outcome_id, governance_subject_ref + FROM session_outcomes WHERE session_id = ?""", + ("complete-session",), + ).fetchone() + assert row is not None + assert row["outcome_id"] == "stable-outcome-id" + assert row["governance_subject_ref"] == migrated._subject_ref_for_user_id( + "complete-user" + ) diff --git a/tests/server/services/storage/test_sqlite_storage.py b/tests/server/services/storage/test_sqlite_storage.py index 81584e44a..e187523cc 100644 --- a/tests/server/services/storage/test_sqlite_storage.py +++ b/tests/server/services/storage/test_sqlite_storage.py @@ -30,6 +30,7 @@ _true_rrf_merge, _vector_rank_rows, ) +from reflexio.server.services.storage.sqlite_storage import _base as sqlite_storage_base from reflexio.server.services.storage.sqlite_storage._base import ( _epoch_to_iso, _iso_to_epoch, @@ -57,6 +58,42 @@ def storage(): yield SQLiteStorage(org_id="0", db_path=f"{temp_dir}/reflexio.db") +def test_sqlite_storage_rejects_sqlite_before_returning_support( + tmp_path, monkeypatch +) -> None: + monkeypatch.setattr(sqlite_storage_base.sqlite3, "sqlite_version_info", (3, 34, 99)) + db_path = tmp_path / "missing" / "nested" / "old.db" + + with ( + patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512), + pytest.raises( + RuntimeError, + match=r"SQLite 3\.35\.0 or newer is required; detected 3\.34\.99", + ), + ): + SQLiteStorage(org_id="version-check", db_path=str(db_path)) + + assert not db_path.parent.exists() + assert not db_path.exists() + + +def test_sqlite_storage_accepts_sqlite_with_returning_support( + tmp_path, monkeypatch +) -> None: + monkeypatch.setattr(sqlite_storage_base.sqlite3, "sqlite_version_info", (3, 35, 0)) + + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + storage = SQLiteStorage( + org_id="version-check", db_path=str(tmp_path / "supported.db") + ) + try: + assert storage.conn.execute( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'purge_operations'" + ).fetchone() + finally: + storage.conn.close() + + # --------------------------------------------------------------------------- # _sanitize_fts_query tests # --------------------------------------------------------------------------- diff --git a/tests/server/services/storage/test_storage_contract_clear_user_data.py b/tests/server/services/storage/test_storage_contract_clear_user_data.py index cf550ae22..1ef977535 100644 --- a/tests/server/services/storage/test_storage_contract_clear_user_data.py +++ b/tests/server/services/storage/test_storage_contract_clear_user_data.py @@ -24,6 +24,8 @@ UserPlaybook, UserProfile, ) +from reflexio.server.services.governance.config import governance_subject_ref +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.storage_base import BaseStorage pytestmark = pytest.mark.integration @@ -170,6 +172,69 @@ def test_clear_unknown_user_is_noop(self, storage: BaseStorage) -> None: assert len(storage.get_user_playbooks(user_id="userA")) == 1 assert storage.get_request("req_a") is not None + def test_session_outcomes_use_authoritative_user_and_report_stable_zero( + self, storage: BaseStorage, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("REFLEXIO_GOVERNANCE_REF_SECRET", "test-governance-secret") + sqlite_storage = storage + assert isinstance(sqlite_storage, SQLiteStorage) + alice_ref = governance_subject_ref( + sqlite_storage.org_id, "alice", "test-governance-secret" + ) + bob_ref = governance_subject_ref( + sqlite_storage.org_id, "bob", "test-governance-secret" + ) + for outcome_id, user_id, subject_ref in ( + ("alice-stale", "alice", bob_ref), + ("bob-conflict", "bob", alice_ref), + ): + sqlite_storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES (?, 1, ?, ?, 'success', 100, 'test', ?, ?, ?, 101)""", + ( + outcome_id, + user_id, + f"session-{outcome_id}", + "a" * 64, + "b" * 64, + subject_ref, + ), + ) + sqlite_storage.conn.commit() + + counts = storage.clear_user_data("alice") + zero_counts = storage.clear_user_data("missing-user") + + remaining = sqlite_storage.conn.execute( + "SELECT outcome_id FROM session_outcomes ORDER BY outcome_id" + ).fetchall() + assert [row["outcome_id"] for row in remaining] == ["bob-conflict"] + assert counts["session_outcomes"] == 1 + assert zero_counts["session_outcomes"] == 0 + + def test_default_clear_user_data_preserves_session_outcome_count( + self, storage: BaseStorage + ) -> None: + sqlite_storage = storage + assert isinstance(sqlite_storage, SQLiteStorage) + sqlite_storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES ('default-clear', 1, 'alice', 'default-clear-session', + 'success', 100, 'test', ?, ?, 'subref_v1_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', 101)""", + ("a" * 64, "b" * 64), + ) + sqlite_storage.conn.commit() + + counts = BaseStorage.clear_user_data(sqlite_storage, "alice") + + assert counts["session_outcomes"] == 1 + def test_returned_counts_match_seeded_rows(self, storage: BaseStorage) -> None: """Per-entity counts must reflect actual seeded row counts for the user.""" # Seed userA with two of each. diff --git a/tests/server/services/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index b07a2223b..6716214b4 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -1,11 +1,18 @@ """Session outcome storage contract.""" +from typing import cast + from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, Request, + SessionOutcomeFailureReason, SessionOutcomeKind, SetSessionOutcomeRequest, ) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._base import ( + _canonical_session_snapshot, +) from reflexio.server.services.storage.storage_base import BaseStorage @@ -39,14 +46,345 @@ def test_first_write_preserves_outcome_fields(storage: BaseStorage) -> None: assert first.recorded is True assert duplicate.recorded is False + assert duplicate.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION assert duplicate.source == "published" records = storage.get_session_outcomes(GetSessionOutcomesRequest(label="booked")) assert len(records) == 1 + assert records[0].outcome_id + assert records[0].outcome_revision == 1 + outcome_contract_digest = records[0].outcome_contract_digest + finalized_trajectory_digest = records[0].finalized_trajectory_digest + assert outcome_contract_digest is not None + assert finalized_trajectory_digest is not None + assert len(outcome_contract_digest) == 64 + assert len(finalized_trajectory_digest) == 64 assert records[0].outcome == SessionOutcomeKind.SUCCESS assert records[0].value == 12.0 assert records[0].metadata == {"crm": "test"} +def test_exact_finalization_retry_is_idempotent(storage: BaseStorage) -> None: + storage.add_request( + Request( + request_id="retry-r1", + user_id="u1", + session_id="exact-retry", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="exact-retry", + outcome=SessionOutcomeKind.UNKNOWN, + occurred_at=101, + metadata={"reason": "not enough information"}, + ) + context = storage.get_session_outcome_context("exact-retry") + + first = storage.record_session_outcome( + request, created_at=102, expected_context=context + ) + retry = storage.record_session_outcome( + request, created_at=103, expected_context=context + ) + + assert first.recorded is True + assert retry.recorded is False + assert retry.reason is None + assert retry.outcome_id == first.outcome_id + assert retry.outcome_revision == first.outcome_revision == 1 + assert retry.outcome_contract_digest == first.outcome_contract_digest + assert retry.finalized_trajectory_digest == first.finalized_trajectory_digest + + +def test_legacy_all_null_identity_exact_retry_uses_available_context( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="legacy-r1", + user_id="legacy-user", + session_id="legacy-retry", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="legacy-retry", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + metadata={"legacy": True}, + ) + first = storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("legacy-retry"), + ) + assert first.recorded is True + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + "CREATE TABLE legacy_session_outcomes AS SELECT * FROM session_outcomes" + ) + sqlite_storage.conn.execute("DROP TABLE session_outcomes") + sqlite_storage.conn.execute( + "ALTER TABLE legacy_session_outcomes RENAME TO session_outcomes" + ) + sqlite_storage.conn.execute( + """UPDATE session_outcomes + SET outcome_id = NULL, outcome_revision = NULL, + outcome_contract_digest = NULL, + finalized_trajectory_digest = NULL + WHERE session_id = ?""", + ("legacy-retry",), + ) + sqlite_storage.conn.commit() + + retry = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context("legacy-retry"), + ) + + assert retry.recorded is False + assert retry.reason is None + assert retry.outcome_id is None + assert retry.outcome_revision is None + assert retry.outcome_contract_digest is None + assert retry.finalized_trajectory_digest is None + + +def test_legacy_all_null_identity_changed_payload_still_conflicts( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="legacy-conflict-r1", + user_id="legacy-user", + session_id="legacy-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="legacy-conflict", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("legacy-conflict"), + ) + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + "CREATE TABLE legacy_session_outcomes AS SELECT * FROM session_outcomes" + ) + sqlite_storage.conn.execute("DROP TABLE session_outcomes") + sqlite_storage.conn.execute( + "ALTER TABLE legacy_session_outcomes RENAME TO session_outcomes" + ) + sqlite_storage.conn.execute( + """UPDATE session_outcomes + SET outcome_id = NULL, outcome_revision = NULL, + outcome_contract_digest = NULL, + finalized_trajectory_digest = NULL + WHERE session_id = ?""", + ("legacy-conflict",), + ) + sqlite_storage.conn.commit() + + retry = storage.record_session_outcome( + request.model_copy(update={"outcome": SessionOutcomeKind.FAILURE}), + created_at=103, + expected_context=storage.get_session_outcome_context("legacy-conflict"), + ) + + assert retry.recorded is False + assert retry.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + + +def test_legacy_all_null_identity_changed_governance_context_conflicts( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="legacy-governance-r1", + user_id="legacy-user", + session_id="legacy-governance-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="legacy-governance-conflict", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context( + "legacy-governance-conflict" + ), + ) + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + "CREATE TABLE legacy_session_outcomes AS SELECT * FROM session_outcomes" + ) + sqlite_storage.conn.execute("DROP TABLE session_outcomes") + sqlite_storage.conn.execute( + "ALTER TABLE legacy_session_outcomes RENAME TO session_outcomes" + ) + sqlite_storage.conn.execute( + """UPDATE session_outcomes + SET outcome_id = NULL, outcome_revision = NULL, + outcome_contract_digest = NULL, + finalized_trajectory_digest = NULL + WHERE session_id = ?""", + ("legacy-governance-conflict",), + ) + sqlite_storage.conn.execute( + "UPDATE requests SET governance_subject_ref = ? WHERE session_id = ?", + ( + sqlite_storage._subject_ref_for_user_id("different-user"), + "legacy-governance-conflict", + ), + ) + sqlite_storage.conn.commit() + + retry = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context( + "legacy-governance-conflict" + ), + ) + + assert retry.recorded is False + assert retry.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + + +def test_sqlite_canonical_snapshot_loads_interactions_in_one_query( + storage: BaseStorage, +) -> None: + sqlite_storage = cast(SQLiteStorage, storage) + for index in range(3): + storage.add_request( + Request( + request_id=f"snapshot-r{index}", + user_id="snapshot-user", + session_id="snapshot-session", + source="published", + created_at=100 + index, + ) + ) + statements: list[str] = [] + sqlite_storage.conn.set_trace_callback(statements.append) + try: + snapshot = _canonical_session_snapshot(sqlite_storage.conn, "snapshot-session") + finally: + sqlite_storage.conn.set_trace_callback(None) + + interaction_queries = [ + statement for statement in statements if "FROM interactions" in statement + ] + assert [item["request"]["request_id"] for item in snapshot["requests"]] == [ + "snapshot-r0", + "snapshot-r1", + "snapshot-r2", + ] + assert len(interaction_queries) == 1 + + +def test_changed_contract_identity_is_conflicting_finalization( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="contract-r1", + user_id="u1", + session_id="contract-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="contract-conflict", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + first = storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("contract-conflict"), + ) + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + """UPDATE session_outcomes SET outcome_contract_digest = ? + WHERE session_id = ?""", + ("0" * 64, "contract-conflict"), + ) + sqlite_storage.conn.commit() + + changed_contract = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context("contract-conflict"), + ) + + assert first.recorded is True + assert changed_contract.recorded is False + assert ( + changed_contract.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + ) + + +def test_changed_session_trajectory_is_conflicting_finalization( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="trajectory-r1", + user_id="u1", + session_id="trajectory-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="trajectory-conflict", + outcome=SessionOutcomeKind.FAILURE, + occurred_at=101, + ) + first = storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("trajectory-conflict"), + ) + storage.add_request( + Request( + request_id="trajectory-r2", + user_id="u1", + session_id="trajectory-conflict", + source="published", + created_at=103, + ) + ) + + changed_trajectory = storage.record_session_outcome( + request, + created_at=104, + expected_context=storage.get_session_outcome_context("trajectory-conflict"), + ) + + assert first.recorded is True + assert changed_trajectory.recorded is False + assert ( + changed_trajectory.reason + == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + ) + + def test_unknown_session_is_rejected(storage: BaseStorage) -> None: context = storage.get_session_outcome_context("missing") result = storage.record_session_outcome( diff --git a/tests/server/services/test_generation_billing_emission.py b/tests/server/services/test_generation_billing_emission.py index 6138f19f4..cdc597054 100644 --- a/tests/server/services/test_generation_billing_emission.py +++ b/tests/server/services/test_generation_billing_emission.py @@ -15,7 +15,14 @@ from typing import Any from unittest.mock import MagicMock, patch -from reflexio.models.api_schema.service_schemas import Interaction, Request +import pytest + +from reflexio.models.api_schema.domain.entities import UserProfile +from reflexio.models.api_schema.service_schemas import ( + Interaction, + Request, + UserPlaybook, +) from reflexio.models.config_schema import ( Config, ProfileExtractorConfig, @@ -28,6 +35,14 @@ BaseGenerationService, PreparedGenerationRun, ) +from reflexio.server.services.deferred_learning_plan import ( + PlaybookWritePlan, + ProfileWritePlan, +) +from reflexio.server.services.playbook.service import ( + PlaybookGenerationService, + PlaybookGenerationServiceConfig, +) from reflexio.server.services.profile.profile_generation_service_utils import ( ProfileGenerationRequest, ) @@ -165,12 +180,123 @@ def test_real_extraction_emits_tokens_and_learnings(tmp_path): assert tok.billing_input_tokens == tok.count_value assert tok.platform_llm is True # no api_key_config in the seeded Config - # learnings_generated.count_value must equal the existing generation_succeeded count. + # This fixture retains every generated output, so billing and success telemetry + # have equal counts here. gen = next(e for e in learning if e.event_name == "learnings_generated") succ = next(e for e in events if e.event_name == "generation_succeeded") assert gen.count_value == succ.count_value +def test_online_learning_bills_survivors_but_telemetry_counts_raw_results(tmp_path): + """Telemetry counts raw output while billing counts retained write-plan items.""" + storage = _build_sqlite_storage(tmp_path) + service = _build_profile_service(storage) + extracted_profiles = [ + UserProfile( + profile_id=profile_id, + user_id=_USER_ID, + content=profile_id, + last_modified_timestamp=1_000, + generated_from_request_id=_REQUEST_ID, + ) + for profile_id in ("retained", "dropped") + ] + write_plan = ProfileWritePlan( + user_id=_USER_ID, + request_id=_REQUEST_ID, + new_profiles=extracted_profiles[:1], + superseded_ids=[], + ) + + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object(service, "_prepare_generation_run", return_value=_prepared()), + patch.object( + service, "_execute_extractor", return_value=extracted_profiles + ), + patch.object(service, "_resolve_write_plan", return_value=write_plan), + patch.object(service, "_finalize_extraction_runs"), + patch.object(service, "_persist_write_plan"), + patch.object(service, "_extraction_input_text", return_value=""), + ): + plan = service.compute_generation(MagicMock()) + assert plan is not None + assert plan.generated_count == 2 + assert plan.billable_count == 1 + service.persist_generation(plan) + service.emit_generation_side_effects(plan) + finally: + configure_usage_event_recorder(None) + + billed = [event for event in events if event.event_name == "learnings_generated"] + assert [event.count_value for event in billed] == [1] + succeeded = [ + event for event in events if event.event_name == "generation_succeeded" + ] + assert [event.count_value for event in succeeded] == [2] + + +def test_online_playbook_bills_only_write_plan_survivors(tmp_path): + """Playbook candidates removed during write-plan resolution are not billable.""" + storage = _build_sqlite_storage(tmp_path) + context = _request_context(storage) + service = PlaybookGenerationService( + llm_client=LiteLLMClient(LiteLLMConfig(model="gpt-4o-mini")), + request_context=context, + ) + service.service_config = PlaybookGenerationServiceConfig( + request_id=_REQUEST_ID, + agent_version="v1", + user_id=_USER_ID, + source="api", + auto_run=True, + ) + extracted_playbooks = [ + UserPlaybook( + user_playbook_id=playbook_id, + agent_version="v1", + request_id=_REQUEST_ID, + content=f"content-{playbook_id}", + trigger=f"trigger-{playbook_id}", + ) + for playbook_id in (1, 2) + ] + write_plan = PlaybookWritePlan( + request_id=_REQUEST_ID, + output_pending_status=False, + skip_aggregation=True, + new_playbooks=extracted_playbooks[:1], + superseded_ids=[], + merge_groups=[], + ) + + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object(service, "_prepare_generation_run", return_value=_prepared()), + patch.object( + service, "_execute_extractor", return_value=extracted_playbooks + ), + patch.object(service, "_resolve_write_plan", return_value=write_plan), + patch.object(service, "_finalize_extraction_runs"), + patch.object(service, "_persist_write_plan"), + patch.object(service, "_extraction_input_text", return_value=""), + patch.object(service, "_dispatch_playbook_schedulers"), + ): + plan = service.compute_generation(MagicMock()) + assert plan is not None + service.persist_generation(plan) + service.emit_generation_side_effects(plan) + finally: + configure_usage_event_recorder(None) + + billed = [event for event in events if event.event_name == "learnings_generated"] + assert [event.count_value for event in billed] == [1] + + def test_should_run_skip_emits_no_learning_billing(tmp_path, monkeypatch): """A should_run-gated skip emits NO extraction_tokens / learnings_generated. @@ -376,6 +502,25 @@ def test_non_learning_service_emits_no_learning_billing_events(): ) +def test_base_finalization_requires_receipt_aware_override_for_run_id(): + """A receipt-less service cannot finalize a resumable extraction run.""" + service = _StubService( + llm_client=LiteLLMClient(LiteLLMConfig(model="gpt-4o-mini")), + request_context=_make_minimal_request_context(), + ) + + with patch.object(service, "_process_results") as process_results: + assert service._finalize_extracted_items(["legacy-item"]) is None + process_results.assert_called_once_with([["legacy-item"]]) + + with pytest.raises(NotImplementedError, match=r"(?i)receipt-aware"): + service._finalize_extracted_items( + ["resumable-item"], finalization_run_id="run-1" + ) + + process_results.assert_called_once() + + # --------------------------------------------------------------------------- # Dedup: billing reuses the should-run gate's already-fetched window instead of # re-querying storage purely to recompute billing_input_tokens. diff --git a/tests/server/services/test_non_extraction_learning_metering.py b/tests/server/services/test_non_extraction_learning_metering.py index 3536b39cc..e261ff3a5 100644 --- a/tests/server/services/test_non_extraction_learning_metering.py +++ b/tests/server/services/test_non_extraction_learning_metering.py @@ -1,12 +1,17 @@ from __future__ import annotations -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest +from reflexio.models.api_schema.service_schemas import AgentPlaybook, UserPlaybook +from reflexio.models.config_schema import PlaybookAggregatorConfig, PlaybookConfig from reflexio.server.api_endpoints.request_context import RequestContext from reflexio.server.services.extraction.resume_worker import ExtractionResumeWorker from reflexio.server.services.playbook.components.aggregator import PlaybookAggregator +from reflexio.server.services.playbook.playbook_service_utils import ( + PlaybookAggregatorRequest, +) from reflexio.server.services.storage.storage_base import ( AgentBinding, AgentRunRecord, @@ -47,11 +52,8 @@ def _agent_run(*, extractor_kind: str) -> AgentRunRecord: ) -def test_resumable_finalization_falls_back_when_items_lack_ids() -> None: - """Items with no durable id (e.g. plain objects) fall back to the - count-based aggregate event -- Task A3's documented fallback, since there - is no safe per-record id to key a dedup event on. - """ +def test_resumable_profile_bills_only_ids_returned_by_finalization() -> None: + """A preassigned candidate ID is not billed when finalization drops it.""" events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -60,21 +62,14 @@ def test_resumable_finalization_falls_back_when_items_lack_ids() -> None: ) run = _agent_run(extractor_kind="profile") - worker._record_finalized_learnings( - run, - [object(), object()], - entity_type="profile", - ) + dropped_candidate = MagicMock(profile_id="dropped-before-persist") + with patch( + "reflexio.server.services.extraction.resume_worker.ProfileGenerationService" + ) as service_class: + service_class.return_value._finalize_extracted_items.return_value = [] + worker._finalize_items(run, [dropped_candidate]) - assert len(events) == 1 - event = events[0] - assert event.event_name == "learnings_generated" - assert event.count_value == 2 - assert event.pipeline == "profile" - assert event.source == "resumable_extraction" - assert event.entity_type == "profile" - assert event.metadata == {"run_id": "run-1", "extractor_kind": "profile"} - assert event.event_key is not None and event.event_key.startswith("learn-batch:") + assert events == [] def test_resumable_fallback_reuses_its_event_key_on_finalization_retry() -> None: @@ -96,10 +91,7 @@ def test_resumable_fallback_reuses_its_event_key_on_finalization_retry() -> None def test_resumable_finalization_emits_one_event_per_profile_id() -> None: - """When every item carries a durable ``profile_id`` (the common case -- - profile ids are assigned by the extractor before finalize runs), emit one - entity-backed event per profile instead of the count-only fallback. - """ + """Finalization survivor IDs emit one entity-backed event per profile.""" events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -108,13 +100,9 @@ def test_resumable_finalization_emits_one_event_per_profile_id() -> None: ) run = _agent_run(extractor_kind="profile") - class _FakeProfile: - def __init__(self, profile_id: str) -> None: - self.profile_id = profile_id - worker._record_finalized_learnings( run, - [_FakeProfile("prof-1"), _FakeProfile("prof-2")], + ["prof-1", "prof-2"], entity_type="profile", ) @@ -134,7 +122,7 @@ def __init__(self, profile_id: str) -> None: def test_resumable_finalization_emits_one_event_per_playbook_id() -> None: - """Same as above for the playbook (``user_playbook_id``) kind.""" + """Finalization survivor IDs emit one entity-backed event per playbook.""" events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -143,13 +131,9 @@ def test_resumable_finalization_emits_one_event_per_playbook_id() -> None: ) run = _agent_run(extractor_kind="playbook") - class _FakePlaybook: - def __init__(self, user_playbook_id: int) -> None: - self.user_playbook_id = user_playbook_id - worker._record_finalized_learnings( run, - [_FakePlaybook(11), _FakePlaybook(12), _FakePlaybook(13)], + ["11", "12", "13"], entity_type="user_playbook", ) @@ -166,11 +150,8 @@ def __init__(self, user_playbook_id: int) -> None: assert event.event_key == f"learn:{event.entity_type}:{event.entity_id}" -def test_resumable_finalization_falls_back_when_a_playbook_id_is_unset() -> None: - """A ``user_playbook_id=0`` (default, unset) mixed in with real ids means - dedup dropped that item before persist -- fall back to the count-based - aggregate rather than emit a colliding ``learn:0`` key or fabricate an id. - """ +def test_resumable_playbook_bills_consolidation_replacement_id() -> None: + """Billing follows the persisted replacement, not its input candidate.""" events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -179,93 +160,87 @@ def test_resumable_finalization_falls_back_when_a_playbook_id_is_unset() -> None ) run = _agent_run(extractor_kind="playbook") - class _FakePlaybook: - def __init__(self, user_playbook_id: int) -> None: - self.user_playbook_id = user_playbook_id - - worker._record_finalized_learnings( - run, - [_FakePlaybook(21), _FakePlaybook(0)], - entity_type="user_playbook", - ) + original_candidate = MagicMock(user_playbook_id=21) + with patch( + "reflexio.server.services.extraction.resume_worker.PlaybookGenerationService" + ) as service_class: + service_class.return_value._finalize_extracted_items.return_value = ["88"] + worker._finalize_items(run, [original_candidate]) assert len(events) == 1 - assert events[0].count_value == 2 # total unchanged vs old count=2 - assert events[0].event_key is not None and events[0].event_key.startswith( - "learn-batch:" - ) - + assert events[0].count_value == 1 + assert events[0].event_key == "learn:user_playbook:88" + assert events[0].entity_id == "88" -def test_aggregation_records_attributed_learnings_generated() -> None: - """Aggregation emits one entity-backed event per generated playbook. - ``saved_playbook_list`` entries always carry a real ``agent_playbook_id`` - (``save_agent_playbooks`` raises rather than - returning a partial row) -- aggregator.py is the one caller with a clean, - always-populated per-record id list, so it uses the entity-backed path - (Task A3) rather than the count-only fallback. - """ +def test_aggregation_emits_no_learnings_generated( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A completed aggregation remains observable but adds no billable learning.""" events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) + request_context = _request_context() + storage = MagicMock() + configurator = MagicMock() + request_context.storage = storage + request_context.configurator = configurator aggregator = PlaybookAggregator( llm_client=MagicMock(), - request_context=_request_context(), + request_context=request_context, agent_version="v1", ) - - aggregator._record_learnings_generated( - learning_ids=["101", "102", "103"], - playbook_name="agent_rules", - request_id="agg-run-1", - metadata={"playbooks_generated": 3}, + config = MagicMock() + configurator.get_config.return_value = config + config.user_playbook_extractor_config = PlaybookConfig( + extractor_name="billing-boundary", + extraction_definition_prompt="Extract user playbooks.", + aggregation_config=PlaybookAggregatorConfig( + min_cluster_size=2, + reaggregation_trigger_count=2, + ), ) - - assert len(events) == 3 - assert sum(e.count_value for e in events) == 3 # total unchanged vs old count=3 - assert {e.event_key for e in events} == { - "learn:agent_playbook:101", - "learn:agent_playbook:102", - "learn:agent_playbook:103", - } - assert {e.entity_id for e in events} == {"101", "102", "103"} - for event in events: - assert event.event_name == "learnings_generated" - assert event.count_value == 1 - assert event.pipeline == "playbook" - assert event.source == "aggregation" - assert event.entity_type == "agent_playbook" - assert event.agent_version == "v1" - assert event.playbook_name == "agent_rules" - # tie key<->entity together so a swapped association can't pass on sets alone - assert event.event_key == f"learn:{event.entity_type}:{event.entity_id}" - - -def test_aggregation_falls_back_when_an_agent_playbook_id_is_falsy() -> None: - """Whole-branch-review finding (2): a falsy/0 ``agent_playbook_id`` mixed - into the run must not mint a colliding ``learn:agent_playbook:0`` key -- - fall back to the count-based aggregate event instead, matching - ``ExtractionResumeWorker``'s guard for the same failure mode. - """ - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) - aggregator = PlaybookAggregator( - llm_client=MagicMock(), - request_context=_request_context(), + user_playbooks = [ + UserPlaybook( + user_playbook_id=1, + agent_version="v1", + request_id="request-1", + playbook_name="user_playbook", + content="Document deployment decisions.", + ), + UserPlaybook( + user_playbook_id=2, + agent_version="v1", + request_id="request-2", + playbook_name="user_playbook", + content="Verify deployment outcomes.", + ), + ] + generated = AgentPlaybook( + agent_playbook_id=101, + playbook_name="user_playbook", agent_version="v1", + content="Deploy changes and verify results.", ) - - aggregator._record_learnings_generated( - learning_ids=["201"], # one id missing relative to total_count=2 - playbook_name="agent_rules", - request_id="agg-run-2", - metadata={"playbooks_generated": 2}, - total_count=2, + storage.count_user_playbooks.return_value = len(user_playbooks) + storage.get_agent_playbooks.return_value = [] + storage.get_user_playbooks.return_value = user_playbooks + storage.save_agent_playbooks.return_value = [generated] + monkeypatch.setattr( + aggregator, + "get_clusters", + lambda *_args: {0: user_playbooks}, ) - - assert len(events) == 1 - assert events[0].count_value == 2 # total unchanged vs old count=2 - assert events[0].event_key is not None and events[0].event_key.startswith( - "learn-batch:" + monkeypatch.setattr( + aggregator, + "_generate_playbooks_with_source_clusters", + lambda *_args, **_kwargs: [(generated, user_playbooks, None)], + ) + monkeypatch.setattr( + aggregator, "_enqueue_playbook_optimization", lambda _items: None ) - assert events[0].entity_type == "agent_playbook" - assert events[0].source == "aggregation" + + stats = aggregator.run(PlaybookAggregatorRequest(agent_version="v1", rerun=True)) + + assert stats["playbooks_generated"] == 1 + assert any(event.event_name == "aggregation_succeeded" for event in events) + assert not any(event.event_name == "learnings_generated" for event in events) diff --git a/tests/server/services/test_search_exposure.py b/tests/server/services/test_search_exposure.py new file mode 100644 index 000000000..96d577def --- /dev/null +++ b/tests/server/services/test_search_exposure.py @@ -0,0 +1,244 @@ +"""Direct behavior tests for user-playbook search exposure identities.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest + +from reflexio.models.api_schema.domain import BlockingIssue, UserPlaybook +from reflexio.models.api_schema.domain.enums import BlockingIssueKind, Status +from reflexio.server.services.search_exposure import ( + SearchExposureBatch, + build_user_playbook_exposure_event, + user_playbook_full_version_fingerprint, +) + + +def _playbook() -> UserPlaybook: + return UserPlaybook( + user_playbook_id=101, + user_id="user-1", + agent_version="agent-v1", + request_id="source-request-1", + playbook_name="Support policy", + created_at=1_700_000_000, + content="Escalate refund requests after verification.", + trigger="refund escalation", + rationale="Historical resolution pattern.", + blocking_issue=BlockingIssue( + kind=BlockingIssueKind.MISSING_TOOL, details="CRM access is absent." + ), + status=Status.ARCHIVED, + source="support-import", + source_interaction_ids=[11, 12], + expanded_terms="refund return escalation", + tags=["support", "refund"], + embedding=[0.25] * 512, + source_span="messages 4-6", + notes="Reviewed by ops.", + reader_angle="customer impact", + merged_into=88, + superseded_by=99, + ) + + +def _batch( + playbook: UserPlaybook, + *, + request_id: str | None = "request-1", + session_id: str | None = "session-1", + interaction_id: int | None = 41, + invocation_id: str = "invocation-1", +) -> SearchExposureBatch: + return SearchExposureBatch( + org_id="org-1", + request_id=request_id, + session_id=session_id, + interaction_id=interaction_id, + user_id="user-1", + user_playbooks=(playbook,), + invocation_id=invocation_id, + ) + + +def _event(batch: SearchExposureBatch, playbook: UserPlaybook): + return build_user_playbook_exposure_event( + batch, + playbook, + exposed_at=1_700_000_100, + ingested_at=1_700_000_101, + governance_subject_ref="user:user-1", + ) + + +def test_correlated_retries_keep_one_exposure_event_id_despite_invocation_id() -> None: + playbook = _playbook() + initial = _event(_batch(playbook, invocation_id="invocation-a"), playbook) + retry = _event(_batch(playbook, invocation_id="invocation-b"), playbook) + + assert initial.exposure_event_id == retry.exposure_event_id + + +def test_correlation_free_invocations_get_distinct_exposure_event_ids() -> None: + playbook = _playbook() + first = _event( + _batch( + playbook, + request_id=None, + session_id=None, + interaction_id=None, + invocation_id="invocation-a", + ), + playbook, + ) + second = _event( + _batch( + playbook, + request_id=None, + session_id=None, + interaction_id=None, + invocation_id="invocation-b", + ), + playbook, + ) + + assert first.exposure_event_id != second.exposure_event_id + + +def test_request_and_session_correlation_ids_normalize_whitespace_consistently() -> ( + None +): + playbook = _playbook() + whitespace = _batch( + playbook, + request_id=" \trequest-1\n", + session_id="\tsession-1 ", + interaction_id=None, + ) + normalized = _batch( + playbook, + request_id="request-1", + session_id="session-1", + interaction_id=None, + ) + blank = _batch( + playbook, + request_id=" \t", + session_id="\n ", + interaction_id=None, + invocation_id="fallback-invocation", + ) + absent = _batch( + playbook, + request_id=None, + session_id=None, + interaction_id=None, + invocation_id="fallback-invocation", + ) + + assert ( + (whitespace.request_id, whitespace.session_id) + == ( + normalized.request_id, + normalized.session_id, + ) + == ("request-1", "session-1") + ) + assert ( + _event(whitespace, playbook).exposure_event_id + == _event(normalized, playbook).exposure_event_id + ) + assert (blank.request_id, blank.session_id) == (None, None) + assert ( + _event(blank, playbook).exposure_event_id + == _event(absent, playbook).exposure_event_id + ) + + +def test_embedding_changes_do_not_change_full_version_fingerprint() -> None: + playbook = _playbook() + reembedded = playbook.model_copy(update={"embedding": [0.5] * 512}) + + assert user_playbook_full_version_fingerprint(playbook) == ( + user_playbook_full_version_fingerprint(reembedded) + ) + + +# Every current UserPlaybook field is persisted except its derived embedding vector. +_PERSISTED_FIELD_CHANGES = [ + ("user_playbook_id", 102), + ("user_id", "user-2"), + ("agent_version", "agent-v2"), + ("request_id", "source-request-2"), + ("playbook_name", "Returns policy"), + ("created_at", 1_700_000_001), + ("content", "Verify returns before escalating."), + ("trigger", "returns escalation"), + ("rationale", "Updated resolution pattern."), + ( + "blocking_issue", + BlockingIssue( + kind=BlockingIssueKind.PERMISSION_DENIED, + details="CRM access was denied.", + ), + ), + ("status", Status.PENDING), + ("source", "returns-import"), + ("source_interaction_ids", [11, 13]), + ("expanded_terms", "return exchange escalation"), + ("tags", ["support", "returns"]), + ("source_span", "messages 7-9"), + ("notes", "Needs legal review."), + ("reader_angle", "policy compliance"), + ("merged_into", 87), + ("superseded_by", 100), +] + + +@pytest.mark.parametrize(("field", "value"), _PERSISTED_FIELD_CHANGES) +def test_full_version_fingerprint_changes_for_each_persisted_non_embedding_field( + field: str, value: object +) -> None: + playbook = _playbook() + changed = playbook.model_copy(update={field: value}) + + assert user_playbook_full_version_fingerprint(playbook) != ( + user_playbook_full_version_fingerprint(changed) + ) + + +def test_full_version_fingerprint_coverage_includes_each_persisted_model_field() -> ( + None +): + assert {field for field, _value in _PERSISTED_FIELD_CHANGES} == ( + set(UserPlaybook.model_fields) - {"embedding"} + ) + + +def test_semantic_digest_and_fallback_identity_are_deterministic_and_domain_separated() -> ( + None +): + playbook = _playbook() + batch = _batch( + playbook, + request_id=None, + session_id=None, + interaction_id=None, + invocation_id="invocation-a", + ) + + first = _event(batch, playbook) + repeated = _event(replace(batch), playbook) + + assert ( + first.exposure_event_id + == repeated.exposure_event_id + == ("80b011b78df4a90e2238a7150d091c4d2f8c0e38d4343ccc7616e3536d40ca49") + ) + assert ( + first.served_semantic_digest + == repeated.served_semantic_digest + == ("d321988fa077b43deb4df4c89753b98c757c7c3167a12ea26af9247e665cc942") + ) + assert first.exposure_event_id != first.served_semantic_digest diff --git a/tests/server/test_billing_meter_events.py b/tests/server/test_billing_meter_events.py index b9ab61401..0012b753d 100644 --- a/tests/server/test_billing_meter_events.py +++ b/tests/server/test_billing_meter_events.py @@ -13,12 +13,11 @@ mint the same ``event_key`` and collapse into one event downstream. The existing count-based ``record_learnings_generated`` / ``emit_learnings_generated`` -remain as the documented FALLBACK for callers that genuinely lack a per-record -id list (e.g. dedup/consolidation can reduce the persisted count below the raw -extracted count, so there is no safe 1:1 id per unit of ``count``). The -fallback path now also carries a synthesized ``event_key=f"learn-batch:{uuid4()}"`` -so every ``learnings_generated`` event -- record-backed or batch -- has a -dedup key. +remain for online extraction callers that have a known billable count but do +not retain per-record ids. Resumable finalization uses only the record-backed +path and skips items without durable ids. The count-based helper carries a +synthesized ``event_key=f"learn-batch:{uuid4()}"`` so every +``learnings_generated`` event -- record-backed or batch -- has a dedup key. Totals are preserved in both paths: the sum of ``count_value`` across the per-record events equals ``len(learning_ids)``; the fallback emits exactly From bfff403ec2ea078fa2dff665fb896cf431ba2d83 Mon Sep 17 00:00:00 2001 From: guangyu-reflexio Date: Mon, 3 Aug 2026 08:00:10 -0800 Subject: [PATCH 02/63] fix: make receipt finalization winner-aware (#408) - make durable finalization receipt insertion atomically identify the sole winning transaction - emit resumable extraction billing only for the receipt winner - keep optimizer, aggregation, and tagging dispatch best-effort and at-most-once - preserve ordered winner IDs, empty receipts, and existing list-returning service wrappers - add an internal `FinalizationResult` carrying learning IDs and receipt ownership - roll back stale loser writes and reuse the committed winner receipt - independently isolate optimizer and aggregation scheduling failures - gate tagging and billing on receipt ownership - add deterministic retry, empty-receipt, identical-ID race, and scheduler coverage - 60 focused shared receipt/resume tests passed - repeated profile/playbook concurrency regressions passed - Ruff format/check and focused Pyright passed - enterprise Supabase/native-Postgres adapter contract is covered in the related enterprise branch * **Bug Fixes** * Improved resumable extraction finalization to prevent duplicate learning records, tags, and billing events. * Added reliable handling for concurrent finalization attempts, preserving the original result and avoiding repeated processing. * Scheduler failures are isolated so one failed follow-up task does not disrupt completed finalization. * Retry flows now maintain durable learning IDs and consistent outcomes. * **Reliability** * Finalization receipts are now immutable and safely reused across retries and competing workers. * Added safeguards for empty or incomplete finalization receipts. --- .../server/services/deferred_learning_plan.py | 18 + .../services/extraction/resume_worker.py | 33 +- reflexio/server/services/playbook/service.py | 99 ++- reflexio/server/services/profile/service.py | 69 +- .../agent_run/_agent_run_store.py | 21 +- .../agent_run/_agent_run_store.py | 4 +- .../services/extraction/test_resume_worker.py | 769 +++++++++++++++--- .../sqlite_storage/test_agent_run_storage.py | 23 +- .../test_non_extraction_learning_metering.py | 13 +- 9 files changed, 867 insertions(+), 182 deletions(-) diff --git a/reflexio/server/services/deferred_learning_plan.py b/reflexio/server/services/deferred_learning_plan.py index a5db43ad1..7e00c2ad9 100644 --- a/reflexio/server/services/deferred_learning_plan.py +++ b/reflexio/server/services/deferred_learning_plan.py @@ -26,6 +26,24 @@ ) +@dataclass(frozen=True) +class FinalizationResult: + """Internal outcome of resumable learning finalization. + + ``won_receipt`` is true only for the caller whose learning writes and + immutable finalization receipt committed together. Receipt-reuse callers + still receive the winner's ordered ids but must not replay billing or other + winner-only side effects. + """ + + learning_ids: list[str] + won_receipt: bool + + +class _FinalizationReceiptAlreadyExistsError(Exception): + """Rollback signal for a finalization transaction that lost receipt ownership.""" + + @dataclass(frozen=True) class ExtractorBookmarkAdvance: """The extractor stride-bookmark advance, deferred out of the extractor (F1). diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index 3d750f36d..00803bbc7 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -20,6 +20,7 @@ from reflexio.server.llm._litellm_types import ModelProvenance from reflexio.server.llm.litellm_client import LiteLLMClient, LiteLLMConfig from reflexio.server.llm.model_defaults import ModelRole, resolve_model_name +from reflexio.server.services.deferred_learning_plan import FinalizationResult from reflexio.server.services.extraction.agent_run_records import build_scope_hash from reflexio.server.services.extraction.pending_tool_call_dispatch import ( PendingToolCallToolContext, @@ -335,8 +336,9 @@ def run_once(self) -> AgentRunRecord | None: try: self.storage.update_agent_run_status(run.id, AgentRunStatus.FINALIZING) - self._finalize_items(run, items, model_provenance=model_provenance) - self._schedule_finalized_tagging(run) + result = self._finalize_items(run, items, model_provenance=model_provenance) + if result.won_receipt: + self._schedule_finalized_tagging(run) self.storage.consume_run_tool_dependencies(run.id) finalized_status = ( AgentRunStatus.FINALIZED_PENDING_TOOL @@ -382,8 +384,9 @@ def _retry_finalization(self, run: AgentRunRecord) -> AgentRunRecord | None: items, pending_tool_call_ids, model_provenance = ( self._items_from_committed_output(run) ) - self._finalize_items(run, items, model_provenance=model_provenance) - self._schedule_finalized_tagging(run) + result = self._finalize_items(run, items, model_provenance=model_provenance) + if result.won_receipt: + self._schedule_finalized_tagging(run) self.storage.consume_run_tool_dependencies(run.id) finalized_status = ( AgentRunStatus.FINALIZED_PENDING_TOOL @@ -896,7 +899,7 @@ def _finalize_items( items: list[Any], *, model_provenance: ModelProvenance | None = None, - ) -> None: + ) -> FinalizationResult: if run.binding.extractor_kind == "profile": service = ProfileGenerationService( llm_client=self.client, @@ -909,13 +912,16 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - learning_ids = service._finalize_extracted_items( + result = service._finalize_extracted_items_with_outcome( items, model_provenance=model_provenance, finalization_run_id=run.id, ) - self._record_finalized_learnings(run, learning_ids, entity_type="profile") - return + if result.won_receipt: + self._record_finalized_learnings( + run, result.learning_ids, entity_type="profile" + ) + return result if run.binding.extractor_kind == "playbook": user_id = run.binding.user_id if not user_id: @@ -932,16 +938,17 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - learning_ids = service._finalize_extracted_items( + result = service._finalize_extracted_items_with_outcome( items, model_provenance=model_provenance, extraction_run=run, finalization_run_id=run.id, ) - self._record_finalized_learnings( - run, learning_ids, entity_type="user_playbook" - ) - return + if result.won_receipt: + self._record_finalized_learnings( + run, result.learning_ids, entity_type="user_playbook" + ) + return result raise ResumeWorkerError( f"Unsupported extractor kind {run.binding.extractor_kind!r}" ) diff --git a/reflexio/server/services/playbook/service.py b/reflexio/server/services/playbook/service.py index 1954278c4..dd18358d2 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -34,7 +34,11 @@ BaseGenerationService, StatusChangeOperation, ) -from reflexio.server.services.deferred_learning_plan import PlaybookWritePlan +from reflexio.server.services.deferred_learning_plan import ( + FinalizationResult, + PlaybookWritePlan, + _FinalizationReceiptAlreadyExistsError, +) from reflexio.server.services.playbook.aggregation_trigger import ( maybe_trigger_user_playbook_aggregation, ) @@ -673,14 +677,27 @@ def _finalize_extracted_items( ) -> list[str]: """Finalize extracted playbooks for synchronous resume/manual callers. - Kept for the synchronous resume/manual callers - (``ExtractionResumeWorker`` calls this directly). Routes them through the - same ``_resolve_write_plan`` (compute) + ``_persist_write_plan`` - (persist) split the durable worker uses. Derived schedulers dispatch - best-effort after the finalization transaction commits. When an existing - finalization receipt is found, the method intentionally returns its - learning ids without replaying those schedulers. + Compatibility surface for synchronous callers that expect ordered + learning ids. Routes them through the same ``_resolve_write_plan`` + (compute) + ``_persist_write_plan`` (persist) split the durable worker + uses. Derived schedulers dispatch best-effort after the finalization + transaction commits. When an existing finalization receipt is found, + the method returns its learning ids without replaying those schedulers. """ + return self._finalize_extracted_items_with_outcome( + all_playbooks, + model_provenance=model_provenance, + finalization_run_id=finalization_run_id, + ).learning_ids + + def _finalize_extracted_items_with_outcome( + self, + all_playbooks: list[UserPlaybook], + *, + model_provenance: ModelProvenance | None = None, + finalization_run_id: str | None = None, + ) -> FinalizationResult: + """Finalize playbooks and expose the atomic receipt winner internally.""" entity_type = "user_playbook" if finalization_run_id is not None: receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] @@ -690,7 +707,7 @@ def _finalize_extracted_items( if receipt is not None: # Receipts make persistence/billing idempotent. Derived schedulers # are best-effort at-most-once and lack durable replay idempotency. - return receipt + return FinalizationResult(receipt, won_receipt=False) if model_provenance is not None: self._last_model_provenance = model_provenance previous_review_run = self._review_run @@ -709,38 +726,54 @@ def _finalize_extracted_items( with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] self._persist_write_plan(plan) self._dispatch_playbook_schedulers(plan) - return ( - [ - str(playbook.user_playbook_id) - for playbook in plan.new_playbooks - if playbook.user_playbook_id - ] - if plan is not None - else [] + return FinalizationResult( + learning_ids=( + [ + str(playbook.user_playbook_id) + for playbook in plan.new_playbooks + if playbook.user_playbook_id + ] + if plan is not None + else [] + ), + won_receipt=False, ) - with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + try: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + if plan is not None: + self._persist_write_plan(plan) + learning_ids = [ + str(playbook.user_playbook_id) + for playbook in plan.new_playbooks + if playbook.user_playbook_id + ] + inserted = self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + learning_ids=learning_ids, + ) + if not inserted: + raise _FinalizationReceiptAlreadyExistsError + except _FinalizationReceiptAlreadyExistsError: receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] run_id=finalization_run_id, entity_type=entity_type, ) - if receipt is not None: - return receipt - if plan is not None: - self._persist_write_plan(plan) - learning_ids = [ - str(playbook.user_playbook_id) - for playbook in plan.new_playbooks - if playbook.user_playbook_id - ] - self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] - run_id=finalization_run_id, - entity_type=entity_type, - learning_ids=learning_ids, - ) + if receipt is None: + raise RuntimeError( + "finalization receipt disappeared after insert conflict" + ) from None + return FinalizationResult(receipt, won_receipt=False) if plan is not None: self._dispatch_playbook_schedulers(plan) - return learning_ids + return FinalizationResult(learning_ids, won_receipt=True) def _apply_consolidation_lineage( self, diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index 691f8d34c..dc847f68e 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -30,7 +30,11 @@ BaseGenerationService, StatusChangeOperation, ) -from reflexio.server.services.deferred_learning_plan import ProfileWritePlan +from reflexio.server.services.deferred_learning_plan import ( + FinalizationResult, + ProfileWritePlan, + _FinalizationReceiptAlreadyExistsError, +) from reflexio.server.services.profile.components.extractor import ProfileExtractor from reflexio.server.services.profile.profile_generation_service_utils import ( ProfileGenerationRequest, @@ -341,12 +345,26 @@ def _finalize_extracted_items( ) -> list[str]: """Permanent V3 wrapper: compute-then-persist together (no external fence). - Kept for the synchronous resume/manual callers - (``ExtractionResumeWorker`` calls this directly). Routes them through the - same ``_resolve_write_plan`` (compute) + ``_persist_write_plan`` - (persist) split the durable worker uses — with no external - ``commit_scope`` — so the result is identical to the pre-split monolith. + Compatibility surface for synchronous callers that expect ordered + learning ids. Routes them through the same ``_resolve_write_plan`` + (compute) + ``_persist_write_plan`` (persist) split the durable worker + uses, with no external ``commit_scope``, so the result is identical to + the pre-split monolith. """ + return self._finalize_extracted_items_with_outcome( + all_new_profiles, + model_provenance=model_provenance, + finalization_run_id=finalization_run_id, + ).learning_ids + + def _finalize_extracted_items_with_outcome( + self, + all_new_profiles: list[UserProfile], + *, + model_provenance: ModelProvenance | None = None, + finalization_run_id: str | None = None, + ) -> FinalizationResult: + """Finalize profiles and expose the atomic receipt winner internally.""" entity_type = "profile" if finalization_run_id is not None: receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] @@ -354,7 +372,7 @@ def _finalize_extracted_items( entity_type=entity_type, ) if receipt is not None: - return receipt + return FinalizationResult(receipt, won_receipt=False) if model_provenance is not None: self._last_model_provenance = model_provenance plan = self._resolve_write_plan([all_new_profiles]) @@ -372,23 +390,36 @@ def _finalize_extracted_items( if plan is not None: with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] self._persist_write_plan(plan) - return learning_ids + return FinalizationResult(learning_ids, won_receipt=False) - with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + try: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + if plan is not None: + self._persist_write_plan(plan) + inserted = self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + learning_ids=learning_ids, + ) + if not inserted: + raise _FinalizationReceiptAlreadyExistsError + except _FinalizationReceiptAlreadyExistsError: receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] run_id=finalization_run_id, entity_type=entity_type, ) - if receipt is not None: - return receipt - if plan is not None: - self._persist_write_plan(plan) - self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] - run_id=finalization_run_id, - entity_type=entity_type, - learning_ids=learning_ids, - ) - return learning_ids + if receipt is None: + raise RuntimeError( + "finalization receipt disappeared after insert conflict" + ) from None + return FinalizationResult(receipt, won_receipt=False) + return FinalizationResult(learning_ids, won_receipt=True) def check_and_update_profiles(self, profiles: list[UserProfile]) -> None: """check if the profiles are expired and update them if they are""" diff --git a/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py b/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py index 299b20871..4ecc0fc05 100644 --- a/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py +++ b/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py @@ -254,7 +254,7 @@ def save_agent_run_finalization_receipt( run_id: str, entity_type: str, learning_ids: list[str], - ) -> None: + ) -> bool: expected_by_extractor = { "profile": "profile", "playbook": "user_playbook", @@ -275,13 +275,16 @@ def save_agent_run_finalization_receipt( raise ValueError( "agent-run finalization receipt entity type is invalid" ) - self.conn.execute( - """ + inserted = ( + self.conn.execute( + """ INSERT OR IGNORE INTO _agent_run_finalization_receipts (run_id, entity_type, learning_ids) VALUES (?, ?, ?) """, - (run_id, entity_type, encoded_ids), + (run_id, entity_type, encoded_ids), + ).rowcount + == 1 ) stored = self.conn.execute( """ @@ -291,14 +294,14 @@ def save_agent_run_finalization_receipt( """, (run_id,), ).fetchone() - if ( - stored is None - or stored["entity_type"] != entity_type - or json.loads(stored["learning_ids"]) != learning_ids - ): + if stored is None or stored["entity_type"] != entity_type: raise ValueError("agent-run finalization receipt is immutable") + stored_ids = json.loads(stored["learning_ids"]) + if not _valid_finalized_learning_ids(stored_ids): + raise ValueError("agent-run finalization receipt is corrupt") if self._own_transaction(): self.conn.commit() + return inserted @SQLiteStorageBase.handle_exceptions def get_latest_finalized_agent_run_for_request( diff --git a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py index 39edaf1fd..3acd3c9fd 100644 --- a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py +++ b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py @@ -39,8 +39,8 @@ def save_agent_run_finalization_receipt( run_id: str, entity_type: str, learning_ids: list[str], - ) -> None: - """Persist an immutable run-to-learning binding with the learning writes.""" + ) -> bool: + """Persist an immutable run-to-learning binding and report insert ownership.""" raise NotImplementedError(f"{type(self).__name__} does not support agent runs") def get_latest_finalized_agent_run_for_request( diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 348919ed7..35ed4daec 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -23,6 +23,7 @@ StorageConfigSQLite, ) from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.services.deferred_learning_plan import FinalizationResult from reflexio.server.services.extraction.resume_worker import ( ExtractionResumeWorker, _run_playbook_contract_selection, @@ -32,7 +33,10 @@ PlaybookConsolidationOutput, UnifyDecision, ) -from reflexio.server.services.playbook.service import PlaybookGenerationService +from reflexio.server.services.playbook.service import ( + PlaybookGenerationService, + PlaybookGenerationServiceConfig, +) from reflexio.server.services.profile.service import ( ProfileGenerationService, ProfileGenerationServiceConfig, @@ -86,6 +90,131 @@ def request_context(storage): return ctx +def _finalization_context(storage: SQLiteStorage) -> RequestContext: + context = RequestContext.__new__(RequestContext) + context.org_id = "org_1" + context.storage = storage + context.storage_base_dir = None + context.configurator = MagicMock() + context.configurator.get_config.return_value = Config( + storage_config=StorageConfigSQLite(), + profile_extractor_config=ProfileExtractorConfig( + extraction_definition_prompt="Extract durable user facts.", + ), + user_playbook_extractor_config=PlaybookConfig( + extraction_definition_prompt="Extract durable operating rules.", + ), + pending_tool_call_config=PendingToolCallConfig(enabled=True), + ) + context.prompt_manager = MagicMock() + context.prompt_manager.get_active_version.return_value = None + return context + + +def _finalizing_run( + *, run_id: str, extractor_kind: str, request_id: str +) -> AgentRunRecord: + return AgentRunRecord( + id=run_id, + binding=AgentBinding( + org_id="org_1", + extractor_kind=extractor_kind, + user_id="user_1", + request_id=request_id, + agent_version="v1", + source="api", + ), + status=AgentRunStatus.FINALIZING, + generation_request_snapshot={"request_id": request_id}, + ) + + +def _gate_initial_receipt_reads( + storages: tuple[SQLiteStorage, SQLiteStorage], +) -> None: + barrier = threading.Barrier(len(storages)) + + def install_gate(storage: SQLiteStorage) -> None: + original = storage.get_agent_run_finalization_receipt + first_read = True + + def gated_read(*, run_id: str, entity_type: str) -> list[str] | None: + nonlocal first_read + receipt = original(run_id=run_id, entity_type=entity_type) + if first_read: + first_read = False + barrier.wait(timeout=5) + return receipt + + storage.get_agent_run_finalization_receipt = MagicMock( # type: ignore[method-assign] + side_effect=gated_read + ) + + for storage in storages: + install_gate(storage) + + +def _hide_next_receipt_reads(storage: SQLiteStorage, *, count: int = 2) -> None: + original = storage.get_agent_run_finalization_receipt + remaining = count + + def stale_read(*, run_id: str, entity_type: str) -> list[str] | None: + nonlocal remaining + if remaining > 0: + remaining -= 1 + return None + return original(run_id=run_id, entity_type=entity_type) + + storage.get_agent_run_finalization_receipt = MagicMock( # type: ignore[method-assign] + side_effect=stale_read + ) + + +def _profile_service( + context: RequestContext, *, request_id: str +) -> ProfileGenerationService: + service = ProfileGenerationService(llm_client=MagicMock(), request_context=context) + service.service_config = ProfileGenerationServiceConfig( + user_id="user_1", + request_id=request_id, + source="api", + auto_run=False, + force_extraction=True, + ) + return service + + +def _playbook_service( + context: RequestContext, *, request_id: str +) -> PlaybookGenerationService: + service = PlaybookGenerationService(llm_client=MagicMock(), request_context=context) + service.service_config = PlaybookGenerationServiceConfig( + request_id=request_id, + agent_version="v1", + user_id="user_1", + source="api", + auto_run=False, + force_extraction=True, + ) + return service + + +def _playbook_candidates(*, request_id: str, prefix: str) -> list[UserPlaybook]: + return [ + UserPlaybook( + user_id="user_1", + agent_version="v1", + request_id=request_id, + content=f"Use deployment procedure {prefix}-{index}.", + trigger=f"when deployment condition {prefix}-{index} occurs", + rationale=f"Procedure {prefix}-{index} is required.", + source="api", + source_interaction_ids=[1, 2], + ) + for index in range(2) + ] + + @pytest.mark.parametrize( ("schema_name", "expected"), [ @@ -387,7 +516,7 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( patch("litellm.completion", side_effect=[response]), patch( "reflexio.server.services.profile.service." - "ProfileGenerationService._finalize_extracted_items", + "ProfileGenerationService._finalize_extracted_items_with_outcome", side_effect=RuntimeError("storage write failed"), ), ): @@ -414,8 +543,8 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( ), patch( "reflexio.server.services.profile.service." - "ProfileGenerationService._finalize_extracted_items", - return_value=None, + "ProfileGenerationService._finalize_extracted_items_with_outcome", + return_value=FinalizationResult([], won_receipt=False), ) as finalize, ): resumed = worker.drain(max_runs=1) @@ -769,122 +898,444 @@ def assert_playbook_scheduler_calls() -> None: } assert observed == { "profile_persisted": 1, - "profile_events": 2, + "profile_events": 1, "profile_distinct_keys": 1, "playbook_persisted": 1, "playbook_lineage_events": 1, - "playbook_events": 2, + "playbook_events": 1, "playbook_distinct_keys": 1, } assert profile_receipt_ids == profile_survivor_ids - assert profile_event_ids == profile_survivor_ids * 2 - assert profile_keys == [f"learn:profile:{profile_survivor_ids[0]}"] * 2 + assert profile_event_ids == profile_survivor_ids + assert profile_keys == [f"learn:profile:{profile_survivor_ids[0]}"] assert playbook_receipt_ids == playbook_survivor_ids - assert playbook_event_ids == playbook_survivor_ids * 2 - assert playbook_keys == [f"learn:user_playbook:{playbook_survivor_ids[0]}"] * 2 + assert playbook_event_ids == playbook_survivor_ids + assert playbook_keys == [f"learn:user_playbook:{playbook_survivor_ids[0]}"] -def test_two_stale_sqlite_workers_share_one_finalization_receipt(tmp_path): - db_path = str(tmp_path / "stale-finalizers.db") - with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): - storage_a = SQLiteStorage(org_id="org_1", db_path=db_path) - storage_b = SQLiteStorage(org_id="org_1", db_path=db_path) - storage_a.create_agent_run( - AgentRunRecord( - id="run_race", - binding=AgentBinding( - org_id="org_1", - extractor_kind="profile", - user_id="user_1", - request_id="request_race", - agent_version="v1", - source="api", - ), - status=AgentRunStatus.FINALIZING, - generation_request_snapshot={"request_id": "request_race"}, - ) +@pytest.mark.parametrize("failing_scheduler", ["optimizer", "aggregation"]) +def test_playbook_scheduler_failure_preserves_billing_and_isolated_retry( + storage, + failing_scheduler, +): + run = _finalizing_run( + run_id=f"run_scheduler_failure_{failing_scheduler}", + extractor_kind="playbook", + request_id=f"request_scheduler_failure_{failing_scheduler}", ) - barrier = threading.Barrier(2) - - def _service(storage_instance: SQLiteStorage) -> ProfileGenerationService: - context = RequestContext.__new__(RequestContext) - context.org_id = "org_1" - context.storage = storage_instance - context.configurator = MagicMock() - context.configurator.get_config.return_value = Config( - storage_config=StorageConfigSQLite(), - profile_extractor_config=ProfileExtractorConfig( - extraction_definition_prompt="Extract durable user facts.", + storage.create_agent_run(run) + worker = ExtractionResumeWorker( + request_context=_finalization_context(storage), + llm_client=MagicMock(), + ) + candidates = _playbook_candidates( + request_id=run.binding.request_id, + prefix=failing_scheduler, + ) + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object( + PlaybookGenerationService, + "_configured_playbook_config", + return_value=None, ), - ) - service = ProfileGenerationService( - llm_client=MagicMock(), - request_context=context, - ) - service.service_config = ProfileGenerationServiceConfig( + patch( + "reflexio.server.services.playbook.components.consolidator." + "PlaybookConsolidator.deduplicate", + side_effect=lambda results, *_args, **_kwargs: ( + [playbook for result in results for playbook in result], + [], + [], + ), + ) as deduplicate, + patch.object( + PlaybookGenerationService, + "_enqueue_user_playbook_optimization", + side_effect=( + RuntimeError("optimizer unavailable") + if failing_scheduler == "optimizer" + else None + ), + ) as optimize, + patch.object( + PlaybookGenerationService, + "_trigger_playbook_aggregation", + side_effect=( + RuntimeError("aggregation unavailable") + if failing_scheduler == "aggregation" + else None + ), + ) as aggregate, + ): + winner = worker._finalize_items(run, candidates) + retry = worker._finalize_items(run, candidates) + finally: + configure_usage_event_recorder(None) + + assert winner.won_receipt is True + assert retry == FinalizationResult(winner.learning_ids, won_receipt=False) + assert deduplicate.call_count == 1 + optimize.assert_called_once() + aggregate.assert_called_once() + billing_ids = [ + event.entity_id + for event in events + if event.event_name == "learnings_generated" + and event.entity_type == "user_playbook" + ] + assert billing_ids == winner.learning_ids + + +def test_empty_profile_receipt_wins_once_without_billing_or_recompute(storage): + run = _finalizing_run( + run_id="run_empty_profile", + extractor_kind="profile", + request_id="request_empty_profile", + ) + storage.create_agent_run(run) + context = _finalization_context(storage) + worker = ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with patch.object( + ProfileGenerationService, + "_resolve_write_plan", + return_value=None, + ) as resolve: + winner = worker._finalize_items(run, []) + retry = worker._finalize_items(run, []) + wrapper_ids = _profile_service( + context, request_id=run.binding.request_id + )._finalize_extracted_items([], finalization_run_id=run.id) + finally: + configure_usage_event_recorder(None) + + assert winner == FinalizationResult([], won_receipt=True) + assert retry == FinalizationResult([], won_receipt=False) + assert type(wrapper_ids) is list + assert wrapper_ids == [] + resolve.assert_called_once() + assert ( + storage.get_agent_run_finalization_receipt(run_id=run.id, entity_type="profile") + == [] + ) + assert [ + event for event in events if event.event_name == "learnings_generated" + ] == [] + + +def test_empty_playbook_receipt_retries_without_redispatch(storage): + run = AgentRunRecord( + id="run_empty_playbook", + binding=AgentBinding( + org_id="org_1", + extractor_kind="playbook", user_id="user_1", - request_id="request_race", + request_id="request_empty_playbook", + agent_version="v1", source="api", - auto_run=False, - force_extraction=True, + ), + status=AgentRunStatus.RESUME_READY, + generation_request_snapshot={"output_schema_name": "StructuredPlaybookList"}, + committed_output={"playbooks": []}, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + storage.create_agent_run(run) + context = _finalization_context(storage) + worker = ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + original_finalize = worker._finalize_items + outcomes: list[FinalizationResult] = [] + + def tracked_finalize(*args, **kwargs) -> FinalizationResult: + outcome = original_finalize(*args, **kwargs) + outcomes.append(outcome) + return outcome + + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object(storage, "claim_ready_agent_run", return_value=run), + patch.object( + worker, "_load_resolved_tool_calls", return_value=[MagicMock()] + ), + patch.object(worker, "_resume_run", return_value=([], [], None)), + patch.object( + worker, "_items_from_committed_output", return_value=([], [], None) + ), + patch.object(worker, "_finalize_items", side_effect=tracked_finalize), + patch.object( + PlaybookGenerationService, + "_resolve_write_plan", + return_value=None, + ) as resolve, + patch.object( + PlaybookGenerationService, + "_enqueue_user_playbook_optimization", + ) as optimize, + patch.object( + PlaybookGenerationService, + "_trigger_playbook_aggregation", + ) as aggregate, + patch( + "reflexio.server.services.extraction.resume_worker.schedule_tagging" + ) as schedule_tagging, + patch.object( + storage, + "consume_run_tool_dependencies", + side_effect=[RuntimeError("failed after finalization"), 0], + ), + ): + failed = worker.run_once() + assert failed is not None + assert failed.status == AgentRunStatus.FINALIZATION_FAILED + schedule_tagging.assert_called_once() + storage.update_agent_run_status( + run.id, + AgentRunStatus.FINALIZATION_FAILED, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + retried = worker.run_once() + schedule_tagging.assert_called_once() + wrapper_ids = _playbook_service( + context, request_id=run.binding.request_id + )._finalize_extracted_items([], finalization_run_id=run.id) + schedule_tagging.assert_called_once() + + assert retried is not None + assert retried.status == AgentRunStatus.FINALIZED + assert outcomes == [ + FinalizationResult([], won_receipt=True), + FinalizationResult([], won_receipt=False), + ] + assert type(wrapper_ids) is list + assert wrapper_ids == [] + resolve.assert_called_once() + optimize.assert_not_called() + aggregate.assert_not_called() + assert ( + storage.get_agent_run_finalization_receipt( + run_id=run.id, entity_type="user_playbook" + ) + == [] ) - return service + assert [ + event for event in events if event.event_name == "learnings_generated" + ] == [] + finally: + configure_usage_event_recorder(None) - def _gate_first_receipt_read(storage_instance: SQLiteStorage) -> None: - original = storage_instance.get_agent_run_finalization_receipt - first_read = True - def gated_read(*, run_id: str, entity_type: str) -> list[str] | None: - nonlocal first_read - receipt = original(run_id=run_id, entity_type=entity_type) - if first_read: - first_read = False - barrier.wait(timeout=5) - return receipt +def test_identical_profile_ids_use_atomic_receipt_owner(tmp_path): + db_path = str(tmp_path / "identical-profile-receipt.db") + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + storage_a = SQLiteStorage(org_id="org_1", db_path=db_path) + storage_b = SQLiteStorage(org_id="org_1", db_path=db_path) + run = _finalizing_run( + run_id="run_identical_profile", + extractor_kind="profile", + request_id="request_identical_profile", + ) + storage_a.create_agent_run(run) + contexts = [_finalization_context(item) for item in (storage_a, storage_b)] + workers = [ + ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + for context in contexts + ] + candidate_batches = [ + [ + UserProfile( + profile_id="profile-shared", + user_id="user_1", + content=f"Profile content from attempt {index}.", + last_modified_timestamp=1_000 + index, + generated_from_request_id=run.binding.request_id, + ) + ] + for index in range(2) + ] + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with patch( + "reflexio.server.services.profile.components.consolidator." + "ProfileConsolidator.deduplicate", + side_effect=lambda profiles, _user_id, _request_id: (profiles, [], []), + ): + winner = workers[0]._finalize_items(run, candidate_batches[0]) + _hide_next_receipt_reads(storage_b) + loser = workers[1]._finalize_items(run, candidate_batches[1]) + finally: + configure_usage_event_recorder(None) - storage_instance.get_agent_run_finalization_receipt = MagicMock( # type: ignore[method-assign] - side_effect=gated_read - ) + assert winner == FinalizationResult(["profile-shared"], won_receipt=True) + assert loser == FinalizationResult(["profile-shared"], won_receipt=False) + assert [profile.content for profile in storage_a.get_user_profile("user_1")] == [ + "Profile content from attempt 0." + ] + assert [ + event.entity_id + for event in events + if event.event_name == "learnings_generated" and event.entity_type == "profile" + ] == ["profile-shared"] - _gate_first_receipt_read(storage_a) - _gate_first_receipt_read(storage_b) - services = [_service(storage_a), _service(storage_b)] - candidates = [ - UserProfile( - profile_id=f"profile-race-{index}", - user_id="user_1", - content="User deployment target is AWS ECS.", - last_modified_timestamp=1_000, - generated_from_request_id="request_race", + +def test_identical_playbook_ids_use_atomic_receipt_owner(tmp_path): + db_path = str(tmp_path / "identical-playbook-receipt.db") + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + storage_a = SQLiteStorage(org_id="org_1", db_path=db_path) + storage_b = SQLiteStorage(org_id="org_1", db_path=db_path) + storage_a.conn.execute( + "CREATE TABLE receipt_race_writes (attempt INTEGER NOT NULL, content TEXT NOT NULL)" + ) + storage_a.conn.commit() + run = _finalizing_run( + run_id="run_identical_playbook", + extractor_kind="playbook", + request_id="request_identical_playbook", + ) + storage_a.create_agent_run(run) + contexts = [_finalization_context(item) for item in (storage_a, storage_b)] + workers = [ + ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + for context in contexts + ] + candidate_batches = [ + _playbook_candidates( + request_id=run.binding.request_id, prefix=f"attempt-{index}" ) for index in range(2) ] - results: list[list[str]] = [] + persist_attempts = iter(range(2)) + + def persist_fixed_ids(service, plan) -> None: + attempt = next(persist_attempts) + for index, playbook in enumerate(plan.new_playbooks): + playbook.user_playbook_id = 88 + index + service.storage.conn.execute( + "INSERT INTO receipt_race_writes (attempt, content) VALUES (?, ?)", + (attempt, playbook.content), + ) + + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object( + PlaybookGenerationService, + "_configured_playbook_config", + return_value=None, + ), + patch( + "reflexio.server.services.playbook.components.consolidator." + "PlaybookConsolidator.deduplicate", + side_effect=lambda results, *_args, **_kwargs: ( + [playbook for result in results for playbook in result], + [], + [], + ), + ), + patch.object( + PlaybookGenerationService, + "_persist_write_plan", + autospec=True, + side_effect=persist_fixed_ids, + ), + patch.object( + PlaybookGenerationService, + "_enqueue_user_playbook_optimization", + ) as optimize, + patch.object( + PlaybookGenerationService, + "_trigger_playbook_aggregation", + ) as aggregate, + ): + winner = workers[0]._finalize_items(run, candidate_batches[0]) + _hide_next_receipt_reads(storage_b) + loser = workers[1]._finalize_items(run, candidate_batches[1]) + finally: + configure_usage_event_recorder(None) + + assert winner == FinalizationResult(["88", "89"], won_receipt=True) + assert loser == FinalizationResult(["88", "89"], won_receipt=False) + persisted_writes = storage_a.conn.execute( + "SELECT attempt, content FROM receipt_race_writes ORDER BY content" + ).fetchall() + assert [(row["attempt"], row["content"]) for row in persisted_writes] == [ + (0, "Use deployment procedure attempt-0-0."), + (0, "Use deployment procedure attempt-0-1."), + ] + optimize.assert_called_once() + aggregate.assert_called_once() + assert [ + event.entity_id + for event in events + if event.event_name == "learnings_generated" + and event.entity_type == "user_playbook" + ] == ["88", "89"] + + +def test_two_stale_profile_workers_preserve_order_and_bill_only_winner(tmp_path): + db_path = str(tmp_path / "stale-finalizers.db") + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + storage_a = SQLiteStorage(org_id="org_1", db_path=db_path) + storage_b = SQLiteStorage(org_id="org_1", db_path=db_path) + run = _finalizing_run( + run_id="run_profile_race", + extractor_kind="profile", + request_id="request_profile_race", + ) + storage_a.create_agent_run(run) + _gate_initial_receipt_reads((storage_a, storage_b)) + contexts = [_finalization_context(storage) for storage in (storage_a, storage_b)] + resume_workers = [ + ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + for context in contexts + ] + candidate_batches = [ + [ + UserProfile( + profile_id=f"profile-{worker_index}-{item_index}", + user_id="user_1", + content=f"Profile candidate {worker_index}-{item_index}.", + last_modified_timestamp=1_000 + item_index, + generated_from_request_id="request_profile_race", + ) + for item_index in range(2) + ] + for worker_index in range(2) + ] + outcomes: list[FinalizationResult] = [] errors: list[BaseException] = [] def finalize(index: int) -> None: try: - results.append( - services[index]._finalize_extracted_items( - [candidates[index]], - finalization_run_id="run_race", - ) + outcomes.append( + resume_workers[index]._finalize_items(run, candidate_batches[index]) ) except BaseException as exc: # noqa: BLE001 - intentional thread error capture errors.append(exc) - with patch( - "reflexio.server.services.profile.components.consolidator." - "ProfileConsolidator.deduplicate", - side_effect=lambda profiles, _user_id, _request_id: (profiles, [], []), - ): - workers = [ - threading.Thread(target=finalize, args=(index,)) for index in range(2) - ] - for worker in workers: - worker.start() - for worker in workers: - worker.join(timeout=10) + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with patch( + "reflexio.server.services.profile.components.consolidator." + "ProfileConsolidator.deduplicate", + side_effect=lambda profiles, _user_id, _request_id: (profiles, [], []), + ): + workers = [ + threading.Thread(target=finalize, args=(index,)) for index in range(2) + ] + for worker in workers: + worker.start() + for worker in workers: + worker.join(timeout=10) + finally: + configure_usage_event_recorder(None) assert all(not worker.is_alive() for worker in workers) assert errors == [] @@ -892,11 +1343,141 @@ def finalize(index: int) -> None: profile.profile_id for profile in storage_a.get_user_profile("user_1") ] receipt_ids = storage_a.get_agent_run_finalization_receipt( - run_id="run_race", entity_type="profile" + run_id=run.id, entity_type="profile" ) - assert len(persisted_ids) == 1 + assert len(persisted_ids) == 2 assert receipt_ids == persisted_ids - assert results == [persisted_ids, persisted_ids] + assert [outcome.won_receipt for outcome in outcomes].count(True) == 1 + assert [outcome.won_receipt for outcome in outcomes].count(False) == 1 + assert all(outcome.learning_ids == receipt_ids for outcome in outcomes) + assert receipt_ids in [ + [profile.profile_id for profile in batch] for batch in candidate_batches + ] + wrapper_ids = _profile_service( + contexts[0], request_id=run.binding.request_id + )._finalize_extracted_items( + candidate_batches[1], + finalization_run_id=run.id, + ) + assert type(wrapper_ids) is list + assert wrapper_ids == receipt_ids + billing_events = [ + event + for event in events + if event.event_name == "learnings_generated" and event.entity_type == "profile" + ] + assert [event.entity_id for event in billing_events] == persisted_ids + + +def test_two_stale_playbook_workers_preserve_order_and_dispatch_once(tmp_path): + db_path = str(tmp_path / "stale-playbook-finalizers.db") + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + storage_a = SQLiteStorage(org_id="org_1", db_path=db_path) + storage_b = SQLiteStorage(org_id="org_1", db_path=db_path) + run = _finalizing_run( + run_id="run_playbook_race", + extractor_kind="playbook", + request_id="request_playbook_race", + ) + storage_a.create_agent_run(run) + _gate_initial_receipt_reads((storage_a, storage_b)) + contexts = [_finalization_context(storage) for storage in (storage_a, storage_b)] + resume_workers = [ + ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + for context in contexts + ] + candidate_batches = [ + _playbook_candidates( + request_id=run.binding.request_id, + prefix=f"worker-{worker_index}", + ) + for worker_index in range(2) + ] + outcomes: list[FinalizationResult] = [] + errors: list[BaseException] = [] + + def finalize(index: int) -> None: + try: + outcomes.append( + resume_workers[index]._finalize_items(run, candidate_batches[index]) + ) + except BaseException as exc: # noqa: BLE001 - intentional thread error capture + errors.append(exc) + + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object( + PlaybookGenerationService, + "_configured_playbook_config", + return_value=None, + ), + patch( + "reflexio.server.services.playbook.components.consolidator." + "PlaybookConsolidator.deduplicate", + side_effect=lambda results, *_args, **_kwargs: ( + [playbook for result in results for playbook in result], + [], + [], + ), + ), + patch.object( + PlaybookGenerationService, + "_enqueue_user_playbook_optimization", + ) as optimize, + patch.object( + PlaybookGenerationService, + "_trigger_playbook_aggregation", + ) as aggregate, + ): + threads = [ + threading.Thread(target=finalize, args=(index,)) for index in range(2) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert all(not thread.is_alive() for thread in threads) + assert errors == [] + receipt_ids = storage_a.get_agent_run_finalization_receipt( + run_id=run.id, + entity_type="user_playbook", + ) + persisted_ids = [ + str(row[0]) + for row in storage_a.conn.execute( + "SELECT user_playbook_id FROM user_playbooks " + "WHERE request_id = ? ORDER BY user_playbook_id ASC", + (run.binding.request_id,), + ).fetchall() + ] + wrapper_ids = _playbook_service( + contexts[0], request_id=run.binding.request_id + )._finalize_extracted_items( + candidate_batches[1], + finalization_run_id=run.id, + ) + + assert len(persisted_ids) == 2 + assert receipt_ids == persisted_ids + assert [outcome.won_receipt for outcome in outcomes].count(True) == 1 + assert [outcome.won_receipt for outcome in outcomes].count(False) == 1 + assert all(outcome.learning_ids == receipt_ids for outcome in outcomes) + assert type(wrapper_ids) is list + assert wrapper_ids == receipt_ids + optimize.assert_called_once() + aggregate.assert_called_once() + billing_ids = [ + event.entity_id + for event in events + if event.event_name == "learnings_generated" + and event.entity_type == "user_playbook" + ] + assert billing_ids == receipt_ids + finally: + configure_usage_event_recorder(None) def test_resume_worker_fails_run_when_step_budget_exhausted( diff --git a/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py b/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py index ce972c2de..cb369c7d4 100644 --- a/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py @@ -102,13 +102,15 @@ def test_sqlite_agent_run_crud_round_trip(storage): def test_finalization_receipt_accepts_empty_ids_idempotently(storage): storage.create_agent_run(_agent_run("run_empty", AgentRunStatus.FINALIZING)) - storage.save_agent_run_finalization_receipt( + inserted = storage.save_agent_run_finalization_receipt( run_id="run_empty", entity_type="profile", learning_ids=[] ) - storage.save_agent_run_finalization_receipt( + reused = storage.save_agent_run_finalization_receipt( run_id="run_empty", entity_type="profile", learning_ids=[] ) + assert inserted is True + assert reused is False assert ( storage.get_agent_run_finalization_receipt( run_id="run_empty", entity_type="profile" @@ -196,21 +198,22 @@ def test_finalization_receipt_rolls_back_with_learning(storage): ) -def test_finalization_receipt_rejects_conflicting_immutable_value(storage): +def test_finalization_receipt_reports_existing_immutable_value(storage): storage.create_agent_run(_agent_run("run_immutable", AgentRunStatus.FINALIZING)) - storage.save_agent_run_finalization_receipt( + inserted = storage.save_agent_run_finalization_receipt( run_id="run_immutable", entity_type="profile", learning_ids=["profile-1"], ) - with pytest.raises(StorageError, match="immutable"): - storage.save_agent_run_finalization_receipt( - run_id="run_immutable", - entity_type="profile", - learning_ids=["profile-2"], - ) + reused = storage.save_agent_run_finalization_receipt( + run_id="run_immutable", + entity_type="profile", + learning_ids=["profile-2"], + ) + assert inserted is True + assert reused is False assert storage.get_agent_run_finalization_receipt( run_id="run_immutable", entity_type="profile" ) == ["profile-1"] diff --git a/tests/server/services/test_non_extraction_learning_metering.py b/tests/server/services/test_non_extraction_learning_metering.py index e261ff3a5..ce43cbeb0 100644 --- a/tests/server/services/test_non_extraction_learning_metering.py +++ b/tests/server/services/test_non_extraction_learning_metering.py @@ -7,6 +7,7 @@ from reflexio.models.api_schema.service_schemas import AgentPlaybook, UserPlaybook from reflexio.models.config_schema import PlaybookAggregatorConfig, PlaybookConfig from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.services.deferred_learning_plan import FinalizationResult from reflexio.server.services.extraction.resume_worker import ExtractionResumeWorker from reflexio.server.services.playbook.components.aggregator import PlaybookAggregator from reflexio.server.services.playbook.playbook_service_utils import ( @@ -66,9 +67,13 @@ def test_resumable_profile_bills_only_ids_returned_by_finalization() -> None: with patch( "reflexio.server.services.extraction.resume_worker.ProfileGenerationService" ) as service_class: - service_class.return_value._finalize_extracted_items.return_value = [] + finalizer = service_class.return_value._finalize_extracted_items_with_outcome + finalizer.return_value = FinalizationResult([], won_receipt=True) worker._finalize_items(run, [dropped_candidate]) + finalizer.assert_called_once_with( + [dropped_candidate], model_provenance=None, finalization_run_id=run.id + ) assert events == [] @@ -164,9 +169,13 @@ def test_resumable_playbook_bills_consolidation_replacement_id() -> None: with patch( "reflexio.server.services.extraction.resume_worker.PlaybookGenerationService" ) as service_class: - service_class.return_value._finalize_extracted_items.return_value = ["88"] + finalizer = service_class.return_value._finalize_extracted_items_with_outcome + finalizer.return_value = FinalizationResult(["88"], won_receipt=True) worker._finalize_items(run, [original_candidate]) + finalizer.assert_called_once_with( + [original_candidate], model_provenance=None, finalization_run_id=run.id + ) assert len(events) == 1 assert events[0].count_value == 1 assert events[0].event_key == "learn:user_playbook:88" From 6ffeb1c4d7b9df301bd98df22132d72b315a0d56 Mon Sep 17 00:00:00 2001 From: guangyu-reflexio Date: Mon, 3 Aug 2026 13:38:41 -0800 Subject: [PATCH 03/63] fix: preserve offline tuner exposure evidence (#409) ## Summary - preserve every user-playbook exposure ingested within the fixed 14-day evidence window, even when the generic row cap is exceeded - separate the optional retrieval subject from the served playbook owner so unscoped searches remain attributable without becoming join-eligible - normalize blank retrieval subjects to the unscoped representation while retaining fail-closed owner checks for scoped searches ## Why The Phase 1 offline-tuner evidence foundation requires complete, reviewable exposure evidence. Generic row-cap retention could delete current-window events, and the original exposure envelope overloaded playbook ownership as the retrieval subject for unscoped searches. Both behaviors could distort later evidence eligibility and exact-user governance. ## Behavior - exposure rows younger than 14 days are protected from row-cap deletion - an exposure may carry a retrieval subject and a playbook owner independently - blank or whitespace retrieval subjects become unscoped - scoped retrieval still rejects a user playbook owned by another user before persistence ## Verification - 46 shared exposure and retention contract tests passed after rebasing onto current `main` - final affected enterprise matrix: 164 passed, 9 adapter-applicability skips - Ruff formatting/lint, Pyright, import, diff, and gitlink checks passed - correctness, security-resilience, architecture, and verification-testing review lenses are clean ## Related PRs - Enterprise Phase 1 stack: https://github.com/ReflexioAI/reflexio-enterprise/pull/936 ## Summary by CodeRabbit * **New Features** * Exposure records now preserve the playbook owner separately from the person retrieving it. * Blank retrieval identifiers are normalized consistently. * Playbook exposure data now includes governance subject references where available. * **Bug Fixes** * Prevented exposure records from being created when a playbook belongs to a different scoped user. * Added safeguards for unscoped exposure scenarios. * **Data Retention** * Open-world evidence, including playbook exposure records, is retained for at least 14 days before cleanup. --- reflexio/server/services/search_exposure.py | 13 ++++- reflexio/server/services/storage/retention.py | 3 ++ .../services/storage/retention_mixin.py | 36 ++++++++++++-- .../storage/sqlite_storage/base/_deletion.py | 9 +++- .../test_storage_contract_retention.py | 46 ++++++++++++++++++ tests/server/services/test_search_exposure.py | 48 +++++++++++++++++++ 6 files changed, 147 insertions(+), 8 deletions(-) diff --git a/reflexio/server/services/search_exposure.py b/reflexio/server/services/search_exposure.py index aae7fc4ff..63e887fff 100644 --- a/reflexio/server/services/search_exposure.py +++ b/reflexio/server/services/search_exposure.py @@ -34,6 +34,7 @@ def __post_init__(self) -> None: object.__setattr__( self, "session_id", _normalize_correlation_id(self.session_id) ) + object.__setattr__(self, "user_id", _normalize_correlation_id(self.user_id)) if self.interaction_id is not None and self.interaction_id <= 0: object.__setattr__(self, "interaction_id", None) @@ -46,12 +47,14 @@ class UserPlaybookExposureEvent: request_id: str | None session_id: str | None user_id: str | None + playbook_owner_user_id: str | None user_playbook_id: int | None served_semantic_digest: str | None served_full_version_fingerprint: str | None exposed_at: int | None ingested_at: int governance_subject_ref: str | None + playbook_owner_governance_subject_ref: str | None @dataclass(frozen=True) @@ -107,8 +110,14 @@ def build_user_playbook_exposure_event( exposed_at: int, ingested_at: int, governance_subject_ref: str | None, + playbook_owner_governance_subject_ref: str | None, ) -> UserPlaybookExposureEvent: """Build one deterministic event identity from retrieval-owned correlation.""" + if batch.user_id is not None and playbook.user_id != batch.user_id: + raise ValueError( + "served playbook owner does not match retrieval subject: " + f"user_playbook_id={playbook.user_playbook_id}" + ) identity: dict[str, object] = { "schema_version": "user-playbook-exposure-event-v1", "org_id": batch.org_id, @@ -128,7 +137,8 @@ def build_user_playbook_exposure_event( exposure_event_id=sha256(canonical_json_bytes(identity)).hexdigest(), request_id=batch.request_id, session_id=batch.session_id, - user_id=playbook.user_id, + user_id=batch.user_id, + playbook_owner_user_id=playbook.user_id, user_playbook_id=playbook.user_playbook_id, served_semantic_digest=incumbent_user_playbook_semantic_digest( content_digest=content_digest, @@ -140,6 +150,7 @@ def build_user_playbook_exposure_event( exposed_at=exposed_at, ingested_at=ingested_at, governance_subject_ref=governance_subject_ref, + playbook_owner_governance_subject_ref=(playbook_owner_governance_subject_ref), ) diff --git a/reflexio/server/services/storage/retention.py b/reflexio/server/services/storage/retention.py index 3a58a9251..9dcb86292 100644 --- a/reflexio/server/services/storage/retention.py +++ b/reflexio/server/services/storage/retention.py @@ -8,6 +8,7 @@ DEFAULT_ROW_RETENTION_LIMIT = 250_000 ROW_RETENTION_DELETE_FRACTION = 0.20 +OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS = 14 * 24 * 60 * 60 TOMBSTONE_STATUSES = ("archived", "merged", "superseded", "expired") @@ -20,6 +21,7 @@ class RetentionTarget: order_column: str id_columns: tuple[str, ...] priority_statuses: tuple[str, ...] = () + minimum_age_seconds: int = 0 @dataclass(frozen=True, slots=True) @@ -141,6 +143,7 @@ class OptimizationRetentionClass: "user_playbook_exposure_events", "ingested_at", ("exposure_event_id",), + minimum_age_seconds=OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS, ), RetentionTarget("skills", "skills", "created_at", ("skill_id",)), ) diff --git a/reflexio/server/services/storage/retention_mixin.py b/reflexio/server/services/storage/retention_mixin.py index 969cab1e2..8a0e4db29 100644 --- a/reflexio/server/services/storage/retention_mixin.py +++ b/reflexio/server/services/storage/retention_mixin.py @@ -9,6 +9,7 @@ from __future__ import annotations +import time from abc import ABC, abstractmethod from collections.abc import Iterator, Sequence from typing import Any @@ -98,7 +99,16 @@ def delete_oldest_retention_target_rows(self, target_name: str, count: int) -> i target = get_retention_target(target_name) if not self._retention_table_exists(target.table_name): return 0 - keys = self._retention_select_keys(target, count) + older_than_epoch = ( + int(time.time()) - target.minimum_age_seconds + if target.minimum_age_seconds > 0 + else None + ) + keys = self._retention_select_keys( + target, + count, + older_than_epoch=older_than_epoch, + ) if not keys: return 0 self._retention_perform_delete(target, keys) @@ -137,7 +147,11 @@ def _retention_gc_retired_optimization_jobs( return 0 def _retention_select_keys( - self, target: RetentionTarget, count: int + self, + target: RetentionTarget, + count: int, + *, + older_than_epoch: int | None, ) -> list[tuple[Any, ...]]: """Select tombstones first, then oldest rows when a target opts in. @@ -147,16 +161,27 @@ def _retention_select_keys( table holds that many rows. """ if not target.priority_statuses: - return self._retention_select_oldest_keys(target, count) + return self._retention_select_oldest_keys( + target, + count, + older_than_epoch=older_than_epoch, + ) keys = self._retention_select_oldest_keys( - target, count, statuses=target.priority_statuses + target, + count, + statuses=target.priority_statuses, + older_than_epoch=older_than_epoch, ) if len(keys) >= count: return keys seen = set(keys) - for key in self._retention_select_oldest_keys(target, count): + for key in self._retention_select_oldest_keys( + target, + count, + older_than_epoch=older_than_epoch, + ): if key not in seen: keys.append(key) seen.add(key) @@ -194,6 +219,7 @@ def _retention_select_oldest_keys( target: RetentionTarget, count: int, statuses: tuple[str, ...] | None = None, + older_than_epoch: int | None = None, ) -> list[tuple[Any, ...]]: """Return up to ``count`` oldest key tuples for ``target``. diff --git a/reflexio/server/services/storage/sqlite_storage/base/_deletion.py b/reflexio/server/services/storage/sqlite_storage/base/_deletion.py index a7227d222..3f3394170 100644 --- a/reflexio/server/services/storage/sqlite_storage/base/_deletion.py +++ b/reflexio/server/services/storage/sqlite_storage/base/_deletion.py @@ -50,16 +50,21 @@ def _retention_select_oldest_keys( target: RetentionTarget, count: int, statuses: tuple[str, ...] | None = None, + older_than_epoch: int | None = None, ) -> list[tuple[Any, ...]]: if statuses is not None and not statuses: return [] id_sql = ", ".join(target.id_columns) - where_sql = "" + predicates: list[str] = [] params: list[Any] = [] if statuses: placeholders = ", ".join("?" for _ in statuses) - where_sql = f"WHERE status IN ({placeholders}) " + predicates.append(f"status IN ({placeholders})") params.extend(statuses) + if older_than_epoch is not None: + predicates.append(f"{target.order_column} < ?") + params.append(older_than_epoch) + where_sql = f"WHERE {' AND '.join(predicates)} " if predicates else "" params.append(count) rows = self._fetchall( f"SELECT {id_sql} FROM {target.table_name} {where_sql}" # noqa: S608 diff --git a/tests/server/services/storage/test_storage_contract_retention.py b/tests/server/services/storage/test_storage_contract_retention.py index 6cdd17083..2612dc1c4 100644 --- a/tests/server/services/storage/test_storage_contract_retention.py +++ b/tests/server/services/storage/test_storage_contract_retention.py @@ -1,6 +1,8 @@ """Contract tests for generic row-retention storage methods.""" from datetime import UTC, datetime +from typing import Any, cast +from unittest.mock import patch import pytest @@ -188,6 +190,50 @@ def test_retention_request_cascade_cleans_interaction_fts( assert len(fts_kept) == 1, "fts row for surviving interaction must remain" +def test_retention_exposure_age_boundary_is_strict(storage: BaseStorage) -> None: + """Only exposure evidence strictly older than 14 days is row-cap eligible.""" + from reflexio.server.services.storage.retention import ( + OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS, + ) + + now = 2_000_000_000 + cutoff = now - OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS + conn = storage.conn # type: ignore[attr-defined] + conn.execute( + """CREATE TABLE user_playbook_exposure_events ( + exposure_event_id TEXT PRIMARY KEY, + ingested_at INTEGER NOT NULL + )""" + ) + conn.executemany( + """INSERT INTO user_playbook_exposure_events + (exposure_event_id, ingested_at) + VALUES (?, ?)""", + [ + ("older", cutoff - 1), + ("exact", cutoff), + ("newer", cutoff + 1), + ], + ) + conn.commit() + + retention_storage = cast(Any, storage) + with patch( + "reflexio.server.services.storage.retention_mixin.time.time", + return_value=now, + ): + deleted = retention_storage.delete_oldest_retention_target_rows( + "user_playbook_exposure_events", 3 + ) + + assert deleted == 1 + remaining = conn.execute( + "SELECT exposure_event_id FROM user_playbook_exposure_events " + "ORDER BY exposure_event_id" + ).fetchall() + assert [row["exposure_event_id"] for row in remaining] == ["exact", "newer"] + + # --------------------------------------------------------------------------- # Playbook retention FTS + vec cleanup (B3h) # --------------------------------------------------------------------------- diff --git a/tests/server/services/test_search_exposure.py b/tests/server/services/test_search_exposure.py index 96d577def..c8e330a80 100644 --- a/tests/server/services/test_search_exposure.py +++ b/tests/server/services/test_search_exposure.py @@ -69,6 +69,7 @@ def _event(batch: SearchExposureBatch, playbook: UserPlaybook): exposed_at=1_700_000_100, ingested_at=1_700_000_101, governance_subject_ref="user:user-1", + playbook_owner_governance_subject_ref="owner:user-1", ) @@ -106,6 +107,53 @@ def test_correlation_free_invocations_get_distinct_exposure_event_ids() -> None: assert first.exposure_event_id != second.exposure_event_id +def test_unscoped_exposure_keeps_unknown_subject_separate_from_playbook_owner() -> None: + playbook = _playbook() + batch = replace(_batch(playbook), user_id=None) + + event = build_user_playbook_exposure_event( + batch, + playbook, + exposed_at=1_700_000_100, + ingested_at=1_700_000_101, + governance_subject_ref=None, + playbook_owner_governance_subject_ref="owner:user-1", + ) + + assert event.user_id is None + assert event.governance_subject_ref is None + assert event.playbook_owner_user_id == "user-1" + assert event.playbook_owner_governance_subject_ref == "owner:user-1" + + +@pytest.mark.parametrize("user_id", ["", " \t\n"], ids=["empty", "whitespace"]) +def test_blank_retrieval_subject_normalizes_to_unscoped(user_id: str) -> None: + playbook = _playbook() + batch = replace(_batch(playbook), user_id=user_id) + + event = build_user_playbook_exposure_event( + batch, + playbook, + exposed_at=1_700_000_100, + ingested_at=1_700_000_101, + governance_subject_ref=None, + playbook_owner_governance_subject_ref="owner:user-1", + ) + + assert batch.user_id is None + assert event.user_id is None + assert event.governance_subject_ref is None + assert event.playbook_owner_user_id == "user-1" + + +def test_scoped_exposure_rejects_a_playbook_owned_by_another_user() -> None: + playbook = _playbook().model_copy(update={"user_id": "user-2"}) + batch = _batch(playbook) + + with pytest.raises(ValueError, match="does not match retrieval subject"): + _event(batch, playbook) + + def test_request_and_session_correlation_ids_normalize_whitespace_consistently() -> ( None ): From b8795bfa3d11d5db89a41a1027dda1f6523590ba Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 02:22:59 +0000 Subject: [PATCH 04/63] fix: fingerprint internal user playbook metadata --- reflexio/models/api_schema/domain/entities.py | 2 ++ reflexio/server/services/search_exposure.py | 6 +++++- tests/server/services/test_search_exposure.py | 11 +++++++++++ 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index f4f280945..1f0971704 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -364,6 +364,8 @@ class UserPlaybook(BaseModel): reader_angle: str | None = None merged_into: int | None = None superseded_by: int | None = None + governance_subject_ref: str | None = Field(default=None, exclude=True) + retired_at: int | None = Field(default=None, exclude=True) class ProfileChangeLog(BaseModel): diff --git a/reflexio/server/services/search_exposure.py b/reflexio/server/services/search_exposure.py index 63e887fff..5e7c13acd 100644 --- a/reflexio/server/services/search_exposure.py +++ b/reflexio/server/services/search_exposure.py @@ -98,7 +98,11 @@ def user_playbook_full_version_fingerprint(playbook: UserPlaybook) -> str: """ payload = { "schema_version": "user-playbook-full-version-v1", - "user_playbook": playbook.model_dump(mode="json", exclude={"embedding"}), + "user_playbook": playbook.model_dump(mode="json", exclude={"embedding"}) + | { + "governance_subject_ref": playbook.governance_subject_ref, + "retired_at": playbook.retired_at, + }, } return sha256(canonical_json_bytes(payload)).hexdigest() diff --git a/tests/server/services/test_search_exposure.py b/tests/server/services/test_search_exposure.py index c8e330a80..546f00b53 100644 --- a/tests/server/services/test_search_exposure.py +++ b/tests/server/services/test_search_exposure.py @@ -40,6 +40,8 @@ def _playbook() -> UserPlaybook: reader_angle="customer impact", merged_into=88, superseded_by=99, + governance_subject_ref="subject:user-1", + retired_at=1_700_000_050, ) @@ -213,6 +215,13 @@ def test_embedding_changes_do_not_change_full_version_fingerprint() -> None: ) +def test_internal_fields_are_excluded_from_user_playbook_serialization() -> None: + serialized = _playbook().model_dump(mode="json") + + assert "governance_subject_ref" not in serialized + assert "retired_at" not in serialized + + # Every current UserPlaybook field is persisted except its derived embedding vector. _PERSISTED_FIELD_CHANGES = [ ("user_playbook_id", 102), @@ -241,6 +250,8 @@ def test_embedding_changes_do_not_change_full_version_fingerprint() -> None: ("reader_angle", "policy compliance"), ("merged_into", 87), ("superseded_by", 100), + ("governance_subject_ref", "subject:user-2"), + ("retired_at", 1_700_000_051), ] From 18683a03915f8edbf96a996bb544fb79e3b59cc2 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 04:27:25 +0000 Subject: [PATCH 05/63] fix: retry receipt-backed learning billing Make resumable finalization replay stable per-learning usage keys until the configured recorder durably accepts or deduplicates them, while preserving fail-open metering for ordinary product calls. --- reflexio/server/billing_meter.py | 110 ++++++- .../services/extraction/resume_worker.py | 20 +- reflexio/server/usage_metrics.py | 166 ++++++++-- .../services/extraction/test_resume_worker.py | 290 ++++++++++++++++-- tests/server/test_usage_metrics.py | 68 +++- 5 files changed, 585 insertions(+), 69 deletions(-) diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py index 9a3633328..beb4f4588 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -13,7 +13,7 @@ from collections.abc import Mapping from typing import Any -from reflexio.server.usage_metrics import record_usage_event +from reflexio.server.usage_metrics import record_usage_event, record_usage_event_strict logger = logging.getLogger(__name__) @@ -193,9 +193,79 @@ def record_learnings_generated_records( entity_type: Optional entity type (e.g. ``"profile"``). metadata: Optional path-specific usage metadata (shared across events). """ + _record_learnings_generated_records( + record_event=record_usage_event, + org_id=org_id, + learning_ids=learning_ids, + platform_llm=platform_llm, + platform_storage=platform_storage, + pipeline=pipeline, + user_id=user_id, + request_id=request_id, + session_id=session_id, + source=source, + agent_version=agent_version, + playbook_name=playbook_name, + entity_type=entity_type, + metadata=metadata, + ) + + +def record_learnings_generated_records_strict( + *, + org_id: str, + learning_ids: list[str], + platform_llm: bool | None, + platform_storage: bool | None, + pipeline: str | None = None, + user_id: str | None = None, + request_id: str | None = None, + session_id: str | None = None, + source: str | None = None, + agent_version: str | None = None, + playbook_name: str | None = None, + entity_type: str | None = None, + metadata: Mapping[str, Any] | None = None, +) -> None: + """Strict per-record emission for receipt-backed finalization only.""" + _record_learnings_generated_records( + record_event=record_usage_event_strict, + org_id=org_id, + learning_ids=learning_ids, + platform_llm=platform_llm, + platform_storage=platform_storage, + pipeline=pipeline, + user_id=user_id, + request_id=request_id, + session_id=session_id, + source=source, + agent_version=agent_version, + playbook_name=playbook_name, + entity_type=entity_type, + metadata=metadata, + ) + + +def _record_learnings_generated_records( + *, + record_event: Any, + org_id: str, + learning_ids: list[str], + platform_llm: bool | None, + platform_storage: bool | None, + pipeline: str | None, + user_id: str | None, + request_id: str | None, + session_id: str | None, + source: str | None, + agent_version: str | None, + playbook_name: str | None, + entity_type: str | None, + metadata: Mapping[str, Any] | None, +) -> None: key_entity_type = entity_type or "_" for learning_id in learning_ids: - record_usage_event( + record_event( org_id=org_id, event_name="learnings_generated", event_category="learning", @@ -356,6 +426,42 @@ def emit_learnings_generated_records( ) +def emit_learnings_generated_records_strict( + *, + org_id: str, + configurator: Any, + learning_ids: list[str], + source: str, + pipeline: str | None = None, + user_id: str | None = None, + request_id: str | None = None, + agent_version: str | None = None, + playbook_name: str | None = None, + entity_type: str | None = None, + metadata: Mapping[str, Any] | None = None, +) -> None: + """Strict receipt-backed counterpart to the ordinary fail-open emitter.""" + if not learning_ids: + return + from reflexio.server.billing_signals import platform_llm_from_config + + config = configurator.get_config() + record_learnings_generated_records_strict( + org_id=org_id, + learning_ids=learning_ids, + platform_llm=platform_llm_from_config(config), + platform_storage=None, + pipeline=pipeline, + user_id=user_id, + request_id=request_id, + source=source, + agent_version=agent_version, + playbook_name=playbook_name, + entity_type=entity_type, + metadata=metadata, + ) + + def record_applied_learnings( *, org_id: str, diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index 00803bbc7..d81e1d6b3 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -917,10 +917,9 @@ def _finalize_items( model_provenance=model_provenance, finalization_run_id=run.id, ) - if result.won_receipt: - self._record_finalized_learnings( - run, result.learning_ids, entity_type="profile" - ) + self._record_finalized_learnings( + run, result.learning_ids, entity_type="profile" + ) return result if run.binding.extractor_kind == "playbook": user_id = run.binding.user_id @@ -944,10 +943,9 @@ def _finalize_items( extraction_run=run, finalization_run_id=run.id, ) - if result.won_receipt: - self._record_finalized_learnings( - run, result.learning_ids, entity_type="user_playbook" - ) + self._record_finalized_learnings( + run, result.learning_ids, entity_type="user_playbook" + ) return result raise ResumeWorkerError( f"Unsupported extractor kind {run.binding.extractor_kind!r}" @@ -961,7 +959,9 @@ def _record_finalized_learnings( Emits one event per durable learning id returned by finalization. Per-record keys make finalization retries idempotent downstream. """ - from reflexio.server.billing_meter import emit_learnings_generated_records + from reflexio.server.billing_meter import ( + emit_learnings_generated_records_strict, + ) if not learning_ids: return @@ -969,7 +969,7 @@ def _record_finalized_learnings( "run_id": run.id, "extractor_kind": run.binding.extractor_kind, } - emit_learnings_generated_records( + emit_learnings_generated_records_strict( org_id=self.request_context.org_id, configurator=self.request_context.configurator, learning_ids=learning_ids, diff --git a/reflexio/server/usage_metrics.py b/reflexio/server/usage_metrics.py index de7f6f1b0..1ebbcc02f 100644 --- a/reflexio/server/usage_metrics.py +++ b/reflexio/server/usage_metrics.py @@ -11,6 +11,7 @@ import time from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from enum import Enum from typing import Any logger = logging.getLogger(__name__) @@ -47,7 +48,25 @@ class UsageEvent: created_at: float = field(default_factory=time.time) -UsageEventRecorder = Callable[[UsageEvent], None] +class UsageEventDeliveryStatus(Enum): + """Recorder-reported durability outcome for one usage event.""" + + ACCEPTED = "accepted" + APPENDED = "appended" + DUPLICATE = "duplicate" + FAILED = "failed" + REJECTED = "rejected" + + +class UsageEventDeliveryError(RuntimeError): + """Raised when strict usage delivery was not durably accepted.""" + + def __init__(self, status: UsageEventDeliveryStatus) -> None: + self.status = status + super().__init__(f"usage event delivery {status.value}") + + +UsageEventRecorder = Callable[[UsageEvent], UsageEventDeliveryStatus | None] _recorder: UsageEventRecorder | None = None @@ -96,39 +115,120 @@ def record_usage_event( The product path must never fail because metrics failed, so this function catches and logs all recorder errors. """ - recorder = _recorder - if recorder is None: - return try: - recorder( - UsageEvent( - org_id=str(org_id), - event_name=event_name, - event_category=event_category, - user_id=user_id, - request_id=request_id, - session_id=session_id, - pipeline=pipeline, - entity_type=entity_type, - entity_id=entity_id, - event_key=event_key, - extractor_name=extractor_name, - playbook_name=playbook_name, - source=source, - agent_version=agent_version, - backend=backend, - outcome=outcome, - count_value=count_value, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - billing_input_tokens=billing_input_tokens, - platform_llm=platform_llm, - platform_storage=platform_storage, - caller_type=caller_type, - duration_ms=duration_ms, - error_kind=error_kind, - metadata=metadata or {}, - ) + record_usage_event_strict( + org_id=org_id, + event_name=event_name, + event_category=event_category, + user_id=user_id, + request_id=request_id, + session_id=session_id, + pipeline=pipeline, + entity_type=entity_type, + entity_id=entity_id, + event_key=event_key, + extractor_name=extractor_name, + playbook_name=playbook_name, + source=source, + agent_version=agent_version, + backend=backend, + outcome=outcome, + count_value=count_value, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + billing_input_tokens=billing_input_tokens, + platform_llm=platform_llm, + platform_storage=platform_storage, + caller_type=caller_type, + duration_ms=duration_ms, + error_kind=error_kind, + metadata=metadata, ) except Exception as exc: # noqa: BLE001 logger.warning("Usage metrics recorder failed: %s", exc) + + +def record_usage_event_strict( + *, + org_id: str, + event_name: str, + event_category: str, + user_id: str | None = None, + request_id: str | None = None, + session_id: str | None = None, + pipeline: str | None = None, + entity_type: str | None = None, + entity_id: str | None = None, + event_key: str | None = None, + extractor_name: str | None = None, + playbook_name: str | None = None, + source: str | None = None, + agent_version: str | None = None, + backend: str | None = None, + outcome: str | None = None, + count_value: int = 1, + prompt_tokens: int | None = None, + completion_tokens: int | None = None, + billing_input_tokens: int | None = None, + platform_llm: bool | None = None, + platform_storage: bool | None = None, + caller_type: str | None = None, + duration_ms: int | None = None, + error_kind: str | None = None, + metadata: Mapping[str, Any] | None = None, +) -> UsageEventDeliveryStatus: + """Deliver one event and fail unless the recorder accepted it durably. + + Recorders predating delivery outcomes return ``None``; that remains an + accepted result so direct test recorders and non-enterprise integrations + keep working. A missing recorder is also accepted because OSS deployments + intentionally have no billing sink. + """ + recorder = _recorder + if recorder is None: + return UsageEventDeliveryStatus.ACCEPTED + delivery_outcome = recorder( + UsageEvent( + org_id=str(org_id), + event_name=event_name, + event_category=event_category, + user_id=user_id, + request_id=request_id, + session_id=session_id, + pipeline=pipeline, + entity_type=entity_type, + entity_id=entity_id, + event_key=event_key, + extractor_name=extractor_name, + playbook_name=playbook_name, + source=source, + agent_version=agent_version, + backend=backend, + outcome=outcome, + count_value=count_value, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + billing_input_tokens=billing_input_tokens, + platform_llm=platform_llm, + platform_storage=platform_storage, + caller_type=caller_type, + duration_ms=duration_ms, + error_kind=error_kind, + metadata=metadata or {}, + ) + ) + status = ( + UsageEventDeliveryStatus.ACCEPTED + if delivery_outcome is None + else delivery_outcome + ) + if not isinstance(status, UsageEventDeliveryStatus): + raise TypeError( + f"usage event recorder returned an invalid delivery status: {status!r}" + ) + if status in { + UsageEventDeliveryStatus.FAILED, + UsageEventDeliveryStatus.REJECTED, + }: + raise UsageEventDeliveryError(status) + return status diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 35ed4daec..412a954f8 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -4,6 +4,7 @@ import tempfile import threading from collections.abc import Callable +from contextlib import ExitStack from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock, patch @@ -55,6 +56,7 @@ ) from reflexio.server.usage_metrics import ( UsageEvent, + UsageEventDeliveryStatus, configure_usage_event_recorder, ) @@ -215,6 +217,22 @@ def _playbook_candidates(*, request_id: str, prefix: str) -> list[UserPlaybook]: ] +class _DeduplicatingUsageRecorder: + def __init__(self) -> None: + self.attempts: list[UsageEvent] = [] + self.events: list[UsageEvent] = [] + self._accepted_keys: set[str] = set() + + def __call__(self, event: UsageEvent) -> UsageEventDeliveryStatus: + self.attempts.append(event) + assert event.event_key is not None + if event.event_key in self._accepted_keys: + return UsageEventDeliveryStatus.DUPLICATE + self._accepted_keys.add(event.event_key) + self.events.append(event) + return UsageEventDeliveryStatus.APPENDED + + @pytest.mark.parametrize( ("schema_name", "expected"), [ @@ -652,7 +670,10 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( learning_ids = ["profile_1"] worker = ExtractionResumeWorker(request_context=request_context) - with patch("reflexio.server.billing_meter.record_usage_event") as record_event: + with patch( + "reflexio.server.billing_meter.record_usage_event_strict", + return_value=UsageEventDeliveryStatus.APPENDED, + ) as record_event: worker._record_finalized_learnings(run, learning_ids, entity_type="profile") worker._record_finalized_learnings(run, learning_ids, entity_type="profile") @@ -670,6 +691,192 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( ] +@pytest.mark.parametrize( + ("extractor_kind", "entity_type"), + [("profile", "profile"), ("playbook", "user_playbook")], +) +def test_delivery_failure_after_receipt_commit_retries_billing_without_recompute( + request_context, + storage, + extractor_kind, + entity_type, +): + _seed_interactions(storage) + request_context.configurator.get_config.return_value = Config( + storage_config=StorageConfigSQLite(), + profile_extractor_config=ProfileExtractorConfig( + extraction_definition_prompt="Extract durable user facts.", + ), + user_playbook_extractor_config=PlaybookConfig( + extraction_definition_prompt="Extract durable operating rules.", + ), + pending_tool_call_config=PendingToolCallConfig(enabled=True), + ) + request_id = f"request_{extractor_kind}_delivery_retry" + run_id = f"run_{extractor_kind}_delivery_retry" + committed_output = ( + { + "profiles": [ + { + "content": "User deployment target is AWS ECS.", + "time_to_live": "infinity", + } + ] + } + if extractor_kind == "profile" + else { + "playbooks": [ + { + "content": "Prefer AWS ECS as the deployment target.", + "trigger": "when selecting a deployment target", + "rationale": "The team standardizes on AWS.", + } + ] + } + ) + storage.create_agent_run( + AgentRunRecord( + id=run_id, + binding=AgentBinding( + org_id="org_1", + extractor_kind=extractor_kind, + user_id="user_1", + request_id=request_id, + agent_version="v1", + source="api", + source_interaction_ids=[1, 2], + ), + status=AgentRunStatus.FINALIZATION_FAILED, + generation_request_snapshot={ + "request_id": request_id, + "output_schema_name": "StructuredPlaybookList", + }, + committed_output=committed_output, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + ) + worker = ExtractionResumeWorker( + request_context=request_context, + llm_client=MagicMock(), + ) + attempts: list[UsageEvent] = [] + accepted: dict[str, UsageEvent] = {} + fail_next = True + + def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: + nonlocal fail_next + attempts.append(event) + if fail_next: + fail_next = False + raise RuntimeError("billing sink unavailable") + assert event.event_key is not None + if event.event_key in accepted: + return UsageEventDeliveryStatus.DUPLICATE + accepted[event.event_key] = event + return UsageEventDeliveryStatus.APPENDED + + configure_usage_event_recorder(recorder) + try: + with ExitStack() as stack: + schedule_tagging = stack.enter_context( + patch.object(worker, "_schedule_finalized_tagging") + ) + if extractor_kind == "profile": + stack.enter_context( + patch( + "reflexio.server.services.profile.components.consolidator." + "ProfileConsolidator.deduplicate", + side_effect=lambda profiles, _user_id, _request_id: ( + profiles, + [], + [], + ), + ) + ) + optimize = aggregate = None + else: + stack.enter_context( + patch.object( + PlaybookGenerationService, + "_configured_playbook_config", + return_value=None, + ) + ) + stack.enter_context( + patch( + "reflexio.server.services.playbook.components.consolidator." + "PlaybookConsolidator.deduplicate", + side_effect=lambda results, *_args, **_kwargs: ( + [playbook for result in results for playbook in result], + [], + [], + ), + ) + ) + optimize = stack.enter_context( + patch.object( + PlaybookGenerationService, + "_enqueue_user_playbook_optimization", + ) + ) + aggregate = stack.enter_context( + patch.object( + PlaybookGenerationService, + "_trigger_playbook_aggregation", + ) + ) + + failed = worker.run_once() + assert failed is not None + assert failed.status == AgentRunStatus.FINALIZATION_FAILED + receipt_after_failure = storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type=entity_type, + ) + assert receipt_after_failure + rows_after_failure = ( + [profile.profile_id for profile in storage.get_user_profile("user_1")] + if extractor_kind == "profile" + else [ + str(row[0]) + for row in storage.conn.execute( + "SELECT user_playbook_id FROM user_playbooks " + "WHERE request_id = ? ORDER BY user_playbook_id", + (request_id,), + ).fetchall() + ] + ) + + storage.update_agent_run_status( + run_id, + AgentRunStatus.FINALIZATION_FAILED, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + retry = worker.run_once() + + assert retry is not None + assert retry.status == AgentRunStatus.FINALIZED + assert ( + storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type=entity_type, + ) + == receipt_after_failure + ) + assert rows_after_failure == receipt_after_failure + assert [event.event_key for event in attempts] == [ + f"learn:{entity_type}:{receipt_after_failure[0]}", + f"learn:{entity_type}:{receipt_after_failure[0]}", + ] + assert list(accepted) == [f"learn:{entity_type}:{receipt_after_failure[0]}"] + schedule_tagging.assert_not_called() + if optimize is not None and aggregate is not None: + optimize.assert_called_once() + aggregate.assert_called_once() + finally: + configure_usage_event_recorder(None) + + def test_retry_after_billing_reuses_ids_without_replaying_playbook_schedulers( request_context, storage, @@ -687,7 +894,19 @@ def test_retry_after_billing_reuses_ids_without_replaying_playbook_schedulers( pending_tool_call_config=PendingToolCallConfig(enabled=True), ) events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + delivery_attempts: list[UsageEvent] = [] + accepted_keys: set[str] = set() + + def deduplicating_recorder(event: UsageEvent) -> UsageEventDeliveryStatus: + delivery_attempts.append(event) + assert event.event_key is not None + if event.event_key in accepted_keys: + return UsageEventDeliveryStatus.DUPLICATE + accepted_keys.add(event.event_key) + events.append(event) + return UsageEventDeliveryStatus.APPENDED + + configure_usage_event_recorder(deduplicating_recorder) def _run_failed_then_retried( run_id: str, @@ -911,6 +1130,12 @@ def assert_playbook_scheduler_calls() -> None: assert playbook_receipt_ids == playbook_survivor_ids assert playbook_event_ids == playbook_survivor_ids assert playbook_keys == [f"learn:user_playbook:{playbook_survivor_ids[0]}"] + assert [event.event_key for event in delivery_attempts] == [ + profile_keys[0], + profile_keys[0], + playbook_keys[0], + playbook_keys[0], + ] @pytest.mark.parametrize("failing_scheduler", ["optimizer", "aggregation"]) @@ -932,8 +1157,8 @@ def test_playbook_scheduler_failure_preserves_billing_and_isolated_retry( request_id=run.binding.request_id, prefix=failing_scheduler, ) - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) try: with ( patch.object( @@ -981,11 +1206,15 @@ def test_playbook_scheduler_failure_preserves_billing_and_isolated_retry( aggregate.assert_called_once() billing_ids = [ event.entity_id - for event in events + for event in recorder.events if event.event_name == "learnings_generated" and event.entity_type == "user_playbook" ] assert billing_ids == winner.learning_ids + assert [event.entity_id for event in recorder.attempts] == [ + *winner.learning_ids, + *winner.learning_ids, + ] def test_empty_profile_receipt_wins_once_without_billing_or_recompute(storage): @@ -997,8 +1226,8 @@ def test_empty_profile_receipt_wins_once_without_billing_or_recompute(storage): storage.create_agent_run(run) context = _finalization_context(storage) worker = ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) try: with patch.object( ProfileGenerationService, @@ -1023,7 +1252,7 @@ def test_empty_profile_receipt_wins_once_without_billing_or_recompute(storage): == [] ) assert [ - event for event in events if event.event_name == "learnings_generated" + event for event in recorder.events if event.event_name == "learnings_generated" ] == [] @@ -1054,8 +1283,8 @@ def tracked_finalize(*args, **kwargs) -> FinalizationResult: outcomes.append(outcome) return outcome - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) try: with ( patch.object(storage, "claim_ready_agent_run", return_value=run), @@ -1123,7 +1352,9 @@ def tracked_finalize(*args, **kwargs) -> FinalizationResult: == [] ) assert [ - event for event in events if event.event_name == "learnings_generated" + event + for event in recorder.events + if event.event_name == "learnings_generated" ] == [] finally: configure_usage_event_recorder(None) @@ -1157,8 +1388,8 @@ def test_identical_profile_ids_use_atomic_receipt_owner(tmp_path): ] for index in range(2) ] - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) try: with patch( "reflexio.server.services.profile.components.consolidator." @@ -1178,9 +1409,13 @@ def test_identical_profile_ids_use_atomic_receipt_owner(tmp_path): ] assert [ event.entity_id - for event in events + for event in recorder.events if event.event_name == "learnings_generated" and event.entity_type == "profile" ] == ["profile-shared"] + assert [event.entity_id for event in recorder.attempts] == [ + "profile-shared", + "profile-shared", + ] def test_identical_playbook_ids_use_atomic_receipt_owner(tmp_path): @@ -1220,8 +1455,8 @@ def persist_fixed_ids(service, plan) -> None: (attempt, playbook.content), ) - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) try: with ( patch.object( @@ -1272,10 +1507,11 @@ def persist_fixed_ids(service, plan) -> None: aggregate.assert_called_once() assert [ event.entity_id - for event in events + for event in recorder.events if event.event_name == "learnings_generated" and event.entity_type == "user_playbook" ] == ["88", "89"] + assert [event.entity_id for event in recorder.attempts] == ["88", "89", "88", "89"] def test_two_stale_profile_workers_preserve_order_and_bill_only_winner(tmp_path): @@ -1319,8 +1555,8 @@ def finalize(index: int) -> None: except BaseException as exc: # noqa: BLE001 - intentional thread error capture errors.append(exc) - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) try: with patch( "reflexio.server.services.profile.components.consolidator." @@ -1363,10 +1599,14 @@ def finalize(index: int) -> None: assert wrapper_ids == receipt_ids billing_events = [ event - for event in events + for event in recorder.events if event.event_name == "learnings_generated" and event.entity_type == "profile" ] assert [event.entity_id for event in billing_events] == persisted_ids + assert [event.entity_id for event in recorder.attempts] == [ + *persisted_ids, + *persisted_ids, + ] def test_two_stale_playbook_workers_preserve_order_and_dispatch_once(tmp_path): @@ -1404,8 +1644,8 @@ def finalize(index: int) -> None: except BaseException as exc: # noqa: BLE001 - intentional thread error capture errors.append(exc) - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) try: with ( patch.object( @@ -1471,11 +1711,15 @@ def finalize(index: int) -> None: aggregate.assert_called_once() billing_ids = [ event.entity_id - for event in events + for event in recorder.events if event.event_name == "learnings_generated" and event.entity_type == "user_playbook" ] assert billing_ids == receipt_ids + assert [event.entity_id for event in recorder.attempts] == [ + *receipt_ids, + *receipt_ids, + ] finally: configure_usage_event_recorder(None) diff --git a/tests/server/test_usage_metrics.py b/tests/server/test_usage_metrics.py index 5bf359332..dff33a96f 100644 --- a/tests/server/test_usage_metrics.py +++ b/tests/server/test_usage_metrics.py @@ -1,5 +1,11 @@ +import pytest + from reflexio.server import usage_metrics -from reflexio.server.usage_metrics import UsageEvent +from reflexio.server.usage_metrics import ( + UsageEvent, + UsageEventDeliveryError, + UsageEventDeliveryStatus, +) def test_usage_event_carries_event_key(): @@ -20,3 +26,63 @@ def test_usage_event_carries_event_key(): def test_event_key_defaults_none(): assert UsageEvent(org_id="7", event_name="x", event_category="y").event_key is None + + +def test_strict_delivery_accepts_legacy_none_recorder(): + captured = [] + usage_metrics.configure_usage_event_recorder(captured.append) + try: + outcome = usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert outcome is UsageEventDeliveryStatus.ACCEPTED + assert [event.event_key for event in captured] == ["learn:profile:1"] + + +@pytest.mark.parametrize( + "outcome", + [UsageEventDeliveryStatus.FAILED, UsageEventDeliveryStatus.REJECTED], +) +def test_strict_delivery_raises_for_unaccepted_outcome(outcome): + usage_metrics.configure_usage_event_recorder(lambda _event: outcome) + try: + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert exc_info.value.status is outcome + + +def test_strict_delivery_propagates_recorder_exception_while_ordinary_stays_fail_open(): + def fail(_event): + raise RuntimeError("sink unavailable") + + usage_metrics.configure_usage_event_recorder(fail) + try: + usage_metrics.record_usage_event( + org_id="7", + event_name="search_request", + event_category="application", + event_key="search:1", + ) + with pytest.raises(RuntimeError, match="sink unavailable"): + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) From 58e0a4d16b36ce6aed279ffbe5624f8039238586 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 05:07:04 +0000 Subject: [PATCH 06/63] fix: keep receipt billing obligations retryable --- reflexio/server/api.py | 13 +++++++ reflexio/server/billing_meter.py | 39 ++++++++++++------- .../services/extraction/resume_worker.py | 31 +++++++++++---- reflexio/server/usage_metrics.py | 26 ++++++++----- .../services/extraction/test_resume_worker.py | 17 +++++++- tests/server/test_usage_metrics.py | 38 ++++++++++++++++-- 6 files changed, 127 insertions(+), 37 deletions(-) diff --git a/reflexio/server/api.py b/reflexio/server/api.py index 5117b2aac..bb3439a9d 100644 --- a/reflexio/server/api.py +++ b/reflexio/server/api.py @@ -474,6 +474,13 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # noqa: ARG001 durable_learning_scheduler = None aggregation_scheduler = None started_caps: list = [] + if capabilities is None: + from reflexio.server.usage_metrics import ( + configure_usage_event_recorder, + exempt_usage_event_recorder, + ) + + configure_usage_event_recorder(exempt_usage_event_recorder) if mounts_data_plane: _log_multi_worker_daemons() log_publish_hardware_capacity() @@ -555,6 +562,12 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # noqa: ARG001 ) stop_search_metering_worker(timeout=5.0) + if capabilities is None: + from reflexio.server.usage_metrics import ( + configure_usage_event_recorder, + ) + + configure_usage_event_recorder(None) for cap in reversed(started_caps): try: await cap.on_shutdown() diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py index beb4f4588..239040262 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -22,6 +22,10 @@ ) +class ReceiptBillingDeliveryError(RuntimeError): + """A durable finalization receipt still has an undelivered billing event.""" + + def record_extraction_tokens( *, org_id: str, @@ -445,21 +449,26 @@ def emit_learnings_generated_records_strict( return from reflexio.server.billing_signals import platform_llm_from_config - config = configurator.get_config() - record_learnings_generated_records_strict( - org_id=org_id, - learning_ids=learning_ids, - platform_llm=platform_llm_from_config(config), - platform_storage=None, - pipeline=pipeline, - user_id=user_id, - request_id=request_id, - source=source, - agent_version=agent_version, - playbook_name=playbook_name, - entity_type=entity_type, - metadata=metadata, - ) + try: + config = configurator.get_config() + record_learnings_generated_records_strict( + org_id=org_id, + learning_ids=learning_ids, + platform_llm=platform_llm_from_config(config), + platform_storage=None, + pipeline=pipeline, + user_id=user_id, + request_id=request_id, + source=source, + agent_version=agent_version, + playbook_name=playbook_name, + entity_type=entity_type, + metadata=metadata, + ) + except Exception as exc: + raise ReceiptBillingDeliveryError( + "receipt-backed learning billing delivery failed" + ) from exc def record_applied_learnings( diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index d81e1d6b3..68d37ae11 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -16,6 +16,7 @@ from reflexio.models.api_schema.service_schemas import Interaction, Request from reflexio.models.config_schema import PlaybookConfig, ProfileExtractorConfig from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.billing_meter import ReceiptBillingDeliveryError from reflexio.server.error_reporting import error_tags from reflexio.server.llm._litellm_types import ModelProvenance from reflexio.server.llm.litellm_client import LiteLLMClient, LiteLLMConfig @@ -92,6 +93,20 @@ def _next_retry_at(attempt_count: int) -> datetime: return datetime.now(UTC) + timedelta(seconds=delay_seconds) +def _finalization_failure_status( + exc: Exception, + *, + next_attempt_count: int, + max_finalization_attempts: int, +) -> AgentRunStatus: + """Keep durable receipt billing obligations retryable without an attempt cap.""" + if isinstance(exc, ReceiptBillingDeliveryError): + return AgentRunStatus.FINALIZATION_FAILED + if next_attempt_count >= max_finalization_attempts: + return AgentRunStatus.FAILED + return AgentRunStatus.FINALIZATION_FAILED + + def _create_llm_client(request_context: RequestContext) -> LiteLLMClient: # The tool loop re-resolves the model per call via ModelRole.EXTRACTION_AGENT # (see run_tool_loop); this client's model name is only a fallback. We resolve @@ -363,10 +378,10 @@ def run_once(self) -> AgentRunRecord | None: run.id, ) next_attempt_count = run.finalization_attempts + 1 - failed_status = ( - AgentRunStatus.FAILED - if next_attempt_count >= pending_config.max_finalization_attempts - else AgentRunStatus.FINALIZATION_FAILED + failed_status = _finalization_failure_status( + exc, + next_attempt_count=next_attempt_count, + max_finalization_attempts=pending_config.max_finalization_attempts, ) return self.storage.update_agent_run_status( run.id, @@ -411,10 +426,10 @@ def _retry_finalization(self, run: AgentRunRecord) -> AgentRunRecord | None: run.id, ) next_attempt_count = run.finalization_attempts + 1 - failed_status = ( - AgentRunStatus.FAILED - if next_attempt_count >= pending_config.max_finalization_attempts - else AgentRunStatus.FINALIZATION_FAILED + failed_status = _finalization_failure_status( + exc, + next_attempt_count=next_attempt_count, + max_finalization_attempts=pending_config.max_finalization_attempts, ) return self.storage.update_agent_run_status( run.id, diff --git a/reflexio/server/usage_metrics.py b/reflexio/server/usage_metrics.py index 1ebbcc02f..0aed99384 100644 --- a/reflexio/server/usage_metrics.py +++ b/reflexio/server/usage_metrics.py @@ -51,9 +51,10 @@ class UsageEvent: class UsageEventDeliveryStatus(Enum): """Recorder-reported durability outcome for one usage event.""" - ACCEPTED = "accepted" APPENDED = "appended" DUPLICATE = "duplicate" + EXEMPT = "exempt" + UNKNOWN = "unknown" FAILED = "failed" REJECTED = "rejected" @@ -81,6 +82,11 @@ def configure_usage_event_recorder(recorder: UsageEventRecorder | None) -> None: _recorder = recorder +def exempt_usage_event_recorder(_event: UsageEvent) -> UsageEventDeliveryStatus: + """Explicitly exempt a deployment from durable usage-event delivery.""" + return UsageEventDeliveryStatus.EXEMPT + + def record_usage_event( *, org_id: str, @@ -179,14 +185,13 @@ def record_usage_event_strict( ) -> UsageEventDeliveryStatus: """Deliver one event and fail unless the recorder accepted it durably. - Recorders predating delivery outcomes return ``None``; that remains an - accepted result so direct test recorders and non-enterprise integrations - keep working. A missing recorder is also accepted because OSS deployments - intentionally have no billing sink. + Only an explicit append, duplicate, or deployment exemption proves the + receipt's billing obligation was handled. Missing and legacy recorders are + unknown delivery, which strict receipt-backed callers must retry. """ recorder = _recorder if recorder is None: - return UsageEventDeliveryStatus.ACCEPTED + raise UsageEventDeliveryError(UsageEventDeliveryStatus.UNKNOWN) delivery_outcome = recorder( UsageEvent( org_id=str(org_id), @@ -218,7 +223,7 @@ def record_usage_event_strict( ) ) status = ( - UsageEventDeliveryStatus.ACCEPTED + UsageEventDeliveryStatus.UNKNOWN if delivery_outcome is None else delivery_outcome ) @@ -226,9 +231,10 @@ def record_usage_event_strict( raise TypeError( f"usage event recorder returned an invalid delivery status: {status!r}" ) - if status in { - UsageEventDeliveryStatus.FAILED, - UsageEventDeliveryStatus.REJECTED, + if status not in { + UsageEventDeliveryStatus.APPENDED, + UsageEventDeliveryStatus.DUPLICATE, + UsageEventDeliveryStatus.EXEMPT, }: raise UsageEventDeliveryError(status) return status diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 412a954f8..d85c252fd 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -58,6 +58,7 @@ UsageEvent, UsageEventDeliveryStatus, configure_usage_event_recorder, + exempt_usage_event_recorder, ) @@ -70,6 +71,15 @@ def storage(): yield SQLiteStorage(org_id="org_1", db_path=f"{temp_dir}/reflexio.db") +@pytest.fixture(autouse=True) +def explicit_test_billing_exemption(): + configure_usage_event_recorder(exempt_usage_event_recorder) + try: + yield + finally: + configure_usage_event_recorder(None) + + @pytest.fixture def request_context(storage): ctx = RequestContext.__new__(RequestContext) @@ -82,7 +92,10 @@ def request_context(storage): profile_extractor_config=ProfileExtractorConfig( extraction_definition_prompt="Extract durable user deployment facts.", ), - pending_tool_call_config=PendingToolCallConfig(enabled=True), + pending_tool_call_config=PendingToolCallConfig( + enabled=True, + max_finalization_attempts=3, + ), ) ctx.configurator.get_agent_context.return_value = "Test agent context" ctx.prompt_manager = MagicMock() @@ -753,6 +766,7 @@ def test_delivery_failure_after_receipt_commit_retries_billing_without_recompute }, committed_output=committed_output, next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + finalization_attempts=2, ) ) worker = ExtractionResumeWorker( @@ -829,6 +843,7 @@ def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: failed = worker.run_once() assert failed is not None assert failed.status == AgentRunStatus.FINALIZATION_FAILED + assert failed.finalization_attempts == 3 receipt_after_failure = storage.get_agent_run_finalization_receipt( run_id=run_id, entity_type=entity_type, diff --git a/tests/server/test_usage_metrics.py b/tests/server/test_usage_metrics.py index dff33a96f..2992ace87 100644 --- a/tests/server/test_usage_metrics.py +++ b/tests/server/test_usage_metrics.py @@ -28,9 +28,42 @@ def test_event_key_defaults_none(): assert UsageEvent(org_id="7", event_name="x", event_category="y").event_key is None -def test_strict_delivery_accepts_legacy_none_recorder(): +def test_strict_delivery_rejects_legacy_none_recorder_as_unknown(): captured = [] usage_metrics.configure_usage_event_recorder(captured.append) + try: + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + assert [event.event_key for event in captured] == ["learn:profile:1"] + + +def test_strict_delivery_rejects_missing_recorder_as_unknown(): + usage_metrics.configure_usage_event_recorder(None) + + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + + +def test_strict_delivery_accepts_explicit_deployment_exemption(): + usage_metrics.configure_usage_event_recorder( + usage_metrics.exempt_usage_event_recorder + ) try: outcome = usage_metrics.record_usage_event_strict( org_id="7", @@ -41,8 +74,7 @@ def test_strict_delivery_accepts_legacy_none_recorder(): finally: usage_metrics.configure_usage_event_recorder(None) - assert outcome is UsageEventDeliveryStatus.ACCEPTED - assert [event.event_key for event in captured] == ["learn:profile:1"] + assert outcome is UsageEventDeliveryStatus.EXEMPT @pytest.mark.parametrize( From e1d7ca0595d8d0b0b1504baf4b6956dcb8349266 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 05:46:47 +0000 Subject: [PATCH 07/63] fix: pin receipt billing event time --- reflexio/server/billing_meter.py | 6 ++++++ reflexio/server/services/extraction/resume_worker.py | 8 ++++++++ reflexio/server/usage_metrics.py | 4 ++++ .../server/services/extraction/test_resume_worker.py | 12 ++++++++++++ 4 files changed, 30 insertions(+) diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py index 239040262..7b0c8eeb2 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -230,6 +230,7 @@ def record_learnings_generated_records_strict( playbook_name: str | None = None, entity_type: str | None = None, metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, ) -> None: """Strict per-record emission for receipt-backed finalization only.""" _record_learnings_generated_records( @@ -247,6 +248,7 @@ def record_learnings_generated_records_strict( playbook_name=playbook_name, entity_type=entity_type, metadata=metadata, + created_at=created_at, ) @@ -266,6 +268,7 @@ def _record_learnings_generated_records( playbook_name: str | None, entity_type: str | None, metadata: Mapping[str, Any] | None, + created_at: float | None = None, ) -> None: key_entity_type = entity_type or "_" for learning_id in learning_ids: @@ -288,6 +291,7 @@ def _record_learnings_generated_records( platform_storage=platform_storage, caller_type=_INTERNAL, metadata=metadata, + created_at=created_at, ) @@ -443,6 +447,7 @@ def emit_learnings_generated_records_strict( playbook_name: str | None = None, entity_type: str | None = None, metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, ) -> None: """Strict receipt-backed counterpart to the ordinary fail-open emitter.""" if not learning_ids: @@ -464,6 +469,7 @@ def emit_learnings_generated_records_strict( playbook_name=playbook_name, entity_type=entity_type, metadata=metadata, + created_at=created_at, ) except Exception as exc: raise ReceiptBillingDeliveryError( diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index 68d37ae11..25e973313 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -980,6 +980,13 @@ def _record_finalized_learnings( if not learning_ids: return + stored_run = self.storage.get_agent_run(run.id) + durable_run = stored_run or run + billing_timestamp = durable_run.agent_completed_at or durable_run.created_at + if billing_timestamp is None: + raise ReceiptBillingDeliveryError( + "receipt-backed learning billing timestamp is not durable" + ) metadata = { "run_id": run.id, "extractor_kind": run.binding.extractor_kind, @@ -995,6 +1002,7 @@ def _record_finalized_learnings( agent_version=run.binding.agent_version, entity_type=entity_type, metadata=metadata, + created_at=billing_timestamp.timestamp(), ) def _schedule_finalized_tagging(self, run: AgentRunRecord) -> None: diff --git a/reflexio/server/usage_metrics.py b/reflexio/server/usage_metrics.py index 0aed99384..642da6287 100644 --- a/reflexio/server/usage_metrics.py +++ b/reflexio/server/usage_metrics.py @@ -115,6 +115,7 @@ def record_usage_event( duration_ms: int | None = None, error_kind: str | None = None, metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, ) -> None: """Record one usage event if a recorder is configured. @@ -149,6 +150,7 @@ def record_usage_event( duration_ms=duration_ms, error_kind=error_kind, metadata=metadata, + created_at=created_at, ) except Exception as exc: # noqa: BLE001 logger.warning("Usage metrics recorder failed: %s", exc) @@ -182,6 +184,7 @@ def record_usage_event_strict( duration_ms: int | None = None, error_kind: str | None = None, metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, ) -> UsageEventDeliveryStatus: """Deliver one event and fail unless the recorder accepted it durably. @@ -220,6 +223,7 @@ def record_usage_event_strict( duration_ms=duration_ms, error_kind=error_kind, metadata=metadata or {}, + created_at=time.time() if created_at is None else created_at, ) ) status = ( diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index d85c252fd..5eab72925 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -667,6 +667,7 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( request_context, ): """A mixed batch charges its persisted profile once across finalization retries.""" + billing_created_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) run = AgentRunRecord( id="run_mixed_billing", binding=AgentBinding( @@ -679,6 +680,7 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( ), status=AgentRunStatus.FINALIZATION_FAILED, generation_request_snapshot={"request_id": "request_1"}, + agent_completed_at=billing_created_at, ) learning_ids = ["profile_1"] worker = ExtractionResumeWorker(request_context=request_context) @@ -702,6 +704,10 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( "profile_1", "profile_1", ] + assert [call.kwargs["created_at"] for call in record_event.call_args_list] == [ + billing_created_at.timestamp(), + billing_created_at.timestamp(), + ] @pytest.mark.parametrize( @@ -747,6 +753,7 @@ def test_delivery_failure_after_receipt_commit_retries_billing_without_recompute ] } ) + billing_created_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) storage.create_agent_run( AgentRunRecord( id=run_id, @@ -767,6 +774,7 @@ def test_delivery_failure_after_receipt_commit_retries_billing_without_recompute committed_output=committed_output, next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), finalization_attempts=2, + agent_completed_at=billing_created_at, ) ) worker = ExtractionResumeWorker( @@ -884,6 +892,10 @@ def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: f"learn:{entity_type}:{receipt_after_failure[0]}", ] assert list(accepted) == [f"learn:{entity_type}:{receipt_after_failure[0]}"] + assert [event.created_at for event in attempts] == [ + billing_created_at.timestamp(), + billing_created_at.timestamp(), + ] schedule_tagging.assert_not_called() if optimize is not None and aggregate is not None: optimize.assert_called_once() From 56287b5ad4bbfc8444affdb3aeb47ac429a4ddd6 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 06:27:02 +0000 Subject: [PATCH 08/63] fix: use loaded run time for receipt billing --- .../services/extraction/resume_worker.py | 4 +- .../services/extraction/test_resume_worker.py | 63 +++++++++++++------ 2 files changed, 45 insertions(+), 22 deletions(-) diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index 25e973313..98fd9cd7e 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -980,9 +980,7 @@ def _record_finalized_learnings( if not learning_ids: return - stored_run = self.storage.get_agent_run(run.id) - durable_run = stored_run or run - billing_timestamp = durable_run.agent_completed_at or durable_run.created_at + billing_timestamp = run.created_at or run.agent_completed_at if billing_timestamp is None: raise ReceiptBillingDeliveryError( "receipt-backed learning billing timestamp is not durable" diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 5eab72925..90cccac5c 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -144,6 +144,16 @@ def _finalizing_run( ) +def _persist_and_load_run( + storage: SQLiteStorage, run: AgentRunRecord +) -> AgentRunRecord: + storage.create_agent_run(run) + persisted = storage.get_agent_run(run.id) + assert persisted is not None + assert persisted.created_at is not None + return persisted + + def _gate_initial_receipt_reads( storages: tuple[SQLiteStorage, SQLiteStorage], ) -> None: @@ -667,7 +677,8 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( request_context, ): """A mixed batch charges its persisted profile once across finalization retries.""" - billing_created_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) + run_created_at = datetime(2026, 8, 1, 12, 0, tzinfo=UTC) + agent_completed_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) run = AgentRunRecord( id="run_mixed_billing", binding=AgentBinding( @@ -680,15 +691,23 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( ), status=AgentRunStatus.FINALIZATION_FAILED, generation_request_snapshot={"request_id": "request_1"}, - agent_completed_at=billing_created_at, + agent_completed_at=agent_completed_at, + created_at=run_created_at, ) learning_ids = ["profile_1"] worker = ExtractionResumeWorker(request_context=request_context) - with patch( - "reflexio.server.billing_meter.record_usage_event_strict", - return_value=UsageEventDeliveryStatus.APPENDED, - ) as record_event: + with ( + patch.object( + worker.storage, + "get_agent_run", + side_effect=AssertionError("billing must not re-read the run"), + ), + patch( + "reflexio.server.billing_meter.record_usage_event_strict", + return_value=UsageEventDeliveryStatus.APPENDED, + ) as record_event, + ): worker._record_finalized_learnings(run, learning_ids, entity_type="profile") worker._record_finalized_learnings(run, learning_ids, entity_type="profile") @@ -705,8 +724,8 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( "profile_1", ] assert [call.kwargs["created_at"] for call in record_event.call_args_list] == [ - billing_created_at.timestamp(), - billing_created_at.timestamp(), + run_created_at.timestamp(), + run_created_at.timestamp(), ] @@ -753,7 +772,8 @@ def test_delivery_failure_after_receipt_commit_retries_billing_without_recompute ] } ) - billing_created_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) + run_created_at = datetime(2026, 8, 1, 12, 0, tzinfo=UTC) + agent_completed_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) storage.create_agent_run( AgentRunRecord( id=run_id, @@ -774,9 +794,14 @@ def test_delivery_failure_after_receipt_commit_retries_billing_without_recompute committed_output=committed_output, next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), finalization_attempts=2, - agent_completed_at=billing_created_at, + agent_completed_at=agent_completed_at, + created_at=run_created_at, ) ) + stored_run = storage.get_agent_run(run_id) + assert stored_run is not None + assert stored_run.created_at is not None + durable_run_created_at = stored_run.created_at worker = ExtractionResumeWorker( request_context=request_context, llm_client=MagicMock(), @@ -893,8 +918,8 @@ def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: ] assert list(accepted) == [f"learn:{entity_type}:{receipt_after_failure[0]}"] assert [event.created_at for event in attempts] == [ - billing_created_at.timestamp(), - billing_created_at.timestamp(), + durable_run_created_at.timestamp(), + durable_run_created_at.timestamp(), ] schedule_tagging.assert_not_called() if optimize is not None and aggregate is not None: @@ -1175,7 +1200,7 @@ def test_playbook_scheduler_failure_preserves_billing_and_isolated_retry( extractor_kind="playbook", request_id=f"request_scheduler_failure_{failing_scheduler}", ) - storage.create_agent_run(run) + run = _persist_and_load_run(storage, run) worker = ExtractionResumeWorker( request_context=_finalization_context(storage), llm_client=MagicMock(), @@ -1250,7 +1275,7 @@ def test_empty_profile_receipt_wins_once_without_billing_or_recompute(storage): extractor_kind="profile", request_id="request_empty_profile", ) - storage.create_agent_run(run) + run = _persist_and_load_run(storage, run) context = _finalization_context(storage) worker = ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) recorder = _DeduplicatingUsageRecorder() @@ -1299,7 +1324,7 @@ def test_empty_playbook_receipt_retries_without_redispatch(storage): committed_output={"playbooks": []}, next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), ) - storage.create_agent_run(run) + run = _persist_and_load_run(storage, run) context = _finalization_context(storage) worker = ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) original_finalize = worker._finalize_items @@ -1397,7 +1422,7 @@ def test_identical_profile_ids_use_atomic_receipt_owner(tmp_path): extractor_kind="profile", request_id="request_identical_profile", ) - storage_a.create_agent_run(run) + run = _persist_and_load_run(storage_a, run) contexts = [_finalization_context(item) for item in (storage_a, storage_b)] workers = [ ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) @@ -1459,7 +1484,7 @@ def test_identical_playbook_ids_use_atomic_receipt_owner(tmp_path): extractor_kind="playbook", request_id="request_identical_playbook", ) - storage_a.create_agent_run(run) + run = _persist_and_load_run(storage_a, run) contexts = [_finalization_context(item) for item in (storage_a, storage_b)] workers = [ ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) @@ -1551,7 +1576,7 @@ def test_two_stale_profile_workers_preserve_order_and_bill_only_winner(tmp_path) extractor_kind="profile", request_id="request_profile_race", ) - storage_a.create_agent_run(run) + run = _persist_and_load_run(storage_a, run) _gate_initial_receipt_reads((storage_a, storage_b)) contexts = [_finalization_context(storage) for storage in (storage_a, storage_b)] resume_workers = [ @@ -1646,7 +1671,7 @@ def test_two_stale_playbook_workers_preserve_order_and_dispatch_once(tmp_path): extractor_kind="playbook", request_id="request_playbook_race", ) - storage_a.create_agent_run(run) + run = _persist_and_load_run(storage_a, run) _gate_initial_receipt_reads((storage_a, storage_b)) contexts = [_finalization_context(storage) for storage in (storage_a, storage_b)] resume_workers = [ From 70bc28e78ff09c6d9cc6cac8e92d9482792e618d Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 10:26:02 +0000 Subject: [PATCH 09/63] fix: require claims for governance rebuild calls --- .../server/services/governance/service.py | 8 ++- .../storage_base/governance/_rebuild_hide.py | 11 +++- .../governance/test_governance_local_e2e.py | 55 +++++++++++++++++++ 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/reflexio/server/services/governance/service.py b/reflexio/server/services/governance/service.py index de3004e7a..ca7db7177 100644 --- a/reflexio/server/services/governance/service.py +++ b/reflexio/server/services/governance/service.py @@ -574,7 +574,12 @@ def _rebuilt_agent_playbook_ids_from_targets(self, purge_id: str) -> list[int]: ) ] - def _rebuild_agent_playbooks(self, purge_id: str) -> list[int]: + def _rebuild_agent_playbooks( + self, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, + ) -> list[int]: rebuilt_ids: list[int] = [] for target in self.storage.list_purge_targets( purge_id, @@ -600,6 +605,7 @@ def _rebuild_agent_playbooks(self, purge_id: str) -> list[int]: blocking_issue=rebuild_fields["blocking_issue"], expanded_terms=rebuild_fields["expanded_terms"], tags=rebuild_fields["tags"], + execution_claim=execution_claim, ) rebuilt_ids.append(agent_playbook_id) return rebuilt_ids diff --git a/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py b/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py index 6020245a2..66f905a92 100644 --- a/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py +++ b/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py @@ -2,6 +2,8 @@ from abc import ABC, abstractmethod +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim + class RebuildHideMixin(ABC): """Backend-neutral governance rebuild-hide store contract. @@ -13,7 +15,12 @@ class RebuildHideMixin(ABC): """ @abstractmethod - def hide_governance_agent_playbooks_for_rebuild(self, purge_id: str) -> list[int]: + def hide_governance_agent_playbooks_for_rebuild( + self, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, + ) -> list[int]: raise NotImplementedError @abstractmethod @@ -28,5 +35,7 @@ def apply_governance_agent_playbook_rebuild( blocking_issue: dict[str, object] | None, expanded_terms: str | None, tags: list[str] | None, + *, + execution_claim: PurgeExecutionClaim, ) -> None: raise NotImplementedError diff --git a/tests/server/services/governance/test_governance_local_e2e.py b/tests/server/services/governance/test_governance_local_e2e.py index 576419903..aebbb587d 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -1510,3 +1510,58 @@ def get_sessions(self, *, user_id: str, top_k: int, offset: int): assert storage.calls == [0, 1000] assert [request.request_id for request in requests] == ["req-1"] assert sessions == [{"session_id": "session-a", "request_ids": ["req-1"]}] + + +def test_rebuild_agent_playbooks_forwards_the_active_execution_claim() -> None: + target = SimpleNamespace( + target_name="agent_playbook", + target_ref="17", + status="running", + detail={"remaining_source_windows": []}, + ) + + class _Storage: + def __init__(self) -> None: + self.applied: list[dict[str, object]] = [] + + def list_purge_targets(self, purge_id: str, *, phase: str): + assert purge_id == "purge_claimed_rebuild" + assert phase == "rebuild_without_erased_sources" + return [target] + + def get_user_playbooks_by_ids_any_user(self, ids: list[int]): + assert ids == [] + return [] + + def apply_governance_agent_playbook_rebuild(self, **kwargs: object) -> None: + self.applied.append(kwargs) + + storage = _Storage() + service = GovernanceService(storage=storage, org_id="org", ref_secret="secret") + claim = PurgeExecutionClaim( + purge_id="purge_claimed_rebuild", + owner="worker-a", + fence=1, + expires_at=2_000_000_000, + ) + + rebuilt_ids = service._rebuild_agent_playbooks( + "purge_claimed_rebuild", + execution_claim=claim, + ) + + assert rebuilt_ids == [17] + assert storage.applied == [ + { + "purge_id": "purge_claimed_rebuild", + "agent_playbook_id": 17, + "remaining_source_windows": [], + "content": None, + "trigger": None, + "rationale": None, + "blocking_issue": None, + "expanded_terms": None, + "tags": None, + "execution_claim": claim, + } + ] From 0d3dd5d060813ccfd2c6005a5fdb8b62739070ae Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 10:45:14 +0000 Subject: [PATCH 10/63] fix: fence sqlite governance rebuild mutations --- .../governance/_rebuild_hide.py | 31 ++- .../sqlite_storage/test_governance_storage.py | 261 +++++++++++++++++- 2 files changed, 277 insertions(+), 15 deletions(-) diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py b/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py index eea1397b3..edb821057 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py @@ -35,6 +35,7 @@ from reflexio.models.api_schema.domain import AgentPlaybookSourceWindow from reflexio.models.api_schema.domain.enums import Status from reflexio.server.services.embedding_text import playbook_trigger_embedding_text +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _canonicalize_governance_windows, _parse_governance_window_list, @@ -81,6 +82,9 @@ class RebuildHideMixin: # Provided via MRO by the co-composed PurgeOperationStoreMixin (purge bucket); # reached here by the cross-bucket rebuild-hide method. _record_purge_target_locked: Callable[..., None] + _assert_purge_operation_execution_claim_locked: Callable[ + [str, PurgeExecutionClaim | None], None + ] def _replace_agent_playbook_source_windows_locked( self, agent_playbook_id: int, windows: list[AgentPlaybookSourceWindow] @@ -137,11 +141,19 @@ def _upsert_agent_playbook_search_rows_locked( (agent_playbook_id, json.dumps(embedding)), ) - def hide_governance_agent_playbooks_for_rebuild(self, purge_id: str) -> list[int]: + def hide_governance_agent_playbooks_for_rebuild( + self, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, + ) -> list[int]: purge_id = _validate_governance_purge_id("purge_id", purge_id) with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) target_rows = self.conn.execute( """SELECT target_ref FROM purge_operation_targets @@ -202,6 +214,8 @@ def apply_governance_agent_playbook_rebuild( blocking_issue: dict[str, object] | None, expanded_terms: str | None, tags: list[str] | None, + *, + execution_claim: PurgeExecutionClaim, ) -> None: purge_id = _validate_governance_purge_id("purge_id", purge_id) windows = _parse_governance_window_list( @@ -210,13 +224,18 @@ def apply_governance_agent_playbook_rebuild( canonical_remaining_windows = [window.model_dump() for window in windows] content_value = content or "" trigger_value = trigger or None - embedding_text = playbook_trigger_embedding_text(trigger_value) - embedding = ( - self._deps()._get_embedding(embedding_text) if embedding_text else [] - ) with self._lock: try: - self.conn.execute("BEGIN") + self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) + embedding_text = playbook_trigger_embedding_text(trigger_value) + embedding = ( + self._deps()._get_embedding(embedding_text) + if embedding_text + else [] + ) rebuild_target_row = self.conn.execute( """SELECT status, detail FROM purge_operation_targets diff --git a/tests/server/services/storage/sqlite_storage/test_governance_storage.py b/tests/server/services/storage/sqlite_storage/test_governance_storage.py index 2d2e45f77..f64d5425a 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -30,6 +30,7 @@ from reflexio.models.api_schema.retriever_schema import SearchAgentPlaybookRequest from reflexio.models.config_schema import GovernanceRetentionConfig from reflexio.server.services.governance.config import governance_subject_ref +from reflexio.server.services.governance.service import GovernanceService from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, @@ -50,6 +51,9 @@ from reflexio.server.services.storage.storage_base.governance._purge import ( PurgeOperationStoreMixin, ) +from reflexio.server.services.storage.storage_base.governance._rebuild_hide import ( + RebuildHideMixin as RebuildHideContractMixin, +) from reflexio.server.services.storage.storage_base.governance._subject_barrier import ( SubjectBarrierMixin, ) @@ -72,6 +76,8 @@ "complete_subject_erasure_barrier_after_empty_check", "fail_subject_erasure_barrier", "apply_governance_user_data_delete", + "hide_governance_agent_playbooks_for_rebuild", + "apply_governance_agent_playbook_rebuild", "complete_purge_operation_with_audit", } @@ -1647,6 +1653,10 @@ def test_claimed_erasure_mutation_signatures_and_callers_require_claim() -> None "apply_governance_user_data_delete", "complete_purge_operation_with_audit", }, + RebuildHideContractMixin: { + "hide_governance_agent_playbooks_for_rebuild", + "apply_governance_agent_playbook_rebuild", + }, SQLiteStorage: CLAIMED_ERASURE_MUTATIONS, } for owner, method_names in method_owners.items(): @@ -2227,6 +2237,195 @@ def test_prepare_governance_erase_targets_records_full_delete_matrix_counts(stor } +def test_sqlite_rebuild_mutations_reject_missing_and_stale_claims_without_state_changes( + storage, +): + purge_id = "purge_rebuild_claim_matrix" + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key="idem_purge_rebuild_claim_matrix", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + agent_playbook_id = _seed_agent_playbook( + storage, + status=None, + source_windows=[ + AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), + AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), + ], + ) + _record_agent_playbook_rebuild_target( + storage, + purge_id=purge_id, + agent_playbook_id=agent_playbook_id, + previous_lifecycle_status=None, + ) + + def state_snapshot() -> tuple[object, ...]: + playbook = storage.get_agent_playbook_by_id( + agent_playbook_id, + include_tombstones=True, + ) + assert playbook is not None + return ( + playbook.model_dump( + mode="json", + include={ + "content", + "trigger", + "rationale", + "blocking_issue", + "expanded_terms", + "tags", + "status", + }, + ), + [ + window.model_dump(mode="json") + for window in storage.get_source_windows_for_agent_playbook( + agent_playbook_id + ) + ], + [ + target.model_dump(mode="json") + for target in storage.list_purge_targets(purge_id) + ], + dict( + storage.conn.execute( + "SELECT * FROM purge_operations WHERE org_id = ? AND purge_id = ?", + (storage.org_id, purge_id), + ).fetchone() + ), + ) + + before_missing_hide = state_snapshot() + _assert_rejects_missing_claim( + lambda: storage.hide_governance_agent_playbooks_for_rebuild(purge_id), + lambda: storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert state_snapshot() == before_missing_hide + + stale_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + before_stale_hide = state_snapshot() + with pytest.raises(ValueError, match="purge execution claim"): + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=stale_claim, + ) + assert state_snapshot() == before_stale_hide + + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=takeover_claim, + ) + apply_kwargs = { + "purge_id": purge_id, + "agent_playbook_id": agent_playbook_id, + "remaining_source_windows": [ + {"user_playbook_id": 9, "source_interaction_ids": [201]} + ], + "content": "rebuilt claimed content", + "trigger": "rebuilt claimed trigger", + "rationale": "rebuilt claimed rationale", + "blocking_issue": None, + "expanded_terms": "rebuilt claimed terms", + "tags": ["rebuilt-claimed"], + } + before_missing_apply = state_snapshot() + _assert_rejects_missing_claim( + lambda: storage.apply_governance_agent_playbook_rebuild(**apply_kwargs), + lambda: storage.apply_governance_agent_playbook_rebuild( + **apply_kwargs, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert state_snapshot() == before_missing_apply + with pytest.raises(ValueError, match="purge execution claim"): + storage.apply_governance_agent_playbook_rebuild( + **apply_kwargs, + execution_claim=stale_claim, + ) + assert state_snapshot() == before_missing_apply + + storage.apply_governance_agent_playbook_rebuild( + **apply_kwargs, + execution_claim=takeover_claim, + ) + rebuilt = storage.get_agent_playbook_by_id(agent_playbook_id) + assert rebuilt is not None + assert rebuilt.content == "rebuilt claimed content" + assert rebuilt.status is None + assert storage.get_source_windows_for_agent_playbook(agent_playbook_id) == [ + AgentPlaybookSourceWindow( + user_playbook_id=9, + source_interaction_ids=[201], + ) + ] + rebuild_target = storage.list_purge_targets( + purge_id, + phase="rebuild_without_erased_sources", + ) + assert len(rebuild_target) == 1 + assert rebuild_target[0].status == "complete" + + +def test_governance_service_rebuilds_through_claimed_sqlite_contract(storage): + purge_id = "purge_service_claimed_rebuild" + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key="idem_purge_service_claimed_rebuild", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + agent_playbook_id = _seed_agent_playbook( + storage, + status=None, + source_windows=[ + AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), + AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), + ], + ) + _record_agent_playbook_rebuild_target( + storage, + purge_id=purge_id, + agent_playbook_id=agent_playbook_id, + previous_lifecycle_status=None, + ) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + + rebuilt_ids = service._rebuild_agent_playbooks( + purge_id, + execution_claim=claim, + ) + + assert rebuilt_ids == [agent_playbook_id] + rebuilt = storage.get_agent_playbook_by_id(agent_playbook_id) + assert rebuilt is not None + assert rebuilt.content == "source-playbook-9" + assert rebuilt.status is None + + def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_and_hide_marker( storage, ): @@ -2266,7 +2465,10 @@ def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_an ], } - hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=_claim_purge(storage, purge_id), + ) assert hidden_ids == [agent_playbook_id] status = storage.conn.execute( @@ -2362,6 +2564,7 @@ def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) rebuild_target = next( @@ -2438,6 +2641,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_ad_hoc_rebuild_without_ blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=_claim_purge(storage, purge_id), ) assert ( @@ -2521,6 +2725,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_rebuild_before_hide_pha blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) assert ( @@ -2584,7 +2789,11 @@ def test_apply_governance_agent_playbook_rebuild_succeeds_after_prepare_and_hide agent_playbook_id=agent_playbook_id, previous_lifecycle_status=None, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, @@ -2598,6 +2807,7 @@ def test_apply_governance_agent_playbook_rebuild_succeeds_after_prepare_and_hide blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) rebuild_target = next( @@ -2708,6 +2918,7 @@ def fail_search_refresh(*args, **kwargs): blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) assert ( @@ -2796,6 +3007,7 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when blocking_issue=None, expanded_terms=None, tags=None, + execution_claim=claim, ) rebuild_target = next( @@ -2907,6 +3119,7 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) rebuilt_status = storage.conn.execute( @@ -2943,7 +3156,11 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl agent_playbook_id=agent_playbook_id, previous_lifecycle_status=Status.ARCHIVED.value, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, agent_playbook_id=agent_playbook_id, @@ -2956,6 +3173,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) before_row = storage.conn.execute( @@ -2994,6 +3212,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl blocking_issue={"issue": "should not persist"}, expanded_terms="mutated terms", tags=["mutated"], + execution_claim=claim, ) after_row = storage.conn.execute( @@ -3056,7 +3275,11 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet agent_playbook_id=agent_playbook_id, previous_lifecycle_status=Status.ARCHIVED.value, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, agent_playbook_id=agent_playbook_id, @@ -3069,6 +3292,7 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) before_status = storage.conn.execute( @@ -3091,7 +3315,10 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet and target.target_ref == str(agent_playbook_id) ) - hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) after_status = storage.conn.execute( "SELECT status FROM agent_playbooks WHERE agent_playbook_id = ?", @@ -3149,7 +3376,11 @@ def test_hide_governance_agent_playbooks_for_rebuild_does_not_reopen_complete_ta agent_playbook_id=agent_playbook_id, previous_lifecycle_status=Status.ARCHIVED.value, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, agent_playbook_id=agent_playbook_id, @@ -3162,6 +3393,7 @@ def test_hide_governance_agent_playbooks_for_rebuild_does_not_reopen_complete_ta blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) before_status = storage.conn.execute( @@ -3194,7 +3426,10 @@ def stale_list_purge_targets(*_args, **_kwargs): monkeypatch.setattr(storage, "list_purge_targets", stale_list_purge_targets) - hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) assert hidden_ids == [] assert ( @@ -4666,7 +4901,10 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( owned_user_playbook_ids=owned_user_playbook_ids, execution_claim=_claim_purge(storage, purge_id), ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=_claim_purge(storage, purge_id), + ) counts = storage.apply_governance_user_data_delete( purge_id=purge_id, @@ -4780,7 +5018,10 @@ def test_apply_governance_user_data_delete_is_failure_atomic(storage, monkeypatc owned_user_playbook_ids=owned_user_playbook_ids, execution_claim=_claim_purge(storage, purge_id), ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=_claim_purge(storage, purge_id), + ) before_counts = _user_scoped_row_counts(storage, user_id=user_id) before_profile_rows = storage.conn.execute( @@ -4885,6 +5126,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_unsafe_purge_id_before_ blocking_issue=None, expanded_terms="updated terms", tags=["updated"], + execution_claim=_typed_test_claim_for_unvalidated_purge_id("request_12345"), ) after_row = storage.conn.execute( @@ -4983,6 +5225,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=_claim_purge(storage, purge_id), ) rebuilt_row = storage.conn.execute( From 46a218a01161ed56871bb72b2a905040362373ae Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 17:41:46 +0000 Subject: [PATCH 11/63] fix: preserve receipt billing delivery outcomes --- reflexio/server/billing_meter.py | 21 ++++++++--- .../services/extraction/resume_worker.py | 8 +++-- .../services/extraction/test_resume_worker.py | 26 ++++++++++++++ tests/server/test_billing_meter.py | 36 ++++++++++++++++++- 4 files changed, 84 insertions(+), 7 deletions(-) diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py index 7b0c8eeb2..9de981eca 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -13,7 +13,12 @@ from collections.abc import Mapping from typing import Any -from reflexio.server.usage_metrics import record_usage_event, record_usage_event_strict +from reflexio.server.usage_metrics import ( + UsageEventDeliveryError, + UsageEventDeliveryStatus, + record_usage_event, + record_usage_event_strict, +) logger = logging.getLogger(__name__) @@ -25,6 +30,14 @@ class ReceiptBillingDeliveryError(RuntimeError): """A durable finalization receipt still has an undelivered billing event.""" + def __init__( + self, + status: UsageEventDeliveryStatus, + message: str = "receipt-backed learning billing delivery failed", + ) -> None: + super().__init__(message) + self.status = status + def record_extraction_tokens( *, @@ -471,10 +484,10 @@ def emit_learnings_generated_records_strict( metadata=metadata, created_at=created_at, ) + except UsageEventDeliveryError as exc: + raise ReceiptBillingDeliveryError(exc.status) from exc except Exception as exc: - raise ReceiptBillingDeliveryError( - "receipt-backed learning billing delivery failed" - ) from exc + raise ReceiptBillingDeliveryError(UsageEventDeliveryStatus.UNKNOWN) from exc def record_applied_learnings( diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index 98fd9cd7e..908d5719e 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -80,6 +80,7 @@ ) from reflexio.server.services.tagging.tagging_scheduler import schedule_tagging from reflexio.server.site_var.site_var_manager import SiteVarManager +from reflexio.server.usage_metrics import UsageEventDeliveryStatus logger = logging.getLogger(__name__) @@ -99,8 +100,10 @@ def _finalization_failure_status( next_attempt_count: int, max_finalization_attempts: int, ) -> AgentRunStatus: - """Keep durable receipt billing obligations retryable without an attempt cap.""" + """Retry transient receipt delivery failures; terminate permanent rejection.""" if isinstance(exc, ReceiptBillingDeliveryError): + if exc.status is UsageEventDeliveryStatus.REJECTED: + return AgentRunStatus.FAILED return AgentRunStatus.FINALIZATION_FAILED if next_attempt_count >= max_finalization_attempts: return AgentRunStatus.FAILED @@ -983,7 +986,8 @@ def _record_finalized_learnings( billing_timestamp = run.created_at or run.agent_completed_at if billing_timestamp is None: raise ReceiptBillingDeliveryError( - "receipt-backed learning billing timestamp is not durable" + UsageEventDeliveryStatus.UNKNOWN, + "receipt-backed learning billing timestamp is not durable", ) metadata = { "run_id": run.id, diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 90cccac5c..b6b33db70 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -24,9 +24,11 @@ StorageConfigSQLite, ) from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.billing_meter import ReceiptBillingDeliveryError from reflexio.server.services.deferred_learning_plan import FinalizationResult from reflexio.server.services.extraction.resume_worker import ( ExtractionResumeWorker, + _finalization_failure_status, _run_playbook_contract_selection, _run_uses_strict_playbook_evidence, ) @@ -929,6 +931,30 @@ def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: configure_usage_event_recorder(None) +@pytest.mark.parametrize( + ("delivery_status", "expected_status"), + [ + (UsageEventDeliveryStatus.FAILED, AgentRunStatus.FINALIZATION_FAILED), + (UsageEventDeliveryStatus.UNKNOWN, AgentRunStatus.FINALIZATION_FAILED), + (UsageEventDeliveryStatus.REJECTED, AgentRunStatus.FAILED), + ], +) +def test_receipt_delivery_failure_status_distinguishes_transient_and_permanent( + delivery_status, + expected_status, +): + error = ReceiptBillingDeliveryError(delivery_status) + + assert ( + _finalization_failure_status( + error, + next_attempt_count=1, + max_finalization_attempts=3, + ) + is expected_status + ) + + def test_retry_after_billing_reuses_ids_without_replaying_playbook_schedulers( request_context, storage, diff --git a/tests/server/test_billing_meter.py b/tests/server/test_billing_meter.py index 740d8d477..416bbbc6e 100644 --- a/tests/server/test_billing_meter.py +++ b/tests/server/test_billing_meter.py @@ -1,10 +1,18 @@ -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +import pytest from reflexio.server.billing_meter import ( + ReceiptBillingDeliveryError, + emit_learnings_generated_records_strict, record_applied_learnings, record_extraction_tokens, record_learnings_generated, ) +from reflexio.server.usage_metrics import ( + UsageEventDeliveryError, + UsageEventDeliveryStatus, +) HOOK = "reflexio.server.billing_meter.record_usage_event" @@ -110,3 +118,29 @@ def test_record_applied_learnings_noop_for_empty_result(): platform_storage=None, ) hook.assert_not_called() + + +@pytest.mark.parametrize( + "status", + [UsageEventDeliveryStatus.FAILED, UsageEventDeliveryStatus.REJECTED], +) +def test_strict_receipt_billing_preserves_delivery_status(status): + configurator = MagicMock() + configurator.get_config.return_value = None + + with ( + patch( + "reflexio.server.billing_meter.record_usage_event_strict", + side_effect=UsageEventDeliveryError(status), + ), + pytest.raises(ReceiptBillingDeliveryError) as exc_info, + ): + emit_learnings_generated_records_strict( + org_id="org1", + configurator=configurator, + learning_ids=["profile-1"], + source="resumable_extraction", + entity_type="profile", + ) + + assert exc_info.value.status is status From ea777b0ff2cc6d6bc9dc1e8b6535fa1a5a21fd26 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 22:31:01 +0000 Subject: [PATCH 12/63] fix: add always-evaluated global scheduler hooks --- .../server/services/lineage/gc_scheduler.py | 42 ++++++++++++++++++- .../lineage/test_gc_scheduler_global_sweep.py | 30 +++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/reflexio/server/services/lineage/gc_scheduler.py b/reflexio/server/services/lineage/gc_scheduler.py index 37d6e9dc4..215ab03ef 100644 --- a/reflexio/server/services/lineage/gc_scheduler.py +++ b/reflexio/server/services/lineage/gc_scheduler.py @@ -126,6 +126,30 @@ def clear_global_sweeps() -> None: _global_sweep_hooks.clear() +# Always-global sweeps run once per elected tick and own their applicability +# checks. Enterprise uses this for global maintenance that must not inherit the +# unrelated expiry-reclamation feature gate. +_always_global_sweep_hooks: list[Callable[[int], int]] = [] + + +def register_always_global_sweep(fn: Callable[[int], int]) -> None: + """Register an always-evaluated global sweep. + + The scheduler's leader gate still controls the tick. The sweep itself must + decide whether its backing service is configured for this deployment. + + Args: + fn (Callable[[int], int]): Called with the current unix epoch; returns + the number of rows it processed. + """ + _always_global_sweep_hooks.append(fn) + + +def clear_always_global_sweeps() -> None: + """Clear all always-evaluated global sweeps (tests restore defaults).""" + _always_global_sweep_hooks.clear() + + # Per-org sweeps run once per org per tick (for per-org reclamation concerns). # Each fn takes (org_id, now) and returns a deleted-row count. Enterprise # registers its closure here at startup so governance retention folds into the @@ -447,6 +471,19 @@ def _run_global_sweeps(self, cfg: object) -> None: capture_anomaly("lineage.global_sweep.failed", sweep=sweep_id) logger.exception("event=global_sweep_failed sweep=%s", sweep_id) + def _run_always_global_sweeps(self) -> None: + """Invoke each applicability-owning global sweep once per elected tick.""" + now = int(time.time()) + for sweep in _always_global_sweep_hooks: + try: + processed = sweep(now) + if processed: + logger.info("event=always_global_sweep processed=%d", processed) + except Exception: + sweep_id = getattr(sweep, "__qualname__", repr(sweep)) + capture_anomaly("lineage.always_global_sweep.failed", sweep=sweep_id) + logger.exception("event=always_global_sweep_failed sweep=%s", sweep_id) + def _run_once(self) -> float: poll_interval = _DEFAULT_POLL_INTERVAL_SECONDS try: @@ -468,6 +505,7 @@ def _run_once(self) -> float: org_ids = self._discover_org_ids(bootstrap_ctx) self._gc_tick(org_ids, max_workers=self._org_fanout_workers(bootstrap_ctx)) self._run_global_sweeps(cfg) + self._run_always_global_sweeps() except Exception: logger.exception("event=lineage_gc_scheduler_tick_failed") return max(poll_interval, _MIN_POLL_SECONDS) @@ -537,7 +575,9 @@ def maybe_start_lineage_gc( # must not silence sweeps that were explicitly registered. Production ran # with 8 registered sweeps (metering force-seal/retention/derive, governance # retention, missing-vector backfill) never firing because this read raised. - has_registered_sweeps = bool(_per_org_sweep_hooks) or bool(_global_sweep_hooks) + has_registered_sweeps = bool( + _per_org_sweep_hooks or _global_sweep_hooks or _always_global_sweep_hooks + ) try: ctx = request_context_factory(bootstrap_org_id) diff --git a/tests/server/services/lineage/test_gc_scheduler_global_sweep.py b/tests/server/services/lineage/test_gc_scheduler_global_sweep.py index 642d9a125..c4466b84e 100644 --- a/tests/server/services/lineage/test_gc_scheduler_global_sweep.py +++ b/tests/server/services/lineage/test_gc_scheduler_global_sweep.py @@ -26,8 +26,13 @@ def _cfg(enabled: bool): @pytest.fixture(autouse=True) def _isolate_hooks(): clear_global_sweeps() + clear_always = getattr(gc_scheduler, "clear_always_global_sweeps", None) + if clear_always is not None: + clear_always() yield clear_global_sweeps() + if clear_always is not None: + clear_always() def test_global_sweep_runs_once_when_enabled(): @@ -97,3 +102,28 @@ def test_run_once_invokes_global_sweeps(monkeypatch): scheduler._run_once() assert len(calls) == 1 + + +def test_run_once_invokes_always_global_sweep_when_expiry_disabled(monkeypatch): + calls: list[int] = [] + register = getattr(gc_scheduler, "register_always_global_sweep", None) + assert callable(register) + register(lambda now: calls.append(now) or 1) + + cfg = types.SimpleNamespace( + lineage_gc=types.SimpleNamespace(poll_interval_seconds=10), + expiry_reclamation=types.SimpleNamespace(enabled=False), + ) + ctx = types.SimpleNamespace( + configurator=types.SimpleNamespace(get_config=lambda: cfg), + ) + scheduler = LineageGCScheduler( + request_context_factory=lambda _org_id: ctx, # type: ignore[arg-type] + bootstrap_org_id="org-boot", + ) + monkeypatch.setattr(scheduler, "_discover_org_ids", lambda _ctx: []) + monkeypatch.setattr(scheduler, "_gc_tick", lambda _org_ids, **_kwargs: None) + + scheduler._run_once() + + assert len(calls) == 1 From dfca8aaf45858a44519ecb6cb009ab69ee164d26 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 4 Aug 2026 23:27:08 +0000 Subject: [PATCH 13/63] fix: bound search exposures and release sqlite locks --- .../models/api_schema/retriever_schema.py | 11 +- reflexio/server/services/search_exposure.py | 13 ++ .../governance/_rebuild_hide.py | 12 +- .../routes/test_search_exposure_boundary.py | 45 ++++- .../sqlite_storage/test_governance_storage.py | 171 ++++++++++++++++++ 5 files changed, 239 insertions(+), 13 deletions(-) diff --git a/reflexio/models/api_schema/retriever_schema.py b/reflexio/models/api_schema/retriever_schema.py index eab732023..f1ed46096 100644 --- a/reflexio/models/api_schema/retriever_schema.py +++ b/reflexio/models/api_schema/retriever_schema.py @@ -794,7 +794,8 @@ class UnifiedSearchRequest(BaseModel): Args: query (str): Search query text - top_k (int, optional): Maximum results per entity type. Defaults to 5 + top_k (int, optional): Maximum results per entity type, up to 100. + Defaults to 5. threshold (float, optional): Similarity threshold for vector search. When omitted, the embedding model's default is used. agent_version (str, optional): Filter by agent version (agent_playbooks, user_playbooks) @@ -812,11 +813,11 @@ class UnifiedSearchRequest(BaseModel): """ query: NonEmptyStr - top_k: int | None = Field(default=5, gt=0) + top_k: int | None = Field(default=5, gt=0, le=100) threshold: float | None = Field(default=None, ge=0.0, le=1.0) agent_version: str | None = None playbook_name: str | None = None - user_id: str | None = None + user_id: str | None = Field(default=None, max_length=255) tags: list[str] | None = None entity_types: list[UnifiedSearchEntityType] | None = None agent_playbook_status_filter: list[PlaybookStatus] | None = None @@ -829,8 +830,8 @@ class UnifiedSearchRequest(BaseModel): # ``session_id`` additionally enables session-scoped result dedup: items # already served to the same (org, session) are skipped and the next-best # matches backfilled (see server/services/retrieval/session_dedup.py). - request_id: str | None = None - session_id: str | None = None + request_id: str | None = Field(default=None, max_length=255) + session_id: str | None = Field(default=None, max_length=255) interaction_id: int | None = Field(default=None, gt=0) diff --git a/reflexio/server/services/search_exposure.py b/reflexio/server/services/search_exposure.py index 5e7c13acd..33270eb4d 100644 --- a/reflexio/server/services/search_exposure.py +++ b/reflexio/server/services/search_exposure.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Sized from dataclasses import dataclass, field from hashlib import sha256 from secrets import token_hex @@ -14,6 +15,8 @@ incumbent_user_playbook_semantic_digest, ) +MAX_EXPOSURE_EVENTS_PER_BATCH = 100 + @dataclass(frozen=True) class SearchExposureBatch: @@ -89,6 +92,14 @@ def record_search_exposures(batch: SearchExposureBatch) -> None: recorder.record(batch) +def validate_exposure_batch_size(events: Sized) -> None: + """Reject exposure batches that exceed the fixed storage safety bound.""" + if len(events) > MAX_EXPOSURE_EVENTS_PER_BATCH: + raise ValueError( + f"exposure batch must contain at most {MAX_EXPOSURE_EVENTS_PER_BATCH} events" + ) + + def user_playbook_full_version_fingerprint(playbook: UserPlaybook) -> str: """Bind every persisted playbook field except its derived embedding vector. @@ -161,10 +172,12 @@ def build_user_playbook_exposure_event( __all__ = [ "SEARCH_EXPOSURE_RECORDER", "ExposureEventWriteResult", + "MAX_EXPOSURE_EVENTS_PER_BATCH", "SearchExposureBatch", "SearchExposureRecorder", "UserPlaybookExposureEvent", "build_user_playbook_exposure_event", "record_search_exposures", "user_playbook_full_version_fingerprint", + "validate_exposure_batch_size", ] diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py b/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py index edb821057..ead650ff6 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py @@ -224,18 +224,18 @@ def apply_governance_agent_playbook_rebuild( canonical_remaining_windows = [window.model_dump() for window in windows] content_value = content or "" trigger_value = trigger or None + embedding: list[float] = [] + if windows: + embedding_text = playbook_trigger_embedding_text(trigger_value) + embedding = ( + self._deps()._get_embedding(embedding_text) if embedding_text else [] + ) with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") self._assert_purge_operation_execution_claim_locked( purge_id, execution_claim ) - embedding_text = playbook_trigger_embedding_text(trigger_value) - embedding = ( - self._deps()._get_embedding(embedding_text) - if embedding_text - else [] - ) rebuild_target_row = self.conn.execute( """SELECT status, detail FROM purge_operation_targets diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py index 94cef8c68..7497e362e 100644 --- a/tests/server/routes/test_search_exposure_boundary.py +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -8,6 +8,7 @@ from typing import Any from unittest.mock import MagicMock, patch +import pytest from fastapi.testclient import TestClient from reflexio.models.api_schema.domain import UserPlaybook @@ -32,7 +33,7 @@ def _playbook(playbook_id: int, content: str) -> UserPlaybook: @contextmanager -def _search_results(playbooks: list[UserPlaybook]) -> Iterator[None]: +def _search_results(playbooks: list[UserPlaybook]) -> Iterator[MagicMock]: reflexio = MagicMock() reflexio.request_context.configurator.get_config.return_value = Config( storage_config=StorageConfigSQLite() @@ -52,7 +53,7 @@ def _search_results(playbooks: list[UserPlaybook]) -> Iterator[None]: "reflexio.server.routes.search.reflexio_cache.get_reflexio", return_value=reflexio, ): - yield + yield reflexio def _client() -> TestClient: @@ -151,3 +152,43 @@ def test_oss_search_succeeds_when_no_recorder_is_registered() -> None: assert [item["user_playbook_id"] for item in response.json()["user_playbooks"]] == [ 11 ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("top_k", 101), + ("request_id", "r" * 256), + ("session_id", "s" * 256), + ("user_id", "u" * 256), + ], +) +def test_unified_search_rejects_oversized_work_before_search_execution( + field: str, + value: object, +) -> None: + with _search_results([]) as reflexio: + response = _client().post( + "/api/search", + json={"query": "answer", field: value}, + ) + + assert response.status_code == 422 + reflexio.unified_search.assert_not_called() + + +def test_unified_search_accepts_exact_workload_and_identifier_limits() -> None: + with _search_results([]) as reflexio: + response = _client().post( + "/api/search", + json={ + "query": "answer", + "top_k": 100, + "request_id": "r" * 255, + "session_id": "s" * 255, + "user_id": "u" * 255, + }, + ) + + assert response.status_code == 200, response.text + reflexio.unified_search.assert_called_once() diff --git a/tests/server/services/storage/sqlite_storage/test_governance_storage.py b/tests/server/services/storage/sqlite_storage/test_governance_storage.py index f64d5425a..33419069a 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -2377,6 +2377,177 @@ def state_snapshot() -> tuple[object, ...]: assert rebuild_target[0].status == "complete" +def _prepare_blocking_embedding_rebuild( + storage: SQLiteStorage, + *, + suffix: str, +) -> tuple[int, PurgeExecutionClaim, dict[str, object]]: + purge_id = f"purge_blocking_embedding_{suffix}" + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + agent_playbook_id = _seed_agent_playbook( + storage, + status=None, + source_windows=[ + AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), + AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), + ], + ) + _record_agent_playbook_rebuild_target( + storage, + purge_id=purge_id, + agent_playbook_id=agent_playbook_id, + previous_lifecycle_status=None, + ) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) + return ( + agent_playbook_id, + claim, + { + "purge_id": purge_id, + "agent_playbook_id": agent_playbook_id, + "remaining_source_windows": [ + {"user_playbook_id": 9, "source_interaction_ids": [201]} + ], + "content": "rebuilt after embedding", + "trigger": "embedding trigger", + "rationale": "embedding rationale", + "blocking_issue": None, + "expanded_terms": "embedding terms", + "tags": ["embedding"], + "execution_claim": claim, + }, + ) + + +def test_sqlite_rebuild_embedding_does_not_hold_shared_storage_lock(storage) -> None: + agent_playbook_id, _, apply_kwargs = _prepare_blocking_embedding_rebuild( + storage, + suffix="ordinary_read", + ) + embedding_started = threading.Event() + release_embedding = threading.Event() + rebuild_errors: list[BaseException] = [] + read_finished = threading.Event() + + def blocking_embedding(_text: str, purpose: str = "document") -> list[float]: + assert purpose == "document" + embedding_started.set() + assert release_embedding.wait(timeout=5) + return [0.0] * 512 + + def rebuild() -> None: + try: + storage.apply_governance_agent_playbook_rebuild(**apply_kwargs) + except BaseException as exc: # pragma: no cover - asserted below + rebuild_errors.append(exc) + + def ordinary_read() -> None: + storage.get_agent_playbook_by_id(agent_playbook_id, include_tombstones=True) + read_finished.set() + + with patch.object(storage, "_get_embedding", side_effect=blocking_embedding): + rebuild_thread = threading.Thread(target=rebuild) + rebuild_thread.start() + read_thread = threading.Thread(target=ordinary_read) + try: + assert embedding_started.wait(timeout=5) + read_thread.start() + read_completed_while_embedding_blocked = read_finished.wait(timeout=0.5) + finally: + release_embedding.set() + rebuild_thread.join(timeout=5) + if read_thread.ident is not None: + read_thread.join(timeout=5) + + assert read_completed_while_embedding_blocked is True + assert rebuild_errors == [] + assert not rebuild_thread.is_alive() + assert not read_thread.is_alive() + + +def test_sqlite_rebuild_revalidates_claim_after_embedding_takeover( + storage_factory, +) -> None: + storage = storage_factory("org1") + peer = storage_factory("org1") + agent_playbook_id, stale_claim, apply_kwargs = _prepare_blocking_embedding_rebuild( + storage, + suffix="claim_takeover", + ) + before = storage.get_agent_playbook_by_id( + agent_playbook_id, + include_tombstones=True, + ) + assert before is not None + embedding_started = threading.Event() + release_embedding = threading.Event() + rebuild_errors: list[BaseException] = [] + + def blocking_embedding(_text: str, purpose: str = "document") -> list[float]: + assert purpose == "document" + embedding_started.set() + assert release_embedding.wait(timeout=5) + return [0.0] * 512 + + def rebuild() -> None: + try: + storage.apply_governance_agent_playbook_rebuild(**apply_kwargs) + except BaseException as exc: + rebuild_errors.append(exc) + + takeover_claim: PurgeExecutionClaim | None = None + with patch.object(storage, "_get_embedding", side_effect=blocking_embedding): + rebuild_thread = threading.Thread(target=rebuild) + rebuild_thread.start() + try: + assert embedding_started.wait(timeout=5) + peer.conn.execute("PRAGMA busy_timeout = 500") + peer.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (peer.org_id, apply_kwargs["purge_id"]), + ) + peer.conn.commit() + takeover_claim = peer.claim_purge_operation_execution( + str(apply_kwargs["purge_id"]), + lease_owner="takeover-owner", + lease_ttl_seconds=30, + ) + finally: + release_embedding.set() + rebuild_thread.join(timeout=5) + + assert takeover_claim is not None + assert takeover_claim.fence == stale_claim.fence + 1 + assert len(rebuild_errors) == 1 + assert isinstance(rebuild_errors[0], ValueError) + assert "purge execution claim" in str(rebuild_errors[0]) + after = storage.get_agent_playbook_by_id( + agent_playbook_id, + include_tombstones=True, + ) + assert after == before + [target] = storage.list_purge_targets( + str(apply_kwargs["purge_id"]), + phase="rebuild_without_erased_sources", + ) + assert target.status == "running" + + def test_governance_service_rebuilds_through_claimed_sqlite_contract(storage): purge_id = "purge_service_claimed_rebuild" _begin_test_purge_operation( From 50480d8e36e55a1fddad1437a891711d4a82e6d6 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 06:27:14 +0000 Subject: [PATCH 14/63] fix: resolve shared rebuild integration Preserve current-main storage and billing contracts while restoring the Phase 1 replay semantics. --- .../playbook/components/aggregator.py | 11 --- reflexio/server/services/playbook/service.py | 2 + .../services/storage/sqlite_storage/_base.py | 84 ++++++++++++++----- .../storage/sqlite_storage/_governance.py | 6 +- .../sqlite_storage/_session_outcomes.py | 6 +- .../governance/_erase_execution.py | 6 +- .../services/extraction/test_resume_worker.py | 12 ++- .../test_subject_write_barrier_sqlite.py | 3 +- .../sqlite_storage/test_governance_storage.py | 30 +++++-- ...est_session_outcome_downgrade_migration.py | 62 +++++++++----- ..._playbook_aggregation_state_integration.py | 16 +--- .../test_non_extraction_learning_metering.py | 44 +++++----- 12 files changed, 168 insertions(+), 114 deletions(-) diff --git a/reflexio/server/services/playbook/components/aggregator.py b/reflexio/server/services/playbook/components/aggregator.py index 11bd366ce..0b53055cd 100644 --- a/reflexio/server/services/playbook/components/aggregator.py +++ b/reflexio/server/services/playbook/components/aggregator.py @@ -1079,17 +1079,6 @@ def save_generated_outcome( duration_ms=int((time.perf_counter() - aggregation_start) * 1000), metadata=stats, ) - self._record_learnings_generated( - learning_ids=[ - str(saved.agent_playbook_id) - for saved in saved_playbooks - if getattr(saved, "agent_playbook_id", None) - ], - playbook_name=SINGLETON_USER_PLAYBOOK_NAME, - request_id=run_id, - metadata=stats, - total_count=len(saved_playbooks), - ) return stats def _stable_aggregation_cluster_id(self, fingerprint: str) -> str: diff --git a/reflexio/server/services/playbook/service.py b/reflexio/server/services/playbook/service.py index dd18358d2..4f79508fb 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -687,6 +687,7 @@ def _finalize_extracted_items( return self._finalize_extracted_items_with_outcome( all_playbooks, model_provenance=model_provenance, + extraction_run=extraction_run, finalization_run_id=finalization_run_id, ).learning_ids @@ -695,6 +696,7 @@ def _finalize_extracted_items_with_outcome( all_playbooks: list[UserPlaybook], *, model_provenance: ModelProvenance | None = None, + extraction_run: AgentRunRecord | None = None, finalization_run_id: str | None = None, ) -> FinalizationResult: """Finalize playbooks and expose the atomic receipt winner internally.""" diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index e6d4ab9f0..25d7da93f 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -1193,24 +1193,46 @@ def _migrate_session_outcomes_schema(self) -> None: with self._lock: self.conn.execute("BEGIN IMMEDIATE") try: - columns = { - row["name"] - for row in self.conn.execute( - "PRAGMA table_info(session_outcomes)" - ).fetchall() - } + table_info = self.conn.execute( + "PRAGMA table_info(session_outcomes)" + ).fetchall() + columns = {str(row["name"]): row for row in table_info} table = self.conn.execute( "SELECT sql FROM sqlite_master WHERE type = 'table' " "AND name = 'session_outcomes'" ).fetchone() table_sql = (table["sql"] if table is not None else "") or "" - required = { + identity_columns = { "outcome_id", "outcome_revision", "outcome_contract_digest", "finalized_trajectory_digest", } - if required.issubset(columns) and "'unknown'" in table_sql: + expected_columns = identity_columns | { + "user_id", + "session_id", + "outcome", + "occurred_at", + "source", + "label", + "value", + "metadata", + "governance_subject_ref", + "created_at", + } + subject_column = columns.get("governance_subject_ref") + has_empty_subject_default = ( + subject_column is not None + and subject_column["dflt_value"] in ("''", '""') + ) + schema_is_canonical = ( + expected_columns.issubset(columns) + and "'unknown'" in table_sql + and subject_column is not None + and int(subject_column["notnull"]) == 1 + and not has_empty_subject_default + ) + if schema_is_canonical: rows_missing_subject_ref = self.conn.execute( """SELECT user_id, session_id FROM session_outcomes WHERE governance_subject_ref IS NULL @@ -1258,9 +1280,26 @@ def _migrate_session_outcomes_schema(self) -> None: ) for row in legacy_rows: source = str(row["source"]) - subject_ref = row["governance_subject_ref"] + subject_ref = ( + row["governance_subject_ref"] + if "governance_subject_ref" in columns + else None + ) if subject_ref is None or not str(subject_ref).strip(): subject_ref = self._subject_ref_for_user_id(str(row["user_id"])) + existing_outcome_id = ( + row["outcome_id"] if "outcome_id" in columns else None + ) + existing_contract_digest = ( + row["outcome_contract_digest"] + if "outcome_contract_digest" in columns + else None + ) + existing_trajectory_digest = ( + row["finalized_trajectory_digest"] + if "finalized_trajectory_digest" in columns + else None + ) self.conn.execute( """INSERT INTO session_outcomes ( outcome_id, outcome_revision, user_id, session_id, outcome, @@ -1269,25 +1308,32 @@ def _migrate_session_outcomes_schema(self) -> None: governance_subject_ref, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( - _legacy_session_outcome_id( + existing_outcome_id + or _legacy_session_outcome_id( str(row["user_id"]), str(row["session_id"]) ), - 1, + ( + row["outcome_revision"] + if "outcome_revision" in columns + else 1 + ), row["user_id"], row["session_id"], row["outcome"], row["occurred_at"], source, - row["label"], - row["value"], - row["metadata"], - outcome_contract_digest( + row["label"] if "label" in columns else None, + row["value"] if "value" in columns else None, + row["metadata"] if "metadata" in columns else None, + existing_contract_digest + or outcome_contract_digest( source=source, schema_version=1, allowed_values=_OUTCOME_ALLOWED_VALUES, finalization_rule="first_write", ), - trajectory_digest( + existing_trajectory_digest + or trajectory_digest( _canonical_session_snapshot( self.conn, str(row["session_id"]) ) @@ -2710,7 +2756,6 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: "SELECT rowid, profile_id FROM profiles WHERE user_id = ?", (user_id,), ).fetchall() - subject_ref = self._subject_ref_for_user_id(user_id) # Build a rowid lookup for FTS/vec cleanup (SQLite-specific need). profile_rowid_by_id: dict[str, int] = { @@ -2761,9 +2806,8 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: "DELETE FROM interactions WHERE user_id = ?", (user_id,) ) session_outcomes_cur = self.conn.execute( - """DELETE FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "DELETE FROM session_outcomes WHERE user_id = ?", + (user_id,), ) requests_cur = self.conn.execute( "DELETE FROM requests WHERE user_id = ?", (user_id,) diff --git a/reflexio/server/services/storage/sqlite_storage/_governance.py b/reflexio/server/services/storage/sqlite_storage/_governance.py index 77f8fec91..922691fc7 100644 --- a/reflexio/server/services/storage/sqlite_storage/_governance.py +++ b/reflexio/server/services/storage/sqlite_storage/_governance.py @@ -431,11 +431,9 @@ def _planned_governance_delete_counts( "SELECT COUNT(DISTINCT session_id) AS cnt FROM requests WHERE user_id = ?", (user_id,), ).fetchone() - subject_ref = self._deps()._subject_ref_for_user_id(user_id) session_outcome_row = self.conn.execute( - """SELECT COUNT(*) AS cnt FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "SELECT COUNT(*) AS cnt FROM session_outcomes WHERE user_id = ?", + (user_id,), ).fetchone() profile_rows = self.conn.execute( "SELECT profile_id FROM profiles WHERE user_id = ?", diff --git a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py index e41fa2fa0..d0657f4eb 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -339,12 +339,10 @@ def get_session_outcomes( @SQLiteStorageBase.handle_exceptions def clear_session_outcomes_for_user(self, user_id: str) -> dict[str, int]: - subject_ref = self._subject_ref_for_user_id(user_id) with self._lock: outcome_cursor = self.conn.execute( - """DELETE FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "DELETE FROM session_outcomes WHERE user_id = ?", + (user_id,), ) self.conn.commit() return { diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py index 6d1cdf726..4333632a3 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py @@ -149,11 +149,9 @@ def _clear_user_data_for_governance_locked( expected_user_playbook_ids: set[int] | None = None, ) -> dict[str, int]: deps = self._deps() - subject_ref = deps._subject_ref_for_user_id(user_id) session_outcomes_cur = self.conn.execute( - """DELETE FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "DELETE FROM session_outcomes WHERE user_id = ?", + (user_id,), ) interaction_ids = [ int(row["interaction_id"]) diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index b6b33db70..519bdaee6 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -606,13 +606,13 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( ( "profile", "reflexio.server.services.profile.service." - "ProfileGenerationService._finalize_extracted_items", + "ProfileGenerationService._finalize_extracted_items_with_outcome", "profile", ), ( "playbook", "reflexio.server.services.playbook.service." - "PlaybookGenerationService._finalize_extracted_items", + "PlaybookGenerationService._finalize_extracted_items_with_outcome", "user_playbook", ), ], @@ -635,15 +635,19 @@ def test_resume_bills_only_items_that_survive_finalization( ) dropped = object() survivor = object() + survivor_id = "durable-survivor-id" worker = ExtractionResumeWorker(request_context=request_context) with ( - patch(finalize_path, return_value=[survivor]) as finalize, + patch( + finalize_path, + return_value=FinalizationResult([survivor_id], won_receipt=True), + ) 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) + record.assert_called_once_with(run, [survivor_id], entity_type=entity_type) if extractor_kind == "playbook": assert finalize.call_args.kwargs["extraction_run"] is run else: diff --git a/tests/server/services/governance/test_subject_write_barrier_sqlite.py b/tests/server/services/governance/test_subject_write_barrier_sqlite.py index a0598f144..c82547119 100644 --- a/tests/server/services/governance/test_subject_write_barrier_sqlite.py +++ b/tests/server/services/governance/test_subject_write_barrier_sqlite.py @@ -490,10 +490,11 @@ def test_barrier_blocks_deferred_evaluation_tag_write( idempotency_key="idem_deferred_tag_write", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_11111111111111111111111111111112", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.update_agent_success_evaluation_result_tags( diff --git a/tests/server/services/storage/sqlite_storage/test_governance_storage.py b/tests/server/services/storage/sqlite_storage/test_governance_storage.py index 33419069a..f4488146c 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -7,8 +7,8 @@ import sqlite3 import threading import time -from concurrent.futures import ThreadPoolExecutor from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Literal, cast from unittest.mock import patch @@ -1036,6 +1036,7 @@ def test_complete_purge_operation_with_audit_begins_immediate_transaction_before def test_apply_governance_delete_begins_immediate_transaction_before_reads(storage): purge_id = _begin_purge(storage, "purge_delete_begin_immediate") + claim = _claim_purge(storage, purge_id) for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -1043,13 +1044,18 @@ def test_apply_governance_delete_begins_immediate_transaction_before_reads(stora target_ref="all", phase="delete", status="pending", + execution_claim=claim, detail={"count": 0}, ) statements: list[str] = [] storage.conn.set_trace_callback(statements.append) try: with pytest.raises(ValueError, match="prepared purge snapshot"): - storage.apply_governance_user_data_delete(purge_id, "empty-user") + storage.apply_governance_user_data_delete( + purge_id, + "alice", + execution_claim=claim, + ) finally: storage.conn.set_trace_callback(None) @@ -4161,6 +4167,7 @@ def test_fail_missing_purge_rolls_back_implicit_transaction(storage): "purge_missing", error_code="PURGE_TARGET_FAILED", error_detail="target_delete_failed", + execution_claim=_typed_test_claim_for_unvalidated_purge_id("purge_missing"), ) assert storage.conn.in_transaction is False @@ -4169,6 +4176,7 @@ def test_fail_missing_purge_rolls_back_implicit_transaction(storage): def test_record_purge_target_rolls_back_after_write_failure(storage, monkeypatch): purge_id = _begin_purge(storage, "purge_target_write_failure") + claim = _claim_purge(storage, purge_id) def _write_then_raise(**_kwargs: object) -> None: storage.conn.execute( @@ -4186,6 +4194,7 @@ def _write_then_raise(**_kwargs: object) -> None: target_ref="all", phase="delete", status="running", + execution_claim=claim, ) assert storage.conn.in_transaction is False @@ -4225,6 +4234,7 @@ def _trace(statement: str) -> None: "user", SUBJECT_REF, REQUEST_REF, + authoritative_user_id="alice", ) assert entered.wait(timeout=1) time.sleep(0.05) @@ -4245,9 +4255,11 @@ def test_prepare_targets_rechecks_snapshot_after_two_connection_write_lock( idempotency_key="idem_two_connection_prepare", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, ) + claim = _claim_purge(first, purge_id) first.conn.execute("BEGIN IMMEDIATE") first._record_purge_target_locked( purge_id=purge_id, @@ -4265,7 +4277,14 @@ def test_prepare_targets_rechecks_snapshot_after_two_connection_write_lock( target_ref="all", phase="prepare_targets", status="complete", - detail={"owned_user_playbook_ids": []}, + detail={ + "authoritative_user_digest": first.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (first.org_id, purge_id), + ).fetchone()["authoritative_user_digest"], + "owned_user_playbook_ids": [], + }, deleted_count=0, error_detail=None, ) @@ -4283,8 +4302,9 @@ def _trace(statement: str) -> None: future = executor.submit( second.prepare_governance_erase_targets, purge_id, - "two-connection-user", - set(), + "alice", + execution_claim=claim, + owned_user_playbook_ids=set(), ) assert entered.wait(timeout=1) time.sleep(0.05) diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py index 54a8d5d13..4e484c818 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py @@ -1,4 +1,4 @@ -"""Regression coverage for reverting the session-outcome identity schema.""" +"""Regression coverage for rebuilding session-outcome identity after rollback.""" from unittest.mock import patch @@ -42,7 +42,7 @@ def _storage(db_path: str) -> SQLiteStorage: return SQLiteStorage(org_id="downgrade-org", db_path=db_path) -def test_identity_schema_is_downgraded_and_current_writes_resume(tmp_path) -> None: +def test_identity_schema_is_preserved_and_current_writes_resume(tmp_path) -> None: db_path = str(tmp_path / "identity-schema.db") storage = _storage(db_path) for session_id in ("kept", "unknown", "new"): @@ -84,6 +84,8 @@ def test_identity_schema_is_downgraded_and_current_writes_resume(tmp_path) -> No row["name"] for row in migrated.conn.execute("PRAGMA table_info(session_outcomes)") } == { + "outcome_id", + "outcome_revision", "user_id", "session_id", "outcome", @@ -92,13 +94,17 @@ def test_identity_schema_is_downgraded_and_current_writes_resume(tmp_path) -> No "label", "value", "metadata", + "outcome_contract_digest", + "finalized_trajectory_digest", "governance_subject_ref", "created_at", } records = migrated.get_session_outcomes(GetSessionOutcomesRequest()) - assert [(record.session_id, record.outcome) for record in records] == [ - ("kept", SessionOutcomeKind.SUCCESS) - ] + records_by_session = {record.session_id: record for record in records} + assert records_by_session["kept"].outcome is SessionOutcomeKind.SUCCESS + assert records_by_session["kept"].outcome_id == "outcome-kept" + assert records_by_session["unknown"].outcome is SessionOutcomeKind.UNKNOWN + assert records_by_session["unknown"].outcome_id == "outcome-unknown" request = SetSessionOutcomeRequest( session_id="new", outcome=SessionOutcomeKind.FAILURE, occurred_at=101 @@ -165,29 +171,28 @@ def test_legacy_schema_backfills_required_subject_ref( ) -def test_legacy_empty_subject_default_is_rebuilt_and_backfilled(tmp_path) -> None: +def test_identity_empty_subject_default_is_rebuilt_without_changing_identity( + tmp_path, +) -> None: db_path = str(tmp_path / "legacy-empty-subject-default.db") storage = _storage(db_path) storage.conn.execute("DROP TABLE session_outcomes") storage.conn.executescript( - """CREATE TABLE session_outcomes ( - user_id TEXT NOT NULL, - session_id TEXT NOT NULL, - outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), - occurred_at INTEGER NOT NULL, - source TEXT NOT NULL, - label TEXT, - value REAL, - metadata TEXT, - governance_subject_ref TEXT NOT NULL DEFAULT '', - created_at INTEGER NOT NULL, - PRIMARY KEY (user_id, session_id) - );""" + _IDENTITY_SCHEMA.replace( + "governance_subject_ref TEXT NOT NULL,", + "governance_subject_ref TEXT NOT NULL DEFAULT '',", + ) ) storage.conn.execute( """INSERT INTO session_outcomes ( - user_id, session_id, outcome, occurred_at, source, created_at - ) VALUES ('legacy-user', 'legacy-session', 'success', 101, 'legacy', 102)""" + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, created_at + ) VALUES ( + 'stable-outcome-id', 3, 'legacy-user', 'legacy-session', 'unknown', + 101, 'legacy', ?, ?, 102 + )""", + ("a" * 64, "b" * 64), ) storage.conn.commit() storage.conn.close() @@ -195,12 +200,25 @@ def test_legacy_empty_subject_default_is_rebuilt_and_backfilled(tmp_path) -> Non migrated = _storage(db_path) row = migrated.conn.execute( - "SELECT governance_subject_ref FROM session_outcomes" + """SELECT outcome_id, outcome_revision, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref + FROM session_outcomes""" ).fetchone() assert row is not None + assert row["outcome_id"] == "stable-outcome-id" + assert row["outcome_revision"] == 3 + assert row["outcome_contract_digest"] == "a" * 64 + assert row["finalized_trajectory_digest"] == "b" * 64 assert row["governance_subject_ref"] == migrated._subject_ref_for_user_id( "legacy-user" ) + subject_column = next( + item + for item in migrated.conn.execute("PRAGMA table_info(session_outcomes)") + if item["name"] == "governance_subject_ref" + ) + assert subject_column["notnull"] == 1 + assert subject_column["dflt_value"] is None def test_sqlite_versions_without_returning_and_drop_column_are_rejected( diff --git a/tests/server/services/storage/test_playbook_aggregation_state_integration.py b/tests/server/services/storage/test_playbook_aggregation_state_integration.py index 557326342..0dd748a1f 100644 --- a/tests/server/services/storage/test_playbook_aggregation_state_integration.py +++ b/tests/server/services/storage/test_playbook_aggregation_state_integration.py @@ -1102,23 +1102,9 @@ def test_incremental_run_refreshes_agent_and_centroid_after_match( request_context=context, agent_version="v1", ) - learning_meter = MagicMock() - monkeypatch.setattr(aggregator, "_record_learnings_generated", learning_meter) - first = aggregator.run( - PlaybookAggregatorRequest(agent_version="v1", operation_key="test-run-1") - ) + first = aggregator.run(PlaybookAggregatorRequest(agent_version="v1")) assert first["playbooks_generated"] == 1 - first_saved_id = store.conn.execute( - "SELECT agent_playbook_id FROM agent_playbooks" - ).fetchone()[0] - learning_meter.assert_called_once_with( - learning_ids=[str(first_saved_id)], - playbook_name="playbook", - request_id="test-run-1", - metadata=first, - total_count=1, - ) assert store.conn.execute("SELECT count(*) FROM agent_playbooks").fetchone()[0] == 1 _insert_current(store, 3, embedding=encoded, trigger="same trigger") diff --git a/tests/server/services/test_non_extraction_learning_metering.py b/tests/server/services/test_non_extraction_learning_metering.py index ce43cbeb0..2c2d2675d 100644 --- a/tests/server/services/test_non_extraction_learning_metering.py +++ b/tests/server/services/test_non_extraction_learning_metering.py @@ -1,5 +1,6 @@ from __future__ import annotations +from datetime import UTC, datetime from unittest.mock import MagicMock, patch import pytest @@ -18,7 +19,11 @@ AgentRunRecord, AgentRunStatus, ) -from reflexio.server.usage_metrics import UsageEvent, configure_usage_event_recorder +from reflexio.server.usage_metrics import ( + UsageEvent, + UsageEventDeliveryStatus, + configure_usage_event_recorder, +) @pytest.fixture(autouse=True) @@ -37,6 +42,14 @@ def _request_context() -> RequestContext: return ctx +def _capture_events(events: list[UsageEvent]) -> None: + def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: + events.append(event) + return UsageEventDeliveryStatus.APPENDED + + configure_usage_event_recorder(recorder) + + def _agent_run(*, extractor_kind: str) -> AgentRunRecord: return AgentRunRecord( id="run-1", @@ -50,13 +63,14 @@ def _agent_run(*, extractor_kind: str) -> AgentRunRecord: ), status=AgentRunStatus.FINALIZING, generation_request_snapshot={}, + created_at=datetime(2026, 8, 1, tzinfo=UTC), ) def test_resumable_profile_bills_only_ids_returned_by_finalization() -> None: """A preassigned candidate ID is not billed when finalization drops it.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), @@ -77,28 +91,10 @@ def test_resumable_profile_bills_only_ids_returned_by_finalization() -> None: assert events == [] -def test_resumable_fallback_reuses_its_event_key_on_finalization_retry() -> None: - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) - worker = ExtractionResumeWorker( - request_context=_request_context(), - llm_client=MagicMock(), - ) - run = _agent_run(extractor_kind="profile") - - worker._record_finalized_learnings(run, [object()], entity_type="profile") - worker._record_finalized_learnings(run, [object()], entity_type="profile") - - assert [event.event_key for event in events] == [ - "learn-batch:resumable:run-1:profile", - "learn-batch:resumable:run-1:profile", - ] - - def test_resumable_finalization_emits_one_event_per_profile_id() -> None: """Finalization survivor IDs emit one entity-backed event per profile.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), @@ -129,7 +125,7 @@ def test_resumable_finalization_emits_one_event_per_profile_id() -> None: def test_resumable_finalization_emits_one_event_per_playbook_id() -> None: """Finalization survivor IDs emit one entity-backed event per playbook.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), @@ -158,7 +154,7 @@ def test_resumable_finalization_emits_one_event_per_playbook_id() -> None: def test_resumable_playbook_bills_consolidation_replacement_id() -> None: """Billing follows the persisted replacement, not its input candidate.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), @@ -187,7 +183,7 @@ def test_aggregation_emits_no_learnings_generated( ) -> None: """A completed aggregation remains observable but adds no billable learning.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) request_context = _request_context() storage = MagicMock() configurator = MagicMock() From b698092927d3a00f68b3122a79449cd747f9323c Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 06:56:13 +0000 Subject: [PATCH 15/63] fix: run always-global sweeps before bootstrap --- .../server/services/lineage/gc_scheduler.py | 2 +- .../lineage/test_gc_scheduler_global_sweep.py | 42 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/reflexio/server/services/lineage/gc_scheduler.py b/reflexio/server/services/lineage/gc_scheduler.py index 215ab03ef..55cd405cb 100644 --- a/reflexio/server/services/lineage/gc_scheduler.py +++ b/reflexio/server/services/lineage/gc_scheduler.py @@ -485,6 +485,7 @@ def _run_always_global_sweeps(self) -> None: logger.exception("event=always_global_sweep_failed sweep=%s", sweep_id) def _run_once(self) -> float: + self._run_always_global_sweeps() poll_interval = _DEFAULT_POLL_INTERVAL_SECONDS try: if ( @@ -505,7 +506,6 @@ def _run_once(self) -> float: org_ids = self._discover_org_ids(bootstrap_ctx) self._gc_tick(org_ids, max_workers=self._org_fanout_workers(bootstrap_ctx)) self._run_global_sweeps(cfg) - self._run_always_global_sweeps() except Exception: logger.exception("event=lineage_gc_scheduler_tick_failed") return max(poll_interval, _MIN_POLL_SECONDS) diff --git a/tests/server/services/lineage/test_gc_scheduler_global_sweep.py b/tests/server/services/lineage/test_gc_scheduler_global_sweep.py index c4466b84e..62ed4cfb6 100644 --- a/tests/server/services/lineage/test_gc_scheduler_global_sweep.py +++ b/tests/server/services/lineage/test_gc_scheduler_global_sweep.py @@ -2,6 +2,7 @@ import pytest +from reflexio.server.auth import DEFAULT_ORG_ID from reflexio.server.services.lineage import gc_scheduler from reflexio.server.services.lineage.gc_scheduler import ( LineageGCScheduler, @@ -127,3 +128,44 @@ def test_run_once_invokes_always_global_sweep_when_expiry_disabled(monkeypatch): scheduler._run_once() assert len(calls) == 1 + + +def test_run_once_invokes_always_global_sweep_during_empty_fleet_retry(): + calls: list[int] = [] + gc_scheduler.register_always_global_sweep(lambda now: calls.append(now) or 1) + factory_calls: list[str] = [] + scheduler = LineageGCScheduler( + request_context_factory=lambda org_id: factory_calls.append(org_id), # type: ignore[arg-type] + bootstrap_org_id=DEFAULT_ORG_ID, + org_id_provider=lambda: [DEFAULT_ORG_ID], + ) + + assert scheduler._run_once() == 5 + assert len(calls) == 1 + assert factory_calls == [] + + +@pytest.mark.parametrize("bootstrap_failure", ["context", "config"]) +def test_run_once_invokes_always_global_sweep_during_bootstrap_failure( + bootstrap_failure: str, +): + calls: list[int] = [] + gc_scheduler.register_always_global_sweep(lambda now: calls.append(now) or 1) + + def get_config(): + raise RuntimeError("config bootstrap failed") + + def request_context_factory(_org_id: str): + if bootstrap_failure == "context": + raise RuntimeError("context bootstrap failed") + return types.SimpleNamespace( + configurator=types.SimpleNamespace(get_config=get_config) + ) + + scheduler = LineageGCScheduler( + request_context_factory=request_context_factory, # type: ignore[arg-type] + bootstrap_org_id="org-boot", + ) + + assert scheduler._run_once() == 86400 + assert len(calls) == 1 From 592b30a0cd824bd3dc3c43df829cc0704dd2cac4 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 07:21:28 +0000 Subject: [PATCH 16/63] fix: bound session outcome migration queries Prefetch request and interaction rows in SQLite-safe chunks while preserving the atomic legacy outcome rebuild and canonical trajectory identities. --- .../services/storage/sqlite_storage/_base.py | 75 ++++++++++++++++++- .../test_session_outcome_migration.py | 61 +++++++++++++++ 2 files changed, 132 insertions(+), 4 deletions(-) diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 25d7da93f..8bb731eb7 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -60,7 +60,7 @@ StorageError, require_non_empty_session_id, ) -from reflexio.server.services.storage.retention_mixin import RetentionMixin +from reflexio.server.services.storage.retention_mixin import RetentionMixin, chunked from reflexio.server.services.storage.session_outcome_identity import ( CanonicalSessionTrajectory, canonical_json_bytes, @@ -147,6 +147,66 @@ def _canonical_session_snapshot( ) +def _prefetch_canonical_session_snapshots( + conn: sqlite3.Connection, session_ids: Sequence[str] +) -> dict[str, CanonicalSessionTrajectory]: + """Bulk-load canonical session state in bounded SQLite-safe chunks.""" + snapshots: dict[str, CanonicalSessionTrajectory] = {} + for session_id_chunk in chunked(list(dict.fromkeys(session_ids))): + placeholders = ",".join("?" for _ in session_id_chunk) + request_rows = conn.execute( + f"""SELECT request_id, user_id, created_at, source, agent_version, + session_id, evaluation_only, retrieval_experiment_id, + retrieval_experiment_arm + FROM requests WHERE session_id IN ({placeholders}) + ORDER BY session_id ASC, created_at ASC, request_id ASC""", # noqa: S608 + session_id_chunk, + ).fetchall() + interaction_rows = conn.execute( + f"""SELECT requests.session_id AS trajectory_session_id, + interactions.interaction_id, interactions.user_id, + interactions.request_id, interactions.created_at, + interactions.content, interactions.role, + interactions.token_count, interactions.user_action, + interactions.user_action_description, + interactions.interacted_image_url, + interactions.image_encoding, interactions.shadow_content, + interactions.expert_content, interactions.tools_used, + interactions.citations, interactions.retrieved_learnings + FROM interactions + JOIN requests ON requests.request_id = interactions.request_id + WHERE requests.session_id IN ({placeholders}) + ORDER BY requests.session_id ASC, interactions.request_id ASC, + interactions.created_at ASC, + interactions.interaction_id ASC""", # noqa: S608 + session_id_chunk, + ).fetchall() + + requests_by_session: dict[str, list[dict[str, object]]] = { + session_id: [] for session_id in session_id_chunk + } + interactions_by_session: dict[str, dict[str, list[dict[str, object]]]] = { + session_id: {} for session_id in session_id_chunk + } + for request in request_rows: + session_id = str(request["session_id"]) + requests_by_session[session_id].append(dict(request)) + interactions_by_session[session_id][str(request["request_id"])] = [] + for interaction in interaction_rows: + interaction_payload = dict(interaction) + session_id = str(interaction_payload.pop("trajectory_session_id")) + interactions_by_session[session_id][ + str(interaction_payload["request_id"]) + ].append(interaction_payload) + for session_id in session_id_chunk: + snapshots[session_id] = canonical_session_trajectory( + session_id, + requests_by_session[session_id], + interactions_by_session[session_id], + ) + return snapshots + + def _legacy_session_outcome_id(user_id: str, session_id: str) -> str: """Return a delimiter-safe immutable identity for a migrated legacy outcome.""" return sha256(canonical_json_bytes([user_id, session_id])).hexdigest() @@ -1256,6 +1316,15 @@ def _migrate_session_outcomes_schema(self) -> None: legacy_rows = self.conn.execute( "SELECT * FROM session_outcomes" ).fetchall() + trajectory_snapshots = _prefetch_canonical_session_snapshots( + self.conn, + [ + str(row["session_id"]) + for row in legacy_rows + if "finalized_trajectory_digest" not in columns + or not row["finalized_trajectory_digest"] + ], + ) self.conn.execute( "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" ) @@ -1334,9 +1403,7 @@ def _migrate_session_outcomes_schema(self) -> None: ), existing_trajectory_digest or trajectory_digest( - _canonical_session_snapshot( - self.conn, str(row["session_id"]) - ) + trajectory_snapshots[str(row["session_id"])] ), subject_ref, row["created_at"], diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index 73439ffc7..4a4f4326f 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -191,6 +191,67 @@ def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( ) +def test_migration_prefetches_trajectory_inputs_in_bounded_chunks(tmp_path) -> None: + db_path = str(tmp_path / "legacy-session-outcome-query-scaling.db") + storage = SQLiteStorage(org_id="legacy-query-scaling", db_path=db_path) + row_count = 501 + storage.conn.executemany( + """INSERT INTO requests ( + request_id, user_id, created_at, source, session_id + ) VALUES (?, ?, ?, ?, ?)""", + [ + ( + f"request-{index}", + f"user-{index}", + str(index), + "legacy-source", + f"session-{index}", + ) + for index in range(row_count) + ], + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + storage.conn.executemany( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?)""", + [ + ( + f"user-{index}", + f"session-{index}", + "success", + index, + "legacy-source", + f"subject-{index}", + index, + ) + for index in range(row_count) + ], + ) + storage.conn.commit() + + statements: list[str] = [] + storage.conn.set_trace_callback(statements.append) + try: + storage._migrate_session_outcomes_schema() + finally: + storage.conn.set_trace_callback(None) + + trajectory_input_queries = [ + statement + for statement in statements + if statement.lstrip().upper().startswith("SELECT") + and ("FROM requests" in statement or "FROM interactions" in statement) + ] + assert len(trajectory_input_queries) == 4 + assert ( + storage.conn.execute("SELECT COUNT(*) FROM session_outcomes").fetchone()[0] + == row_count + ) + + @pytest.mark.parametrize( ("with_subject_column", "stored_subject_ref"), [(False, None), (True, None), (True, " ")], From 421950074ea4025831676131c79ff703af4ed0e2 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 07:31:21 +0000 Subject: [PATCH 17/63] fix: restore OSS E2E billing exemption --- docs/lib/methods/unified-search.ts | 21 ++++++++++++++++--- reflexio/client/client.py | 15 +++++++++----- tests/client/test_search.py | 33 ++++++++++++++++++++++++++++++ tests/e2e_tests/conftest.py | 13 ++++++++++++ 4 files changed, 74 insertions(+), 8 deletions(-) diff --git a/docs/lib/methods/unified-search.ts b/docs/lib/methods/unified-search.ts index 76e2004b9..b63470725 100644 --- a/docs/lib/methods/unified-search.ts +++ b/docs/lib/methods/unified-search.ts @@ -23,7 +23,7 @@ export const unifiedSearchMethods: MethodDef[] = [ type: "number", required: false, default: 5, - description: "Maximum results per entity type", + description: "Maximum results per entity type, from 1 to 100", }, { name: "threshold", @@ -50,7 +50,8 @@ export const unifiedSearchMethods: MethodDef[] = [ name: "user_id", type: "string", required: false, - description: "Filter by user ID (profiles, user_playbooks)", + description: + "Filter by user ID (profiles, user_playbooks), at most 255 characters", }, { name: "entity_types", @@ -90,12 +91,26 @@ export const unifiedSearchMethods: MethodDef[] = [ "Search mode: vector (embedding similarity), fts (full-text search), or hybrid (combined with RRF)", enumValues: ["vector", "fts", "hybrid"], }, + { + name: "request_id", + type: "string", + required: false, + description: + "Caller correlation ID for the search turn, at most 255 characters", + }, { name: "session_id", type: "string", required: false, description: - "Agent session this search serves. When set, results already returned to the same session are skipped and next-best matches backfilled; searches without it neither read nor record session dedup state", + "Agent session this search serves, at most 255 characters. When set, results already returned to the same session are skipped and next-best matches backfilled; searches without it neither read nor record session dedup state", + }, + { + name: "interaction_id", + type: "number", + required: false, + description: + "Caller interaction ID for the search turn; must be a positive integer (minimum 1)", }, ], }, diff --git a/reflexio/client/client.py b/reflexio/client/client.py index 684c8e728..fe3fd08ea 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -2829,12 +2829,14 @@ def search( Args: request (Optional[UnifiedSearchRequest]): The search request object (alternative to kwargs) query (str): Search query text - top_k (Optional[int]): Maximum results per entity type (default: 5) + top_k (Optional[int]): Maximum results per entity type, from 1 to 100 + (default: 5). threshold (Optional[float]): Similarity threshold for vector search. When omitted, the embedding model's default is used. agent_version (Optional[str]): Filter by agent version (agent_playbooks, user_playbooks) playbook_name (Optional[str]): Filter by playbook name (agent_playbooks, user_playbooks) - user_id (Optional[str]): Filter by user ID (profiles, user_playbooks) + user_id (Optional[str]): Filter by user ID (profiles, user_playbooks), + at most 255 characters. tags (Optional[list[str]]): Match entities having any requested tag. entity_types (Optional[list[str]]): Entity types to search. Valid values: "profiles", "user_playbooks", "agent_playbooks". @@ -2850,9 +2852,12 @@ def search( the configured search backend supports it (default: False). conversation_history (Optional[list[ConversationTurn] | list[dict]]): Prior conversation turns for context-aware query reformulation. Accepts ConversationTurn objects or dicts with "role" and "content" keys. search_mode (Optional[SearchMode | str]): Search mode to use. Accepts SearchMode enum or string value ("vector", "fts", "hybrid"). - request_id (Optional[str]): Caller correlation id for the search turn. - session_id (Optional[str]): Caller session id for the search turn. - interaction_id (Optional[int]): Caller interaction id for the search turn. + request_id (Optional[str]): Caller correlation ID for the search turn, + at most 255 characters. + session_id (Optional[str]): Caller session ID for the search turn, + at most 255 characters. Also enables session-scoped result deduplication. + interaction_id (Optional[int]): Caller interaction ID for the search + turn; must be a positive integer (minimum 1). Returns: UnifiedSearchViewResponse: Combined search results from all entity types diff --git a/tests/client/test_search.py b/tests/client/test_search.py index 285abc8a9..e48776b8d 100644 --- a/tests/client/test_search.py +++ b/tests/client/test_search.py @@ -1,8 +1,20 @@ +import inspect +from pathlib import Path from typing import Any import pytest from reflexio import ReflexioClient +from reflexio.models.api_schema.retriever_schema import UnifiedSearchRequest + + +def _non_null_schema(field_name: str) -> dict[str, Any]: + field_schema = UnifiedSearchRequest.model_json_schema()["properties"][field_name] + return next( + option + for option in field_schema.get("anyOf", [field_schema]) + if option.get("type") != "null" + ) def test_unified_search_serializes_tag_filter(monkeypatch) -> None: @@ -69,3 +81,24 @@ async def fake_make_async_request( assert captured["json"]["user_id"] == "u1" assert captured["json"]["agent_version"] == "a1" assert captured["json"]["top_k"] == 4 + + +def test_unified_search_docs_track_schema_contract() -> None: + top_k_schema = _non_null_schema("top_k") + identifier_limit = _non_null_schema("request_id")["maxLength"] + interaction_minimum = _non_null_schema("interaction_id")["exclusiveMinimum"] + 1 + client_docs = inspect.getdoc(ReflexioClient.search) or "" + registry_docs = ( + Path(__file__).parents[2] / "docs/lib/methods/unified-search.ts" + ).read_text(encoding="utf-8") + + assert f"1 to {top_k_schema['maximum']}" in client_docs + assert client_docs.count(f"at most {identifier_limit} characters") >= 3 + assert f"positive integer (minimum {interaction_minimum})" in client_docs + + assert f"1 to {top_k_schema['maximum']}" in registry_docs + for field_name in ("user_id", "request_id", "session_id"): + assert f'name: "{field_name}"' in registry_docs + assert registry_docs.count(f"at most {identifier_limit} characters") >= 3 + assert 'name: "interaction_id"' in registry_docs + assert f"positive integer (minimum {interaction_minimum})" in registry_docs diff --git a/tests/e2e_tests/conftest.py b/tests/e2e_tests/conftest.py index 4085fbc39..af76368ae 100644 --- a/tests/e2e_tests/conftest.py +++ b/tests/e2e_tests/conftest.py @@ -24,6 +24,7 @@ StorageConfigSQLite, ToolUseConfig, ) +from reflexio.server import usage_metrics from reflexio.server.services.configurator.configurator import DefaultConfigurator from reflexio.server.services.tagging.tagging_scheduler import drain_tagging from reflexio.test_support.llm_mock import ( @@ -58,6 +59,18 @@ """ +@pytest.fixture(autouse=True) +def _install_oss_usage_event_exemption() -> Iterator[None]: + """Match the explicit billing exemption installed by the OSS app lifespan.""" + usage_metrics.configure_usage_event_recorder( + usage_metrics.exempt_usage_event_recorder + ) + try: + yield + finally: + usage_metrics.configure_usage_event_recorder(None) + + @pytest.fixture(autouse=True) def mock_llm(request: pytest.FixtureRequest) -> Iterator[None]: """Keep the standard E2E tier deterministic and credential-free. From e85fa246f2e52805baf5caaef3cbe208dfbbf98f Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 07:34:40 +0000 Subject: [PATCH 18/63] test: cover session outcome migration rollback --- .../test_session_outcome_migration.py | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index 4a4f4326f..d932fa11a 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -1,6 +1,7 @@ """Regression coverage for the SQLite session-outcome identity migration.""" import json +import sqlite3 from hashlib import sha256 import pytest @@ -252,6 +253,126 @@ def test_migration_prefetches_trajectory_inputs_in_bounded_chunks(tmp_path) -> N ) +def test_session_outcome_rebuild_failure_rolls_back_renamed_legacy_table( + tmp_path, monkeypatch +) -> None: + db_path = str(tmp_path / "legacy-session-outcome-atomicity.db") + storage = SQLiteStorage(org_id="legacy-atomicity", db_path=db_path) + storage.add_request( + Request( + request_id="request-atomicity-session", + user_id="atomicity-user", + session_id="atomicity-session", + source="atomicity-source", + created_at=100, + ) + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + legacy_rows = [ + ( + "atomicity-user", + "atomicity-session", + "success", + 101, + "atomicity-source", + "atomicity-label", + 4.5, + json.dumps({"kept": True}, sort_keys=True), + "subject:atomicity-user", + 102, + ), + ( + "atomicity-user-2", + "atomicity-session-2", + "failure", + 201, + "atomicity-source-2", + None, + None, + None, + "subject:atomicity-user-2", + 202, + ), + ] + storage.conn.executemany( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, label, value, + metadata, governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + legacy_rows, + ) + storage.conn.commit() + storage.conn.close() + + with sqlite3.connect(db_path) as probe: + original_schema = probe.execute( + """SELECT sql FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes'""" + ).fetchone()[0] + original_rows = probe.execute( + "SELECT * FROM session_outcomes ORDER BY user_id, session_id" + ).fetchall() + + real_connect = sqlite3.connect + failing_connections = [] + + class _FailAfterSessionOutcomeRename(sqlite3.Connection): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.saw_session_outcome_rename = False + self.injected_failure = False + + def execute(self, sql, parameters=(), /): + normalized = " ".join(str(sql).split()) + if ( + normalized + == "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" + ): + self.saw_session_outcome_rename = True + if self.saw_session_outcome_rename and normalized.startswith( + "CREATE TABLE session_outcomes (" + ): + self.injected_failure = True + raise RuntimeError("injected session_outcomes rebuild failure") + return super().execute(sql, parameters) + + def _connect_with_failure(*args, **kwargs): + kwargs["factory"] = _FailAfterSessionOutcomeRename + conn = real_connect(*args, **kwargs) + failing_connections.append(conn) + return conn + + monkeypatch.setattr(sqlite3, "connect", _connect_with_failure) + + with pytest.raises(RuntimeError, match="injected session_outcomes rebuild failure"): + SQLiteStorage(org_id="legacy-atomicity", db_path=db_path) + + assert any( + conn.saw_session_outcome_rename and conn.injected_failure + for conn in failing_connections + ) + for conn in failing_connections: + conn.close() + + with sqlite3.connect(db_path) as probe: + restored_schema = probe.execute( + """SELECT sql FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes'""" + ).fetchone()[0] + restored_rows = probe.execute( + "SELECT * FROM session_outcomes ORDER BY user_id, session_id" + ).fetchall() + stranded_legacy_table = probe.execute( + """SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes_legacy'""" + ).fetchone() + + assert restored_schema == original_schema + assert restored_rows == original_rows + assert stranded_legacy_table is None + + @pytest.mark.parametrize( ("with_subject_column", "stored_subject_ref"), [(False, None), (True, None), (True, " ")], From 38eb8d3146608b2b0bb80c6da50ce77c48b1b8e9 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 07:49:15 +0000 Subject: [PATCH 19/63] test: isolate generation billing from E2E mode --- .../services/test_generation_billing_emission.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/server/services/test_generation_billing_emission.py b/tests/server/services/test_generation_billing_emission.py index cdc597054..ec622bc1c 100644 --- a/tests/server/services/test_generation_billing_emission.py +++ b/tests/server/services/test_generation_billing_emission.py @@ -53,6 +53,7 @@ from reflexio.server.services.service_utils import format_sessions_to_history_string from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.usage_metrics import UsageEvent, configure_usage_event_recorder +from reflexio.test_support.llm_mock import patched_litellm # --------------------------------------------------------------------------- # Helpers: build a RequestContext backed by a real SQLiteStorage, mirroring @@ -152,9 +153,9 @@ def _run_profile_generation(storage: SQLiteStorage, *, auto_run: bool = True) -> def test_real_extraction_emits_tokens_and_learnings(tmp_path): """A successful extraction emits extraction_tokens + learnings_generated. - MOCK_LLM_RESPONSE=true is active (autouse conftest), so the profile extractor - takes the deterministic mock path and we verify the billing events it triggers. - auto_run=False bypasses the should_run gate so extraction always fires. + The local LiteLLM patch keeps this deterministic even when the test is selected + alongside E2E paths. auto_run=False bypasses the should_run gate so extraction + always fires. """ events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) @@ -164,7 +165,8 @@ def test_real_extraction_emits_tokens_and_learnings(tmp_path): org_id=_ORG_ID, db_path=str(tmp_path / "reflexio.db") ) _seed_interactions(storage) - _run_profile_generation(storage, auto_run=False) + with patched_litellm(): + _run_profile_generation(storage, auto_run=False) finally: configure_usage_event_recorder(None) From a4dd97d345b0f3ad2431a23d1c5c267c4e47b4b5 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 08:04:12 +0000 Subject: [PATCH 20/63] test: isolate playbook retry from E2E mode --- tests/server/services/extraction/test_resume_worker.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 519bdaee6..7e23160c0 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -962,8 +962,10 @@ def test_receipt_delivery_failure_status_distinguishes_transient_and_permanent( def test_retry_after_billing_reuses_ids_without_replaying_playbook_schedulers( request_context, storage, + monkeypatch, ): """A post-billing retry reuses durable IDs and skips derived schedulers.""" + monkeypatch.setenv("MOCK_LLM_RESPONSE", "true") _seed_interactions(storage) request_context.configurator.get_config.return_value = Config( storage_config=StorageConfigSQLite(), From fb8983717e7f70d34ad92d46ef347bf693f94282 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 08:25:18 +0000 Subject: [PATCH 21/63] fix: bound session outcome migration memory --- .../services/storage/sqlite_storage/_base.py | 26 ++++++++--------- .../test_session_outcome_migration.py | 28 +++++++++++++++++++ 2 files changed, 41 insertions(+), 13 deletions(-) diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 8bb731eb7..cb8b79adb 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -147,11 +147,11 @@ def _canonical_session_snapshot( ) -def _prefetch_canonical_session_snapshots( +def _prefetch_canonical_session_trajectory_digests( conn: sqlite3.Connection, session_ids: Sequence[str] -) -> dict[str, CanonicalSessionTrajectory]: - """Bulk-load canonical session state in bounded SQLite-safe chunks.""" - snapshots: dict[str, CanonicalSessionTrajectory] = {} +) -> dict[str, str]: + """Derive trajectory digests while retaining only one payload chunk.""" + digests: dict[str, str] = {} for session_id_chunk in chunked(list(dict.fromkeys(session_ids))): placeholders = ",".join("?" for _ in session_id_chunk) request_rows = conn.execute( @@ -199,12 +199,14 @@ def _prefetch_canonical_session_snapshots( str(interaction_payload["request_id"]) ].append(interaction_payload) for session_id in session_id_chunk: - snapshots[session_id] = canonical_session_trajectory( - session_id, - requests_by_session[session_id], - interactions_by_session[session_id], + digests[session_id] = trajectory_digest( + canonical_session_trajectory( + session_id, + requests_by_session[session_id], + interactions_by_session[session_id], + ) ) - return snapshots + return digests def _legacy_session_outcome_id(user_id: str, session_id: str) -> str: @@ -1316,7 +1318,7 @@ def _migrate_session_outcomes_schema(self) -> None: legacy_rows = self.conn.execute( "SELECT * FROM session_outcomes" ).fetchall() - trajectory_snapshots = _prefetch_canonical_session_snapshots( + trajectory_digests = _prefetch_canonical_session_trajectory_digests( self.conn, [ str(row["session_id"]) @@ -1402,9 +1404,7 @@ def _migrate_session_outcomes_schema(self) -> None: finalization_rule="first_write", ), existing_trajectory_digest - or trajectory_digest( - trajectory_snapshots[str(row["session_id"])] - ), + or trajectory_digests[str(row["session_id"])], subject_ref, row["created_at"], ), diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index d932fa11a..b8ef00f17 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -18,6 +18,7 @@ from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.sqlite_storage._base import ( _canonical_session_snapshot, + _prefetch_canonical_session_trajectory_digests, ) pytestmark = pytest.mark.integration @@ -253,6 +254,33 @@ def test_migration_prefetches_trajectory_inputs_in_bounded_chunks(tmp_path) -> N ) +def test_migration_prefetch_retains_only_trajectory_digests(tmp_path) -> None: + storage = SQLiteStorage( + org_id="legacy-digest-retention", + db_path=str(tmp_path / "legacy-digest-retention.db"), + ) + storage.add_request( + Request( + request_id="request-digest-retention", + user_id="digest-user", + session_id="digest-session", + source="legacy-source", + created_at=100, + ) + ) + + digests = _prefetch_canonical_session_trajectory_digests( + storage.conn, ["digest-session"] + ) + + assert digests == { + "digest-session": trajectory_digest( + _canonical_session_snapshot(storage.conn, "digest-session") + ) + } + assert all(isinstance(digest, str) for digest in digests.values()) + + def test_session_outcome_rebuild_failure_rolls_back_renamed_legacy_table( tmp_path, monkeypatch ) -> None: From ac40f9d2d125c793fefa0570c089a29a87d16c11 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 09:16:11 +0000 Subject: [PATCH 22/63] fix: harden session outcome identity contracts --- reflexio/server/services/search_exposure.py | 2 +- .../storage/session_outcome_identity.py | 21 +++-- .../services/storage/sqlite_storage/_base.py | 1 + .../agent_run/_agent_run_store.py | 6 +- tests/models/test_session_outcome_identity.py | 42 +++++++++ .../test_session_outcomes_integration.py | 2 +- .../test_session_outcome_migration.py | 88 +++++++++++++++++++ 7 files changed, 154 insertions(+), 8 deletions(-) diff --git a/reflexio/server/services/search_exposure.py b/reflexio/server/services/search_exposure.py index 33270eb4d..131d98631 100644 --- a/reflexio/server/services/search_exposure.py +++ b/reflexio/server/services/search_exposure.py @@ -170,9 +170,9 @@ def build_user_playbook_exposure_event( __all__ = [ + "MAX_EXPOSURE_EVENTS_PER_BATCH", "SEARCH_EXPOSURE_RECORDER", "ExposureEventWriteResult", - "MAX_EXPOSURE_EVENTS_PER_BATCH", "SearchExposureBatch", "SearchExposureRecorder", "UserPlaybookExposureEvent", diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py index 273b0195a..ccf39edf6 100644 --- a/reflexio/server/services/storage/session_outcome_identity.py +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -10,12 +10,14 @@ __all__ = [ "CanonicalSessionTrajectory", - "canonical_session_trajectory", "canonical_json_bytes", + "canonical_session_trajectory", "outcome_contract_digest", "trajectory_digest", ] +MAX_CANONICAL_TRAJECTORY_JSON_DEPTH = 100 + class CanonicalRequest(TypedDict): request_id: str @@ -162,12 +164,20 @@ def outcome_contract_digest( return sha256(canonical_json_bytes(payload)).hexdigest() -def _canonical_trajectory_json(value: object) -> str: - """Encode trajectory JSON, including finite floats, deterministically.""" +def _canonical_trajectory_json(value: object, *, depth: int = 0) -> str: + """Encode trajectory JSON deterministically within a bounded nesting depth.""" + if depth > MAX_CANONICAL_TRAJECTORY_JSON_DEPTH: + raise ValueError("canonical trajectory JSON exceeds maximum depth") if isinstance(value, float): return json.dumps(value, allow_nan=False, separators=(",", ":")) if isinstance(value, tuple | list): - return "[" + ",".join(_canonical_trajectory_json(item) for item in value) + "]" + return ( + "[" + + ",".join( + _canonical_trajectory_json(item, depth=depth + 1) for item in value + ) + + "]" + ) if isinstance(value, Mapping): if not all(isinstance(key, str) for key in value): raise TypeError("canonical trajectory object keys must be strings") @@ -175,7 +185,8 @@ def _canonical_trajectory_json(value: object) -> str: return ( "{" + ",".join( - f"{canonical_json_bytes(key).decode()}:{_canonical_trajectory_json(value[key])}" + f"{canonical_json_bytes(key).decode()}:" + f"{_canonical_trajectory_json(value[key], depth=depth + 1)}" for key in keys ) + "}" diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index cb8b79adb..10cacd5c1 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -1386,6 +1386,7 @@ def _migrate_session_outcomes_schema(self) -> None: ( row["outcome_revision"] if "outcome_revision" in columns + and row["outcome_revision"] is not None else 1 ), row["user_id"], diff --git a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py index 3acd3c9fd..d5f2eec6c 100644 --- a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py +++ b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py @@ -30,7 +30,11 @@ def get_agent_run_finalization_receipt( run_id: str, entity_type: str, ) -> list[str] | None: - """Return persisted learning ids for a completed run finalization.""" + """Return the committed learning ids, or ``None`` when no receipt exists. + + A committed empty receipt returns ``[]`` and must remain distinguishable + from a missing receipt so finalization stays idempotent. + """ raise NotImplementedError(f"{type(self).__name__} does not support agent runs") def save_agent_run_finalization_receipt( diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index 41e55df8f..7a5e68df4 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -12,6 +12,7 @@ SetSessionOutcomeResponse, ) from reflexio.models.api_schema.domain.enums import SessionOutcomeKind +from reflexio.server.services.storage import session_outcome_identity from reflexio.server.services.storage.session_outcome_identity import ( canonical_json_bytes, canonical_session_trajectory, @@ -141,6 +142,47 @@ def test_trajectory_digest_rejects_non_finite_nested_floats(value: float) -> Non trajectory_digest({"nested": [{"value": value}]}) +@pytest.mark.parametrize( + "container_factory", + [ + lambda value: {"nested": value}, + lambda value: [value], + lambda value: (value,), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_accepts_maximum_canonical_json_depth( + container_factory, +) -> None: + value: object = "leaf" + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH): + value = container_factory(value) + + assert trajectory_digest(value) + + +@pytest.mark.parametrize( + "container_factory", + [ + lambda value: {"nested": value}, + lambda value: [value], + lambda value: (value,), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_rejects_over_maximum_canonical_json_depth( + container_factory, +) -> None: + value: object = "leaf" + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH + 1): + value = container_factory(value) + + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + trajectory_digest(value) + + def test_session_outcome_record_accepts_unknown_and_serializes_identities() -> None: record = SessionOutcomeRecord( outcome_id="outcome-1", diff --git a/tests/server/api_endpoints/test_session_outcomes_integration.py b/tests/server/api_endpoints/test_session_outcomes_integration.py index 1b3d408a9..83f2fa33e 100644 --- a/tests/server/api_endpoints/test_session_outcomes_integration.py +++ b/tests/server/api_endpoints/test_session_outcomes_integration.py @@ -185,7 +185,7 @@ def test_retry_compares_metadata_as_json_values( assert retry.status_code == 200 assert retry.json()["success"] is True assert retry.json()["recorded"] is False - assert "reason" not in retry.json() + assert retry.json().get("reason") is None def test_retry_rejects_metadata_with_different_json_value_types( diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index b8ef00f17..c2602c6d0 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -82,6 +82,55 @@ def _identity_complete_session_outcomes_ddl(*, with_subject_column: bool) -> str """ +_IDENTITY_COMPLETE_UNCONSTRAINED_REVISION_DDL = """ +CREATE TABLE session_outcomes ( + outcome_id TEXT NOT NULL UNIQUE, + outcome_revision INTEGER, + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), + occurred_at INTEGER NOT NULL, + source TEXT NOT NULL, + label TEXT, + value REAL, + metadata TEXT, + outcome_contract_digest TEXT NOT NULL, + finalized_trajectory_digest TEXT NOT NULL, + governance_subject_ref TEXT NOT NULL, + created_at INTEGER NOT NULL, + PRIMARY KEY (user_id, session_id) +); +""" + + +def _replace_session_outcomes_with_unconstrained_revision( + storage: SQLiteStorage, revision: int | None +) -> None: + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_IDENTITY_COMPLETE_UNCONSTRAINED_REVISION_DDL) + storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + "legacy-outcome-id", + revision, + "legacy-user", + "legacy-session", + "success", + 101, + "legacy-source", + "a" * 64, + "b" * 64, + "subject:legacy-user", + 102, + ), + ) + storage.conn.commit() + + def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( tmp_path, ) -> None: @@ -193,6 +242,45 @@ def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( ) +def test_migration_defaults_null_legacy_outcome_revision_to_one(tmp_path) -> None: + db_path = str(tmp_path / "legacy-null-revision.db") + storage = SQLiteStorage(org_id="legacy-null-revision", db_path=db_path) + _replace_session_outcomes_with_unconstrained_revision(storage, None) + storage.conn.close() + + migrated = SQLiteStorage(org_id="legacy-null-revision", db_path=db_path) + + row = migrated.conn.execute( + "SELECT outcome_revision FROM session_outcomes WHERE session_id = ?", + ("legacy-session",), + ).fetchone() + assert row is not None + assert row["outcome_revision"] == 1 + + +def test_migration_rejects_zero_legacy_outcome_revision_atomically(tmp_path) -> None: + db_path = str(tmp_path / "legacy-zero-revision.db") + storage = SQLiteStorage(org_id="legacy-zero-revision", db_path=db_path) + _replace_session_outcomes_with_unconstrained_revision(storage, 0) + storage.conn.close() + + with pytest.raises(sqlite3.IntegrityError, match="outcome_revision"): + SQLiteStorage(org_id="legacy-zero-revision", db_path=db_path) + + with sqlite3.connect(db_path) as probe: + stored_revision = probe.execute( + "SELECT outcome_revision FROM session_outcomes WHERE session_id = ?", + ("legacy-session",), + ).fetchone()[0] + stranded_legacy_table = probe.execute( + """SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes_legacy'""" + ).fetchone() + + assert stored_revision == 0 + assert stranded_legacy_table is None + + def test_migration_prefetches_trajectory_inputs_in_bounded_chunks(tmp_path) -> None: db_path = str(tmp_path / "legacy-session-outcome-query-scaling.db") storage = SQLiteStorage(org_id="legacy-query-scaling", db_path=db_path) From 0b7578bc88f482676ba3c64fb2410d22d28dc8f1 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 09:21:29 +0000 Subject: [PATCH 23/63] fix: restore fail-open usage lifecycle --- reflexio/server/api.py | 114 ++++++++++--------- reflexio/server/billing_meter.py | 21 ++-- reflexio/server/usage_metrics.py | 12 +- tests/server/test_create_app_capabilities.py | 29 +++++ tests/server/test_usage_metrics.py | 32 ++++++ 5 files changed, 138 insertions(+), 70 deletions(-) diff --git a/reflexio/server/api.py b/reflexio/server/api.py index bb3439a9d..ed860fe15 100644 --- a/reflexio/server/api.py +++ b/reflexio/server/api.py @@ -481,64 +481,66 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # noqa: ARG001 ) configure_usage_event_recorder(exempt_usage_event_recorder) - if mounts_data_plane: - _log_multi_worker_daemons() - log_publish_hardware_capacity() - validate_llm_availability() - # The scheduler discovers every org with resumable work each tick and - # drives a per-org worker with org-scoped claims, so it is not limited - # to the bootstrap org. The bootstrap org is only used to read config - # and to seed cross-org discovery. - bootstrap_org_id = _resolve_lifespan_org_id(get_org_id) - scheduler = maybe_start_resume_scheduler( - lambda org_id: RequestContext(org_id=org_id), - bootstrap_org_id=bootstrap_org_id, - org_id_provider=resume_org_ids_provider, - ) - # Register the missing-vector backfill sweep (opt-in via - # REFLEXIO_MISSING_VECTOR_BACKFILL_ENABLED) BEFORE starting the GC - # scheduler, so its per-org hook is visible when maybe_start_lineage_gc - # evaluates its start conditions. No-op when the flag is off. - install_missing_vector_backfill_sweep() - gc_scheduler = maybe_start_lineage_gc( - lambda org_id: RequestContext(org_id=org_id), - bootstrap_org_id=bootstrap_org_id, - ) - # Durable learning drains the learning_jobs queue per org. Gated on - # REFLEXIO_DURABLE_LEARNING_QUEUE; the default provider discovers - # orgs-with-work via the bootstrap storage (single-ref). A deployment - # may inject its own discovery via ``durable_org_ids_provider`` (e.g. - # enterprise cross-ref fan-out); when None the single-ref default runs. - durable_learning_scheduler = maybe_start_durable_learning( - lambda org_id: RequestContext(org_id=org_id), - bootstrap_org_id=bootstrap_org_id, - org_ids_provider=durable_org_ids_provider, - ) - # Enterprise owns cross-org aggregation through its capability. - # OSS still needs a startup-owned scheduler so pending SQLite work - # drains after a restart even when no new publish arrives. - if capabilities is None: - from reflexio.server.services.playbook.aggregation_scheduler import ( - ensure_local_playbook_aggregation_scheduler, + try: + if mounts_data_plane: + _log_multi_worker_daemons() + log_publish_hardware_capacity() + validate_llm_availability() + # The scheduler discovers every org with resumable work each tick and + # drives a per-org worker with org-scoped claims, so it is not limited + # to the bootstrap org. The bootstrap org is only used to read config + # and to seed cross-org discovery. + bootstrap_org_id = _resolve_lifespan_org_id(get_org_id) + scheduler = maybe_start_resume_scheduler( + lambda org_id: RequestContext(org_id=org_id), + bootstrap_org_id=bootstrap_org_id, + org_id_provider=resume_org_ids_provider, ) - - try: - aggregation_context = RequestContext(org_id=bootstrap_org_id) - aggregation_scheduler = ensure_local_playbook_aggregation_scheduler( - aggregation_context - ) - except Exception: # noqa: BLE001 - # A composition root may intentionally install an enterprise - # configurator while exercising the OSS app shape without - # enterprise deployment settings. The post-persist trigger - # still starts the same scheduler lazily once a real context - # exists. - logger.warning( - "playbook aggregation scheduler startup deferred: " - "bootstrap context is unavailable", - exc_info=True, + # Register the missing-vector backfill sweep (opt-in via + # REFLEXIO_MISSING_VECTOR_BACKFILL_ENABLED) BEFORE starting the GC + # scheduler, so its per-org hook is visible when maybe_start_lineage_gc + # evaluates its start conditions. No-op when the flag is off. + install_missing_vector_backfill_sweep() + gc_scheduler = maybe_start_lineage_gc( + lambda org_id: RequestContext(org_id=org_id), + bootstrap_org_id=bootstrap_org_id, + ) + # Durable learning drains the learning_jobs queue per org. Gated on + # REFLEXIO_DURABLE_LEARNING_QUEUE; the default provider discovers + # orgs-with-work via the bootstrap storage (single-ref). A deployment + # may inject its own discovery via ``durable_org_ids_provider`` (e.g. + # enterprise cross-ref fan-out); when None the single-ref default runs. + durable_learning_scheduler = maybe_start_durable_learning( + lambda org_id: RequestContext(org_id=org_id), + bootstrap_org_id=bootstrap_org_id, + org_ids_provider=durable_org_ids_provider, + ) + # Enterprise owns cross-org aggregation through its capability. + # OSS still needs a startup-owned scheduler so pending SQLite work + # drains after a restart even when no new publish arrives. + if capabilities is None: + from reflexio.server.services.playbook.aggregation_scheduler import ( + ensure_local_playbook_aggregation_scheduler, ) - try: + + try: + aggregation_context = RequestContext(org_id=bootstrap_org_id) + aggregation_scheduler = ( + ensure_local_playbook_aggregation_scheduler( + aggregation_context + ) + ) + except Exception: # noqa: BLE001 + # A composition root may intentionally install an enterprise + # configurator while exercising the OSS app shape without + # enterprise deployment settings. The post-persist trigger + # still starts the same scheduler lazily once a real context + # exists. + logger.warning( + "playbook aggregation scheduler startup deferred: " + "bootstrap context is unavailable", + exc_info=True, + ) if capabilities is not None: ctx = ( app_context_factory() diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py index 9de981eca..4a5cdf6a4 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -110,7 +110,7 @@ def record_learnings_generated( Intended for online extraction paths that have a known billable count but do not retain a complete per-record id list. Resumable finalization must - use :func:`record_learnings_generated_records` and skip items without + use :func:`record_learnings_generated_records_strict` and skip items without durable ids. No-op when ``count <= 0``. Emits a single event carrying ``event_key`` when the caller has a durable @@ -331,7 +331,8 @@ def emit_learnings_generated( product path must never fail because metering failed, so config resolution and emission are wrapped and any exception is logged and swallowed (mirroring the extraction path's ``_record_billing_learning_events``). - Resumable finalization must use :func:`emit_learnings_generated_records`. + Resumable finalization must use + :func:`emit_learnings_generated_records_strict`. No-op when ``count <= 0``. Args: @@ -395,14 +396,14 @@ def emit_learnings_generated_records( ) -> None: """Resolve ``platform_llm`` from config and emit one event per learning id. - Entity-backed counterpart to :func:`emit_learnings_generated`, used by - resumable-extraction finalization for every created user learning with a - durable id. Items without ids are not billable on that path. Online - extraction uses the count-based :func:`record_learnings_generated` helper - because it does not retain a safe 1:1 id per generated unit. Same guard - semantics: config resolution and emission are wrapped and any exception is - logged and swallowed — the product path must never fail because metering - failed. No-op when ``learning_ids`` is empty. + Ordinary fail-open per-record counterpart to + :func:`emit_learnings_generated`. Receipt-backed resumable finalization uses + :func:`emit_learnings_generated_records_strict`; items without durable ids + are not billable on that path. Online extraction uses the count-based + :func:`record_learnings_generated` helper because it does not retain a safe + 1:1 id per generated unit. Config resolution and emission are wrapped and + any exception is logged and swallowed — the product path must never fail + because metering failed. No-op when ``learning_ids`` is empty. Args: org_id: Organisation identifier. diff --git a/reflexio/server/usage_metrics.py b/reflexio/server/usage_metrics.py index 642da6287..d505f0916 100644 --- a/reflexio/server/usage_metrics.py +++ b/reflexio/server/usage_metrics.py @@ -119,8 +119,9 @@ def record_usage_event( ) -> None: """Record one usage event if a recorder is configured. - The product path must never fail because metrics failed, so this function - catches and logs all recorder errors. + The product path must never fail because metrics failed. Missing and legacy + recorders are silent on this ordinary fail-open path; explicit delivery + failures and recorder exceptions are logged. """ try: record_usage_event_strict( @@ -152,8 +153,11 @@ def record_usage_event( metadata=metadata, created_at=created_at, ) - except Exception as exc: # noqa: BLE001 - logger.warning("Usage metrics recorder failed: %s", exc) + except UsageEventDeliveryError as exc: + if exc.status is not UsageEventDeliveryStatus.UNKNOWN: + logger.warning("Usage metrics recorder failed: %s", exc) + except Exception: # noqa: BLE001 + logger.warning("Usage metrics recorder failed", exc_info=True) def record_usage_event_strict( diff --git a/tests/server/test_create_app_capabilities.py b/tests/server/test_create_app_capabilities.py index 45a1e23da..05aa9226b 100644 --- a/tests/server/test_create_app_capabilities.py +++ b/tests/server/test_create_app_capabilities.py @@ -15,8 +15,13 @@ from fastapi import APIRouter from fastapi.testclient import TestClient +from reflexio.server import usage_metrics from reflexio.server.api import create_app from reflexio.server.extensions import AppContext, Capability, CapabilityRegistry +from reflexio.server.usage_metrics import ( + UsageEventDeliveryError, + UsageEventDeliveryStatus, +) class RouterCap(Capability): @@ -92,6 +97,30 @@ def test_capabilities_none_is_unchanged() -> None: assert app is not None +def test_oss_startup_failure_clears_exempt_usage_recorder(monkeypatch) -> None: + usage_metrics.configure_usage_event_recorder(None) + + def fail_startup_guards() -> None: + raise RuntimeError("startup guard failed") + + monkeypatch.setattr( + "reflexio.server.llm.providers.embedder_warmup.run_startup_config_guards", + fail_startup_guards, + ) + + app = create_app(capabilities=None, mount_data_plane=False) + with pytest.raises(RuntimeError, match="startup guard failed"), TestClient(app): + pass + + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + ) + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + + def test_partial_cleanup_invariant() -> None: """Already-started caps must be shut down; caps that never started must not be.""" LifecycleCap.started = LifecycleCap.stopped = False diff --git a/tests/server/test_usage_metrics.py b/tests/server/test_usage_metrics.py index 2992ace87..823503404 100644 --- a/tests/server/test_usage_metrics.py +++ b/tests/server/test_usage_metrics.py @@ -1,3 +1,5 @@ +import logging + import pytest from reflexio.server import usage_metrics @@ -24,6 +26,36 @@ def test_usage_event_carries_event_key(): assert captured and captured[0].event_key == "search:abc" +def test_ordinary_delivery_is_silent_without_recorder(caplog): + usage_metrics.configure_usage_event_recorder(None) + + with caplog.at_level(logging.WARNING, logger=usage_metrics.__name__): + usage_metrics.record_usage_event( + org_id="7", + event_name="search_request", + event_category="application", + ) + + assert caplog.records == [] + + +def test_ordinary_delivery_is_silent_for_legacy_recorder(caplog): + captured = [] + usage_metrics.configure_usage_event_recorder(captured.append) + try: + with caplog.at_level(logging.WARNING, logger=usage_metrics.__name__): + usage_metrics.record_usage_event( + org_id="7", + event_name="search_request", + event_category="application", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert len(captured) == 1 + assert caplog.records == [] + + def test_event_key_defaults_none(): assert UsageEvent(org_id="7", event_name="x", event_category="y").event_key is None From f6d8c0c602cf155428af551af6579b92364972c3 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 09:24:43 +0000 Subject: [PATCH 24/63] docs: clarify receipt billing retry ownership --- .../server/services/deferred_learning_plan.py | 6 ++++-- .../server/services/extraction/resume_worker.py | 17 +++++++++++------ .../services/extraction/test_resume_worker.py | 16 +++++++++++++++- 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/reflexio/server/services/deferred_learning_plan.py b/reflexio/server/services/deferred_learning_plan.py index 7e00c2ad9..e540357e6 100644 --- a/reflexio/server/services/deferred_learning_plan.py +++ b/reflexio/server/services/deferred_learning_plan.py @@ -32,8 +32,10 @@ class FinalizationResult: ``won_receipt`` is true only for the caller whose learning writes and immutable finalization receipt committed together. Receipt-reuse callers - still receive the winner's ordered ids but must not replay billing or other - winner-only side effects. + still receive the winner's ordered ids and may replay retry-safe durable + obligations, such as billing records keyed by those ids. They must not + replay winner-only derived work such as optimization, aggregation, or + tagging. """ learning_ids: list[str] diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index 908d5719e..fc68e8de7 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -16,7 +16,10 @@ from reflexio.models.api_schema.service_schemas import Interaction, Request from reflexio.models.config_schema import PlaybookConfig, ProfileExtractorConfig from reflexio.server.api_endpoints.request_context import RequestContext -from reflexio.server.billing_meter import ReceiptBillingDeliveryError +from reflexio.server.billing_meter import ( + ReceiptBillingDeliveryError, + emit_learnings_generated_records_strict, +) from reflexio.server.error_reporting import error_tags from reflexio.server.llm._litellm_types import ModelProvenance from reflexio.server.llm.litellm_client import LiteLLMClient, LiteLLMConfig @@ -100,7 +103,13 @@ def _finalization_failure_status( next_attempt_count: int, max_finalization_attempts: int, ) -> AgentRunStatus: - """Retry transient receipt delivery failures; terminate permanent rejection.""" + """Classify finalization failures under the receipt-delivery contract. + + Transient receipt delivery failures remain retryable beyond the ordinary + finalization-attempt ceiling because their committed billing obligation must + eventually be delivered. Permanent receipt rejection is terminal, as are + ordinary finalization failures at or above the configured ceiling. + """ if isinstance(exc, ReceiptBillingDeliveryError): if exc.status is UsageEventDeliveryStatus.REJECTED: return AgentRunStatus.FAILED @@ -977,10 +986,6 @@ def _record_finalized_learnings( Emits one event per durable learning id returned by finalization. Per-record keys make finalization retries idempotent downstream. """ - from reflexio.server.billing_meter import ( - emit_learnings_generated_records_strict, - ) - if not learning_ids: return billing_timestamp = run.created_at or run.agent_completed_at diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 7e23160c0..274ad8f4b 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -935,6 +935,7 @@ def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: configure_usage_event_recorder(None) +@pytest.mark.parametrize("next_attempt_count", [1, 3, 4]) @pytest.mark.parametrize( ("delivery_status", "expected_status"), [ @@ -945,6 +946,7 @@ def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: ) def test_receipt_delivery_failure_status_distinguishes_transient_and_permanent( delivery_status, + next_attempt_count, expected_status, ): error = ReceiptBillingDeliveryError(delivery_status) @@ -952,13 +954,25 @@ def test_receipt_delivery_failure_status_distinguishes_transient_and_permanent( assert ( _finalization_failure_status( error, - next_attempt_count=1, + next_attempt_count=next_attempt_count, max_finalization_attempts=3, ) is expected_status ) +@pytest.mark.parametrize("next_attempt_count", [3, 4]) +def test_ordinary_finalization_failure_stops_at_attempt_ceiling(next_attempt_count): + assert ( + _finalization_failure_status( + RuntimeError("ordinary finalization failed"), + next_attempt_count=next_attempt_count, + max_finalization_attempts=3, + ) + is AgentRunStatus.FAILED + ) + + def test_retry_after_billing_reuses_ids_without_replaying_playbook_schedulers( request_context, storage, From af287fae983e17458226d2759245649f1efbc02a Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 10:03:15 +0000 Subject: [PATCH 25/63] fix(storage): align trajectory depth boundary --- .../storage/session_outcome_identity.py | 5 ++- tests/models/test_session_outcome_identity.py | 43 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py index ccf39edf6..36179cee2 100644 --- a/reflexio/server/services/storage/session_outcome_identity.py +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -166,7 +166,10 @@ def outcome_contract_digest( def _canonical_trajectory_json(value: object, *, depth: int = 0) -> str: """Encode trajectory JSON deterministically within a bounded nesting depth.""" - if depth > MAX_CANONICAL_TRAJECTORY_JSON_DEPTH: + if ( + isinstance(value, tuple | list | Mapping) + and depth >= MAX_CANONICAL_TRAJECTORY_JSON_DEPTH + ): raise ValueError("canonical trajectory JSON exceeds maximum depth") if isinstance(value, float): return json.dumps(value, allow_nan=False, separators=(",", ":")) diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index 7a5e68df4..aba7e74d0 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -183,6 +183,49 @@ def test_trajectory_digest_rejects_over_maximum_canonical_json_depth( trajectory_digest(value) +@pytest.mark.parametrize( + ("empty_container", "container_factory"), + [ + ({}, lambda value: {"nested": value}), + ([], lambda value: [value]), + ((), lambda value: (value,)), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_accepts_maximum_empty_container_depth( + empty_container, + container_factory, +) -> None: + value: object = empty_container + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH - 1): + value = container_factory(value) + + assert trajectory_digest(value) + + +@pytest.mark.parametrize( + ("empty_container", "container_factory"), + [ + ({}, lambda value: {"nested": value}), + ([], lambda value: [value]), + ((), lambda value: (value,)), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_rejects_over_maximum_empty_container_depth( + empty_container, + container_factory, +) -> None: + value: object = empty_container + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH): + value = container_factory(value) + + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + trajectory_digest(value) + + def test_session_outcome_record_accepts_unknown_and_serializes_identities() -> None: record = SessionOutcomeRecord( outcome_id="outcome-1", From 343414f7f9c2224dcba8ce421d15f5341dc837a6 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 12:07:15 +0000 Subject: [PATCH 26/63] fix(search): record direct playbook exposures Persist non-empty direct user-playbook search results through the canonical synchronous exposure boundary before metering or response release. --- reflexio/server/README.md | 2 + reflexio/server/routes/search.py | 11 +++ reflexio/server/services/README.md | 1 + reflexio/server/services/search_exposure.py | 2 +- .../routes/test_search_exposure_boundary.py | 91 +++++++++++++++++++ 5 files changed, 106 insertions(+), 1 deletion(-) diff --git a/reflexio/server/README.md b/reflexio/server/README.md index 35707f65c..9eeff5907 100644 --- a/reflexio/server/README.md +++ b/reflexio/server/README.md @@ -100,6 +100,8 @@ Description: FastAPI backend server that processes user interactions to generate **Authentication Pattern**: The open-source app uses `default_get_org_id` and `DEFAULT_ORG_ID` for local/no-auth starts. Enterprise deployments wrap `create_app()` with authenticated org resolution, additional account routers, admin checks, observability hooks, and usage metrics. +**User-playbook exposure boundary**: Every authenticated search boundary that serves user playbooks (`POST /api/search` and `POST /api/search_user_playbooks`) synchronously records the final served set through `services/search_exposure.py` before metering or releasing a successful response. Recorder failures fail closed so a served result cannot escape without exposure evidence; the direct route skips recording when its final result set is empty. + **Pattern**: Core route handlers call `Reflexio` through `get_reflexio(org_id)`; endpoint helper files should not instantiate `Reflexio` directly. ## Extension Registry diff --git a/reflexio/server/routes/search.py b/reflexio/server/routes/search.py index ad098cebe..6834faee0 100644 --- a/reflexio/server/routes/search.py +++ b/reflexio/server/routes/search.py @@ -251,6 +251,17 @@ def search_user_playbooks_endpoint( msg=response.msg, experiment=assignment, ) + if response.user_playbooks: + record_search_exposures( + SearchExposureBatch( + org_id=org_id, + request_id=payload.request_id, + session_id=payload.session_id, + interaction_id=None, + user_id=payload.user_id, + user_playbooks=tuple(response.user_playbooks), + ) + ) enqueue_search_metering( org_id=org_id, caller_type=caller_type, diff --git a/reflexio/server/services/README.md b/reflexio/server/services/README.md index 2e2062aa0..c503670cd 100644 --- a/reflexio/server/services/README.md +++ b/reflexio/server/services/README.md @@ -58,6 +58,7 @@ strings before deleting old import paths in the same PR. | `pre_retrieval/` | `QueryReformulator` (`_query_reformulator.py`) + `DocumentExpander` (`_document_expander.py`) - query rewrite and doc expansion for recall. Compact by design; see [README](pre_retrieval/README.md). | | `tagging/` | `TaggingService` (`service.py`) + deferred `tagging_scheduler.py` - post-generation profile/playbook tagging. Compact by design; see [README](tagging/README.md). | | `unified_search_service.py` | `run_unified_search()` — two-phase parallel search across profiles / agent playbooks / user playbooks. | +| `search_exposure.py` | Optional synchronous recorder contract for the final user-playbook sets served by authenticated unified and direct user-playbook search routes. Both boundaries record before metering/response and fail closed on recorder errors; direct search omits empty batches. | | `retrieval/` | `relevance_floor.py` — result relevance thresholding. `temporal.py` — temporal post-processing driven by reformulation signals: query time windows → per-arm SQL filters, near-duplicate freshness collapse for current-value questions, timestamp ordering for latest-value questions. `user_context_guard.py` — high-precision detection of explicit personalization opt-outs, including Simplified and Traditional Chinese, before user-context retrieval. (Superseded/expired rows are already excluded by storage search SQL.) | ## Persistence & Config diff --git a/reflexio/server/services/search_exposure.py b/reflexio/server/services/search_exposure.py index 131d98631..e2ab47a4c 100644 --- a/reflexio/server/services/search_exposure.py +++ b/reflexio/server/services/search_exposure.py @@ -20,7 +20,7 @@ @dataclass(frozen=True) class SearchExposureBatch: - """The final user-playbook set returned by one unified search.""" + """The final user-playbook set returned by an authenticated search.""" org_id: str request_id: str | None diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py index 7497e362e..cea0a3864 100644 --- a/tests/server/routes/test_search_exposure_boundary.py +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -49,6 +49,7 @@ def _search_results(playbooks: list[UserPlaybook]) -> Iterator[MagicMock]: rehydrated_text=None, ) reflexio.unified_search.return_value = result + reflexio.search_user_playbooks.return_value = result with patch( "reflexio.server.routes.search.reflexio_cache.get_reflexio", return_value=reflexio, @@ -70,8 +71,11 @@ def _client() -> TestClient: class _Recorder: batches: list[Any] = field(default_factory=list) completed: bool = False + order: list[str] | None = None def record(self, batch: Any) -> None: + if self.order is not None: + self.order.append("record") self.batches.append(batch) self.completed = True @@ -141,6 +145,93 @@ def test_no_user_playbook_results_record_one_empty_synchronous_batch() -> None: assert recorder.batches[0].user_playbooks == () +def test_direct_user_playbook_search_records_final_results_before_metering() -> None: + playbooks = [_playbook(21, "Direct first"), _playbook(22, "Direct second")] + order: list[str] = [] + recorder = _Recorder(order=order) + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with ( + _search_results(playbooks), + patch( + "reflexio.server.routes.search._meter_search_request", + side_effect=lambda **_kwargs: order.append("meter_search"), + ), + patch( + "reflexio.server.routes.search._meter_applied_learnings", + side_effect=lambda **_kwargs: order.append("meter_applied"), + ), + ): + response = _client().post( + "/api/search_user_playbooks", + json={ + "query": "direct answer", + "user_id": "user-1", + "request_id": "request-direct-1", + "session_id": "session-direct-1", + }, + ) + + assert response.status_code == 200, response.text + assert order == ["record", "meter_search", "meter_applied"] + assert len(recorder.batches) == 1 + batch = recorder.batches[0] + assert batch.org_id == "org-1" + assert batch.request_id == "request-direct-1" + assert batch.session_id == "session-direct-1" + assert batch.interaction_id is None + assert batch.user_id == "user-1" + assert batch.user_playbooks == tuple(playbooks) + + +def test_direct_user_playbook_recorder_failure_prevents_metering_and_success() -> None: + order: list[str] = [] + + class _FailingRecorder: + def record(self, _batch: Any) -> None: + order.append("record") + raise RuntimeError("ledger unavailable") + + register_service(SEARCH_EXPOSURE_RECORDER, _FailingRecorder()) + + with ( + _search_results([_playbook(21, "Direct first")]), + patch( + "reflexio.server.routes.search._meter_search_request", + side_effect=lambda **_kwargs: order.append("meter_search"), + ), + patch( + "reflexio.server.routes.search._meter_applied_learnings", + side_effect=lambda **_kwargs: order.append("meter_applied"), + ), + ): + response = _client().post( + "/api/search_user_playbooks", + json={ + "query": "direct answer", + "user_id": "user-1", + "request_id": "request-direct-1", + }, + ) + + assert response.status_code == 500 + assert order == ["record"] + + +def test_direct_user_playbook_search_does_not_record_empty_results() -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([]): + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", "user_id": "user-1"}, + ) + + assert response.status_code == 200, response.text + assert recorder.batches == [] + + def test_oss_search_succeeds_when_no_recorder_is_registered() -> None: with _search_results([_playbook(11, "First")]): response = _client().post( From e92bdca08de7e77f3323b6be0ef03678ec7f91f7 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 12:27:40 +0000 Subject: [PATCH 27/63] fix(search): cap direct playbook result batches Align direct user-playbook search admission with the 100-event exposure contract and publish the bound through shared client metadata. --- docs/lib/methods/user-playbooks.ts | 2 +- reflexio/client/client.py | 2 +- .../models/api_schema/retriever_schema.py | 4 +- tests/client/test_search.py | 24 +++++++++-- .../routes/test_search_exposure_boundary.py | 40 +++++++++++++++++++ 5 files changed, 65 insertions(+), 7 deletions(-) diff --git a/docs/lib/methods/user-playbooks.ts b/docs/lib/methods/user-playbooks.ts index 36333c5f5..ef1d55e40 100644 --- a/docs/lib/methods/user-playbooks.ts +++ b/docs/lib/methods/user-playbooks.ts @@ -105,7 +105,7 @@ export const userPlaybookMethods: MethodDef[] = [ type: "number", required: false, default: 10, - description: "Maximum number of results to return", + description: "Maximum number of results to return, from 1 to 100", }, { name: "threshold", diff --git a/reflexio/client/client.py b/reflexio/client/client.py index fe3fd08ea..c2d1fc1b4 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -906,7 +906,7 @@ def search_user_playbooks( end_time (Optional[datetime]): End time for created_at filter status_filter (Optional[list[Optional[Status]]]): Filter by status (None for CURRENT, PENDING, ARCHIVED) tags (Optional[list[str]]): Match playbooks having any of these tags. - top_k (Optional[int]): Maximum number of results to return (default: 10) + top_k (Optional[int]): Maximum results to return, from 1 to 100 (default: 10) threshold (Optional[float]): Similarity threshold for vector search. When omitted, the embedding model's default is used. enable_reformulation (Optional[bool]): Enable LLM query reformulation (default: False) diff --git a/reflexio/models/api_schema/retriever_schema.py b/reflexio/models/api_schema/retriever_schema.py index f1ed46096..46d1ab5da 100644 --- a/reflexio/models/api_schema/retriever_schema.py +++ b/reflexio/models/api_schema/retriever_schema.py @@ -318,7 +318,7 @@ class SearchUserPlaybookRequest(BaseModel): start_time (datetime, optional): Start time for created_at filter end_time (datetime, optional): End time for created_at filter status_filter (list[Optional[Status]], optional): Filter by status (None for CURRENT, PENDING, ARCHIVED) - top_k (int, optional): Maximum number of results to return. Defaults to 10 + top_k (int, optional): Maximum results to return, up to 100. Defaults to 10 threshold (float, optional): Similarity threshold for vector search. When omitted, the embedding model's default is used. """ @@ -331,7 +331,7 @@ class SearchUserPlaybookRequest(BaseModel): end_time: datetime | None = None status_filter: list[Status | None] | None = None tags: list[str] | None = None - top_k: int | None = Field(default=10, gt=0) + top_k: int | None = Field(default=10, gt=0, le=100) threshold: float | None = Field(default=None, ge=0.0, le=1.0) enable_reformulation: bool | None = False search_mode: SearchMode = SearchMode.HYBRID diff --git a/tests/client/test_search.py b/tests/client/test_search.py index e48776b8d..a1e644b98 100644 --- a/tests/client/test_search.py +++ b/tests/client/test_search.py @@ -5,11 +5,16 @@ import pytest from reflexio import ReflexioClient -from reflexio.models.api_schema.retriever_schema import UnifiedSearchRequest +from reflexio.models.api_schema.retriever_schema import ( + SearchUserPlaybookRequest, + UnifiedSearchRequest, +) -def _non_null_schema(field_name: str) -> dict[str, Any]: - field_schema = UnifiedSearchRequest.model_json_schema()["properties"][field_name] +def _non_null_schema( + field_name: str, *, request_model: type[Any] = UnifiedSearchRequest +) -> dict[str, Any]: + field_schema = request_model.model_json_schema()["properties"][field_name] return next( option for option in field_schema.get("anyOf", [field_schema]) @@ -102,3 +107,16 @@ def test_unified_search_docs_track_schema_contract() -> None: assert registry_docs.count(f"at most {identifier_limit} characters") >= 3 assert 'name: "interaction_id"' in registry_docs assert f"positive integer (minimum {interaction_minimum})" in registry_docs + + +def test_user_playbook_search_docs_track_top_k_schema_contract() -> None: + top_k_limit = _non_null_schema("top_k", request_model=SearchUserPlaybookRequest)[ + "maximum" + ] + client_docs = inspect.getdoc(ReflexioClient.search_user_playbooks) or "" + registry_docs = ( + Path(__file__).parents[2] / "docs/lib/methods/user-playbooks.ts" + ).read_text(encoding="utf-8") + + assert f"1 to {top_k_limit}" in client_docs + assert f"1 to {top_k_limit}" in registry_docs diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py index cea0a3864..69dba82a7 100644 --- a/tests/server/routes/test_search_exposure_boundary.py +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -232,6 +232,46 @@ def test_direct_user_playbook_search_does_not_record_empty_results() -> None: assert recorder.batches == [] +def test_direct_user_playbook_search_rejects_101_before_search_or_recording() -> None: + playbooks = [ + _playbook(playbook_id, f"Playbook {playbook_id}") + for playbook_id in range(1, 102) + ] + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results(playbooks) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", "user_id": "user-1", "top_k": 101}, + ) + + assert response.status_code == 422 + reflexio.search_user_playbooks.assert_not_called() + assert recorder.batches == [] + + +def test_direct_user_playbook_search_returns_and_records_exactly_100() -> None: + playbooks = [ + _playbook(playbook_id, f"Playbook {playbook_id}") + for playbook_id in range(1, 101) + ] + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results(playbooks) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", "user_id": "user-1", "top_k": 100}, + ) + + assert response.status_code == 200, response.text + reflexio.search_user_playbooks.assert_called_once() + assert len(response.json()["user_playbooks"]) == 100 + assert len(recorder.batches) == 1 + assert recorder.batches[0].user_playbooks == tuple(playbooks) + + def test_oss_search_succeeds_when_no_recorder_is_registered() -> None: with _search_results([_playbook(11, "First")]): response = _client().post( From 064bbc6d7f7cae98f561d33a85e7c5b294837d6e Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 12:51:11 +0000 Subject: [PATCH 28/63] docs: scope search exposure guarantee Clarify that authenticated enterprise capability wiring supplies the fail-closed recorder while shared create_app intentionally has no default implementation. --- reflexio/server/README.md | 2 +- reflexio/server/services/README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/reflexio/server/README.md b/reflexio/server/README.md index 9eeff5907..b4b8d9572 100644 --- a/reflexio/server/README.md +++ b/reflexio/server/README.md @@ -100,7 +100,7 @@ Description: FastAPI backend server that processes user interactions to generate **Authentication Pattern**: The open-source app uses `default_get_org_id` and `DEFAULT_ORG_ID` for local/no-auth starts. Enterprise deployments wrap `create_app()` with authenticated org resolution, additional account routers, admin checks, observability hooks, and usage metrics. -**User-playbook exposure boundary**: Every authenticated search boundary that serves user playbooks (`POST /api/search` and `POST /api/search_user_playbooks`) synchronously records the final served set through `services/search_exposure.py` before metering or releasing a successful response. Recorder failures fail closed so a served result cannot escape without exposure evidence; the direct route skips recording when its final result set is empty. +**User-playbook exposure boundary**: Authenticated enterprise deployments install `ExposureLedgerRecorder` through their capability registry. In those deployments, `POST /api/search` and `POST /api/search_user_playbooks` synchronously record the final served user-playbook set through `services/search_exposure.py` before metering or releasing success, and recorder failures fail closed. Shared `create_app()` intentionally installs no default recorder; local/no-auth OSS and custom shared app constructions persist exposures only when their deployment registers `SEARCH_EXPOSURE_RECORDER`. The direct route skips recording when its final result set is empty. **Pattern**: Core route handlers call `Reflexio` through `get_reflexio(org_id)`; endpoint helper files should not instantiate `Reflexio` directly. diff --git a/reflexio/server/services/README.md b/reflexio/server/services/README.md index c503670cd..cce700102 100644 --- a/reflexio/server/services/README.md +++ b/reflexio/server/services/README.md @@ -58,7 +58,7 @@ strings before deleting old import paths in the same PR. | `pre_retrieval/` | `QueryReformulator` (`_query_reformulator.py`) + `DocumentExpander` (`_document_expander.py`) - query rewrite and doc expansion for recall. Compact by design; see [README](pre_retrieval/README.md). | | `tagging/` | `TaggingService` (`service.py`) + deferred `tagging_scheduler.py` - post-generation profile/playbook tagging. Compact by design; see [README](tagging/README.md). | | `unified_search_service.py` | `run_unified_search()` — two-phase parallel search across profiles / agent playbooks / user playbooks. | -| `search_exposure.py` | Optional synchronous recorder contract for the final user-playbook sets served by authenticated unified and direct user-playbook search routes. Both boundaries record before metering/response and fail closed on recorder errors; direct search omits empty batches. | +| `search_exposure.py` | Optional synchronous recorder contract for final user-playbook result sets. Enterprise capability registration installs the recorder and makes authenticated unified/direct search fail closed before metering/response. Shared `create_app()` has no default recorder; other constructions persist exposures only when they register one. Direct search omits empty batches. | | `retrieval/` | `relevance_floor.py` — result relevance thresholding. `temporal.py` — temporal post-processing driven by reformulation signals: query time windows → per-arm SQL filters, near-duplicate freshness collapse for current-value questions, timestamp ordering for latest-value questions. `user_context_guard.py` — high-precision detection of explicit personalization opt-outs, including Simplified and Traditional Chinese, before user-context retrieval. (Superseded/expired rows are already excluded by storage search SQL.) | ## Persistence & Config From cf20c2a4324511b8b7c2d8dc981db69e8ed360a5 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 20:54:48 +0000 Subject: [PATCH 29/63] fix(outcomes): constrain source labels Require non-empty session outcome producer labels to use the approved lowercase machine-label contract while preserving empty and optional source semantics. --- reflexio/client/client.py | 4 +- reflexio/models/api_schema/domain/entities.py | 14 +-- reflexio/models/api_schema/validators.py | 16 +++ .../api_schema/test_input_size_bounds.py | 11 +- tests/models/test_session_outcome_identity.py | 101 +++++++++++++++++- 5 files changed, 133 insertions(+), 13 deletions(-) diff --git a/reflexio/client/client.py b/reflexio/client/client.py index c2d1fc1b4..272efe918 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -528,7 +528,9 @@ def publish_interaction( Args: user_id: The user ID. interactions: List of interaction data. - source: The source of the interaction. + source: Non-sensitive producer/workflow label. A non-empty value + must match ``^[a-z0-9][a-z0-9._:-]{0,127}$`` and must not + contain user identifiers or PII. agent_version: The agent version. session_id: Required non-empty session ID for grouping requests. wait_for_response: If True, the **server** waits for diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 1f0971704..9747b705e 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -26,6 +26,7 @@ from ..validators import ( EmbeddingVector, NonEmptyStr, + SessionOutcomeSource, TimeRangeValidatorMixin, _validate_image_url, ) @@ -278,7 +279,8 @@ class Request(BaseModel): user_id (str): Owner of the request. created_at (int): Unix epoch seconds at request creation. Defaults to the current UTC time. - source (str): Free-form origin tag (integration name, etc.). + source (str): Non-sensitive producer/workflow label. Non-empty values + use the session outcome source contract. agent_version (str): The agent version that handled this request. session_id (str): Non-empty session this request belongs to. evaluation_only (bool): Whether this request is stored for @@ -293,7 +295,7 @@ class Request(BaseModel): request_id: str user_id: str created_at: int = Field(default_factory=lambda: int(datetime.now(UTC).timestamp())) - source: str = "" + source: SessionOutcomeSource = "" agent_version: str = "" session_id: NonEmptyStr evaluation_only: bool = False @@ -833,7 +835,7 @@ class SessionOutcomeRecord(BaseModel): session_id: NonEmptyStr outcome: SessionOutcomeKind occurred_at: int = Field(ge=0) - source: str + source: SessionOutcomeSource label: str | None = Field(default=None, max_length=128) value: float | None = Field(default=None, allow_inf_nan=False) metadata: dict[str, Any] | None = None @@ -910,7 +912,7 @@ class SetSessionOutcomeResponse(BaseModel): reason: SessionOutcomeFailureReason | None = None message: str = "" user_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None outcome_id: NonEmptyStr | None = None outcome_revision: int | None = Field(default=None, ge=1) outcome_contract_digest: Sha256Digest | None = None @@ -945,7 +947,7 @@ def validate_identity_shape(self) -> Self: class GetSessionOutcomesRequest(CapturesUnknownFields): session_ids: list[NonEmptyStr] | None = Field(default=None, max_length=100) user_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None outcome: SessionOutcomeKind | None = None label: str | None = None start_time: int | None = Field(default=None, ge=0) @@ -1242,7 +1244,7 @@ class PublishUserInteractionRequest(CapturesUnknownFields): request_id: NonEmptyStr | None = None user_id: NonEmptyStr interaction_data_list: list[InteractionData] = Field(min_length=1, max_length=1_000) - source: str = Field(default="", max_length=1_000) + source: SessionOutcomeSource = "" # this is used for aggregating interactions for generating agent playbooks agent_version: str = Field(default="", max_length=1_000) session_id: NonEmptyStr # used for grouping requests together diff --git a/reflexio/models/api_schema/validators.py b/reflexio/models/api_schema/validators.py index 50bf9d9d6..994365b68 100644 --- a/reflexio/models/api_schema/validators.py +++ b/reflexio/models/api_schema/validators.py @@ -104,6 +104,22 @@ def _check_embedding_dimensions(v: list[float]) -> list[float]: """Embedding vector that must be either empty or exactly EMBEDDING_DIMENSIONS (512) floats.""" +SESSION_OUTCOME_SOURCE_PATTERN = r"^[a-z0-9][a-z0-9._:-]{0,127}$" + + +def _check_session_outcome_source(v: str) -> str: + """Validate a non-empty source as a non-sensitive machine label.""" + if v and re.fullmatch(SESSION_OUTCOME_SOURCE_PATTERN, v, flags=re.ASCII) is None: + raise ValueError( + f"non-empty outcome source must match {SESSION_OUTCOME_SOURCE_PATTERN}" + ) + return v + + +SessionOutcomeSource = Annotated[str, AfterValidator(_check_session_outcome_source)] +"""Outcome producer/workflow label; empty preserves the existing absent-source value.""" + + # ============================================================================= # Security Validators — SSRF Prevention # ============================================================================= diff --git a/tests/models/api_schema/test_input_size_bounds.py b/tests/models/api_schema/test_input_size_bounds.py index 911294744..ee35fe46c 100644 --- a/tests/models/api_schema/test_input_size_bounds.py +++ b/tests/models/api_schema/test_input_size_bounds.py @@ -138,10 +138,11 @@ def test_interaction_data_nested_list_bounds() -> None: def test_publish_request_string_bounds() -> None: - cap = 1_000 - _publish(source="s" * cap) - _publish(agent_version="v" * cap) + source_cap = 128 + agent_version_cap = 1_000 + _publish(source="s" * source_cap) + _publish(agent_version="v" * agent_version_cap) with pytest.raises(ValidationError): - _publish(source="s" * (cap + 1)) + _publish(source="s" * (source_cap + 1)) with pytest.raises(ValidationError): - _publish(agent_version="v" * (cap + 1)) + _publish(agent_version="v" * (agent_version_cap + 1)) diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index aba7e74d0..d948d4e93 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -1,13 +1,18 @@ +from collections.abc import Callable from datetime import UTC, datetime from hashlib import sha256 from itertools import combinations -from typing import Any +from typing import Any, Protocol import pytest from pydantic import ValidationError from reflexio.models.api_schema.domain.entities import ( + GetSessionOutcomesRequest, GetSessionOutcomesResponse, + InteractionData, + PublishUserInteractionRequest, + Request, SessionOutcomeRecord, SetSessionOutcomeResponse, ) @@ -45,6 +50,100 @@ def test_outcome_contract_digest_changes_when_source_changes() -> None: ) != _outcome_contract_digest(source="customer_batch") +class _HasSource(Protocol): + @property + def source(self) -> str | None: ... + + +_OUTCOME_SOURCE_MODEL_FACTORIES: tuple[Callable[[str], _HasSource], ...] = ( + lambda source: Request( + request_id="request-1", + user_id="user-1", + session_id="session-1", + source=source, + ), + lambda source: PublishUserInteractionRequest( + user_id="user-1", + session_id="session-1", + interaction_data_list=[InteractionData(content="hello")], + source=source, + ), + lambda source: SessionOutcomeRecord( + user_id="user-1", + session_id="session-1", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source=source, + created_at=2, + ), + lambda source: SetSessionOutcomeResponse(success=True, source=source), + lambda source: GetSessionOutcomesRequest(source=source), +) +_OUTCOME_SOURCE_MODEL_IDS = ("request", "publish", "record", "set-response", "get") + + +@pytest.mark.parametrize( + "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS +) +def test_outcome_source_models_accept_machine_label( + factory: Callable[[str], _HasSource], +) -> None: + assert factory("support-agent:v2").source == "support-agent:v2" + assert factory("a" * 128).source == "a" * 128 + + +@pytest.mark.parametrize( + "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS +) +@pytest.mark.parametrize( + "source", + [ + "alice@example.com", + "https://example.com/hook", + "support agent", + " support-agent", + "support/agent", + "Support-Agent", + "support-agént", + "a" * 129, + ], +) +def test_outcome_source_models_reject_sensitive_or_free_form_values( + factory: Callable[[str], _HasSource], + source: str, +) -> None: + with pytest.raises(ValidationError, match="outcome source"): + factory(source) + + +@pytest.mark.parametrize( + "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS +) +def test_outcome_source_models_preserve_empty_source( + factory: Callable[[str], _HasSource], +) -> None: + assert factory("").source == "" + + +def test_optional_outcome_source_models_preserve_absence() -> None: + assert SetSessionOutcomeResponse(success=False).source is None + assert GetSessionOutcomesRequest().source is None + + +def test_outcome_contract_digest_is_stable_for_valid_machine_label() -> None: + payload = { + "allowed_values": ["failure", "success", "unknown"], + "finalization_rule": "first_write", + "schema_version": 1, + "source": "support-agent:v2", + } + + assert ( + _outcome_contract_digest(source="support-agent:v2") + == sha256(canonical_json_bytes(payload)).hexdigest() + ) + + def test_outcome_contract_digest_changes_when_schema_version_changes() -> None: assert _outcome_contract_digest(schema_version=1) != _outcome_contract_digest( schema_version=2 From 923534cf2094eef00dd83e8a4ae78bed868cbbd8 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 21:26:18 +0000 Subject: [PATCH 30/63] fix(outcomes): apply source contract to filters Use the shared request-source authority across read and generation filters, and derive runtime plus JSON Schema constraints from one constrained type. --- reflexio/cli/README.md | 2 +- reflexio/models/api_schema/domain/entities.py | 8 +-- .../models/api_schema/retriever_schema.py | 7 +-- reflexio/models/api_schema/validators.py | 22 ++++---- tests/models/test_session_outcome_identity.py | 52 ++++++++++++++++++- 5 files changed, 68 insertions(+), 23 deletions(-) diff --git a/reflexio/cli/README.md b/reflexio/cli/README.md index c0680c50a..eb420b8b8 100644 --- a/reflexio/cli/README.md +++ b/reflexio/cli/README.md @@ -144,7 +144,7 @@ Apply to all three modes: | `--wait` | Block until server-side extraction finishes (returns real counts). | | `--session-id` | Required unless each payload includes `session_id`; group publishes into one session. | | `--agent-version` | Tag the interaction with an agent version (used by playbook filtering). | -| `--source` | Free-form source tag (defaults to `cli`). | +| `--source` | Non-sensitive producer/workflow label matching `^[a-z0-9][a-z0-9._:-]{0,127}$`; never use user identifiers or PII (defaults to `cli`). | | `--skip-aggregation`| Extract profiles/playbooks but skip playbook aggregation. | | `--force-extraction`| Bypass all extraction gates (`stride_size`, cheap pre-filter, LLM `should_run`) and always run extractors. | | `--evaluation-only` | Store the request for session-level evaluation only; requires `--session-id` and skips profile/playbook extraction. | diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 9747b705e..35403f647 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -1732,7 +1732,7 @@ class RerunProfileGenerationRequest(BaseModel): user_id: str | None = None start_time: datetime | None = None end_time: datetime | None = None - source: str | None = None + source: SessionOutcomeSource | None = None extractor_names: list[str] | None = ( None # Deprecated compatibility field; ignored for selection. ) @@ -1759,7 +1759,7 @@ class ManualProfileGenerationRequest(BaseModel): """ user_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None extractor_names: list[str] | None = None @@ -1779,7 +1779,7 @@ class ManualPlaybookGenerationRequest(BaseModel): """ agent_version: str = DEFAULT_AGENT_VERSION - source: str | None = None + source: SessionOutcomeSource | None = None playbook_name: str | None = ( None # Deprecated compatibility field; ignored for selection. ) @@ -1805,7 +1805,7 @@ class RerunPlaybookGenerationRequest(BaseModel): playbook_name: str | None = ( None # Deprecated compatibility field; ignored for selection. ) - source: str | None = None + source: SessionOutcomeSource | None = None @field_validator("agent_version") @classmethod diff --git a/reflexio/models/api_schema/retriever_schema.py b/reflexio/models/api_schema/retriever_schema.py index 46d1ab5da..f1d048a50 100644 --- a/reflexio/models/api_schema/retriever_schema.py +++ b/reflexio/models/api_schema/retriever_schema.py @@ -29,6 +29,7 @@ ) from .validators import ( NonEmptyStr, + SessionOutcomeSource, TimeRangeValidatorMixin, ) @@ -63,7 +64,7 @@ class SearchUserProfileRequest(BaseModel): start_time: datetime | None = None end_time: datetime | None = None top_k: int | None = Field(default=10, gt=0) - source: str | None = None + source: SessionOutcomeSource | None = None custom_feature: str | None = None extractor_name: str | None = ( None # Deprecated compatibility field; accepted but ignored. @@ -196,7 +197,7 @@ class GetUserProfilesRequest(BaseModel): start_time: datetime | None = None end_time: datetime | None = None top_k: int | None = Field(default=30, gt=0) - source: str | None = None + source: SessionOutcomeSource | None = None profile_time_to_live: str | None = None status_filter: list[Status | None] | None = None tags: list[str] | None = None @@ -478,7 +479,7 @@ class GetRequestsRequest(BaseModel): user_id: str | None = None request_id: str | None = None session_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None start_time: datetime | None = None end_time: datetime | None = None top_k: int | None = Field( diff --git a/reflexio/models/api_schema/validators.py b/reflexio/models/api_schema/validators.py index 994365b68..f4a875f53 100644 --- a/reflexio/models/api_schema/validators.py +++ b/reflexio/models/api_schema/validators.py @@ -18,10 +18,10 @@ import os import re import socket -from typing import Annotated, Any +from typing import Annotated, Any, Literal from urllib.parse import urlparse -from pydantic import AfterValidator, HttpUrl +from pydantic import AfterValidator, HttpUrl, StringConstraints # Embedding vector dimensions — must match config_schema.EMBEDDING_DIMENSIONS. # Duplicated here to avoid circular imports (config_schema imports from this module). @@ -106,17 +106,13 @@ def _check_embedding_dimensions(v: list[float]) -> list[float]: SESSION_OUTCOME_SOURCE_PATTERN = r"^[a-z0-9][a-z0-9._:-]{0,127}$" - -def _check_session_outcome_source(v: str) -> str: - """Validate a non-empty source as a non-sensitive machine label.""" - if v and re.fullmatch(SESSION_OUTCOME_SOURCE_PATTERN, v, flags=re.ASCII) is None: - raise ValueError( - f"non-empty outcome source must match {SESSION_OUTCOME_SOURCE_PATTERN}" - ) - return v - - -SessionOutcomeSource = Annotated[str, AfterValidator(_check_session_outcome_source)] +SessionOutcomeSource = ( + Literal[""] + | Annotated[ + str, + StringConstraints(max_length=128, pattern=SESSION_OUTCOME_SOURCE_PATTERN), + ] +) """Outcome producer/workflow label; empty preserves the existing absent-source value.""" diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index d948d4e93..cf4bed58b 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -11,12 +11,21 @@ GetSessionOutcomesRequest, GetSessionOutcomesResponse, InteractionData, + ManualPlaybookGenerationRequest, + ManualProfileGenerationRequest, PublishUserInteractionRequest, Request, + RerunPlaybookGenerationRequest, + RerunProfileGenerationRequest, SessionOutcomeRecord, SetSessionOutcomeResponse, ) from reflexio.models.api_schema.domain.enums import SessionOutcomeKind +from reflexio.models.api_schema.retriever_schema import ( + GetRequestsRequest, + GetUserProfilesRequest, + SearchUserProfileRequest, +) from reflexio.server.services.storage import session_outcome_identity from reflexio.server.services.storage.session_outcome_identity import ( canonical_json_bytes, @@ -78,8 +87,28 @@ def source(self) -> str | None: ... ), lambda source: SetSessionOutcomeResponse(success=True, source=source), lambda source: GetSessionOutcomesRequest(source=source), + lambda source: GetRequestsRequest(source=source), + lambda source: SearchUserProfileRequest(user_id="user-1", source=source), + lambda source: GetUserProfilesRequest(user_id="user-1", source=source), + lambda source: RerunProfileGenerationRequest(source=source), + lambda source: ManualProfileGenerationRequest(source=source), + lambda source: ManualPlaybookGenerationRequest(source=source), + lambda source: RerunPlaybookGenerationRequest(source=source), +) +_OUTCOME_SOURCE_MODEL_IDS = ( + "request", + "publish", + "record", + "set-response", + "get-outcomes", + "get-requests", + "search-profiles", + "get-profiles", + "rerun-profiles", + "manual-profiles", + "manual-playbooks", + "rerun-playbooks", ) -_OUTCOME_SOURCE_MODEL_IDS = ("request", "publish", "record", "set-response", "get") @pytest.mark.parametrize( @@ -112,7 +141,7 @@ def test_outcome_source_models_reject_sensitive_or_free_form_values( factory: Callable[[str], _HasSource], source: str, ) -> None: - with pytest.raises(ValidationError, match="outcome source"): + with pytest.raises(ValidationError): factory(source) @@ -128,6 +157,25 @@ def test_outcome_source_models_preserve_empty_source( def test_optional_outcome_source_models_preserve_absence() -> None: assert SetSessionOutcomeResponse(success=False).source is None assert GetSessionOutcomesRequest().source is None + assert GetRequestsRequest().source is None + assert SearchUserProfileRequest(user_id="user-1").source is None + assert GetUserProfilesRequest(user_id="user-1").source is None + assert RerunProfileGenerationRequest().source is None + assert ManualProfileGenerationRequest().source is None + assert ManualPlaybookGenerationRequest().source is None + assert RerunPlaybookGenerationRequest().source is None + + +def test_outcome_source_json_schema_exposes_exact_constraints() -> None: + source_schema = PublishUserInteractionRequest.model_json_schema()["properties"][ + "source" + ] + + assert { + "maxLength": 128, + "pattern": "^[a-z0-9][a-z0-9._:-]{0,127}$", + "type": "string", + } in source_schema["anyOf"] def test_outcome_contract_digest_is_stable_for_valid_machine_label() -> None: From 1c3c1446701633597b0333e2fd0c36f76f021211 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 5 Aug 2026 23:51:28 +0000 Subject: [PATCH 31/63] docs: define source label contract --- AI_AGENT_INTEGRATION.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/AI_AGENT_INTEGRATION.md b/AI_AGENT_INTEGRATION.md index 4807eb2aa..5952c0a85 100644 --- a/AI_AGENT_INTEGRATION.md +++ b/AI_AGENT_INTEGRATION.md @@ -95,7 +95,7 @@ learns, what stays private to a user scope, and what transfers to other users. | `user_id` | Scope for profiles and user playbooks | Use the human user, tenant, workspace, repo, or project whose preferences should be isolated. For example, use a project id when repo-specific rules should not leak into unrelated repos. | | `agent_version` | Scope for shared agent playbooks | Use a stable agent name plus major behavior version, for example `my-agent-v1`. Keep it stable if learnings should transfer across users/projects. If you omit it, the SDK uses `DEFAULT_AGENT_VERSION` (`"agent-v0"`) — fine for a single agent, but set an explicit value before you run more than one. | | `session_id` | Group turns for one conversation | Use the host session/conversation id. Generate a UUID if the host does not provide one. | -| `source` | Audit label | Use the integration name, for example `my-agent-plugin`. | +| `source` | Producer/workflow label | Use a non-sensitive machine label such as `support-agent:v2`. Where the API permits an empty value, empty means the source is absent. Every non-empty value must match `^[a-z0-9][a-z0-9._:-]{0,127}$` and is limited to 128 ASCII characters. Do not include user identifiers, email addresses, or other PII. | `user_id` and `agent_version` work together: @@ -234,7 +234,7 @@ Publish request fields: | --- | --- | --- | --- | | `user_id` | Yes | The user, tenant, workspace, repo, or project scope whose private profiles and user playbooks should be isolated. | `"alice"`, `"tenant-acme"`, `"repo-reflexio"` | | `interactions` / `interaction_data_list` | Yes | Ordered conversation turns to publish. Include at least one turn; multi-turn correction examples are best for learning. | `[{"role": "User", "content": "Use pnpm here."}, {"role": "Assistant", "content": "Got it, I will use pnpm."}]` | -| `source` | No | Integration label for debugging and filtering. Use a stable name for the plugin, framework, or adapter. | `"my-agent-plugin"`, `"vscode-assistant"`, `"support-chatbot"` | +| `source` | No | Non-sensitive producer/workflow machine label for debugging and filtering. Empty means the source is absent where the API permits it. Every non-empty value must match `^[a-z0-9][a-z0-9._:-]{0,127}$` and is limited to 128 ASCII characters. Do not include user identifiers, email addresses, or other PII. | `"support-agent:v2"` | | `agent_version` | Strongly recommended | The shared-agent learning boundary. Use the same value when playbooks should transfer across users. Change it when old playbooks should not transfer. | `"support-agent-v1"`, `"coding-agent-2026-05"` | | `session_id` | Yes | Host conversation/session id. Generate a UUID if the host has no session id, and reuse it for all turns in that conversation. | `"sess_01HX8Y..."`, `"3f02b7f8-..."` | | `skip_aggregation` | No | `False` when user playbooks should be eligible to roll up into shared agent playbooks. `True` when you want user-level extraction only. | `false` | @@ -315,7 +315,7 @@ def publish_turns( response = client.publish_interaction( user_id=user_id, interactions=interactions, - source="my-agent-plugin", + source="support-agent:v2", agent_version=agent_version, session_id=session_id, wait_for_response=False, From d84e1b7d24fadcaa0857f61d1ef95027be6f0036 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Thu, 6 Aug 2026 02:30:55 +0000 Subject: [PATCH 32/63] fix: serialize purge heartbeat renewals --- .../server/services/governance/service.py | 34 ++++---- .../governance/test_governance_local_e2e.py | 80 +++++++++++++++++++ 2 files changed, 98 insertions(+), 16 deletions(-) diff --git a/reflexio/server/services/governance/service.py b/reflexio/server/services/governance/service.py index ca7db7177..5b60a9be4 100644 --- a/reflexio/server/services/governance/service.py +++ b/reflexio/server/services/governance/service.py @@ -85,6 +85,7 @@ def __init__( self._purge_id = purge_id self._claim = execution_claim self._lock = threading.Lock() + self._renewal_lock = threading.Lock() self._renewal_error: Exception | None = None self._stop = threading.Event() self._thread = threading.Thread(target=self._run, daemon=True) @@ -107,23 +108,24 @@ def claim(self) -> PurgeExecutionClaim: return self._claim def renew_now(self) -> PurgeExecutionClaim: - try: - renewed = self._storage.renew_purge_operation_execution_claim( - self._purge_id, - self.claim(), - lease_ttl_seconds=_PURGE_EXECUTION_LEASE_SECONDS, - ) - except _PurgeExecutionHeartbeatLostError: - raise - except Exception as exc: + with self._renewal_lock: + try: + renewed = self._storage.renew_purge_operation_execution_claim( + self._purge_id, + self.claim(), + lease_ttl_seconds=_PURGE_EXECUTION_LEASE_SECONDS, + ) + except _PurgeExecutionHeartbeatLostError: + raise + except Exception as exc: + with self._lock: + self._renewal_error = exc + raise _PurgeExecutionHeartbeatLostError( + "purge execution heartbeat renewal was lost" + ) from exc with self._lock: - self._renewal_error = exc - raise _PurgeExecutionHeartbeatLostError( - "purge execution heartbeat renewal was lost" - ) from exc - with self._lock: - self._claim = renewed - return renewed + self._claim = renewed + return renewed def _run(self) -> None: while not self._stop.wait(_PURGE_EXECUTION_HEARTBEAT_SECONDS): diff --git a/tests/server/services/governance/test_governance_local_e2e.py b/tests/server/services/governance/test_governance_local_e2e.py index aebbb587d..eb68e97b3 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -142,6 +142,86 @@ def _insert_session_outcome( storage.conn.commit() +def test_purge_execution_heartbeat_serializes_concurrent_renewals() -> None: + initial_claim = PurgeExecutionClaim( + purge_id="purge-1", + owner="owner-1", + fence=1, + expires_at=300, + ) + first_renewed_claim = PurgeExecutionClaim( + purge_id="purge-1", + owner="owner-1", + fence=1, + expires_at=400, + ) + second_renewed_claim = PurgeExecutionClaim( + purge_id="purge-1", + owner="owner-1", + fence=1, + expires_at=401, + ) + + class BlockingRenewalStorage: + def __init__(self) -> None: + self.claims: list[PurgeExecutionClaim] = [] + self._calls_lock = threading.Lock() + self.first_renewal_started = threading.Event() + self.release_first_renewal = threading.Event() + self.second_renewal_started = threading.Event() + + def renew_purge_operation_execution_claim( + self, + _purge_id: str, + claim: PurgeExecutionClaim, + *, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim: + assert lease_ttl_seconds == 300 + with self._calls_lock: + self.claims.append(claim) + call_number = len(self.claims) + if call_number == 1: + self.first_renewal_started.set() + assert self.release_first_renewal.wait(timeout=5) + return first_renewed_claim + self.second_renewal_started.set() + return second_renewed_claim + + storage = BlockingRenewalStorage() + heartbeat = governance_service_module._PurgeExecutionHeartbeat( + storage=storage, + purge_id="purge-1", + execution_claim=initial_claim, + ) + errors: list[BaseException] = [] + + def renew() -> None: + try: + heartbeat.renew_now() + except BaseException as exc: + errors.append(exc) + + first = threading.Thread(target=renew) + second = threading.Thread(target=renew) + first.start() + assert storage.first_renewal_started.wait(timeout=5) + second.start() + second_started_while_first_was_blocked = storage.second_renewal_started.wait( + timeout=0.2 + ) + storage.release_first_renewal.set() + first.join(timeout=5) + second.join(timeout=5) + + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert not second_started_while_first_was_blocked + assert storage.claims == [initial_claim, first_renewed_claim] + assert heartbeat.claim() == second_renewed_claim + + @pytest.fixture def storage( tmp_path: Path, From 64b5ac273c76e509d621d6d40071343263f1f644 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Thu, 6 Aug 2026 04:38:27 +0000 Subject: [PATCH 33/63] test: make heartbeat renewal regression deterministic --- .../governance/test_governance_local_e2e.py | 28 ++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/tests/server/services/governance/test_governance_local_e2e.py b/tests/server/services/governance/test_governance_local_e2e.py index eb68e97b3..c32099686 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -6,6 +6,7 @@ from datetime import UTC, datetime from pathlib import Path from types import SimpleNamespace +from typing import Any from unittest.mock import patch import pytest @@ -188,12 +189,33 @@ def renew_purge_operation_execution_claim( self.second_renewal_started.set() return second_renewed_claim + class TrackingRenewalLock: + def __init__(self) -> None: + self._lock = threading.Lock() + self._attempts_lock = threading.Lock() + self._attempts = 0 + self.second_renewal_attempted = threading.Event() + + def __enter__(self) -> TrackingRenewalLock: + with self._attempts_lock: + self._attempts += 1 + if self._attempts == 2: + self.second_renewal_attempted.set() + self._lock.acquire() + return self + + def __exit__(self, *_exc: object) -> None: + self._lock.release() + storage = BlockingRenewalStorage() heartbeat = governance_service_module._PurgeExecutionHeartbeat( storage=storage, purge_id="purge-1", execution_claim=initial_claim, ) + renewal_lock = TrackingRenewalLock() + cast_heartbeat: Any = heartbeat + cast_heartbeat._renewal_lock = renewal_lock errors: list[BaseException] = [] def renew() -> None: @@ -207,9 +229,8 @@ def renew() -> None: first.start() assert storage.first_renewal_started.wait(timeout=5) second.start() - second_started_while_first_was_blocked = storage.second_renewal_started.wait( - timeout=0.2 - ) + assert renewal_lock.second_renewal_attempted.wait(timeout=5) + assert not storage.second_renewal_started.is_set() storage.release_first_renewal.set() first.join(timeout=5) second.join(timeout=5) @@ -217,7 +238,6 @@ def renew() -> None: assert not first.is_alive() assert not second.is_alive() assert errors == [] - assert not second_started_while_first_was_blocked assert storage.claims == [initial_claim, first_renewed_claim] assert heartbeat.claim() == second_renewed_claim From 1435c630a7140d3bb8558567d463fb0cf6bd9b8d Mon Sep 17 00:00:00 2001 From: Guangyu Date: Thu, 6 Aug 2026 04:54:27 +0000 Subject: [PATCH 34/63] test: prove heartbeat renewal overlap deterministically --- .../governance/test_governance_local_e2e.py | 52 ++++++++++++++----- 1 file changed, 40 insertions(+), 12 deletions(-) diff --git a/tests/server/services/governance/test_governance_local_e2e.py b/tests/server/services/governance/test_governance_local_e2e.py index c32099686..53791890b 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -163,8 +163,22 @@ def test_purge_execution_heartbeat_serializes_concurrent_renewals() -> None: expires_at=401, ) - class BlockingRenewalStorage: + class RenewalProgress: def __init__(self) -> None: + self.event = threading.Event() + self._lock = threading.Lock() + self.source: str | None = None + + def record(self, source: str) -> None: + with self._lock: + if self.event.is_set(): + return + self.source = source + self.event.set() + + class BlockingRenewalStorage: + def __init__(self, progress: RenewalProgress) -> None: + self._progress = progress self.claims: list[PurgeExecutionClaim] = [] self._calls_lock = threading.Lock() self.first_renewal_started = threading.Event() @@ -186,11 +200,13 @@ def renew_purge_operation_execution_claim( self.first_renewal_started.set() assert self.release_first_renewal.wait(timeout=5) return first_renewed_claim + self._progress.record("storage") self.second_renewal_started.set() return second_renewed_claim class TrackingRenewalLock: - def __init__(self) -> None: + def __init__(self, progress: RenewalProgress) -> None: + self._progress = progress self._lock = threading.Lock() self._attempts_lock = threading.Lock() self._attempts = 0 @@ -200,20 +216,21 @@ def __enter__(self) -> TrackingRenewalLock: with self._attempts_lock: self._attempts += 1 if self._attempts == 2: - self.second_renewal_attempted.set() + self._progress.record("renewal lock") self._lock.acquire() return self def __exit__(self, *_exc: object) -> None: self._lock.release() - storage = BlockingRenewalStorage() + progress = RenewalProgress() + storage = BlockingRenewalStorage(progress) heartbeat = governance_service_module._PurgeExecutionHeartbeat( storage=storage, purge_id="purge-1", execution_claim=initial_claim, ) - renewal_lock = TrackingRenewalLock() + renewal_lock = TrackingRenewalLock(progress) cast_heartbeat: Any = heartbeat cast_heartbeat._renewal_lock = renewal_lock errors: list[BaseException] = [] @@ -224,16 +241,27 @@ def renew() -> None: except BaseException as exc: errors.append(exc) + allow_second_renewal = threading.Event() + + def renew_second() -> None: + assert allow_second_renewal.wait(timeout=5) + renew() + first = threading.Thread(target=renew) - second = threading.Thread(target=renew) + second = threading.Thread(target=renew_second) first.start() - assert storage.first_renewal_started.wait(timeout=5) second.start() - assert renewal_lock.second_renewal_attempted.wait(timeout=5) - assert not storage.second_renewal_started.is_set() - storage.release_first_renewal.set() - first.join(timeout=5) - second.join(timeout=5) + try: + assert storage.first_renewal_started.wait(timeout=5) + allow_second_renewal.set() + assert progress.event.wait(timeout=5) + assert progress.source == "renewal lock" + assert not storage.second_renewal_started.is_set() + finally: + allow_second_renewal.set() + storage.release_first_renewal.set() + first.join(timeout=5) + second.join(timeout=5) assert not first.is_alive() assert not second.is_alive() From 67a02da38791c0506ecdb6728c3e81695dd3e9c1 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Thu, 6 Aug 2026 06:01:54 +0000 Subject: [PATCH 35/63] test: narrow heartbeat renewal exception capture --- tests/server/services/governance/test_governance_local_e2e.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/server/services/governance/test_governance_local_e2e.py b/tests/server/services/governance/test_governance_local_e2e.py index 53791890b..e717a807d 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -233,12 +233,12 @@ def __exit__(self, *_exc: object) -> None: renewal_lock = TrackingRenewalLock(progress) cast_heartbeat: Any = heartbeat cast_heartbeat._renewal_lock = renewal_lock - errors: list[BaseException] = [] + errors: list[Exception] = [] def renew() -> None: try: heartbeat.renew_now() - except BaseException as exc: + except Exception as exc: errors.append(exc) allow_second_renewal = threading.Event() From 963c2840e042864cd3f36f5501ae4c0d9cfdfcad Mon Sep 17 00:00:00 2001 From: Guangyu Date: Thu, 6 Aug 2026 19:16:23 +0000 Subject: [PATCH 36/63] test: cover unregistered session outcome acceptance --- .../test_session_outcomes_integration.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/server/api_endpoints/test_session_outcomes_integration.py b/tests/server/api_endpoints/test_session_outcomes_integration.py index 83f2fa33e..645c0ba26 100644 --- a/tests/server/api_endpoints/test_session_outcomes_integration.py +++ b/tests/server/api_endpoints/test_session_outcomes_integration.py @@ -409,6 +409,41 @@ def test_outcome_warning_values_are_sanitized_and_bounded( assert source_warning.endswith("unsafe?session") +def test_unregistered_acceptance_provider_accepts_long_lived_session_marker( + client_with_org: tuple[TestClient, str], +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + session_started_at = int(time.time()) - 90 * 86400 + storage.add_request( + Request( + request_id="long-lived-r1", + user_id="u1", + session_id="long-lived-session", + source="published", + created_at=session_started_at, + ) + ) + + response = client.post( + "/api/session_outcome", + json={ + "session_id": "long-lived-session", + "outcome": "success", + "occurred_at": session_started_at + 1, + }, + ) + + assert response.status_code == 200 + assert response.json()["success"] is True + assert response.json()["recorded"] is True + outcomes = storage.get_session_outcomes( + GetSessionOutcomesRequest(session_ids=["long-lived-session"]) + ) + assert len(outcomes) == 1 + assert outcomes[0].occurred_at == session_started_at + 1 + + def test_acceptance_hook_runs_before_persistence_at_exact_deadline( client_with_org: tuple[TestClient, str], monkeypatch ) -> None: From 2dcd5f523882c1a6f811eff0f97fccab47165934 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Thu, 6 Aug 2026 22:31:43 +0000 Subject: [PATCH 37/63] Validate direct search identifier lengths --- .../models/api_schema/retriever_schema.py | 6 +-- .../routes/test_search_exposure_boundary.py | 46 +++++++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/reflexio/models/api_schema/retriever_schema.py b/reflexio/models/api_schema/retriever_schema.py index f1d048a50..b34be97a6 100644 --- a/reflexio/models/api_schema/retriever_schema.py +++ b/reflexio/models/api_schema/retriever_schema.py @@ -325,7 +325,7 @@ class SearchUserPlaybookRequest(BaseModel): """ query: str | None = None - user_id: str | None = None + user_id: str | None = Field(default=None, max_length=255) agent_version: str | None = None playbook_name: str | None = None start_time: datetime | None = None @@ -338,8 +338,8 @@ class SearchUserPlaybookRequest(BaseModel): search_mode: SearchMode = SearchMode.HYBRID # Caller correlation IDs for billing attribution on the Application line. # Optional; consumed by _meter_applied_learnings in server/api.py. - request_id: str | None = None - session_id: str | None = None + request_id: str | None = Field(default=None, max_length=255) + session_id: str | None = Field(default=None, max_length=255) @model_validator(mode="after") def check_time_range(self) -> Self: diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py index 69dba82a7..73575cdce 100644 --- a/tests/server/routes/test_search_exposure_boundary.py +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -272,6 +272,52 @@ def test_direct_user_playbook_search_returns_and_records_exactly_100() -> None: assert recorder.batches[0].user_playbooks == tuple(playbooks) +@pytest.mark.parametrize( + ("field", "value"), + [ + ("user_id", "u" * 256), + ("request_id", "r" * 256), + ("session_id", "s" * 256), + ], +) +def test_direct_user_playbook_search_rejects_oversized_identifiers_before_search_or_recording( + field: str, + value: str, +) -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([]) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", field: value}, + ) + + assert response.status_code == 422 + reflexio.search_user_playbooks.assert_not_called() + assert recorder.batches == [] + + +def test_direct_user_playbook_search_accepts_255_character_identifiers() -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([]) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={ + "query": "direct answer", + "user_id": "u" * 255, + "request_id": "r" * 255, + "session_id": "s" * 255, + }, + ) + + assert response.status_code == 200, response.text + reflexio.search_user_playbooks.assert_called_once() + assert len(recorder.batches) == 0 + + def test_oss_search_succeeds_when_no_recorder_is_registered() -> None: with _search_results([_playbook(11, "First")]): response = _client().post( From 7e024e3d160c16ce5a8f12523a5f54173d53fee1 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Fri, 7 Aug 2026 02:24:32 +0000 Subject: [PATCH 38/63] test: harden concurrent test assertions --- .../services/extraction/test_resume_worker.py | 31 ++++++++++--------- .../sqlite_storage/test_governance_storage.py | 25 ++++++++++----- 2 files changed, 33 insertions(+), 23 deletions(-) diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index 274ad8f4b..40cead5fa 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -3,6 +3,7 @@ import json import tempfile import threading +from collections import Counter from collections.abc import Callable from contextlib import ExitStack from datetime import UTC, datetime, timedelta @@ -247,15 +248,17 @@ def __init__(self) -> None: self.attempts: list[UsageEvent] = [] self.events: list[UsageEvent] = [] self._accepted_keys: set[str] = set() + self._lock = threading.Lock() def __call__(self, event: UsageEvent) -> UsageEventDeliveryStatus: - self.attempts.append(event) - assert event.event_key is not None - if event.event_key in self._accepted_keys: - return UsageEventDeliveryStatus.DUPLICATE - self._accepted_keys.add(event.event_key) - self.events.append(event) - return UsageEventDeliveryStatus.APPENDED + with self._lock: + self.attempts.append(event) + assert event.event_key is not None + if event.event_key in self._accepted_keys: + return UsageEventDeliveryStatus.DUPLICATE + self._accepted_keys.add(event.event_key) + self.events.append(event) + return UsageEventDeliveryStatus.APPENDED @pytest.mark.parametrize( @@ -1701,10 +1704,9 @@ def finalize(index: int) -> None: if event.event_name == "learnings_generated" and event.entity_type == "profile" ] assert [event.entity_id for event in billing_events] == persisted_ids - assert [event.entity_id for event in recorder.attempts] == [ - *persisted_ids, - *persisted_ids, - ] + assert Counter(event.entity_id for event in recorder.attempts) == Counter( + persisted_ids * 2 + ) def test_two_stale_playbook_workers_preserve_order_and_dispatch_once(tmp_path): @@ -1814,10 +1816,9 @@ def finalize(index: int) -> None: and event.entity_type == "user_playbook" ] assert billing_ids == receipt_ids - assert [event.entity_id for event in recorder.attempts] == [ - *receipt_ids, - *receipt_ids, - ] + assert Counter(event.entity_id for event in recorder.attempts) == Counter( + receipt_ids * 2 + ) finally: configure_usage_event_recorder(None) diff --git a/tests/server/services/storage/sqlite_storage/test_governance_storage.py b/tests/server/services/storage/sqlite_storage/test_governance_storage.py index f4488146c..8ed1b9f87 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -1526,16 +1526,20 @@ def test_shared_file_claim_takeover_fences_independent_storage_instance( ) ready = threading.Barrier(2) claims: list[PurgeExecutionClaim | None] = [] + errors: list[BaseException] = [] def claim(storage_instance: SQLiteStorage, owner: str) -> None: - ready.wait(timeout=5) - claims.append( - storage_instance.claim_purge_operation_execution( - purge.purge_id, - lease_owner=owner, - lease_ttl_seconds=30, + try: + ready.wait(timeout=5) + claims.append( + storage_instance.claim_purge_operation_execution( + purge.purge_id, + lease_owner=owner, + lease_ttl_seconds=30, + ) ) - ) + except BaseException as exc: # noqa: BLE001 - intentional thread error capture + errors.append(exc) callers = [ threading.Thread(target=claim, args=(storage_a, "owner-a")), @@ -1547,6 +1551,8 @@ def claim(storage_instance: SQLiteStorage, owner: str) -> None: caller.join(timeout=5) assert all(not caller.is_alive() for caller in callers) + assert errors == [] + assert len(claims) == len(callers) live_claims = [claim for claim in claims if claim is not None] assert len(live_claims) == 1 first_claim = live_claims[0] @@ -1675,8 +1681,11 @@ def test_claimed_erasure_mutation_signatures_and_callers_require_claim() -> None assert parameter.annotation in {"PurgeExecutionClaim", PurgeExecutionClaim} production_root = Path(__file__).resolve().parents[5] / "reflexio" + assert production_root.exists() + python_modules = list(production_root.rglob("*.py")) + assert python_modules violations: list[str] = [] - for path in production_root.rglob("*.py"): + for path in python_modules: tree = ast.parse(path.read_text(), filename=str(path)) for node in ast.walk(tree): if not isinstance(node, ast.Call): From 6b6e6c24887fa7aa0465446c30fae5beb259dc9e Mon Sep 17 00:00:00 2001 From: Guangyu Date: Fri, 7 Aug 2026 14:00:09 +0000 Subject: [PATCH 39/63] feat: add open-world optimizer identity --- reflexio/models/api_schema/domain/entities.py | 2 ++ .../services/storage/sqlite_storage/_base.py | 6 ++++- .../test_open_world_optimization_identity.py | 24 +++++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 tests/models/test_open_world_optimization_identity.py diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 35403f647..59bc3f27a 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -403,6 +403,7 @@ class AgentPlaybook(BaseModel): OptimizerKind = Literal[ "gepa", "offline_tuner_replay", + "offline_tuner_open_world", "offline_tuner_legacy", "optimizer_legacy_unknown", ] @@ -443,6 +444,7 @@ class AgentPlaybook(BaseModel): "replay_manifest", "candidate", "candidate_search_projection", + "open_world_evidence_bundle", ] Sha256Digest = str diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 10cacd5c1..3f7d1ceb5 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -2086,6 +2086,7 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: "CHECK (stage IS NULL OR stage IN", "CHECK (terminal_outcome IS NULL OR terminal_outcome IN", "'governance_erased'", + "'offline_tuner_open_world'", ) if all(check in table_sql for check in required_checks): return @@ -2106,6 +2107,7 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: CHECK (optimizer_kind IN ( 'gepa', 'offline_tuner_replay', + 'offline_tuner_open_world', 'offline_tuner_legacy', 'optimizer_legacy_unknown' )), @@ -3149,6 +3151,7 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: CHECK (optimizer_kind IN ( 'gepa', 'offline_tuner_replay', + 'offline_tuner_open_world', 'offline_tuner_legacy', 'optimizer_legacy_unknown' )), @@ -3249,7 +3252,8 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: 'generation_selection', 'replay_manifest', 'candidate', - 'candidate_search_projection' + 'candidate_search_projection', + 'open_world_evidence_bundle' )), content_json TEXT NOT NULL, content_digest TEXT NOT NULL, diff --git a/tests/models/test_open_world_optimization_identity.py b/tests/models/test_open_world_optimization_identity.py new file mode 100644 index 000000000..09c58d9d1 --- /dev/null +++ b/tests/models/test_open_world_optimization_identity.py @@ -0,0 +1,24 @@ +from hashlib import sha256 + +from reflexio.models.api_schema.domain.entities import ( + PlaybookOptimizationArtifact, + PlaybookOptimizationJob, +) + + +def test_open_world_optimization_identity_is_accepted() -> None: + content_json = '{"schema_version":"offline-tuner-open-world-evidence-v1"}' + job = PlaybookOptimizationJob( + optimizer_kind="offline_tuner_open_world", + target_kind="user_playbook", + target_id=7, + ) + artifact = PlaybookOptimizationArtifact( + job_id=1, + artifact_kind="open_world_evidence_bundle", + content_json=content_json, + content_digest=sha256(content_json.encode()).hexdigest(), + ) + + assert job.optimizer_kind == "offline_tuner_open_world" + assert artifact.artifact_kind == "open_world_evidence_bundle" From c9708638685f283a3ab62ab3398e1d42a36155a9 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Sat, 8 Aug 2026 23:32:15 +0000 Subject: [PATCH 40/63] fix: upgrade generic evidence artifact storage --- .../storage/session_outcome_identity.py | 8 +- .../services/storage/sqlite_storage/_base.py | 90 +++++++++++++ tests/models/test_session_outcome_identity.py | 52 +++++++ ...ptimization_replay_contract_integration.py | 127 ++++++++++++++++++ 4 files changed, 276 insertions(+), 1 deletion(-) diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py index 36179cee2..28e0e0064 100644 --- a/reflexio/server/services/storage/session_outcome_identity.py +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -12,6 +12,7 @@ "CanonicalSessionTrajectory", "canonical_json_bytes", "canonical_session_trajectory", + "canonical_trajectory_bytes", "outcome_contract_digest", "trajectory_digest", ] @@ -197,6 +198,11 @@ def _canonical_trajectory_json(value: object, *, depth: int = 0) -> str: return canonical_json_bytes(value).decode() +def canonical_trajectory_bytes(trajectory: object) -> bytes: + """Encode a finalized session trajectory to its canonical UTF-8 bytes.""" + return _canonical_trajectory_json(trajectory).encode("utf-8") + + def trajectory_digest(trajectory: object) -> str: """Hash the canonical finalized session trajectory.""" - return sha256(_canonical_trajectory_json(trajectory).encode()).hexdigest() + return sha256(canonical_trajectory_bytes(trajectory)).hexdigest() diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 3f7d1ceb5..763926994 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -1075,6 +1075,7 @@ def migrate(self) -> bool: self._migrate_user_playbook_publication_staging_columns() self._classify_legacy_playbook_optimization_jobs() self._enforce_playbook_optimization_job_constraints() + self._enforce_playbook_optimization_artifact_constraints() self._migrate_retire_profile_change_logs() self._migrate_retire_playbook_aggregation_change_logs() init_stall_state_table(self.conn) @@ -2226,6 +2227,95 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: f"PRAGMA foreign_keys={'ON' if foreign_keys_enabled else 'OFF'}" ) + def _enforce_playbook_optimization_artifact_constraints(self) -> None: + """Rebuild legacy artifact tables with the current kind allowlist.""" + table_sql_row = self.conn.execute( + """SELECT sql FROM sqlite_master + WHERE type = 'table' AND name = 'playbook_optimization_artifacts'""" + ).fetchone() + if table_sql_row is None: + return + table_sql = table_sql_row["sql"] + artifact_kinds = ( + "'expected_population_manifest'", + "'generation_selection'", + "'replay_manifest'", + "'candidate'", + "'candidate_search_projection'", + "'open_world_evidence_bundle'", + ) + if all(artifact_kind in table_sql for artifact_kind in artifact_kinds): + return + + foreign_keys_enabled = bool( + self.conn.execute("PRAGMA foreign_keys").fetchone()[0] + ) + self.conn.commit() + if foreign_keys_enabled: + self.conn.execute("PRAGMA foreign_keys=OFF") + try: + self.conn.execute("BEGIN IMMEDIATE") + self.conn.execute( + "DROP TABLE IF EXISTS playbook_optimization_artifacts_new" + ) + self.conn.execute( + """ + CREATE TABLE playbook_optimization_artifacts_new ( + artifact_id INTEGER PRIMARY KEY AUTOINCREMENT, + job_id INTEGER NOT NULL, + artifact_kind TEXT NOT NULL CHECK (artifact_kind IN ( + 'expected_population_manifest', + 'generation_selection', + 'replay_manifest', + 'candidate', + 'candidate_search_projection', + 'open_world_evidence_bundle' + )), + content_json TEXT NOT NULL, + content_digest TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (job_id, artifact_kind), + FOREIGN KEY (job_id) + REFERENCES playbook_optimization_jobs(job_id) + ON DELETE CASCADE + ) + """ + ) + self.conn.execute( + """ + INSERT INTO playbook_optimization_artifacts_new ( + artifact_id, job_id, artifact_kind, content_json, + content_digest, created_at, updated_at + ) SELECT + artifact_id, job_id, artifact_kind, content_json, + content_digest, created_at, updated_at + FROM playbook_optimization_artifacts + """ + ) + self.conn.execute("DROP TABLE playbook_optimization_artifacts") + self.conn.execute( + "ALTER TABLE playbook_optimization_artifacts_new " + "RENAME TO playbook_optimization_artifacts" + ) + self.conn.execute( + "CREATE INDEX idx_poa_job ON playbook_optimization_artifacts(job_id)" + ) + violations = self.conn.execute("PRAGMA foreign_key_check").fetchall() + if violations: + raise sqlite3.IntegrityError( + "foreign key check failed after optimizer artifact migration: " + f"{violations}" + ) + self.conn.commit() + except Exception: + self.conn.rollback() + raise + finally: + self.conn.execute( + f"PRAGMA foreign_keys={'ON' if foreign_keys_enabled else 'OFF'}" + ) + def _migrate_retire_profile_change_logs(self) -> None: """Retire the frozen ``profile_change_logs`` table via a reversible RENAME. diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index cf4bed58b..cb77b8d35 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -30,6 +30,7 @@ from reflexio.server.services.storage.session_outcome_identity import ( canonical_json_bytes, canonical_session_trajectory, + canonical_trajectory_bytes, outcome_contract_digest, trajectory_digest, ) @@ -222,6 +223,57 @@ def test_trajectory_digest_changes_when_trajectory_data_changes() -> None: ) +def test_canonical_trajectory_bytes_returns_exact_utf8_representation() -> None: + assert canonical_trajectory_bytes( + {"value": 1.5, "message": "caf\N{LATIN SMALL LETTER E WITH ACUTE}"} + ) == (b'{"message":"caf\xc3\xa9","value":1.5}') + + +def test_canonical_trajectory_bytes_ignores_mapping_insertion_order() -> None: + first = {"messages": [{"content": "hello", "role": "user"}], "session": "s1"} + second = {"session": "s1", "messages": [{"role": "user", "content": "hello"}]} + + assert canonical_trajectory_bytes(first) == canonical_trajectory_bytes(second) + + +@pytest.mark.parametrize( + "value", + [ + {"value": float("nan")}, + {"value": 2**53}, + {"value": object()}, + {1: "non-string-key"}, + {"value": "\ud800"}, + ], + ids=[ + "non-finite-float", + "non-ijson-integer", + "unsupported-object", + "key", + "surrogate", + ], +) +def test_canonical_trajectory_bytes_rejects_the_same_values_as_digest( + value: object, +) -> None: + with pytest.raises((TypeError, ValueError)) as digest_error: + trajectory_digest(value) + + with pytest.raises(type(digest_error.value)) as bytes_error: + canonical_trajectory_bytes(value) + + assert str(bytes_error.value) == str(digest_error.value) + + +def test_trajectory_digest_hashes_canonical_trajectory_bytes() -> None: + trajectory = {"messages": [{"role": "user", "content": "hello"}]} + + assert ( + trajectory_digest(trajectory) + == sha256(canonical_trajectory_bytes(trajectory)).hexdigest() + ) + + def test_canonical_session_trajectory_normalizes_sqlite_and_postgres_rows() -> None: sqlite_request = { "request_id": "parity-request", diff --git a/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py b/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py index d6439ec30..c53060c92 100644 --- a/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py +++ b/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py @@ -429,6 +429,133 @@ def test_artifact_model_uses_publication_numeric_contract() -> None: ) +def test_previous_artifact_schema_is_upgraded_without_losing_constraints( + tmp_path: Path, +) -> None: + db_path = tmp_path / "previous-artifact-schema.db" + initial_store = SQLiteStorage( + org_id="previous-artifact-schema", db_path=str(db_path) + ) + initial_store.conn.close() + + legacy_content = '{"source":"legacy"}' + legacy_digest = sha256(legacy_content.encode()).hexdigest() + conn = sqlite3.connect(db_path) + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("DROP INDEX idx_poa_job") + conn.execute("DROP TABLE playbook_optimization_artifacts") + conn.executescript( + """ + CREATE TABLE playbook_optimization_artifacts ( + artifact_id INTEGER PRIMARY KEY AUTOINCREMENT, + job_id INTEGER NOT NULL, + artifact_kind TEXT NOT NULL CHECK (artifact_kind IN ( + 'expected_population_manifest', + 'generation_selection', + 'replay_manifest', + 'candidate', + 'candidate_search_projection' + )), + content_json TEXT NOT NULL, + content_digest TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (job_id, artifact_kind), + FOREIGN KEY (job_id) REFERENCES playbook_optimization_jobs(job_id) + ON DELETE CASCADE + ); + CREATE INDEX idx_poa_job ON playbook_optimization_artifacts(job_id); + """ + ) + conn.execute( + """INSERT INTO playbook_optimization_jobs ( + job_id, optimizer_kind, target_kind, target_id, status, + lease_owner, lease_fence, lease_expires_at, created_at, updated_at + ) VALUES ( + 41, 'offline_tuner_open_world', 'user_playbook', 9, 'running', + 'worker-a', 3, 1000, 101, 102 + )""" + ) + conn.execute( + """INSERT INTO playbook_optimization_artifacts ( + artifact_id, job_id, artifact_kind, content_json, content_digest, + created_at, updated_at + ) VALUES (77, 41, 'candidate', ?, ?, 103, 104)""", + (legacy_content, legacy_digest), + ) + conn.commit() + conn.close() + + store = SQLiteStorage(org_id="previous-artifact-schema", db_path=str(db_path)) + try: + legacy_row = store.conn.execute( + "SELECT * FROM playbook_optimization_artifacts WHERE artifact_id = 77" + ).fetchone() + assert dict(legacy_row) == { + "artifact_id": 77, + "job_id": 41, + "artifact_kind": "candidate", + "content_json": legacy_content, + "content_digest": legacy_digest, + "created_at": 103, + "updated_at": 104, + } + assert store.conn.execute("PRAGMA foreign_key_check").fetchall() == [] + assert store.conn.execute("PRAGMA foreign_keys").fetchone()[0] == 1 + + with pytest.raises(sqlite3.IntegrityError): + store.conn.execute( + """INSERT INTO playbook_optimization_artifacts VALUES + (78, 41, 'candidate', '{}', ?, 105, 106)""", + (sha256(b"{}").hexdigest(),), + ) + store.conn.rollback() + assert ( + store.conn.execute( + """SELECT 1 FROM sqlite_master + WHERE type = 'index' AND name = 'idx_poa_job'""" + ).fetchone() + is not None + ) + with pytest.raises(sqlite3.IntegrityError): + store.conn.execute( + """INSERT INTO playbook_optimization_artifacts VALUES + (79, 41, 'unknown_kind', '{}', ?, 105, 106)""", + (sha256(b"{}").hexdigest(),), + ) + store.conn.rollback() + + evidence_json = '{"cases":[1]}' + evidence = schemas.PlaybookOptimizationArtifact( + job_id=41, + artifact_kind="open_world_evidence_bundle", + content_json=evidence_json, + content_digest=sha256(evidence_json.encode()).hexdigest(), + created_at=107, + updated_at=108, + ) + saved = store.upsert_playbook_optimization_artifact(evidence, fence=3, now=500) + assert ( + store.get_playbook_optimization_artifact(41, "open_world_evidence_bundle") + == saved + ) + + assert store.migrate() is True + store.conn.execute("DELETE FROM playbook_optimization_jobs WHERE job_id = 41") + store.conn.commit() + assert ( + store.conn.execute( + "SELECT COUNT(*) FROM playbook_optimization_artifacts WHERE job_id = 41" + ).fetchone()[0] + == 0 + ) + finally: + store.conn.close() + + assert not hasattr(BaseStorage, "load_open_world_evidence_snapshot") + assert not hasattr(SQLiteStorage, "load_open_world_evidence_snapshot") + + @pytest.mark.parametrize( "current_stage", [ From 063be9f06c3e7c4a6d6619d04b47c0ca11d08f0e Mon Sep 17 00:00:00 2001 From: Guangyu Date: Sun, 9 Aug 2026 07:25:04 +0000 Subject: [PATCH 41/63] fix: harden SQLite optimization storage invariants --- reflexio/server/services/storage/error.py | 8 + .../services/storage/sqlite_storage/_base.py | 45 ++++- .../sqlite_storage/playbook/_optimization.py | 47 +++-- ...ptimization_replay_contract_integration.py | 175 +++++++++++++++++- 4 files changed, 254 insertions(+), 21 deletions(-) diff --git a/reflexio/server/services/storage/error.py b/reflexio/server/services/storage/error.py index cf696ed79..fe14734bd 100644 --- a/reflexio/server/services/storage/error.py +++ b/reflexio/server/services/storage/error.py @@ -22,6 +22,14 @@ class OptimizationJobLeaseLiveError(StorageError): """Raised when an optimizer recovery attempt finds a non-expired lease.""" +class OptimizationJobIdentityConflictError(StorageError): + """Raised when one durable optimizer identity resolves to conflicting jobs.""" + + +class OptimizationArtifactIntegrityError(StorageError): + """Raised when a durable optimizer artifact is malformed or conflicts.""" + + def require_non_empty_session_id(value: Any) -> str: """Return a stripped, non-empty request ``session_id`` or raise ``StorageError``. diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 763926994..dfb471f59 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -2087,9 +2087,10 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: "CHECK (stage IS NULL OR stage IN", "CHECK (terminal_outcome IS NULL OR terminal_outcome IN", "'governance_erased'", - "'offline_tuner_open_world'", ) - if all(check in table_sql for check in required_checks): + if all(check in table_sql for check in required_checks) and ( + "'offline_tuner_open_world'" not in table_sql + ): return foreign_keys_enabled = bool( self.conn.execute("PRAGMA foreign_keys").fetchone()[0] @@ -2099,6 +2100,10 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: self.conn.execute("PRAGMA foreign_keys=OFF") try: self.conn.execute("BEGIN IMMEDIATE") + sequence_high_water = self._autoincrement_high_water( + "playbook_optimization_jobs", + "job_id", + ) self.conn.execute("DROP TABLE IF EXISTS playbook_optimization_jobs_new") self.conn.execute( """ @@ -2108,7 +2113,6 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: CHECK (optimizer_kind IN ( 'gepa', 'offline_tuner_replay', - 'offline_tuner_open_world', 'offline_tuner_legacy', 'optimizer_legacy_unknown' )), @@ -2191,6 +2195,10 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: "ALTER TABLE playbook_optimization_jobs_new " "RENAME TO playbook_optimization_jobs" ) + self._restore_autoincrement_high_water( + "playbook_optimization_jobs", + sequence_high_water, + ) self.conn.execute( "CREATE INDEX idx_poj_target " "ON playbook_optimization_jobs(target_kind, target_id)" @@ -2255,6 +2263,10 @@ def _enforce_playbook_optimization_artifact_constraints(self) -> None: self.conn.execute("PRAGMA foreign_keys=OFF") try: self.conn.execute("BEGIN IMMEDIATE") + sequence_high_water = self._autoincrement_high_water( + "playbook_optimization_artifacts", + "artifact_id", + ) self.conn.execute( "DROP TABLE IF EXISTS playbook_optimization_artifacts_new" ) @@ -2298,6 +2310,10 @@ def _enforce_playbook_optimization_artifact_constraints(self) -> None: "ALTER TABLE playbook_optimization_artifacts_new " "RENAME TO playbook_optimization_artifacts" ) + self._restore_autoincrement_high_water( + "playbook_optimization_artifacts", + sequence_high_water, + ) self.conn.execute( "CREATE INDEX idx_poa_job ON playbook_optimization_artifacts(job_id)" ) @@ -2316,6 +2332,28 @@ def _enforce_playbook_optimization_artifact_constraints(self) -> None: f"PRAGMA foreign_keys={'ON' if foreign_keys_enabled else 'OFF'}" ) + def _autoincrement_high_water(self, table_name: str, id_column: str) -> int: + sequence_row = self.conn.execute( + "SELECT MAX(seq) FROM sqlite_sequence WHERE name = ?", + (table_name,), + ).fetchone() + maximum_row = self.conn.execute( + f"SELECT MAX({id_column}) FROM {table_name}" # noqa: S608 + ).fetchone() + return max(sequence_row[0] or 0, maximum_row[0] or 0) + + def _restore_autoincrement_high_water( + self, + table_name: str, + high_water: int, + ) -> None: + self.conn.execute("DELETE FROM sqlite_sequence WHERE name = ?", (table_name,)) + if high_water: + self.conn.execute( + "INSERT INTO sqlite_sequence(name, seq) VALUES (?, ?)", + (table_name, high_water), + ) + def _migrate_retire_profile_change_logs(self) -> None: """Retire the frozen ``profile_change_logs`` table via a reversible RENAME. @@ -3241,7 +3279,6 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: CHECK (optimizer_kind IN ( 'gepa', 'offline_tuner_replay', - 'offline_tuner_open_world', 'offline_tuner_legacy', 'optimizer_legacy_unknown' )), diff --git a/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py b/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py index d97fd6f24..41b938d95 100644 --- a/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py +++ b/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py @@ -6,6 +6,8 @@ from hashlib import sha256 from typing import Any +from pydantic import ValidationError + from reflexio.models.api_schema.domain.entities import canonicalize_artifact_json from reflexio.models.api_schema.service_schemas import ( OptimizationArtifactKind, @@ -25,6 +27,8 @@ canonical_json_bytes, ) from reflexio.server.services.storage.error import ( + OptimizationArtifactIntegrityError, + OptimizationJobIdentityConflictError, OptimizationJobLeaseLiveError, StorageError, ) @@ -86,15 +90,20 @@ def _row_to_playbook_optimization_job(row: sqlite3.Row) -> PlaybookOptimizationJ def _row_to_playbook_optimization_artifact( row: sqlite3.Row, ) -> PlaybookOptimizationArtifact: - return PlaybookOptimizationArtifact( - artifact_id=row["artifact_id"], - job_id=row["job_id"], - artifact_kind=row["artifact_kind"], - content_json=row["content_json"], - content_digest=row["content_digest"], - created_at=row["created_at"], - updated_at=row["updated_at"], - ) + try: + return PlaybookOptimizationArtifact( + artifact_id=row["artifact_id"], + job_id=row["job_id"], + artifact_kind=row["artifact_kind"], + content_json=row["content_json"], + content_digest=row["content_digest"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + except (ValidationError, ValueError) as exc: + raise OptimizationArtifactIntegrityError( + "optimizer artifact row is malformed" + ) from exc def _job_insert_values(job: PlaybookOptimizationJob) -> tuple[Any, ...]: @@ -531,7 +540,9 @@ def create_or_get_playbook_optimization_job( and by_attempt is not None and by_discovery["job_id"] != by_attempt["job_id"] ): - raise ValueError("conflicting immutable optimizer job identity") + raise OptimizationJobIdentityConflictError( + "conflicting immutable optimizer job identity" + ) existing = by_discovery or by_attempt if existing is not None: if ( @@ -542,7 +553,9 @@ def create_or_get_playbook_optimization_job( and existing["attempt_key"] != job.attempt_key ) ): - raise ValueError("conflicting immutable optimizer job identity") + raise OptimizationJobIdentityConflictError( + "conflicting immutable optimizer job identity" + ) result = _row_to_playbook_optimization_job(existing) else: cur = self.conn.execute(_JOB_INSERT_SQL, _job_insert_values(job)) @@ -801,7 +814,9 @@ def upsert_playbook_optimization_artifact( sha256(artifact_content_json.encode()).hexdigest() != artifact.content_digest ): - raise ValueError("optimizer artifact digest does not match content") + raise OptimizationArtifactIntegrityError( + "optimizer artifact digest does not match content" + ) written_at = self._lease_now(now) with self._lock: owns_transaction = self._own_transaction() @@ -825,9 +840,13 @@ def upsert_playbook_optimization_artifact( ).fetchone() if existing is not None: if existing["content_digest"] != artifact.content_digest: - raise ValueError("optimizer artifact digest conflict") + raise OptimizationArtifactIntegrityError( + "optimizer artifact digest conflict" + ) if existing["content_json"] != artifact_content_json: - raise ValueError("optimizer artifact content conflict") + raise OptimizationArtifactIntegrityError( + "optimizer artifact content conflict" + ) result = _row_to_playbook_optimization_artifact(existing) else: cur = self.conn.execute( diff --git a/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py b/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py index c53060c92..0a6c64f16 100644 --- a/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py +++ b/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py @@ -14,6 +14,8 @@ from reflexio.models.api_schema import service_schemas as schemas from reflexio.server.services.playbook.publication import canonical_json_bytes from reflexio.server.services.storage.error import ( + OptimizationArtifactIntegrityError, + OptimizationJobIdentityConflictError, OptimizationJobLeaseLiveError, StorageError, ) @@ -113,10 +115,21 @@ def test_same_attempt_key_returns_one_active_job(storage: BaseStorage) -> None: def test_conflicting_active_identity_is_rejected(storage: BaseStorage) -> None: storage.create_or_get_playbook_optimization_job(_replay_job("d1", "a1")) - with pytest.raises(StorageError, match="immutable optimizer job identity"): + with pytest.raises( + OptimizationJobIdentityConflictError, + match="immutable optimizer job identity", + ): storage.create_or_get_playbook_optimization_job(_replay_job("d1", "a2")) +def test_sqlite_rejects_open_world_optimizer_jobs(storage: BaseStorage) -> None: + open_world_job = _replay_job("open-world-discovery", "open-world-attempt") + open_world_job.optimizer_kind = "offline_tuner_open_world" + + with pytest.raises(StorageError, match="CHECK constraint failed"): + storage.create_or_get_playbook_optimization_job(open_world_job) + + def test_gepa_publication_reclaim_contract_has_none_live_and_reclaimed_outcomes( storage: BaseStorage, ) -> None: @@ -348,7 +361,10 @@ def test_artifact_upsert_canonicalizes_equivalent_json_and_requires_digest_and_c fence=claim.fence, now=5_001, ) - with pytest.raises(StorageError, match="artifact digest conflict"): + with pytest.raises( + OptimizationArtifactIntegrityError, + match="artifact digest conflict", + ): storage.upsert_playbook_optimization_artifact( _artifact( job_id=job.job_id, @@ -359,6 +375,52 @@ def test_artifact_upsert_canonicalizes_equivalent_json_and_requires_digest_and_c ) +def test_malformed_persisted_artifact_raises_typed_integrity_error( + storage: BaseStorage, +) -> None: + job = storage.create_or_get_playbook_optimization_job(_replay_job("d1", "a1")) + claim = storage.claim_playbook_optimization_job( + job_id=job.job_id, + owner="worker-a", + lease_seconds=60, + now=5_000, + ) + saved = storage.upsert_playbook_optimization_artifact( + _artifact(job_id=job.job_id), + fence=claim.fence, + now=5_001, + ) + assert isinstance(storage, SQLiteStorage) + storage.conn.execute( + "UPDATE playbook_optimization_artifacts SET content_digest = ? " + "WHERE artifact_id = ?", + ("f" * 64, saved.artifact_id), + ) + storage.conn.commit() + + with pytest.raises( + OptimizationArtifactIntegrityError, + match="optimizer artifact row is malformed", + ): + storage.get_playbook_optimization_artifact( + job.job_id, + "expected_population_manifest", + ) + + +def test_artifact_storage_failures_remain_generic(storage: BaseStorage) -> None: + assert isinstance(storage, SQLiteStorage) + storage.conn.execute("DROP TABLE playbook_optimization_artifacts") + + with pytest.raises(StorageError) as raised: + storage.get_playbook_optimization_artifact( + 1, + "expected_population_manifest", + ) + + assert type(raised.value) is StorageError + + def test_artifact_model_rejects_malformed_json() -> None: with pytest.raises( ValidationError, match="artifact content_json must be valid JSON" @@ -472,7 +534,7 @@ def test_previous_artifact_schema_is_upgraded_without_losing_constraints( job_id, optimizer_kind, target_kind, target_id, status, lease_owner, lease_fence, lease_expires_at, created_at, updated_at ) VALUES ( - 41, 'offline_tuner_open_world', 'user_playbook', 9, 'running', + 41, 'offline_tuner_replay', 'user_playbook', 9, 'running', 'worker-a', 3, 1000, 101, 102 )""" ) @@ -556,6 +618,113 @@ def test_previous_artifact_schema_is_upgraded_without_losing_constraints( assert not hasattr(SQLiteStorage, "load_open_world_evidence_snapshot") +def test_optimizer_job_rebuild_preserves_deleted_id_high_water_and_repeats( + tmp_path: Path, +) -> None: + db_path = tmp_path / "legacy-job-sequence.db" + _create_legacy_optimizer_schema(db_path) + conn = sqlite3.connect(db_path) + conn.execute("DELETE FROM playbook_optimization_jobs WHERE job_id IN (5, 6)") + conn.commit() + conn.close() + + first_store = SQLiteStorage(org_id="job-sequence-first", db_path=str(db_path)) + first = first_store.create_playbook_optimization_job(_replay_job("d1", "a1")) + first_store.conn.close() + second_store = SQLiteStorage(org_id="job-sequence-second", db_path=str(db_path)) + second = second_store.create_playbook_optimization_job( + _replay_job("d2", "a2").model_copy(update={"target_id": 42}) + ) + second_store.conn.close() + + assert first.job_id == 7 + assert second.job_id == 8 + + +def test_empty_optimizer_job_rebuild_preserves_deleted_id_high_water( + tmp_path: Path, +) -> None: + db_path = tmp_path / "empty-legacy-job-sequence.db" + _create_legacy_optimizer_schema(db_path) + conn = sqlite3.connect(db_path) + conn.execute("DELETE FROM playbook_optimization_jobs") + conn.commit() + conn.close() + + store = SQLiteStorage(org_id="empty-job-sequence", db_path=str(db_path)) + job = store.create_playbook_optimization_job(_replay_job("d1", "a1")) + store.conn.close() + + assert job.job_id == 7 + + +def test_artifact_rebuild_preserves_deleted_id_high_water_and_repeats( + tmp_path: Path, +) -> None: + db_path = tmp_path / "legacy-artifact-sequence.db" + store = SQLiteStorage(org_id="artifact-sequence-setup", db_path=str(db_path)) + parent = store.create_playbook_optimization_job(_replay_job("d1", "a1")) + store.conn.execute("DROP INDEX idx_poa_job") + store.conn.execute("DROP TABLE playbook_optimization_artifacts") + store.conn.executescript( + """ + CREATE TABLE playbook_optimization_artifacts ( + artifact_id INTEGER PRIMARY KEY AUTOINCREMENT, + job_id INTEGER NOT NULL, + artifact_kind TEXT NOT NULL CHECK (artifact_kind IN ('candidate')), + content_json TEXT NOT NULL, + content_digest TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (job_id, artifact_kind), + FOREIGN KEY (job_id) REFERENCES playbook_optimization_jobs(job_id) + ON DELETE CASCADE + ); + CREATE INDEX idx_poa_job ON playbook_optimization_artifacts(job_id); + """ + ) + digest = sha256(b"{}").hexdigest() + store.conn.execute( + "INSERT INTO playbook_optimization_artifacts VALUES " + "(9, ?, 'candidate', '{}', ?, 1, 1)", + (parent.job_id, digest), + ) + store.conn.execute("DELETE FROM playbook_optimization_artifacts") + store.conn.commit() + store.conn.close() + + first_store = SQLiteStorage(org_id="artifact-sequence-first", db_path=str(db_path)) + first_store.conn.execute( + "INSERT INTO playbook_optimization_artifacts " + "(job_id, artifact_kind, content_json, content_digest, created_at, updated_at) " + "VALUES (?, 'candidate', '{}', ?, 1, 1)", + (parent.job_id, digest), + ) + first_id = first_store.conn.execute( + "SELECT artifact_id FROM playbook_optimization_artifacts" + ).fetchone()[0] + first_store.conn.execute("DELETE FROM playbook_optimization_artifacts") + first_store.conn.commit() + first_store.conn.close() + + second_store = SQLiteStorage( + org_id="artifact-sequence-second", db_path=str(db_path) + ) + second_store.conn.execute( + "INSERT INTO playbook_optimization_artifacts " + "(job_id, artifact_kind, content_json, content_digest, created_at, updated_at) " + "VALUES (?, 'candidate', '{}', ?, 1, 1)", + (parent.job_id, digest), + ) + second_id = second_store.conn.execute( + "SELECT artifact_id FROM playbook_optimization_artifacts" + ).fetchone()[0] + second_store.conn.close() + + assert first_id == 10 + assert second_id == 11 + + @pytest.mark.parametrize( "current_stage", [ From fe4bfdab52ecbfcb306bcda4ab911a9dea5e88a3 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 04:03:51 +0000 Subject: [PATCH 42/63] test: align exposure boundary with queued metering --- .../routes/test_search_exposure_boundary.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py index 73575cdce..1530aef7f 100644 --- a/tests/server/routes/test_search_exposure_boundary.py +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -154,12 +154,8 @@ def test_direct_user_playbook_search_records_final_results_before_metering() -> with ( _search_results(playbooks), patch( - "reflexio.server.routes.search._meter_search_request", - side_effect=lambda **_kwargs: order.append("meter_search"), - ), - patch( - "reflexio.server.routes.search._meter_applied_learnings", - side_effect=lambda **_kwargs: order.append("meter_applied"), + "reflexio.server.routes.search.enqueue_search_metering", + side_effect=lambda **_kwargs: order.append("meter_enqueue"), ), ): response = _client().post( @@ -173,7 +169,7 @@ def test_direct_user_playbook_search_records_final_results_before_metering() -> ) assert response.status_code == 200, response.text - assert order == ["record", "meter_search", "meter_applied"] + assert order == ["record", "meter_enqueue"] assert len(recorder.batches) == 1 batch = recorder.batches[0] assert batch.org_id == "org-1" @@ -197,12 +193,8 @@ def record(self, _batch: Any) -> None: with ( _search_results([_playbook(21, "Direct first")]), patch( - "reflexio.server.routes.search._meter_search_request", - side_effect=lambda **_kwargs: order.append("meter_search"), - ), - patch( - "reflexio.server.routes.search._meter_applied_learnings", - side_effect=lambda **_kwargs: order.append("meter_applied"), + "reflexio.server.routes.search.enqueue_search_metering", + side_effect=lambda **_kwargs: order.append("meter_enqueue"), ), ): response = _client().post( From 3d4186b9f342c18c41688dea7b5af905c298f7e6 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 06:43:26 +0000 Subject: [PATCH 43/63] fix: align Phase 1 tests with current runtime --- reflexio/server/routes/search.py | 2 +- .../services/test_non_extraction_learning_metering.py | 5 ++++- tests/server/test_create_app_capabilities.py | 8 ++++---- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/reflexio/server/routes/search.py b/reflexio/server/routes/search.py index 6834faee0..1adb23301 100644 --- a/reflexio/server/routes/search.py +++ b/reflexio/server/routes/search.py @@ -48,11 +48,11 @@ from reflexio.server.services.retrieval_experiment import ( active_retrieval_experiment_assignment, ) -from reflexio.server.services.search_metering_worker import enqueue_search_metering from reflexio.server.services.search_exposure import ( SearchExposureBatch, record_search_exposures, ) +from reflexio.server.services.search_metering_worker import enqueue_search_metering from reflexio.server.tracing import profile_step logger = logging.getLogger(__name__) diff --git a/tests/server/services/test_non_extraction_learning_metering.py b/tests/server/services/test_non_extraction_learning_metering.py index 2c2d2675d..2b2c76ad6 100644 --- a/tests/server/services/test_non_extraction_learning_metering.py +++ b/tests/server/services/test_non_extraction_learning_metering.py @@ -170,7 +170,10 @@ def test_resumable_playbook_bills_consolidation_replacement_id() -> None: worker._finalize_items(run, [original_candidate]) finalizer.assert_called_once_with( - [original_candidate], model_provenance=None, finalization_run_id=run.id + [original_candidate], + model_provenance=None, + extraction_run=run, + finalization_run_id=run.id, ) assert len(events) == 1 assert events[0].count_value == 1 diff --git a/tests/server/test_create_app_capabilities.py b/tests/server/test_create_app_capabilities.py index 05aa9226b..95ae4c542 100644 --- a/tests/server/test_create_app_capabilities.py +++ b/tests/server/test_create_app_capabilities.py @@ -100,15 +100,15 @@ def test_capabilities_none_is_unchanged() -> None: def test_oss_startup_failure_clears_exempt_usage_recorder(monkeypatch) -> None: usage_metrics.configure_usage_event_recorder(None) - def fail_startup_guards() -> None: + def fail_startup_guard() -> None: raise RuntimeError("startup guard failed") monkeypatch.setattr( - "reflexio.server.llm.providers.embedder_warmup.run_startup_config_guards", - fail_startup_guards, + "reflexio.server.llm.model_defaults.validate_llm_availability", + fail_startup_guard, ) - app = create_app(capabilities=None, mount_data_plane=False) + app = create_app(capabilities=None, mount_data_plane=True) with pytest.raises(RuntimeError, match="startup guard failed"), TestClient(app): pass From 7e42b17552fe85f0ea73ed7cc4c34031a9d82c6b Mon Sep 17 00:00:00 2001 From: Guangyu Date: Mon, 10 Aug 2026 01:22:01 +0000 Subject: [PATCH 44/63] test: align lifecycle coverage with queued metering --- tests/server/test_create_app_capabilities.py | 26 ++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/server/test_create_app_capabilities.py b/tests/server/test_create_app_capabilities.py index 95ae4c542..eb9c2f2b8 100644 --- a/tests/server/test_create_app_capabilities.py +++ b/tests/server/test_create_app_capabilities.py @@ -121,6 +121,32 @@ def fail_startup_guard() -> None: assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN +def test_oss_metering_startup_failure_clears_exempt_usage_recorder( + monkeypatch, +) -> None: + usage_metrics.configure_usage_event_recorder(None) + + def fail_start_search_metering_worker() -> None: + raise RuntimeError("startup guard failed") + + monkeypatch.setattr( + "reflexio.server.services.search_metering_worker.start_search_metering_worker", + fail_start_search_metering_worker, + ) + + app = create_app(capabilities=None, mount_data_plane=False) + with pytest.raises(RuntimeError, match="startup guard failed"), TestClient(app): + pass + + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + ) + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + + def test_partial_cleanup_invariant() -> None: """Already-started caps must be shut down; caps that never started must not be.""" LifecycleCap.started = LifecycleCap.stopped = False From 0d2f77dd3be7f036ef53c4d41950203bd9da9144 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Mon, 10 Aug 2026 02:41:00 +0000 Subject: [PATCH 45/63] fix: clean up partial metering worker starts --- .../server/services/search_metering_worker.py | 17 +++++++-- .../services/test_search_metering_worker.py | 37 +++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/reflexio/server/services/search_metering_worker.py b/reflexio/server/services/search_metering_worker.py index 8b879af50..4dc19d6f1 100644 --- a/reflexio/server/services/search_metering_worker.py +++ b/reflexio/server/services/search_metering_worker.py @@ -65,7 +65,7 @@ def start(self) -> bool: return False self._stop_event.clear() self._abort_event.clear() - self._threads = [ + threads = [ threading.Thread( target=self._worker_loop, name=f"search-metering-worker-{index}", @@ -73,8 +73,19 @@ def start(self) -> bool: ) for index in range(self.worker_count) ] - for thread in self._threads: - thread.start() + started_threads: list[threading.Thread] = [] + try: + for thread in threads: + thread.start() + started_threads.append(thread) + except BaseException: + self._stop_event.set() + for thread in started_threads: + thread.join() + self._stop_event.clear() + self._threads = [] + raise + self._threads = threads self._started = True logger.info( "event=search_metering_worker_started workers=%d queue_capacity=%d", diff --git a/tests/server/services/test_search_metering_worker.py b/tests/server/services/test_search_metering_worker.py index 6c4bf9104..37d793cf1 100644 --- a/tests/server/services/test_search_metering_worker.py +++ b/tests/server/services/test_search_metering_worker.py @@ -125,6 +125,43 @@ def blocked_process(_job: SearchMeteringJob) -> None: assert worker.stop() == 0 +def test_partial_start_failure_stops_started_threads_and_allows_retry( + monkeypatch, +) -> None: + worker = SearchMeteringWorker(worker_count=2) + original_start = threading.Thread.start + start_calls = 0 + started_threads: list[threading.Thread] = [] + + def fail_second_start(thread: threading.Thread) -> None: + nonlocal start_calls + start_calls += 1 + if start_calls == 2: + raise RuntimeError("thread start failed") + original_start(thread) + started_threads.append(thread) + + try: + with monkeypatch.context() as patch: + patch.setattr(threading.Thread, "start", fail_second_start) + with pytest.raises(RuntimeError, match="thread start failed"): + worker.start() + + assert len(started_threads) == 1 + assert not started_threads[0].is_alive() + assert worker._threads == [] + assert worker._started is False + + assert worker.start() is True + assert worker.stop() == 0 + assert worker.start() is False + finally: + worker._stop_event.set() + for thread in worker._threads: + if thread.ident is not None: + thread.join(timeout=1.0) + + def test_full_queue_drops_without_backpressure_and_reports(monkeypatch) -> None: started = threading.Event() release = threading.Event() From 00b8bfd0c8b0081b63bc9bb520a64615558e4083 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Mon, 10 Aug 2026 05:21:53 +0000 Subject: [PATCH 46/63] fix: preserve finalized session outcomes --- reflexio/server/services/storage/retention.py | 6 ---- .../services/storage/retention_mixin.py | 2 ++ .../test_storage_contract_session_outcomes.py | 33 +++++++++++++++++++ 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/reflexio/server/services/storage/retention.py b/reflexio/server/services/storage/retention.py index 9dcb86292..9e0dd73fd 100644 --- a/reflexio/server/services/storage/retention.py +++ b/reflexio/server/services/storage/retention.py @@ -50,12 +50,6 @@ class OptimizationRetentionClass: ), RetentionTarget("interactions", "interactions", "created_at", ("interaction_id",)), RetentionTarget("requests", "requests", "created_at", ("request_id",)), - RetentionTarget( - "session_outcomes", - "session_outcomes", - "created_at", - ("user_id", "session_id"), - ), RetentionTarget( "user_playbooks", "user_playbooks", diff --git a/reflexio/server/services/storage/retention_mixin.py b/reflexio/server/services/storage/retention_mixin.py index 8a0e4db29..f7462b757 100644 --- a/reflexio/server/services/storage/retention_mixin.py +++ b/reflexio/server/services/storage/retention_mixin.py @@ -233,6 +233,8 @@ def _retention_select_oldest_keys( statuses (tuple[str, ...] | None): When not None, restrict the select to rows whose ``status`` is one of these values. An empty tuple matches nothing (never every row). + older_than_epoch (int | None): When set, restrict the select to + rows whose target ordering column is strictly older. """ raise NotImplementedError diff --git a/tests/server/services/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index 6716214b4..43fa70ec8 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -2,6 +2,8 @@ from typing import cast +import pytest + from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, Request, @@ -63,6 +65,37 @@ def test_first_write_preserves_outcome_fields(storage: BaseStorage) -> None: assert records[0].metadata == {"crm": "test"} +def test_generic_retention_cannot_delete_finalized_session_outcomes( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="retention-r1", + user_id="u1", + session_id="retention-session", + source="published", + created_at=100, + ) + ) + storage.record_session_outcome( + SetSessionOutcomeRequest( + session_id="retention-session", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ), + created_at=102, + expected_context=storage.get_session_outcome_context("retention-session"), + ) + + with pytest.raises(ValueError, match="Unknown retention target: session_outcomes"): + storage.delete_oldest_retention_target_rows("session_outcomes", 1) # type: ignore[attr-defined] + + [record] = storage.get_session_outcomes( + GetSessionOutcomesRequest(session_ids=["retention-session"]) + ) + assert record.outcome == SessionOutcomeKind.SUCCESS + + def test_exact_finalization_retry_is_idempotent(storage: BaseStorage) -> None: storage.add_request( Request( From 3577a6f5b6c6371fdcf879d068f47fdc4cdc61d4 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Mon, 10 Aug 2026 06:10:55 +0000 Subject: [PATCH 47/63] fix: preserve legacy request sources on read --- AI_AGENT_INTEGRATION.md | 4 + reflexio/models/api_schema/domain/entities.py | 13 ++-- reflexio/models/api_schema/validators.py | 12 ++- .../storage/sqlite_storage/_requests.py | 4 +- tests/models/test_session_outcome_identity.py | 76 ++++++++++++------- .../test_session_outcome_migration.py | 70 ++++++++++++++--- .../storage/test_storage_contract_requests.py | 9 +++ 7 files changed, 145 insertions(+), 43 deletions(-) diff --git a/AI_AGENT_INTEGRATION.md b/AI_AGENT_INTEGRATION.md index 5952c0a85..e9538ec42 100644 --- a/AI_AGENT_INTEGRATION.md +++ b/AI_AGENT_INTEGRATION.md @@ -241,6 +241,10 @@ Publish request fields: | `force_extraction` | No | `False` for normal background publishing. `True` for manual learn-now, tests, or final flushes where you intentionally want extraction to run immediately. | `false` | | `wait_for_response` | SDK/query option | `False` on interactive paths. `True` only when the caller is prepared to wait for extraction results. | `false` | +The source contract applies to new publish and filter inputs. Read responses may +return historical source values created before this contract; Reflexio preserves +those values exactly because source participates in session-outcome identity. + Each interaction row should resemble Reflexio's `InteractionData` shape: ```json diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 59bc3f27a..9e61513f4 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -26,6 +26,7 @@ from ..validators import ( EmbeddingVector, NonEmptyStr, + PersistedSessionOutcomeSource, SessionOutcomeSource, TimeRangeValidatorMixin, _validate_image_url, @@ -279,8 +280,8 @@ class Request(BaseModel): user_id (str): Owner of the request. created_at (int): Unix epoch seconds at request creation. Defaults to the current UTC time. - source (str): Non-sensitive producer/workflow label. Non-empty values - use the session outcome source contract. + source (str): Producer/workflow label. Persisted reads preserve legacy + values verbatim; new publish inputs use the strict source contract. agent_version (str): The agent version that handled this request. session_id (str): Non-empty session this request belongs to. evaluation_only (bool): Whether this request is stored for @@ -295,7 +296,7 @@ class Request(BaseModel): request_id: str user_id: str created_at: int = Field(default_factory=lambda: int(datetime.now(UTC).timestamp())) - source: SessionOutcomeSource = "" + source: PersistedSessionOutcomeSource = "" agent_version: str = "" session_id: NonEmptyStr evaluation_only: bool = False @@ -831,13 +832,15 @@ class DeleteSessionResponse(BaseModel): class SessionOutcomeRecord(BaseModel): + """Persisted outcome row, including its immutable historical source.""" + outcome_id: NonEmptyStr | None = None outcome_revision: int | None = Field(default=None, ge=1) user_id: str session_id: NonEmptyStr outcome: SessionOutcomeKind occurred_at: int = Field(ge=0) - source: SessionOutcomeSource + source: PersistedSessionOutcomeSource label: str | None = Field(default=None, max_length=128) value: float | None = Field(default=None, allow_inf_nan=False) metadata: dict[str, Any] | None = None @@ -914,7 +917,7 @@ class SetSessionOutcomeResponse(BaseModel): reason: SessionOutcomeFailureReason | None = None message: str = "" user_id: str | None = None - source: SessionOutcomeSource | None = None + source: PersistedSessionOutcomeSource | None = None outcome_id: NonEmptyStr | None = None outcome_revision: int | None = Field(default=None, ge=1) outcome_contract_digest: Sha256Digest | None = None diff --git a/reflexio/models/api_schema/validators.py b/reflexio/models/api_schema/validators.py index f4a875f53..008acc91f 100644 --- a/reflexio/models/api_schema/validators.py +++ b/reflexio/models/api_schema/validators.py @@ -21,7 +21,7 @@ from typing import Annotated, Any, Literal from urllib.parse import urlparse -from pydantic import AfterValidator, HttpUrl, StringConstraints +from pydantic import AfterValidator, HttpUrl, StringConstraints, TypeAdapter # Embedding vector dimensions — must match config_schema.EMBEDDING_DIMENSIONS. # Duplicated here to avoid circular imports (config_schema imports from this module). @@ -115,6 +115,16 @@ def _check_embedding_dimensions(v: list[float]) -> list[float]: ) """Outcome producer/workflow label; empty preserves the existing absent-source value.""" +PersistedSessionOutcomeSource = str +"""Historical outcome source returned exactly as stored, without new-input validation.""" + +_SESSION_OUTCOME_SOURCE_ADAPTER = TypeAdapter(SessionOutcomeSource) + + +def validate_session_outcome_source(value: str) -> SessionOutcomeSource: + """Validate an outcome source before writing a new request.""" + return _SESSION_OUTCOME_SOURCE_ADAPTER.validate_python(value) + # ============================================================================= # Security Validators — SSRF Prevention diff --git a/reflexio/server/services/storage/sqlite_storage/_requests.py b/reflexio/server/services/storage/sqlite_storage/_requests.py index 46b557ee3..9f9143b0a 100644 --- a/reflexio/server/services/storage/sqlite_storage/_requests.py +++ b/reflexio/server/services/storage/sqlite_storage/_requests.py @@ -11,6 +11,7 @@ from reflexio.models.api_schema.service_schemas import ( Request, ) +from reflexio.models.api_schema.validators import validate_session_outcome_source from ._base import ( SQLiteStorageBase, @@ -40,6 +41,7 @@ class RequestMixin: @SQLiteStorageBase.handle_exceptions def add_request(self, request: Request) -> None: + source = validate_session_outcome_source(request.source) created_at_iso = _epoch_to_iso(request.created_at) subject_ref = self._subject_ref_for_user_id(request.user_id) with self._lock: @@ -58,7 +60,7 @@ def add_request(self, request: Request) -> None: request.request_id, request.user_id, created_at_iso, - request.source, + source, request.agent_version, request.session_id, 1 if request.evaluation_only else 0, diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index cb77b8d35..e2c7bd1c7 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -65,28 +65,13 @@ class _HasSource(Protocol): def source(self) -> str | None: ... -_OUTCOME_SOURCE_MODEL_FACTORIES: tuple[Callable[[str], _HasSource], ...] = ( - lambda source: Request( - request_id="request-1", - user_id="user-1", - session_id="session-1", - source=source, - ), +_STRICT_OUTCOME_SOURCE_INPUT_FACTORIES: tuple[Callable[[str], _HasSource], ...] = ( lambda source: PublishUserInteractionRequest( user_id="user-1", session_id="session-1", interaction_data_list=[InteractionData(content="hello")], source=source, ), - lambda source: SessionOutcomeRecord( - user_id="user-1", - session_id="session-1", - outcome=SessionOutcomeKind.SUCCESS, - occurred_at=1, - source=source, - created_at=2, - ), - lambda source: SetSessionOutcomeResponse(success=True, source=source), lambda source: GetSessionOutcomesRequest(source=source), lambda source: GetRequestsRequest(source=source), lambda source: SearchUserProfileRequest(user_id="user-1", source=source), @@ -96,11 +81,8 @@ def source(self) -> str | None: ... lambda source: ManualPlaybookGenerationRequest(source=source), lambda source: RerunPlaybookGenerationRequest(source=source), ) -_OUTCOME_SOURCE_MODEL_IDS = ( - "request", +_STRICT_OUTCOME_SOURCE_INPUT_IDS = ( "publish", - "record", - "set-response", "get-outcomes", "get-requests", "search-profiles", @@ -113,9 +95,11 @@ def source(self) -> str | None: ... @pytest.mark.parametrize( - "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS + "factory", + _STRICT_OUTCOME_SOURCE_INPUT_FACTORIES, + ids=_STRICT_OUTCOME_SOURCE_INPUT_IDS, ) -def test_outcome_source_models_accept_machine_label( +def test_outcome_source_inputs_accept_machine_label( factory: Callable[[str], _HasSource], ) -> None: assert factory("support-agent:v2").source == "support-agent:v2" @@ -123,11 +107,14 @@ def test_outcome_source_models_accept_machine_label( @pytest.mark.parametrize( - "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS + "factory", + _STRICT_OUTCOME_SOURCE_INPUT_FACTORIES, + ids=_STRICT_OUTCOME_SOURCE_INPUT_IDS, ) @pytest.mark.parametrize( "source", [ + "Legacy Source", "alice@example.com", "https://example.com/hook", "support agent", @@ -138,7 +125,7 @@ def test_outcome_source_models_accept_machine_label( "a" * 129, ], ) -def test_outcome_source_models_reject_sensitive_or_free_form_values( +def test_outcome_source_inputs_reject_sensitive_or_free_form_values( factory: Callable[[str], _HasSource], source: str, ) -> None: @@ -147,14 +134,51 @@ def test_outcome_source_models_reject_sensitive_or_free_form_values( @pytest.mark.parametrize( - "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS + "factory", + _STRICT_OUTCOME_SOURCE_INPUT_FACTORIES, + ids=_STRICT_OUTCOME_SOURCE_INPUT_IDS, ) -def test_outcome_source_models_preserve_empty_source( +def test_outcome_source_inputs_preserve_empty_source( factory: Callable[[str], _HasSource], ) -> None: assert factory("").source == "" +_PERSISTED_OUTCOME_SOURCE_MODEL_FACTORIES: tuple[Callable[[str], _HasSource], ...] = ( + lambda source: Request( + request_id="request-1", + user_id="user-1", + session_id="session-1", + source=source, + ), + lambda source: SessionOutcomeRecord( + user_id="user-1", + session_id="session-1", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source=source, + created_at=2, + ), + lambda source: SetSessionOutcomeResponse(success=True, source=source), +) + + +@pytest.mark.parametrize( + "factory", + _PERSISTED_OUTCOME_SOURCE_MODEL_FACTORIES, + ids=("request", "outcome", "finalization-response"), +) +@pytest.mark.parametrize( + "source", + ["Legacy Source", "legacy/source", "legacy-sourcé", "x" * 256], +) +def test_persisted_outcome_source_models_preserve_legacy_values( + factory: Callable[[str], _HasSource], + source: str, +) -> None: + assert factory(source).source == source + + def test_optional_outcome_source_models_preserve_absence() -> None: assert SetSessionOutcomeResponse(success=False).source is None assert GetSessionOutcomesRequest().source is None diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index c2602c6d0..36b38831f 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -10,6 +10,8 @@ GetSessionOutcomesRequest, Request, SessionOutcomeKind, + SetSessionOutcomeRequest, + SetSessionOutcomeResponse, ) from reflexio.server.services.storage.session_outcome_identity import ( outcome_contract_digest, @@ -18,6 +20,7 @@ from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.sqlite_storage._base import ( _canonical_session_snapshot, + _epoch_to_iso, _prefetch_canonical_session_trajectory_digests, ) @@ -142,7 +145,7 @@ def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( "session_id": "c", "outcome": "success", "occurred_at": 101, - "source": "legacy-source-one", + "source": "Legacy Source", "label": "resolved:one", "value": 2.5, "metadata": {"nested": {"a": 1}, "legacy": True}, @@ -162,16 +165,23 @@ def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( "created_at": 202, }, ] - for row in legacy_rows: - storage.add_request( - Request( - request_id=f"request-{row['session_id']}", - user_id=row["user_id"], - session_id=row["session_id"], - source=row["source"], - created_at=row["occurred_at"] - 1, + storage.conn.executemany( + """INSERT INTO requests ( + request_id, user_id, created_at, source, session_id, + governance_subject_ref + ) VALUES (?, ?, ?, ?, ?, ?)""", + [ + ( + f"request-{row['session_id']}", + row["user_id"], + _epoch_to_iso(row["occurred_at"] - 1), + row["source"], + row["session_id"], + row["governance_subject_ref"], ) - ) + for row in legacy_rows + ], + ) storage.conn.execute("DROP TABLE session_outcomes") storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) for row in legacy_rows: @@ -214,6 +224,9 @@ def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( assert records_by_session["c"].outcome_id == sha256(b'["a:b","c"]').hexdigest() assert records_by_session["b:c"].outcome_id == sha256(b'["a","b:c"]').hexdigest() assert records_by_session["c"].outcome_id != records_by_session["b:c"].outcome_id + persisted_request = migrated.get_request("request-c") + assert persisted_request is not None + assert persisted_request.source == "Legacy Source" for row in legacy_rows: record = records_by_session[row["session_id"]] @@ -241,6 +254,43 @@ def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( _canonical_session_snapshot(migrated.conn, row["session_id"]) ) + legacy_identity_before = dict(rows_by_session["c"]) + retry = migrated.record_session_outcome( + SetSessionOutcomeRequest( + session_id="c", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + label="resolved:one", + value=2.5, + metadata={"nested": {"a": 1}, "legacy": True}, + ), + created_at=103, + expected_context=migrated.get_session_outcome_context("c"), + ) + response = SetSessionOutcomeResponse( + success=True, + recorded=retry.recorded, + user_id=retry.user_id, + source=retry.source, + outcome_id=retry.outcome_id, + outcome_revision=retry.outcome_revision, + outcome_contract_digest=retry.outcome_contract_digest, + finalized_trajectory_digest=retry.finalized_trajectory_digest, + ) + legacy_identity_after = dict( + migrated.conn.execute( + "SELECT * FROM session_outcomes WHERE session_id = ?", ("c",) + ).fetchone() + ) + + assert retry.recorded is False + assert retry.reason is None + assert response.source == "Legacy Source" + assert legacy_identity_after == legacy_identity_before + assert legacy_identity_after["outcome_contract_digest"] == ( + "a3faaa8073272084ffdcb3d1b12c410676829d8376495062b9600e55bfcc4dd4" + ) + def test_migration_defaults_null_legacy_outcome_revision_to_one(tmp_path) -> None: db_path = str(tmp_path / "legacy-null-revision.db") diff --git a/tests/server/services/storage/test_storage_contract_requests.py b/tests/server/services/storage/test_storage_contract_requests.py index 8bf182bfc..bdfc77476 100644 --- a/tests/server/services/storage/test_storage_contract_requests.py +++ b/tests/server/services/storage/test_storage_contract_requests.py @@ -9,6 +9,7 @@ Request, UserActionType, ) +from reflexio.server.services.storage.error import StorageError from reflexio.server.services.storage.storage_base import BaseStorage pytestmark = pytest.mark.integration @@ -31,6 +32,14 @@ def _make_request( class TestRequestCRUD: + def test_add_request_rejects_legacy_source(self, storage: BaseStorage) -> None: + request = _make_request("legacy-source-write", "u1", source="Legacy Source") + + with pytest.raises(StorageError): + storage.add_request(request) + + assert storage.get_request(request.request_id) is None + def test_add_and_get_request(self, storage: BaseStorage) -> None: req = _make_request("r1", "u1") storage.add_request(req) From 6a584015584e0b663a4e2fc65476c8db9e7acd5f Mon Sep 17 00:00:00 2001 From: Guangyu Date: Mon, 10 Aug 2026 09:11:43 +0000 Subject: [PATCH 48/63] fix: stream session outcome trajectory digests --- .../storage/session_outcome_identity.py | 152 ++++++++++++-- .../services/storage/sqlite_storage/_base.py | 114 +++++++---- .../sqlite_storage/_session_outcomes.py | 106 ++++++---- tests/models/test_session_outcome_identity.py | 193 ++++++++++++++++++ .../test_session_outcome_migration.py | 11 +- .../test_storage_contract_session_outcomes.py | 163 ++++++++++++--- 6 files changed, 611 insertions(+), 128 deletions(-) diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py index 28e0e0064..076cfc535 100644 --- a/reflexio/server/services/storage/session_outcome_identity.py +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -2,6 +2,7 @@ import json from collections.abc import Collection, Mapping, Sequence +from dataclasses import dataclass from datetime import UTC, datetime from hashlib import sha256 from typing import TypedDict @@ -10,6 +11,11 @@ __all__ = [ "CanonicalSessionTrajectory", + "CanonicalTrajectoryDigestResult", + "CanonicalTrajectoryDigestAccumulator", + "OUTCOME_ALLOWED_VALUES", + "OUTCOME_FINALIZATION_RULE", + "OUTCOME_SCHEMA_VERSION", "canonical_json_bytes", "canonical_session_trajectory", "canonical_trajectory_bytes", @@ -18,6 +24,9 @@ ] MAX_CANONICAL_TRAJECTORY_JSON_DEPTH = 100 +OUTCOME_SCHEMA_VERSION = 1 +OUTCOME_FINALIZATION_RULE = "first_write" +OUTCOME_ALLOWED_VALUES = ("success", "failure", "unknown") class CanonicalRequest(TypedDict): @@ -61,6 +70,15 @@ class CanonicalSessionTrajectory(TypedDict): requests: list[CanonicalTrajectoryRequest] +@dataclass(frozen=True) +class CanonicalTrajectoryDigestResult: + """Digest and bounded context derived from one ordered trajectory stream.""" + + digest: str + first_request: dict[str, object] | None + request_count: int + + def _canonical_timestamp(value: object) -> str: if isinstance(value, datetime): if value.tzinfo is None: @@ -112,6 +130,28 @@ def _canonical_interaction( } +def _canonical_request(row: Mapping[str, object]) -> CanonicalRequest: + return { + "request_id": str(row["request_id"]), + "user_id": str(row["user_id"]), + "created_at": _canonical_timestamp(row["created_at"]), + "source": str(row["source"] or ""), + "agent_version": str(row["agent_version"] or ""), + "session_id": str(row["session_id"]), + "evaluation_only": bool(row["evaluation_only"]), + "retrieval_experiment_id": ( + str(row["retrieval_experiment_id"]) + if row["retrieval_experiment_id"] is not None + else None + ), + "retrieval_experiment_arm": ( + str(row["retrieval_experiment_arm"]) + if row["retrieval_experiment_arm"] is not None + else None + ), + } + + def canonical_session_trajectory( session_id: str, request_rows: Sequence[Mapping[str, object]], @@ -121,25 +161,7 @@ def canonical_session_trajectory( requests: list[CanonicalTrajectoryRequest] = [] for row in request_rows: request_id = str(row["request_id"]) - request: CanonicalRequest = { - "request_id": request_id, - "user_id": str(row["user_id"]), - "created_at": _canonical_timestamp(row["created_at"]), - "source": str(row["source"] or ""), - "agent_version": str(row["agent_version"] or ""), - "session_id": str(row["session_id"]), - "evaluation_only": bool(row["evaluation_only"]), - "retrieval_experiment_id": ( - str(row["retrieval_experiment_id"]) - if row["retrieval_experiment_id"] is not None - else None - ), - "retrieval_experiment_arm": ( - str(row["retrieval_experiment_arm"]) - if row["retrieval_experiment_arm"] is not None - else None - ), - } + request = _canonical_request(row) interactions = [ _canonical_interaction(interaction) for interaction in interactions_by_request.get(request_id, ()) @@ -203,6 +225,98 @@ def canonical_trajectory_bytes(trajectory: object) -> bytes: return _canonical_trajectory_json(trajectory).encode("utf-8") +def _canonical_trajectory_bytes_at_depth(value: object, *, depth: int) -> bytes: + return _canonical_trajectory_json(value, depth=depth).encode("utf-8") + + +class CanonicalTrajectoryDigestAccumulator: + """Hash one complete canonical trajectory while retaining only its current row.""" + + def __init__(self, session_id: str) -> None: + self._digest = sha256() + self._digest.update(b'{"requests":[') + self._session_id = session_id + self._request: CanonicalRequest | None = None + self._request_count = 0 + self._interaction_count = 0 + self._hexdigest: str | None = None + self._poisoned = False + + def _raise_if_poisoned(self) -> None: + if self._poisoned: + raise RuntimeError("canonical trajectory digest accumulator is invalid") + + def start_request(self, row: Mapping[str, object]) -> None: + self._raise_if_poisoned() + try: + if self._hexdigest is not None: + raise RuntimeError("canonical trajectory digest is already finalized") + if self._request is not None: + raise RuntimeError("previous canonical request is not finished") + request = _canonical_request(row) + if self._request_count: + self._digest.update(b",") + self._digest.update(b'{"interactions":[') + self._request = request + self._interaction_count = 0 + except Exception: + self._poisoned = True + raise + + def add_interaction(self, row: Mapping[str, object]) -> None: + self._raise_if_poisoned() + try: + if self._request is None: + raise RuntimeError("canonical interaction has no active request") + if str(row["request_id"]) != self._request["request_id"]: + raise ValueError( + "canonical interaction does not belong to active request" + ) + encoded = _canonical_trajectory_bytes_at_depth( + _canonical_interaction(row), depth=4 + ) + if self._interaction_count: + self._digest.update(b",") + self._digest.update(encoded) + self._interaction_count += 1 + except Exception: + self._poisoned = True + raise + + def finish_request(self) -> None: + self._raise_if_poisoned() + try: + if self._request is None: + raise RuntimeError("canonical trajectory has no active request") + encoded = _canonical_trajectory_bytes_at_depth(self._request, depth=3) + self._digest.update(b'],"request":') + self._digest.update(encoded) + self._digest.update(b"}") + self._request = None + self._request_count += 1 + except Exception: + self._poisoned = True + raise + + def hexdigest(self) -> str: + self._raise_if_poisoned() + try: + if self._request is not None: + raise RuntimeError("canonical trajectory has an unfinished request") + if self._hexdigest is None: + encoded_session_id = _canonical_trajectory_bytes_at_depth( + self._session_id, depth=1 + ) + self._digest.update(b'],"session_id":') + self._digest.update(encoded_session_id) + self._digest.update(b"}") + self._hexdigest = self._digest.hexdigest() + return self._hexdigest + except Exception: + self._poisoned = True + raise + + def trajectory_digest(trajectory: object) -> str: """Hash the canonical finalized session trajectory.""" return sha256(canonical_trajectory_bytes(trajectory)).hexdigest() diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index dfb471f59..be0956d7b 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -23,7 +23,6 @@ from typing import Any, ClassVar, Literal from reflexio.models.api_schema.common import BlockingIssue -from reflexio.models.api_schema.domain import SessionOutcomeKind from reflexio.models.api_schema.service_schemas import ( AgentPlaybook, AgentSuccessEvaluationResult, @@ -62,7 +61,9 @@ ) from reflexio.server.services.storage.retention_mixin import RetentionMixin, chunked from reflexio.server.services.storage.session_outcome_identity import ( - CanonicalSessionTrajectory, + OUTCOME_ALLOWED_VALUES, + CanonicalTrajectoryDigestAccumulator, + CanonicalTrajectoryDigestResult, canonical_json_bytes, canonical_session_trajectory, outcome_contract_digest, @@ -76,7 +77,7 @@ logger = logging.getLogger(__name__) -_OUTCOME_ALLOWED_VALUES = tuple(kind.value for kind in SessionOutcomeKind) +_TRAJECTORY_FETCH_SIZE = 256 _MINIMUM_SQLITE_VERSION = (3, 35, 0) _SQLITE_INITIALIZATION_LOCK_STRIPES = 64 _sqlite_initialization_locks = tuple( @@ -111,40 +112,83 @@ def _json_loads(text: str | None) -> Any: return json.loads(text) -def _canonical_session_snapshot( +def _canonical_session_trajectory_snapshot( conn: sqlite3.Connection, session_id: str -) -> CanonicalSessionTrajectory: - """Return the durable session state used to bind an outcome finalization.""" - request_rows = conn.execute( - """SELECT request_id, user_id, created_at, source, agent_version, session_id, - evaluation_only, retrieval_experiment_id, retrieval_experiment_arm - FROM requests WHERE session_id = ? - ORDER BY created_at ASC, request_id ASC""", +) -> CanonicalTrajectoryDigestResult: + """Hash and describe a complete bounded, ordered row stream.""" + cursor = conn.execute( + """SELECT requests.request_id, requests.user_id, requests.created_at, + requests.source, requests.agent_version, requests.session_id, + requests.evaluation_only, requests.retrieval_experiment_id, + requests.retrieval_experiment_arm, + requests.governance_subject_ref, + interactions.interaction_id, + interactions.user_id AS interaction_user_id, + interactions.request_id AS interaction_request_id, + interactions.created_at AS interaction_created_at, + interactions.content, interactions.role, interactions.token_count, + interactions.user_action, interactions.user_action_description, + interactions.interacted_image_url, interactions.image_encoding, + interactions.shadow_content, interactions.expert_content, + interactions.tools_used, interactions.citations, + interactions.retrieved_learnings + FROM requests + LEFT JOIN interactions + ON interactions.request_id = requests.request_id + WHERE requests.session_id = ? + ORDER BY requests.created_at ASC, requests.request_id ASC, + interactions.created_at ASC, interactions.interaction_id ASC""", (session_id,), - ).fetchall() - request_payloads = [dict(request) for request in request_rows] - interactions_by_request: dict[str, list[dict[str, object]]] = { - str(request["request_id"]): [] for request in request_rows - } - interaction_rows = conn.execute( - """SELECT interaction_id, user_id, request_id, created_at, content, role, - token_count, user_action, user_action_description, - interacted_image_url, image_encoding, shadow_content, - expert_content, tools_used, citations, retrieved_learnings - FROM interactions - WHERE request_id IN ( - SELECT request_id FROM requests WHERE session_id = ? - ) - ORDER BY request_id ASC, created_at ASC, interaction_id ASC""", - (session_id,), - ).fetchall() - for interaction in interaction_rows: - interactions_by_request[str(interaction["request_id"])].append( - dict(interaction) - ) - return canonical_session_trajectory( - session_id, request_payloads, interactions_by_request ) + accumulator = CanonicalTrajectoryDigestAccumulator(session_id) + active_request_id: str | None = None + first_request: dict[str, object] | None = None + request_count = 0 + while rows := cursor.fetchmany(_TRAJECTORY_FETCH_SIZE): + for row in rows: + request_id = str(row["request_id"]) + if request_id != active_request_id: + if active_request_id is not None: + accumulator.finish_request() + accumulator.start_request(row) + active_request_id = request_id + request_count += 1 + if first_request is None: + first_request = { + key: row[key] + for key in ( + "request_id", + "user_id", + "created_at", + "source", + "agent_version", + "session_id", + "evaluation_only", + "retrieval_experiment_id", + "retrieval_experiment_arm", + "governance_subject_ref", + ) + } + if row["interaction_id"] is None: + continue + interaction = dict(row) + interaction["user_id"] = row["interaction_user_id"] + interaction["request_id"] = row["interaction_request_id"] + interaction["created_at"] = row["interaction_created_at"] + accumulator.add_interaction(interaction) + if active_request_id is not None: + accumulator.finish_request() + return CanonicalTrajectoryDigestResult( + digest=accumulator.hexdigest(), + first_request=first_request, + request_count=request_count, + ) + + +def _canonical_session_trajectory_digest( + conn: sqlite3.Connection, session_id: str +) -> str: + return _canonical_session_trajectory_snapshot(conn, session_id).digest def _prefetch_canonical_session_trajectory_digests( @@ -1402,7 +1446,7 @@ def _migrate_session_outcomes_schema(self) -> None: or outcome_contract_digest( source=source, schema_version=1, - allowed_values=_OUTCOME_ALLOWED_VALUES, + allowed_values=OUTCOME_ALLOWED_VALUES, finalization_rule="first_write", ), existing_trajectory_digest diff --git a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py index d0657f4eb..a965c0eaf 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -2,7 +2,7 @@ import json import sqlite3 -from typing import Any +from typing import Any, cast from uuid import uuid4 from reflexio.models.api_schema.domain import ( @@ -13,8 +13,10 @@ ) from reflexio.server.services.storage.error import SubjectWriteBarrierError from reflexio.server.services.storage.session_outcome_identity import ( + OUTCOME_ALLOWED_VALUES, + OUTCOME_FINALIZATION_RULE, + OUTCOME_SCHEMA_VERSION, outcome_contract_digest, - trajectory_digest, ) from reflexio.server.services.storage.storage_base._session_outcomes import ( SessionOutcomeContext, @@ -22,15 +24,11 @@ ) from ._base import ( - _OUTCOME_ALLOWED_VALUES, SQLiteStorageBase, - _canonical_session_snapshot, + _canonical_session_trajectory_snapshot, _iso_to_epoch, ) -_OUTCOME_SCHEMA_VERSION = 1 -_OUTCOME_FINALIZATION_RULE = "first_write" - def _canonical_metadata_json(metadata: object) -> str | None: if metadata is None: @@ -78,21 +76,26 @@ def get_session_outcome_context(self, session_id: str) -> SessionOutcomeContext: source=str(existing["source"]), existing=True, ) - rows = self.conn.execute( + first = self.conn.execute( """SELECT user_id, source, created_at, request_id FROM requests WHERE session_id = ? - ORDER BY created_at ASC, request_id ASC""", + ORDER BY created_at ASC, request_id ASC LIMIT 1""", (session_id,), - ).fetchall() - if not rows: + ).fetchone() + if first is None: return SessionOutcomeContext() - first = rows[0] + counts = self.conn.execute( + """SELECT COUNT(DISTINCT user_id) AS user_count, + COUNT(DISTINCT COALESCE(source, '')) AS source_count + FROM requests WHERE session_id = ?""", + (session_id,), + ).fetchone() return SessionOutcomeContext( user_id=str(first["user_id"]), source=str(first["source"]), first_request_at=_iso_to_epoch(first["created_at"]), - user_contract_violation=len({str(row["user_id"]) for row in rows}) > 1, - source_contract_violation=len({str(row["source"]) for row in rows}) > 1, + user_contract_violation=int(counts["user_count"]) > 1, + source_contract_violation=int(counts["source_count"]) > 1, ) @SQLiteStorageBase.handle_exceptions @@ -111,12 +114,20 @@ def record_session_outcome( (request.session_id,), ).fetchone() if existing is not None: - first = self.conn.execute( + early_first = self.conn.execute( """SELECT user_id, source, governance_subject_ref FROM requests WHERE session_id = ? ORDER BY created_at ASC, request_id ASC LIMIT 1""", (request.session_id,), ).fetchone() + snapshot = ( + _canonical_session_trajectory_snapshot( + self.conn, request.session_id + ) + if early_first is not None + else None + ) + first = snapshot.first_request if snapshot is not None else None source = ( str(first["source"]) if first is not None @@ -132,11 +143,7 @@ def record_session_outcome( ) contract_digest = self._outcome_contract_digest(source) current_snapshot_digest = ( - trajectory_digest( - _canonical_session_snapshot(self.conn, request.session_id) - ) - if first is not None - else None + snapshot.digest if snapshot is not None else None ) stored_contract_digest = existing["outcome_contract_digest"] stored_snapshot_digest = existing["finalized_trajectory_digest"] @@ -196,21 +203,46 @@ def record_session_outcome( else None ), ) - first = self.conn.execute( - """SELECT user_id, source, created_at, request_id + barrier_first = self.conn.execute( + """SELECT user_id, source, created_at, request_id, + governance_subject_ref FROM requests WHERE session_id = ? ORDER BY created_at ASC, request_id ASC LIMIT 1""", (request.session_id,), ).fetchone() - if first is None: + if barrier_first is None: + self.conn.rollback() + return SessionOutcomeWriteResult( + recorded=False, + reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, + ) + barrier_user_id = str(barrier_first["user_id"]) + subject_ref = str( + barrier_first["governance_subject_ref"] + or self._subject_ref_for_user_id(barrier_user_id) + ) + try: + self._assert_subject_writable_locked(subject_ref) + except SubjectWriteBarrierError: + self.conn.rollback() + return SessionOutcomeWriteResult( + recorded=False, + user_id=barrier_user_id, + reason=SessionOutcomeFailureReason.SUBJECT_NOT_WRITABLE, + ) + snapshot = _canonical_session_trajectory_snapshot( + self.conn, request.session_id + ) + if snapshot.request_count == 0 or snapshot.first_request is None: self.conn.rollback() return SessionOutcomeWriteResult( recorded=False, reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, ) + first = snapshot.first_request user_id = str(first["user_id"]) - source = str(first["source"]) - first_request_at = _iso_to_epoch(first["created_at"]) + source = str(first["source"] or "") + first_request_at = _iso_to_epoch(cast(str, first["created_at"])) if ( expected_context.user_id != user_id or expected_context.source != source @@ -231,20 +263,12 @@ def record_session_outcome( source=source, reason=SessionOutcomeFailureReason.OCCURRED_BEFORE_SESSION, ) - subject_ref = self._subject_ref_for_user_id(user_id) - try: - self._assert_subject_writable_locked(subject_ref) - except SubjectWriteBarrierError: - self.conn.rollback() - return SessionOutcomeWriteResult( - recorded=False, - user_id=user_id, - reason=SessionOutcomeFailureReason.SUBJECT_NOT_WRITABLE, - ) - contract_digest = self._outcome_contract_digest(source) - snapshot_digest = trajectory_digest( - _canonical_session_snapshot(self.conn, request.session_id) + subject_ref = str( + first["governance_subject_ref"] + or self._subject_ref_for_user_id(user_id) ) + contract_digest = self._outcome_contract_digest(source) + snapshot_digest = snapshot.digest outcome_id = uuid4().hex self.conn.execute( """INSERT INTO session_outcomes @@ -357,7 +381,7 @@ def _metadata_json(request: SetSessionOutcomeRequest) -> str | None: def _outcome_contract_digest(source: str) -> str: return outcome_contract_digest( source=source, - schema_version=_OUTCOME_SCHEMA_VERSION, - allowed_values=_OUTCOME_ALLOWED_VALUES, - finalization_rule=_OUTCOME_FINALIZATION_RULE, + schema_version=OUTCOME_SCHEMA_VERSION, + allowed_values=OUTCOME_ALLOWED_VALUES, + finalization_rule=OUTCOME_FINALIZATION_RULE, ) diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index e2c7bd1c7..a6cae955f 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -28,6 +28,7 @@ ) from reflexio.server.services.storage import session_outcome_identity from reflexio.server.services.storage.session_outcome_identity import ( + CanonicalTrajectoryDigestAccumulator, canonical_json_bytes, canonical_session_trajectory, canonical_trajectory_bytes, @@ -642,3 +643,195 @@ def test_session_outcome_record_rejects_non_sha256_identity_digests( def test_trajectory_digest_matches_sha256_of_canonical_json() -> None: assert trajectory_digest({"b": 2, "a": 1}) == sha256(b'{"a":1,"b":2}').hexdigest() + + +def test_streaming_trajectory_digest_matches_materialized_multi_request_json() -> None: + request_rows = [ + { + "request_id": "request-b", + "user_id": "user-1", + "created_at": datetime(2026, 8, 10, 9, 0, tzinfo=UTC), + "source": "workflow:v2", + "agent_version": "agent-2", + "session_id": "session-1", + "evaluation_only": False, + "retrieval_experiment_id": "experiment-1", + "retrieval_experiment_arm": "treatment", + }, + { + "request_id": "request-a", + "user_id": "user-1", + "created_at": "2026-08-10T09:01:00+00:00", + "source": "workflow:v2", + "agent_version": "agent-2", + "session_id": "session-1", + "evaluation_only": True, + "retrieval_experiment_id": None, + "retrieval_experiment_arm": None, + }, + ] + interactions_by_request = { + "request-b": [ + { + "interaction_id": 11, + "user_id": "user-1", + "request_id": "request-b", + "created_at": datetime(2026, 8, 10, 9, 0, 1, tzinfo=UTC), + "content": "nested payload", + "role": "User", + "token_count": 3, + "user_action": "none", + "user_action_description": None, + "interacted_image_url": "", + "image_encoding": None, + "shadow_content": "", + "expert_content": None, + "tools_used": '[{"input":{"z":1,"a":[true,null,{"x":"y"}]} }]', + "citations": [{"metadata": {"beta": 2, "alpha": 1}}], + "retrieved_learnings": {"items": [[1, 2], {"nested": [3.5]}]}, + }, + { + "interaction_id": "12", + "user_id": "user-1", + "request_id": "request-b", + "created_at": "2026-08-10T09:00:02+00:00", + "content": "second", + "role": "Assistant", + "token_count": None, + "user_action": "none", + "user_action_description": "", + "interacted_image_url": None, + "image_encoding": "", + "shadow_content": None, + "expert_content": "", + "tools_used": [], + "citations": "[]", + "retrieved_learnings": None, + }, + ], + "request-a": [], + } + materialized = canonical_session_trajectory( + "session-1", request_rows, interactions_by_request + ) + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + + for request_row in request_rows: + accumulator.start_request(request_row) + for interaction_row in interactions_by_request[request_row["request_id"]]: + accumulator.add_interaction(interaction_row) + accumulator.finish_request() + + assert accumulator.hexdigest() == trajectory_digest(materialized) + + +def _streaming_request_row() -> dict[str, object]: + return { + "request_id": "request-1", + "user_id": "user-1", + "created_at": "2026-08-10T09:00:00+00:00", + "source": "workflow:v2", + "agent_version": "agent-2", + "session_id": "session-1", + "evaluation_only": False, + "retrieval_experiment_id": None, + "retrieval_experiment_arm": None, + } + + +def _streaming_interaction_row(*, tools_used: object = ()) -> dict[str, object]: + return { + "interaction_id": 1, + "user_id": "user-1", + "request_id": "request-1", + "created_at": "2026-08-10T09:00:01+00:00", + "content": "nested payload", + "role": "User", + "token_count": 3, + "user_action": "none", + "user_action_description": "", + "interacted_image_url": "", + "image_encoding": "", + "shadow_content": "", + "expert_content": "", + "tools_used": tools_used, + "citations": [], + "retrieved_learnings": [], + } + + +@pytest.mark.parametrize( + ("nested_container_count", "raises"), + [(95, False), (96, True)], + ids=["exact-boundary", "over-boundary"], +) +def test_streaming_trajectory_digest_has_materialized_depth_parity( + nested_container_count: int, + raises: bool, +) -> None: + tools_used: object = "leaf" + for _ in range(nested_container_count): + tools_used = [tools_used] + request_row = _streaming_request_row() + interaction_row = _streaming_interaction_row(tools_used=tools_used) + materialized = canonical_session_trajectory( + "session-1", [request_row], {"request-1": [interaction_row]} + ) + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + accumulator.start_request(request_row) + + if raises: + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + trajectory_digest(materialized) + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + accumulator.add_interaction(interaction_row) + else: + accumulator.add_interaction(interaction_row) + accumulator.finish_request() + assert accumulator.hexdigest() == trajectory_digest(materialized) + + +@pytest.mark.parametrize( + "first_error", + ["bad-request", "no-request", "wrong-request", "bad-json", "unfinished"], +) +def test_streaming_trajectory_digest_is_permanently_poisoned_after_error( + first_error: str, +) -> None: + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + request_row = _streaming_request_row() + + with pytest.raises((KeyError, RuntimeError, ValueError)): + if first_error == "bad-request": + accumulator.start_request({}) + elif first_error == "no-request": + accumulator.add_interaction(_streaming_interaction_row()) + else: + accumulator.start_request(request_row) + if first_error == "wrong-request": + accumulator.add_interaction( + { + **_streaming_interaction_row(), + "request_id": "request-2", + } + ) + elif first_error == "bad-json": + accumulator.add_interaction(_streaming_interaction_row(tools_used="{")) + else: + accumulator.hexdigest() + + for operation in ( + lambda: accumulator.start_request(request_row), + lambda: accumulator.add_interaction(_streaming_interaction_row()), + accumulator.finish_request, + accumulator.hexdigest, + ): + with pytest.raises( + RuntimeError, + match="canonical trajectory digest accumulator is invalid$", + ): + operation() diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index 36b38831f..08710f2cd 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -15,11 +15,10 @@ ) from reflexio.server.services.storage.session_outcome_identity import ( outcome_contract_digest, - trajectory_digest, ) from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.sqlite_storage._base import ( - _canonical_session_snapshot, + _canonical_session_trajectory_digest, _epoch_to_iso, _prefetch_canonical_session_trajectory_digests, ) @@ -250,8 +249,8 @@ def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( allowed_values={"success", "failure", "unknown"}, finalization_rule="first_write", ) - assert record.finalized_trajectory_digest == trajectory_digest( - _canonical_session_snapshot(migrated.conn, row["session_id"]) + assert record.finalized_trajectory_digest == ( + _canonical_session_trajectory_digest(migrated.conn, row["session_id"]) ) legacy_identity_before = dict(rows_by_session["c"]) @@ -412,8 +411,8 @@ def test_migration_prefetch_retains_only_trajectory_digests(tmp_path) -> None: ) assert digests == { - "digest-session": trajectory_digest( - _canonical_session_snapshot(storage.conn, "digest-session") + "digest-session": _canonical_session_trajectory_digest( + storage.conn, "digest-session" ) } assert all(isinstance(digest, str) for digest in digests.values()) diff --git a/tests/server/services/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index 43fa70ec8..d88c5c569 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -1,20 +1,22 @@ """Session outcome storage contract.""" -from typing import cast +from typing import Any, cast import pytest from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, + Interaction, Request, SessionOutcomeFailureReason, SessionOutcomeKind, SetSessionOutcomeRequest, ) -from reflexio.server.services.storage.sqlite_storage import SQLiteStorage -from reflexio.server.services.storage.sqlite_storage._base import ( - _canonical_session_snapshot, +from reflexio.server.services.storage.session_outcome_identity import ( + canonical_session_trajectory, + trajectory_digest, ) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.storage_base import BaseStorage @@ -297,36 +299,143 @@ def test_legacy_all_null_identity_changed_governance_context_conflicts( assert retry.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION -def test_sqlite_canonical_snapshot_loads_interactions_in_one_query( +class _NoFetchAllCursor: + def __init__(self, cursor: Any, fetch_sizes: list[int]) -> None: + self._cursor = cursor + self._fetch_sizes = fetch_sizes + + def fetchall(self) -> Any: + raise AssertionError("session outcome finalization must not call fetchall") + + def fetchmany(self, size: int) -> Any: + self._fetch_sizes.append(size) + return self._cursor.fetchmany(size) + + def __getattr__(self, name: str) -> Any: + return getattr(self._cursor, name) + + +class _NoFetchAllConnection: + def __init__(self, connection: Any) -> None: + self._connection = connection + self.fetch_sizes: list[int] = [] + + def execute(self, *args: Any, **kwargs: Any) -> _NoFetchAllCursor: + return _NoFetchAllCursor( + self._connection.execute(*args, **kwargs), self.fetch_sizes + ) + + def __getattr__(self, name: str) -> Any: + return getattr(self._connection, name) + + +def test_sqlite_large_finalization_streams_complete_digest_and_preserves_retry_contract( storage: BaseStorage, ) -> None: sqlite_storage = cast(SQLiteStorage, storage) - for index in range(3): - storage.add_request( - Request( - request_id=f"snapshot-r{index}", - user_id="snapshot-user", - session_id="snapshot-session", - source="published", - created_at=100 + index, - ) + session_id = "large-streamed-session" + request_id = "large-streamed-request" + storage.add_request( + Request( + request_id=request_id, + user_id="stream-user", + session_id=session_id, + source="published", + created_at=100, + ) + ) + interactions = [ + Interaction( + interaction_id=index + 1, + user_id="stream-user", + request_id=request_id, + created_at=101 + index, + content=f"complete-row-{index}", + token_count=index, + ) + for index in range(300) + ] + sqlite_storage.add_user_interactions_bulk( + "stream-user", interactions, embeddings_prepared=True + ) + request_rows = sqlite_storage.conn.execute( + """SELECT request_id, user_id, created_at, source, agent_version, session_id, + evaluation_only, retrieval_experiment_id, retrieval_experiment_arm + FROM requests WHERE session_id = ? + ORDER BY created_at ASC, request_id ASC""", + (session_id,), + ).fetchall() + interaction_rows = sqlite_storage.conn.execute( + """SELECT interaction_id, user_id, request_id, created_at, content, role, + token_count, user_action, user_action_description, + interacted_image_url, image_encoding, shadow_content, + expert_content, tools_used, citations, retrieved_learnings + FROM interactions WHERE request_id = ? + ORDER BY created_at ASC, interaction_id ASC""", + (request_id,), + ).fetchall() + expected_digest = trajectory_digest( + canonical_session_trajectory( + session_id, + [dict(row) for row in request_rows], + {request_id: [dict(row) for row in interaction_rows]}, + ) + ) + outcome = SetSessionOutcomeRequest( + session_id=session_id, + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=500, + metadata={"complete": True}, + ) + raw_connection = sqlite_storage.conn + guarded_connection = _NoFetchAllConnection(raw_connection) + cast(Any, sqlite_storage).conn = guarded_connection + try: + context = storage.get_session_outcome_context(session_id) + first = storage.record_session_outcome( + outcome, created_at=501, expected_context=context + ) + retry = storage.record_session_outcome( + outcome, + created_at=502, + expected_context=storage.get_session_outcome_context(session_id), ) - statements: list[str] = [] - sqlite_storage.conn.set_trace_callback(statements.append) + finally: + cast(Any, sqlite_storage).conn = raw_connection + + sqlite_storage.add_user_interactions_bulk( + "stream-user", + [ + Interaction( + interaction_id=301, + user_id="stream-user", + request_id=request_id, + created_at=401, + content="post-finalization-row", + ) + ], + embeddings_prepared=True, + ) + guarded_connection = _NoFetchAllConnection(raw_connection) + cast(Any, sqlite_storage).conn = guarded_connection try: - snapshot = _canonical_session_snapshot(sqlite_storage.conn, "snapshot-session") + conflict = storage.record_session_outcome( + outcome, + created_at=503, + expected_context=storage.get_session_outcome_context(session_id), + ) finally: - sqlite_storage.conn.set_trace_callback(None) + cast(Any, sqlite_storage).conn = raw_connection - interaction_queries = [ - statement for statement in statements if "FROM interactions" in statement - ] - assert [item["request"]["request_id"] for item in snapshot["requests"]] == [ - "snapshot-r0", - "snapshot-r1", - "snapshot-r2", - ] - assert len(interaction_queries) == 1 + assert first.recorded is True + assert first.finalized_trajectory_digest == expected_digest + assert retry.recorded is False + assert retry.reason is None + assert retry.finalized_trajectory_digest == expected_digest + assert conflict.recorded is False + assert conflict.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + assert guarded_connection.fetch_sizes + assert len(set(guarded_connection.fetch_sizes)) == 1 def test_changed_contract_identity_is_conflicting_finalization( From 81447e04ab3eddd45389148ccd71c26edcb0055a Mon Sep 17 00:00:00 2001 From: Guangyu Date: Mon, 10 Aug 2026 13:18:51 +0000 Subject: [PATCH 49/63] fix: make exposure retention policy fixed --- reflexio/server/services/storage/retention.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/reflexio/server/services/storage/retention.py b/reflexio/server/services/storage/retention.py index 9e0dd73fd..1296545d5 100644 --- a/reflexio/server/services/storage/retention.py +++ b/reflexio/server/services/storage/retention.py @@ -22,6 +22,7 @@ class RetentionTarget: id_columns: tuple[str, ...] priority_statuses: tuple[str, ...] = () minimum_age_seconds: int = 0 + fixed_row_limit: int | None = None @dataclass(frozen=True, slots=True) @@ -138,6 +139,7 @@ class OptimizationRetentionClass: "ingested_at", ("exposure_event_id",), minimum_age_seconds=OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS, + fixed_row_limit=DEFAULT_ROW_RETENTION_LIMIT, ), RetentionTarget("skills", "skills", "created_at", ("skill_id",)), ) @@ -195,12 +197,16 @@ class CascadeRef: def get_row_retention_limits() -> dict[str, int]: """Return per-target row limits from env with code defaults. - ``REFLEXIO_ROW_LIMIT_`` takes precedence for every target. + ``REFLEXIO_ROW_LIMIT_`` takes precedence for targets without a + ``fixed_row_limit``. Fixed targets explicitly reject that override path. ``INTERACTION_CLEANUP_THRESHOLD`` remains the legacy override for interactions when the new variable is not present. """ limits: dict[str, int] = {} for target in RETENTION_TARGETS: + if target.fixed_row_limit is not None: + limits[target.name] = target.fixed_row_limit + continue env_name = f"REFLEXIO_ROW_LIMIT_{target.name.upper()}" default = DEFAULT_ROW_RETENTION_LIMIT if target.name == "interactions": From d3e88f531d1e9dbe786bc42c6ea587e8080f7c6f Mon Sep 17 00:00:00 2001 From: Guangyu Date: Mon, 10 Aug 2026 13:58:07 +0000 Subject: [PATCH 50/63] fix: bound shared maintenance lifecycle --- .../server/services/lineage/gc_scheduler.py | 9 +- .../services/storage/sqlite_storage/_base.py | 265 ++++++++++-------- .../test_session_outcome_migration.py | 21 +- 3 files changed, 177 insertions(+), 118 deletions(-) diff --git a/reflexio/server/services/lineage/gc_scheduler.py b/reflexio/server/services/lineage/gc_scheduler.py index 55cd405cb..8e33701aa 100644 --- a/reflexio/server/services/lineage/gc_scheduler.py +++ b/reflexio/server/services/lineage/gc_scheduler.py @@ -118,7 +118,8 @@ def register_global_sweep(fn: Callable[[int], int]) -> None: fn (Callable[[int], int]): Called with the current unix epoch; returns the number of rows it deleted. """ - _global_sweep_hooks.append(fn) + if fn not in _global_sweep_hooks: + _global_sweep_hooks.append(fn) def clear_global_sweeps() -> None: @@ -142,7 +143,8 @@ def register_always_global_sweep(fn: Callable[[int], int]) -> None: fn (Callable[[int], int]): Called with the current unix epoch; returns the number of rows it processed. """ - _always_global_sweep_hooks.append(fn) + if fn not in _always_global_sweep_hooks: + _always_global_sweep_hooks.append(fn) def clear_always_global_sweeps() -> None: @@ -172,7 +174,8 @@ def register_per_org_sweep(fn: Callable[[str, int], int]) -> None: fn (Callable[[str, int], int]): Called with ``(org_id, now)`` where ``now`` is the current unix epoch; returns the number of rows deleted. """ - _per_org_sweep_hooks.append(fn) + if fn not in _per_org_sweep_hooks: + _per_org_sweep_hooks.append(fn) def clear_per_org_sweeps() -> None: diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index be0956d7b..c1b25ae47 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -62,12 +62,11 @@ from reflexio.server.services.storage.retention_mixin import RetentionMixin, chunked from reflexio.server.services.storage.session_outcome_identity import ( OUTCOME_ALLOWED_VALUES, + OUTCOME_FINALIZATION_RULE, CanonicalTrajectoryDigestAccumulator, CanonicalTrajectoryDigestResult, canonical_json_bytes, - canonical_session_trajectory, outcome_contract_digest, - trajectory_digest, ) from reflexio.server.services.storage.storage_base import BaseStorage from reflexio.server.site_var.site_var_manager import SiteVarManager @@ -78,6 +77,7 @@ logger = logging.getLogger(__name__) _TRAJECTORY_FETCH_SIZE = 256 +_SESSION_OUTCOME_MIGRATION_BATCH_SIZE = 256 _MINIMUM_SQLITE_VERSION = (3, 35, 0) _SQLITE_INITIALIZATION_LOCK_STRIPES = 64 _sqlite_initialization_locks = tuple( @@ -194,22 +194,20 @@ def _canonical_session_trajectory_digest( def _prefetch_canonical_session_trajectory_digests( conn: sqlite3.Connection, session_ids: Sequence[str] ) -> dict[str, str]: - """Derive trajectory digests while retaining only one payload chunk.""" + """Derive trajectory digests from ordered joined rows in bounded chunks.""" digests: dict[str, str] = {} for session_id_chunk in chunked(list(dict.fromkeys(session_ids))): placeholders = ",".join("?" for _ in session_id_chunk) - request_rows = conn.execute( - f"""SELECT request_id, user_id, created_at, source, agent_version, - session_id, evaluation_only, retrieval_experiment_id, - retrieval_experiment_arm - FROM requests WHERE session_id IN ({placeholders}) - ORDER BY session_id ASC, created_at ASC, request_id ASC""", # noqa: S608 - session_id_chunk, - ).fetchall() - interaction_rows = conn.execute( - f"""SELECT requests.session_id AS trajectory_session_id, - interactions.interaction_id, interactions.user_id, - interactions.request_id, interactions.created_at, + cursor = conn.execute( + f"""SELECT requests.request_id, requests.user_id, requests.created_at, + requests.source, requests.agent_version, requests.session_id, + requests.evaluation_only, requests.retrieval_experiment_id, + requests.retrieval_experiment_arm, + requests.governance_subject_ref, + interactions.interaction_id, + interactions.user_id AS interaction_user_id, + interactions.request_id AS interaction_request_id, + interactions.created_at AS interaction_created_at, interactions.content, interactions.role, interactions.token_count, interactions.user_action, interactions.user_action_description, @@ -217,38 +215,52 @@ def _prefetch_canonical_session_trajectory_digests( interactions.image_encoding, interactions.shadow_content, interactions.expert_content, interactions.tools_used, interactions.citations, interactions.retrieved_learnings - FROM interactions - JOIN requests ON requests.request_id = interactions.request_id + FROM requests + LEFT JOIN interactions + ON interactions.request_id = requests.request_id WHERE requests.session_id IN ({placeholders}) - ORDER BY requests.session_id ASC, interactions.request_id ASC, + ORDER BY requests.session_id ASC, requests.created_at ASC, + requests.request_id ASC, interactions.created_at ASC, interactions.interaction_id ASC""", # noqa: S608 session_id_chunk, - ).fetchall() - - requests_by_session: dict[str, list[dict[str, object]]] = { - session_id: [] for session_id in session_id_chunk - } - interactions_by_session: dict[str, dict[str, list[dict[str, object]]]] = { - session_id: {} for session_id in session_id_chunk - } - for request in request_rows: - session_id = str(request["session_id"]) - requests_by_session[session_id].append(dict(request)) - interactions_by_session[session_id][str(request["request_id"])] = [] - for interaction in interaction_rows: - interaction_payload = dict(interaction) - session_id = str(interaction_payload.pop("trajectory_session_id")) - interactions_by_session[session_id][ - str(interaction_payload["request_id"]) - ].append(interaction_payload) + ) + active_session_id: str | None = None + active_request_id: str | None = None + accumulator: CanonicalTrajectoryDigestAccumulator | None = None + while rows := cursor.fetchmany(_TRAJECTORY_FETCH_SIZE): + for row in rows: + session_id = str(row["session_id"]) + request_id = str(row["request_id"]) + if accumulator is None or session_id != active_session_id: + if accumulator is not None: + if active_request_id is not None: + accumulator.finish_request() + assert active_session_id is not None # noqa: S101 + digests[active_session_id] = accumulator.hexdigest() + accumulator = CanonicalTrajectoryDigestAccumulator(session_id) + active_session_id = session_id + active_request_id = None + if request_id != active_request_id: + if active_request_id is not None: + accumulator.finish_request() + accumulator.start_request(row) + active_request_id = request_id + if row["interaction_id"] is not None: + interaction = dict(row) + interaction["user_id"] = row["interaction_user_id"] + interaction["request_id"] = row["interaction_request_id"] + interaction["created_at"] = row["interaction_created_at"] + accumulator.add_interaction(interaction) + if accumulator is not None: + if active_request_id is not None: + accumulator.finish_request() + assert active_session_id is not None # noqa: S101 + digests[active_session_id] = accumulator.hexdigest() for session_id in session_id_chunk: - digests[session_id] = trajectory_digest( - canonical_session_trajectory( - session_id, - requests_by_session[session_id], - interactions_by_session[session_id], - ) + digests.setdefault( + session_id, + CanonicalTrajectoryDigestAccumulator(session_id).hexdigest(), ) return digests @@ -1360,18 +1372,6 @@ def _migrate_session_outcomes_schema(self) -> None: self.conn.commit() return - legacy_rows = self.conn.execute( - "SELECT * FROM session_outcomes" - ).fetchall() - trajectory_digests = _prefetch_canonical_session_trajectory_digests( - self.conn, - [ - str(row["session_id"]) - for row in legacy_rows - if "finalized_trajectory_digest" not in columns - or not row["finalized_trajectory_digest"] - ], - ) self.conn.execute( "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" ) @@ -1394,67 +1394,108 @@ def _migrate_session_outcomes_schema(self) -> None: PRIMARY KEY (user_id, session_id) )""" ) - for row in legacy_rows: - source = str(row["source"]) - subject_ref = ( - row["governance_subject_ref"] - if "governance_subject_ref" in columns - else None - ) - if subject_ref is None or not str(subject_ref).strip(): - subject_ref = self._subject_ref_for_user_id(str(row["user_id"])) - existing_outcome_id = ( - row["outcome_id"] if "outcome_id" in columns else None - ) - existing_contract_digest = ( - row["outcome_contract_digest"] - if "outcome_contract_digest" in columns - else None - ) - existing_trajectory_digest = ( - row["finalized_trajectory_digest"] - if "finalized_trajectory_digest" in columns - else None - ) - self.conn.execute( - """INSERT INTO session_outcomes ( - outcome_id, outcome_revision, user_id, session_id, outcome, - occurred_at, source, label, value, metadata, - outcome_contract_digest, finalized_trajectory_digest, - governance_subject_ref, created_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - existing_outcome_id - or _legacy_session_outcome_id( - str(row["user_id"]), str(row["session_id"]) - ), + last_user_id: str | None = None + last_session_id: str | None = None + while True: + if last_user_id is None: + legacy_rows = self.conn.execute( + """SELECT * FROM session_outcomes_legacy + ORDER BY user_id, session_id + LIMIT ?""", + (_SESSION_OUTCOME_MIGRATION_BATCH_SIZE,), + ).fetchall() + else: + legacy_rows = self.conn.execute( + """SELECT * FROM session_outcomes_legacy + WHERE user_id > ? + OR (user_id = ? AND session_id > ?) + ORDER BY user_id, session_id + LIMIT ?""", ( - row["outcome_revision"] - if "outcome_revision" in columns - and row["outcome_revision"] is not None - else 1 + last_user_id, + last_user_id, + last_session_id, + _SESSION_OUTCOME_MIGRATION_BATCH_SIZE, ), - row["user_id"], - row["session_id"], - row["outcome"], - row["occurred_at"], - source, - row["label"] if "label" in columns else None, - row["value"] if "value" in columns else None, - row["metadata"] if "metadata" in columns else None, - existing_contract_digest - or outcome_contract_digest( - source=source, - schema_version=1, - allowed_values=OUTCOME_ALLOWED_VALUES, - finalization_rule="first_write", - ), - existing_trajectory_digest - or trajectory_digests[str(row["session_id"])], - subject_ref, - row["created_at"], - ), + ).fetchall() + if not legacy_rows: + break + trajectory_digests = _prefetch_canonical_session_trajectory_digests( + self.conn, + [ + str(row["session_id"]) + for row in legacy_rows + if "finalized_trajectory_digest" not in columns + or not row["finalized_trajectory_digest"] + ], ) + for row in legacy_rows: + source = str(row["source"]) + subject_ref = ( + row["governance_subject_ref"] + if "governance_subject_ref" in columns + else None + ) + if subject_ref is None or not str(subject_ref).strip(): + subject_ref = self._subject_ref_for_user_id( + str(row["user_id"]) + ) + existing_outcome_id = ( + row["outcome_id"] if "outcome_id" in columns else None + ) + existing_contract_digest = ( + row["outcome_contract_digest"] + if "outcome_contract_digest" in columns + else None + ) + existing_trajectory_digest = ( + row["finalized_trajectory_digest"] + if "finalized_trajectory_digest" in columns + else None + ) + self.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, + outcome, occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + existing_outcome_id + or _legacy_session_outcome_id( + str(row["user_id"]), str(row["session_id"]) + ), + ( + row["outcome_revision"] + if "outcome_revision" in columns + and row["outcome_revision"] is not None + else 1 + ), + row["user_id"], + row["session_id"], + row["outcome"], + row["occurred_at"], + source, + row["label"] if "label" in columns else None, + row["value"] if "value" in columns else None, + row["metadata"] if "metadata" in columns else None, + existing_contract_digest + or outcome_contract_digest( + source=source, + schema_version=1, + allowed_values=OUTCOME_ALLOWED_VALUES, + finalization_rule=OUTCOME_FINALIZATION_RULE, + ), + existing_trajectory_digest + or trajectory_digests[str(row["session_id"])], + subject_ref, + row["created_at"], + ), + ) + last_user_id = str(legacy_rows[-1]["user_id"]) + last_session_id = str(legacy_rows[-1]["session_id"]) + if len(legacy_rows) < _SESSION_OUTCOME_MIGRATION_BATCH_SIZE: + break self.conn.execute("DROP TABLE session_outcomes_legacy") self.conn.execute( "CREATE INDEX idx_session_outcomes_occurred_at " diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index 08710f2cd..7f93b52c9 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -330,7 +330,9 @@ def test_migration_rejects_zero_legacy_outcome_revision_atomically(tmp_path) -> assert stranded_legacy_table is None -def test_migration_prefetches_trajectory_inputs_in_bounded_chunks(tmp_path) -> None: +def test_migration_streams_legacy_outcomes_and_trajectories_in_bounded_batches( + tmp_path, +) -> None: db_path = str(tmp_path / "legacy-session-outcome-query-scaling.db") storage = SQLiteStorage(org_id="legacy-query-scaling", db_path=db_path) row_count = 501 @@ -378,13 +380,26 @@ def test_migration_prefetches_trajectory_inputs_in_bounded_chunks(tmp_path) -> N finally: storage.conn.set_trace_callback(None) + legacy_outcome_queries = [ + statement + for statement in statements + if statement.lstrip().upper().startswith("SELECT") + and "FROM session_outcomes_legacy" in statement + ] + assert len(legacy_outcome_queries) == 2 + assert all("LIMIT 256" in statement for statement in legacy_outcome_queries) + assert not any( + statement.strip() == "SELECT * FROM session_outcomes" + for statement in statements + ) trajectory_input_queries = [ statement for statement in statements if statement.lstrip().upper().startswith("SELECT") - and ("FROM requests" in statement or "FROM interactions" in statement) + and "FROM requests" in statement ] - assert len(trajectory_input_queries) == 4 + assert len(trajectory_input_queries) == 2 + assert all("LEFT JOIN interactions" in query for query in trajectory_input_queries) assert ( storage.conn.execute("SELECT COUNT(*) FROM session_outcomes").fetchone()[0] == row_count From 09b7c60ea669a75b59e9acda373771a232f53509 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 08:13:16 +0000 Subject: [PATCH 51/63] fix: stream canonical trajectory sizing --- .../storage/session_outcome_identity.py | 47 +++++++++++++----- tests/models/test_session_outcome_identity.py | 48 +++++++++++++++++++ 2 files changed, 84 insertions(+), 11 deletions(-) diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py index 076cfc535..ff0820d28 100644 --- a/reflexio/server/services/storage/session_outcome_identity.py +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -234,7 +234,8 @@ class CanonicalTrajectoryDigestAccumulator: def __init__(self, session_id: str) -> None: self._digest = sha256() - self._digest.update(b'{"requests":[') + self._byte_count = 0 + self._update(b'{"requests":[') self._session_id = session_id self._request: CanonicalRequest | None = None self._request_count = 0 @@ -246,6 +247,10 @@ def _raise_if_poisoned(self) -> None: if self._poisoned: raise RuntimeError("canonical trajectory digest accumulator is invalid") + def _update(self, encoded: bytes) -> None: + self._digest.update(encoded) + self._byte_count += len(encoded) + def start_request(self, row: Mapping[str, object]) -> None: self._raise_if_poisoned() try: @@ -255,8 +260,8 @@ def start_request(self, row: Mapping[str, object]) -> None: raise RuntimeError("previous canonical request is not finished") request = _canonical_request(row) if self._request_count: - self._digest.update(b",") - self._digest.update(b'{"interactions":[') + self._update(b",") + self._update(b'{"interactions":[') self._request = request self._interaction_count = 0 except Exception: @@ -276,8 +281,8 @@ def add_interaction(self, row: Mapping[str, object]) -> None: _canonical_interaction(row), depth=4 ) if self._interaction_count: - self._digest.update(b",") - self._digest.update(encoded) + self._update(b",") + self._update(encoded) self._interaction_count += 1 except Exception: self._poisoned = True @@ -289,15 +294,35 @@ def finish_request(self) -> None: if self._request is None: raise RuntimeError("canonical trajectory has no active request") encoded = _canonical_trajectory_bytes_at_depth(self._request, depth=3) - self._digest.update(b'],"request":') - self._digest.update(encoded) - self._digest.update(b"}") + self._update(b'],"request":') + self._update(encoded) + self._update(b"}") self._request = None self._request_count += 1 except Exception: self._poisoned = True raise + def byte_count_if_finalized(self) -> int: + """Return exact canonical bytes if the current stream ended now.""" + self._raise_if_poisoned() + try: + byte_count = self._byte_count + if self._hexdigest is not None: + return byte_count + if self._request is not None: + encoded_request = _canonical_trajectory_bytes_at_depth( + self._request, depth=3 + ) + byte_count += len(b'],"request":') + len(encoded_request) + 1 + encoded_session_id = _canonical_trajectory_bytes_at_depth( + self._session_id, depth=1 + ) + return byte_count + len(b'],"session_id":') + len(encoded_session_id) + 1 + except Exception: + self._poisoned = True + raise + def hexdigest(self) -> str: self._raise_if_poisoned() try: @@ -307,9 +332,9 @@ def hexdigest(self) -> str: encoded_session_id = _canonical_trajectory_bytes_at_depth( self._session_id, depth=1 ) - self._digest.update(b'],"session_id":') - self._digest.update(encoded_session_id) - self._digest.update(b"}") + self._update(b'],"session_id":') + self._update(encoded_session_id) + self._update(b"}") self._hexdigest = self._digest.hexdigest() return self._hexdigest except Exception: diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index a6cae955f..63a8ee581 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -725,6 +725,54 @@ def test_streaming_trajectory_digest_matches_materialized_multi_request_json() - assert accumulator.hexdigest() == trajectory_digest(materialized) +def test_streaming_trajectory_byte_count_matches_each_materialized_prefix() -> None: + request_row = _streaming_request_row() + interaction_row = _streaming_interaction_row() + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(canonical_session_trajectory("session-1", [], {})) + ) + + accumulator.start_request(request_row) + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes( + canonical_session_trajectory("session-1", [request_row], {}) + ) + ) + + accumulator.add_interaction(interaction_row) + materialized = canonical_session_trajectory( + "session-1", [request_row], {"request-1": [interaction_row]} + ) + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(materialized) + ) + + accumulator.finish_request() + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(materialized) + ) + assert accumulator.hexdigest() == trajectory_digest(materialized) + + +def test_streaming_trajectory_byte_count_handles_one_huge_interaction() -> None: + request_row = _streaming_request_row() + interaction_row = _streaming_interaction_row() + interaction_row["content"] = "x" * 100_000 + materialized = canonical_session_trajectory( + "session-1", [request_row], {"request-1": [interaction_row]} + ) + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + + accumulator.start_request(request_row) + accumulator.add_interaction(interaction_row) + + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(materialized) + ) + + def _streaming_request_row() -> dict[str, object]: return { "request_id": "request-1", From 5c020483c04f7e0376ea972da820a73f8102c5f5 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 10:58:37 +0000 Subject: [PATCH 52/63] fix: allow synchronous playbook aggregation to finish --- reflexio/server/middleware.py | 6 ++-- tests/server/test_api_security_middleware.py | 34 ++++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/reflexio/server/middleware.py b/reflexio/server/middleware.py index 86331792b..40bbec92c 100644 --- a/reflexio/server/middleware.py +++ b/reflexio/server/middleware.py @@ -23,9 +23,11 @@ # Bot protection configuration REQUEST_TIMEOUT_SECONDS = 60 SYNC_REQUEST_TIMEOUT_SECONDS = ( - 600 # Longer timeout for synchronous processing (wait_for_response=true) + 600 # Longer timeout for synchronous long-running processing. +) +SYNC_REQUEST_PATHS = frozenset( + {"/api/review_user_playbooks", "/api/run_playbook_aggregation"} ) -SYNC_REQUEST_PATHS = frozenset({"/api/review_user_playbooks"}) SUSPICIOUS_USER_AGENTS = ["bot", "crawler", "spider", "scraper", "curl", "wget"] ALLOWED_EMPTY_UA_PATHS = ["/health", "/"] # Paths that allow empty user agents DEFAULT_MAX_BODY_BYTES = 10 * 1024 * 1024 diff --git a/tests/server/test_api_security_middleware.py b/tests/server/test_api_security_middleware.py index 900ba9fc7..e8ca7a2b8 100644 --- a/tests/server/test_api_security_middleware.py +++ b/tests/server/test_api_security_middleware.py @@ -174,6 +174,40 @@ async def call_next(_request): assert observed["timeout"] == SYNC_REQUEST_TIMEOUT_SECONDS +def test_playbook_aggregation_post_uses_synchronous_request_timeout_without_wait_query( + monkeypatch, +): + observed: dict[str, float | None] = {} + + async def fake_wait_for(awaitable, *, timeout=None): + observed["timeout"] = timeout + return await awaitable + + async def call_next(_request): + from starlette.responses import Response + + return Response() + + monkeypatch.setattr(asyncio, "wait_for", fake_wait_for) + request = Request( + { + "type": "http", + "method": "POST", + "scheme": "http", + "path": "/api/run_playbook_aggregation", + "raw_path": b"/api/run_playbook_aggregation", + "query_string": b"", + "headers": [], + "client": ("testclient", 50000), + "server": ("testserver", 80), + } + ) + + asyncio.run(TimeoutMiddleware(FastAPI()).dispatch(request, call_next)) + + assert observed["timeout"] == SYNC_REQUEST_TIMEOUT_SECONDS + + def test_security_headers_are_added(monkeypatch): monkeypatch.delenv("REFLEXIO_ALLOWED_ORIGINS", raising=False) From 531b18c292a65a2a93df205ed5da5265736fdc28 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 11:24:09 +0000 Subject: [PATCH 53/63] fix: address Phase 2 storage review findings --- .../sqlite_storage/_session_outcomes.py | 2 +- tests/models/test_session_outcome_identity.py | 2 +- .../test_session_outcome_migration.py | 20 +++++++-- .../test_storage_contract_session_outcomes.py | 44 ++++++++++++++++--- 4 files changed, 57 insertions(+), 11 deletions(-) diff --git a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py index a965c0eaf..dc7d4bf82 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -92,7 +92,7 @@ def get_session_outcome_context(self, session_id: str) -> SessionOutcomeContext: ).fetchone() return SessionOutcomeContext( user_id=str(first["user_id"]), - source=str(first["source"]), + source=str(first["source"] or ""), first_request_at=_iso_to_epoch(first["created_at"]), user_contract_violation=int(counts["user_count"]) > 1, source_contract_violation=int(counts["source_count"]) > 1, diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py index 63a8ee581..614ea6a89 100644 --- a/tests/models/test_session_outcome_identity.py +++ b/tests/models/test_session_outcome_identity.py @@ -880,6 +880,6 @@ def test_streaming_trajectory_digest_is_permanently_poisoned_after_error( ): with pytest.raises( RuntimeError, - match="canonical trajectory digest accumulator is invalid$", + match=r"canonical trajectory digest accumulator is invalid$", ): operation() diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index 7f93b52c9..0d34505b2 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -3,6 +3,7 @@ import json import sqlite3 from hashlib import sha256 +from math import ceil import pytest @@ -13,11 +14,13 @@ SetSessionOutcomeRequest, SetSessionOutcomeResponse, ) +from reflexio.server.services.storage.retention_mixin import RETENTION_DELETE_CHUNK from reflexio.server.services.storage.session_outcome_identity import ( outcome_contract_digest, ) from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.sqlite_storage._base import ( + _SESSION_OUTCOME_MIGRATION_BATCH_SIZE, _canonical_session_trajectory_digest, _epoch_to_iso, _prefetch_canonical_session_trajectory_digests, @@ -386,8 +389,12 @@ def test_migration_streams_legacy_outcomes_and_trajectories_in_bounded_batches( if statement.lstrip().upper().startswith("SELECT") and "FROM session_outcomes_legacy" in statement ] - assert len(legacy_outcome_queries) == 2 - assert all("LIMIT 256" in statement for statement in legacy_outcome_queries) + migration_batch_count = ceil(row_count / _SESSION_OUTCOME_MIGRATION_BATCH_SIZE) + assert len(legacy_outcome_queries) == migration_batch_count + assert all( + f"LIMIT {_SESSION_OUTCOME_MIGRATION_BATCH_SIZE}" in statement + for statement in legacy_outcome_queries + ) assert not any( statement.strip() == "SELECT * FROM session_outcomes" for statement in statements @@ -398,7 +405,14 @@ def test_migration_streams_legacy_outcomes_and_trajectories_in_bounded_batches( if statement.lstrip().upper().startswith("SELECT") and "FROM requests" in statement ] - assert len(trajectory_input_queries) == 2 + expected_trajectory_query_count = sum( + ceil( + min(_SESSION_OUTCOME_MIGRATION_BATCH_SIZE, row_count - batch_start) + / RETENTION_DELETE_CHUNK + ) + for batch_start in range(0, row_count, _SESSION_OUTCOME_MIGRATION_BATCH_SIZE) + ) + assert len(trajectory_input_queries) == expected_trajectory_query_count assert all("LEFT JOIN interactions" in query for query in trajectory_input_queries) assert ( storage.conn.execute("SELECT COUNT(*) FROM session_outcomes").fetchone()[0] diff --git a/tests/server/services/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index d88c5c569..52d8e27e8 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -1,6 +1,8 @@ """Session outcome storage contract.""" +from threading import RLock from typing import Any, cast +from unittest.mock import MagicMock import pytest @@ -17,6 +19,9 @@ trajectory_digest, ) from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._session_outcomes import ( + SessionOutcomeStoreMixin, +) from reflexio.server.services.storage.storage_base import BaseStorage @@ -388,8 +393,8 @@ def test_sqlite_large_finalization_streams_complete_digest_and_preserves_retry_c metadata={"complete": True}, ) raw_connection = sqlite_storage.conn - guarded_connection = _NoFetchAllConnection(raw_connection) - cast(Any, sqlite_storage).conn = guarded_connection + initial_retry_connection = _NoFetchAllConnection(raw_connection) + cast(Any, sqlite_storage).conn = initial_retry_connection try: context = storage.get_session_outcome_context(session_id) first = storage.record_session_outcome( @@ -416,8 +421,8 @@ def test_sqlite_large_finalization_streams_complete_digest_and_preserves_retry_c ], embeddings_prepared=True, ) - guarded_connection = _NoFetchAllConnection(raw_connection) - cast(Any, sqlite_storage).conn = guarded_connection + conflict_connection = _NoFetchAllConnection(raw_connection) + cast(Any, sqlite_storage).conn = conflict_connection try: conflict = storage.record_session_outcome( outcome, @@ -434,8 +439,12 @@ def test_sqlite_large_finalization_streams_complete_digest_and_preserves_retry_c assert retry.finalized_trajectory_digest == expected_digest assert conflict.recorded is False assert conflict.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION - assert guarded_connection.fetch_sizes - assert len(set(guarded_connection.fetch_sizes)) == 1 + all_fetch_sizes = [ + *initial_retry_connection.fetch_sizes, + *conflict_connection.fetch_sizes, + ] + assert all_fetch_sizes + assert len(set(all_fetch_sizes)) == 1 def test_changed_contract_identity_is_conflicting_finalization( @@ -607,6 +616,29 @@ def test_empty_request_source_is_preserved(storage: BaseStorage) -> None: assert record.source == "" +def test_sqlite_context_normalizes_nullable_request_source() -> None: + existing_cursor = MagicMock() + existing_cursor.fetchone.return_value = None + first_cursor = MagicMock() + first_cursor.fetchone.return_value = { + "user_id": "u1", + "source": None, + "created_at": "1970-01-01T00:01:40+00:00", + } + counts_cursor = MagicMock() + counts_cursor.fetchone.return_value = {"user_count": 1, "source_count": 1} + connection = MagicMock() + connection.execute.side_effect = [existing_cursor, first_cursor, counts_cursor] + reader = SessionOutcomeStoreMixin() + cast(Any, reader).conn = connection + cast(Any, reader)._lock = RLock() + + context = reader.get_session_outcome_context("nullable-source") + + assert context.source == "" + assert context.first_request_at == 100 + + def test_clear_outcomes_survives_governance_secret_rotation( storage: BaseStorage, monkeypatch ) -> None: From 9c54640cdf8a9f069db908e8320080a704e4a2f8 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 11:37:29 +0000 Subject: [PATCH 54/63] test: cover streamed migration trajectories --- .../test_session_outcome_migration.py | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index 0d34505b2..692bed112 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -4,6 +4,7 @@ import sqlite3 from hashlib import sha256 from math import ceil +from typing import Any, cast import pytest @@ -21,6 +22,7 @@ from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.sqlite_storage._base import ( _SESSION_OUTCOME_MIGRATION_BATCH_SIZE, + _TRAJECTORY_FETCH_SIZE, _canonical_session_trajectory_digest, _epoch_to_iso, _prefetch_canonical_session_trajectory_digests, @@ -108,6 +110,38 @@ def _identity_complete_session_outcomes_ddl(*, with_subject_column: bool) -> str """ +class _NoTrajectoryFetchAllCursor: + def __init__(self, cursor: Any, fetch_sizes: list[int]) -> None: + self._cursor = cursor + self._fetch_sizes = fetch_sizes + + def fetchall(self) -> Any: + raise AssertionError("trajectory migration must not call fetchall") + + def fetchmany(self, size: int) -> Any: + self._fetch_sizes.append(size) + return self._cursor.fetchmany(size) + + def __getattr__(self, name: str) -> Any: + return getattr(self._cursor, name) + + +class _NoTrajectoryFetchAllConnection: + def __init__(self, connection: sqlite3.Connection) -> None: + self._connection = connection + self.fetch_sizes: list[int] = [] + + def execute(self, *args: Any, **kwargs: Any) -> Any: + cursor = self._connection.execute(*args, **kwargs) + statement = str(args[0]) if args else str(kwargs.get("sql", "")) + if "FROM requests" in statement and "LEFT JOIN interactions" in statement: + return _NoTrajectoryFetchAllCursor(cursor, self.fetch_sizes) + return cursor + + def __getattr__(self, name: str) -> Any: + return getattr(self._connection, name) + + def _replace_session_outcomes_with_unconstrained_revision( storage: SQLiteStorage, revision: int | None ) -> None: @@ -420,6 +454,75 @@ def test_migration_streams_legacy_outcomes_and_trajectories_in_bounded_batches( ) +def test_migration_streams_complete_digest_for_trajectory_over_fetch_window( + tmp_path, +) -> None: + storage = SQLiteStorage( + org_id="legacy-large-trajectory", + db_path=str(tmp_path / "legacy-large-trajectory.db"), + ) + session_id = "large-trajectory-session" + request_id = "large-trajectory-request" + user_id = "large-trajectory-user" + source = "legacy-source" + storage.add_request( + Request( + request_id=request_id, + user_id=user_id, + session_id=session_id, + source=source, + created_at=100, + ) + ) + interaction_count = _TRAJECTORY_FETCH_SIZE + 1 + storage.conn.executemany( + """INSERT INTO interactions ( + interaction_id, user_id, request_id, created_at, content + ) VALUES (?, ?, ?, ?, ?)""", + [ + ( + index + 1, + user_id, + request_id, + _epoch_to_iso(101 + index), + f"trajectory-row-{index}", + ) + for index in range(interaction_count) + ], + ) + expected_digest = _canonical_session_trajectory_digest(storage.conn, session_id) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + storage.conn.execute( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?)""", + (user_id, session_id, "success", 500, source, "subject-large", 501), + ) + storage.conn.commit() + + raw_connection = storage.conn + guarded_connection = _NoTrajectoryFetchAllConnection(raw_connection) + cast(Any, storage).conn = guarded_connection + try: + storage._migrate_session_outcomes_schema() + finally: + cast(Any, storage).conn = raw_connection + + migrated = raw_connection.execute( + """SELECT finalized_trajectory_digest FROM session_outcomes + WHERE session_id = ?""", + (session_id,), + ).fetchone() + assert migrated is not None + assert migrated["finalized_trajectory_digest"] == expected_digest + assert len(guarded_connection.fetch_sizes) == ( + ceil(interaction_count / _TRAJECTORY_FETCH_SIZE) + 1 + ) + assert set(guarded_connection.fetch_sizes) == {_TRAJECTORY_FETCH_SIZE} + + def test_migration_prefetch_retains_only_trajectory_digests(tmp_path) -> None: storage = SQLiteStorage( org_id="legacy-digest-retention", From 0892a97edc9318904770c7b20995f9608ffa582c Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 12:41:14 +0000 Subject: [PATCH 55/63] fix: normalize nullable source on outcome retries --- .../sqlite_storage/_session_outcomes.py | 4 +- .../test_storage_contract_session_outcomes.py | 73 +++++++++++++++++-- 2 files changed, 69 insertions(+), 8 deletions(-) diff --git a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py index dc7d4bf82..aaec427a7 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -129,7 +129,7 @@ def record_session_outcome( ) first = snapshot.first_request if snapshot is not None else None source = ( - str(first["source"]) + str(first["source"] or "") if first is not None else str(existing["source"]) ) @@ -149,7 +149,7 @@ def record_session_outcome( stored_snapshot_digest = existing["finalized_trajectory_digest"] server_context_matches = first is None or ( str(existing["user_id"]) == str(first["user_id"]) - and str(existing["source"]) == str(first["source"]) + and str(existing["source"]) == str(first["source"] or "") and str(existing["governance_subject_ref"]) == subject_ref ) exact_retry = ( diff --git a/tests/server/services/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index 52d8e27e8..675dbeafe 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -19,6 +19,9 @@ trajectory_digest, ) from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._base import ( + _TRAJECTORY_FETCH_SIZE, +) from reflexio.server.services.storage.sqlite_storage._session_outcomes import ( SessionOutcomeStoreMixin, ) @@ -137,6 +140,62 @@ def test_exact_finalization_retry_is_idempotent(storage: BaseStorage) -> None: assert retry.finalized_trajectory_digest == first.finalized_trajectory_digest +def test_sqlite_legacy_null_source_finalization_retry_is_idempotent( + storage: BaseStorage, +) -> None: + sqlite_storage = cast(SQLiteStorage, storage) + storage.add_request( + Request( + request_id="legacy-null-source-r1", + user_id="u1", + session_id="legacy-null-source", + source="", + created_at=100, + ) + ) + schema_sql = sqlite_storage.conn.execute( + "SELECT sql FROM sqlite_schema WHERE type = 'table' AND name = 'requests'" + ).fetchone()["sql"] + nullable_schema_sql = schema_sql.replace( + "source TEXT NOT NULL DEFAULT ''", "source TEXT DEFAULT ''", 1 + ) + sqlite_storage.conn.execute("PRAGMA writable_schema = ON") + sqlite_storage.conn.execute( + "UPDATE sqlite_schema SET sql = ? WHERE type = 'table' AND name = 'requests'", + (nullable_schema_sql,), + ) + schema_version = sqlite_storage.conn.execute("PRAGMA schema_version").fetchone()[0] + sqlite_storage.conn.execute(f"PRAGMA schema_version = {schema_version + 1}") + sqlite_storage.conn.execute("PRAGMA writable_schema = OFF") + sqlite_storage.conn.execute( + "UPDATE requests SET source = NULL WHERE request_id = ?", + ("legacy-null-source-r1",), + ) + sqlite_storage.conn.commit() + + request = SetSessionOutcomeRequest( + session_id="legacy-null-source", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + context = storage.get_session_outcome_context("legacy-null-source") + first = storage.record_session_outcome( + request, created_at=102, expected_context=context + ) + retry = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context("legacy-null-source"), + ) + + assert context.source == "" + assert first.recorded is True + assert first.source == "" + assert retry.recorded is False + assert retry.reason is None + assert retry.source == "" + + def test_legacy_all_null_identity_exact_retry_uses_available_context( storage: BaseStorage, ) -> None: @@ -439,12 +498,14 @@ def test_sqlite_large_finalization_streams_complete_digest_and_preserves_retry_c assert retry.finalized_trajectory_digest == expected_digest assert conflict.recorded is False assert conflict.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION - all_fetch_sizes = [ - *initial_retry_connection.fetch_sizes, - *conflict_connection.fetch_sizes, - ] - assert all_fetch_sizes - assert len(set(all_fetch_sizes)) == 1 + assert initial_retry_connection.fetch_sizes + assert conflict_connection.fetch_sizes + assert all( + size == _TRAJECTORY_FETCH_SIZE for size in initial_retry_connection.fetch_sizes + ) + assert all( + size == _TRAJECTORY_FETCH_SIZE for size in conflict_connection.fetch_sizes + ) def test_changed_contract_identity_is_conflicting_finalization( From 823d0a4a66f1736e72e988955ff281295507cf0e Mon Sep 17 00:00:00 2001 From: Guangyu Date: Tue, 11 Aug 2026 13:34:40 +0000 Subject: [PATCH 56/63] test: harden session outcome storage guards --- .../test_session_outcome_migration.py | 5 +++++ .../test_storage_contract_session_outcomes.py | 14 +++++++++++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py index 692bed112..70d411613 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -118,6 +118,11 @@ def __init__(self, cursor: Any, fetch_sizes: list[int]) -> None: def fetchall(self) -> Any: raise AssertionError("trajectory migration must not call fetchall") + def __iter__(self) -> Any: + raise AssertionError( + "trajectory migration must not iterate the cursor directly" + ) + def fetchmany(self, size: int) -> Any: self._fetch_sizes.append(size) return self._cursor.fetchmany(size) diff --git a/tests/server/services/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index 675dbeafe..b1d1a746f 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -689,7 +689,19 @@ def test_sqlite_context_normalizes_nullable_request_source() -> None: counts_cursor = MagicMock() counts_cursor.fetchone.return_value = {"user_count": 1, "source_count": 1} connection = MagicMock() - connection.execute.side_effect = [existing_cursor, first_cursor, counts_cursor] + + def execute(statement: str, _parameters: tuple[str]) -> MagicMock: + if "FROM session_outcomes" in statement: + return existing_cursor + if "ORDER BY created_at ASC" in statement: + return first_cursor + if "COUNT(DISTINCT user_id)" in statement: + return counts_cursor + raise AssertionError( + f"unexpected SQL in nullable-source context test: {statement}" + ) + + connection.execute.side_effect = execute reader = SessionOutcomeStoreMixin() cast(Any, reader).conn = connection cast(Any, reader)._lock = RLock() From 056fc97adaa66aad1882c51a5482ac2b7a482f52 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Wed, 12 Aug 2026 07:46:43 +0000 Subject: [PATCH 57/63] fix: wait for local embedding before backend startup --- reflexio/cli/utils.py | 16 ++++++++++ tests/cli/test_utils.py | 66 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/reflexio/cli/utils.py b/reflexio/cli/utils.py index 27d32ca04..97ef3b9fd 100644 --- a/reflexio/cli/utils.py +++ b/reflexio/cli/utils.py @@ -504,8 +504,24 @@ def run_services( for path in supervisor.stop_request_paths.values(): remove_pidfile(path) + service_names = {svc.name for svc in services} + gate_local_embedding = {"embedding", "backend"}.issubset(service_names) + try: + if gate_local_embedding: + embedding = next(svc for svc in services if svc.name == "embedding") + supervisor.start_service(embedding) + if not _wait_for_all_ready( + {"embedding": supervisor.ready_events["embedding"]}, + {"embedding": supervisor.processes["embedding"]}, + ): + raise RuntimeError( + "embedding service did not become ready before backend startup" + ) + for svc in services: + if gate_local_embedding and svc.name == "embedding": + continue supervisor.start_service(svc) supervisor.write_current_pidfile() except (OSError, RuntimeError): diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index 10414c673..300d46653 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -240,6 +240,7 @@ def fake_popen(*_args, **_kwargs) -> _FakeProcess: return proc monkeypatch.setattr(utils.subprocess, "Popen", fake_popen) + monkeypatch.setattr(utils, "_wait_for_all_ready", lambda *_args: True) with pytest.raises(OSError, match="simulated initial spawn failure"): utils.run_services( @@ -293,6 +294,71 @@ def fake_popen(*_args, **_kwargs) -> _FakeProcess: assert ready_calls == [{"backend": 8071}] +@pytest.mark.unit +def test_run_services_waits_for_local_embedding_before_starting_backend( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Do not start backend lifespan work before local inference is ready.""" + _patch_run_services_environment(monkeypatch, tmp_path) + lifecycle: list[str] = [] + + def fake_popen(command, **_kwargs) -> _FakeProcess: + lifecycle.append(f"start:{command[0]}") + return _FakeProcess(pid=1000 + len(lifecycle), polls_to_exit=1) + + def fake_wait_for_ready(ready_events, processes) -> bool: + assert set(ready_events) == {"embedding"} + assert set(processes) == {"embedding"} + lifecycle.append("ready:embedding") + return True + + monkeypatch.setattr(utils.subprocess, "Popen", fake_popen) + monkeypatch.setattr(utils, "_wait_for_all_ready", fake_wait_for_ready) + + utils.run_services( + [ + utils.ServiceConfig(name="embedding", command=["embedding"]), + utils.ServiceConfig(name="backend", command=["backend"]), + ], + {"embedding": 8072, "backend": 8071}, + ) + + assert lifecycle[:3] == [ + "start:embedding", + "ready:embedding", + "start:backend", + ] + + +@pytest.mark.unit +def test_run_services_aborts_when_local_embedding_is_not_ready( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Clean up local inference and never start backend when readiness fails.""" + _patch_run_services_environment(monkeypatch, tmp_path) + started: list[tuple[str, _FakeProcess]] = [] + + def fake_popen(command, **_kwargs) -> _FakeProcess: + proc = _FakeProcess(pid=1000 + len(started), polls_to_exit=1) + started.append((command[0], proc)) + return proc + + monkeypatch.setattr(utils.subprocess, "Popen", fake_popen) + monkeypatch.setattr(utils, "_wait_for_all_ready", lambda *_args: False) + + with pytest.raises(RuntimeError, match="embedding.*ready"): + utils.run_services( + [ + utils.ServiceConfig(name="embedding", command=["embedding"]), + utils.ServiceConfig(name="backend", command=["backend"]), + ], + {"embedding": 8072, "backend": 8071}, + ) + + assert [name for name, _proc in started] == ["embedding"] + assert started[0][1].terminated is True + + @pytest.mark.unit def test_run_services_does_not_respawn_an_explicitly_stopped_service( monkeypatch: pytest.MonkeyPatch, tmp_path: Path From 6eab375eabfec9af663a552bfae933a899d351e4 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Thu, 13 Aug 2026 23:05:21 +0000 Subject: [PATCH 58/63] fix: persist inline finalization receipts --- .../services/base_generation_service.py | 65 ++- .../server/services/deferred_learning_plan.py | 13 +- .../services/durable_learning/worker.py | 42 +- reflexio/server/services/playbook/service.py | 49 +- reflexio/server/services/profile/service.py | 40 +- .../services/durable_learning/test_worker.py | 290 +++++++++++ .../services/test_base_generation_service.py | 462 +++++++++++++++++- 7 files changed, 932 insertions(+), 29 deletions(-) diff --git a/reflexio/server/services/base_generation_service.py b/reflexio/server/services/base_generation_service.py index 9fd3f9fe6..89a46cbdb 100644 --- a/reflexio/server/services/base_generation_service.py +++ b/reflexio/server/services/base_generation_service.py @@ -29,6 +29,7 @@ ) from reflexio.server.services.deferred_learning_plan import ( ExtractorBookmarkAdvance, + FinalizationResult, GenerationComputePlan, ) from reflexio.server.services.extractor_config_utils import ( @@ -587,7 +588,11 @@ def _run_generation(self, request: TRequest) -> None: plan = self.compute_generation(request) if plan is None: return - self.persist_generation(plan) + try: + self.persist_generation(plan) + except Exception as exc: + self._mark_extraction_runs_finalization_failed(exc) + raise self.emit_generation_side_effects(plan) except Exception as e: self._record_generation_event( @@ -608,11 +613,11 @@ def _run_generation(self, request: TRequest) -> None: def compute_generation(self, request: TRequest) -> GenerationComputePlan | None: """Compute half of one generation run — NO learning DB write, NO fence. - Runs the prepare gate, the extractor (thread-pool LLM tool-loop), dedup - + embedding resolution (``_resolve_write_plan``), and drives the - ``agent_run`` rows to their terminal state (``_finalize_extraction_runs`` - — agent_run only, §4.3). Snapshots the billing inputs onto the returned - plan so ``emit_generation_side_effects`` never depends on the reused + Runs the prepare gate, the extractor (thread-pool LLM tool-loop), and + dedup + embedding resolution (``_resolve_write_plan``). Agent runs stay + non-terminal until receipt-aware persistence commits and + ``emit_generation_side_effects`` finalizes them. Snapshots the billing + inputs onto the returned plan so emit does not depend on the reused instance's mutable ``_last_*``. Returns a resolved ``GenerationComputePlan`` for persist/emit, or @@ -646,7 +651,6 @@ def compute_generation(self, request: TRequest) -> GenerationComputePlan | None: try: write_plan = self._resolve_write_plan([result]) if result else None - self._finalize_extraction_runs() except Exception as exc: self._mark_extraction_runs_finalization_failed(exc) raise @@ -683,6 +687,31 @@ def persist_generation(self, plan: GenerationComputePlan) -> None: atomically with the row writes it corresponds to (F1) — and it is applied in exactly one place, so it is never double-applied. """ + if self.EMITS_LEARNING_BILLING and plan.extraction_run_ids: + if len(plan.extraction_run_ids) != 1: + raise RuntimeError( + "Receipt-aware generation requires exactly one extraction run" + ) + finalization_run_id = plan.extraction_run_ids[0] + extraction_run = ( + self.storage.get_agent_run(finalization_run_id) + if self.storage is not None + else None + ) + if extraction_run is None: + raise RuntimeError( + "Receipt-aware generation extraction run could not be resolved" + ) + # A pending-tool run has not produced its final learning set yet. + # Its resume worker must remain free to persist that set and create + # the immutable receipt after the tool dependency resolves. + if not extraction_run.pending_tool_call_ids: + plan.finalization_result = self._finalize_write_plan_with_outcome( + plan.write_plan, + finalization_run_id=finalization_run_id, + bookmark_advance=plan.bookmark_advance, + ) + return if plan.write_plan is not None: self._persist_write_plan(plan.write_plan) self._apply_bookmark_advance(plan.bookmark_advance) @@ -710,6 +739,16 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: is exactly as compute left it when emit runs. The plan still snapshots the billing inputs so a future durable reuse can re-plumb locally. """ + try: + self._finalize_extraction_runs() + except Exception as exc: + self._mark_extraction_runs_finalization_failed(exc) + raise + if ( + plan.finalization_result is not None + and not plan.finalization_result.won_receipt + ): + return self._record_generation_event( event_name="generation_succeeded", outcome="success", @@ -728,6 +767,18 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: prepared=plan.prepared, generated_count=plan.billable_count ) + def _finalize_write_plan_with_outcome( + self, + write_plan: Any, + *, + finalization_run_id: str, + bookmark_advance: ExtractorBookmarkAdvance | None, + ) -> FinalizationResult: + """Persist one precomputed inline plan through a receipt-aware finalizer.""" + raise NotImplementedError( + "Receipt-aware plan finalization must be implemented by learning services" + ) + @abstractmethod def _resolve_write_plan(self, results: list) -> Any | None: """Compute-half of item finalization — resolve a write-plan, NO DB write. diff --git a/reflexio/server/services/deferred_learning_plan.py b/reflexio/server/services/deferred_learning_plan.py index e540357e6..1896ff11b 100644 --- a/reflexio/server/services/deferred_learning_plan.py +++ b/reflexio/server/services/deferred_learning_plan.py @@ -147,11 +147,11 @@ class GenerationComputePlan: """Resolved compute output of one ``BaseGenerationService`` run (gate b). ``compute_generation`` runs the prepare gate + extractor + dedup/embedding - resolution (``_resolve_write_plan``) and drives the ``agent_run`` rows to - their terminal state (``_finalize_extraction_runs`` — agent_run only, §4.3), - issuing **no** learning DB write. ``persist_generation`` applies - ``write_plan`` + the extractor bookmark advance inside the fence; - ``emit_generation_side_effects`` fires the post-commit telemetry + billing. + resolution (``_resolve_write_plan``), issuing **no** learning DB write and + leaving receipt-aware ``agent_run`` rows non-terminal. ``persist_generation`` + applies ``write_plan`` + the extractor bookmark advance + immutable receipt + inside the fence; ``emit_generation_side_effects`` terminalizes the run and + fires post-commit telemetry + billing. The billing inputs (``extraction_run_ids`` / ``token_totals`` / ``billable_count`` / ``prepared``) and telemetry's ``generated_count`` are @@ -179,6 +179,8 @@ class GenerationComputePlan: extraction_run_ids: Snapshot of the run's ``agent_run`` ids. token_totals: Snapshot of the run's LLM token totals (billing cost facet), or ``None`` when the extractor reported none. + finalization_result: Receipt ownership recorded by persistence. ``None`` + for services that do not use receipt-aware inline finalization. """ prepared: PreparedGenerationRun[Any] @@ -189,6 +191,7 @@ class GenerationComputePlan: generation_start: float extraction_run_ids: list[str] token_totals: RunTokenTotals | None + finalization_result: FinalizationResult | None = None @dataclass diff --git a/reflexio/server/services/durable_learning/worker.py b/reflexio/server/services/durable_learning/worker.py index fed223419..9f783cfeb 100644 --- a/reflexio/server/services/durable_learning/worker.py +++ b/reflexio/server/services/durable_learning/worker.py @@ -24,7 +24,9 @@ from reflexio.server.api_endpoints.request_context import RequestContext from reflexio.server.cache.reflexio_cache import get_reflexio +from reflexio.server.services.deferred_learning_plan import DeferredLearningPlan from reflexio.server.services.generation_service import GenerationService +from reflexio.server.services.storage.storage_base import AgentRunStatus, BaseStorage from reflexio.server.services.storage.storage_base._learning_jobs import LearningJob logger = logging.getLogger(__name__) @@ -151,6 +153,7 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: return False gen: GenerationService | None = None + plan: DeferredLearningPlan | None = None try: reflexio = get_reflexio( org_id=ctx.org_id, storage_base_dir=ctx.storage_base_dir @@ -202,7 +205,16 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: # POST-COMMIT — billing / telemetry / tagging / off-thread schedulers # + the per-user lock release, only for the winning worker. - gen.emit_deferred_learning_side_effects(plan) + try: + gen.emit_deferred_learning_side_effects(plan) + except Exception: + logger.exception( + "event=learning_job_side_effects_failed " + "job_id=%s org_id=%s user_id=%s", + job.job_id, + job.org_id, + job.user_id, + ) logger.info( "event=learning_job_done job_id=%s org_id=%s user_id=%s", job.job_id, @@ -210,23 +222,25 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: job.user_id, ) return True - except _SupersededError: + except _SupersededError as exc: logger.info( "event=learning_job_superseded job_id=%s org_id=%s", job.job_id, job.org_id, ) + self._abandon_computed_agent_runs(storage, plan, exc) # The persist rolled back and emit never ran, so the per-user lock is # still held by this compute — release it so the reclaim isn't blocked. self._release_user_lock(gen, job) return False - except Exception: + except Exception as exc: logger.exception( "event=learning_job_failed job_id=%s org_id=%s user_id=%s", job.job_id, job.org_id, job.user_id, ) + self._abandon_computed_agent_runs(storage, plan, exc) # emit (which releases the lock) never ran on this path — release the # per-user lock so a failed job doesn't strand it. self._release_user_lock(gen, job) @@ -237,6 +251,28 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: ) return False + @staticmethod + def _abandon_computed_agent_runs( + storage: BaseStorage, + plan: DeferredLearningPlan | None, + exc: Exception, + ) -> None: + """Make runs from a rolled-back durable attempt non-resumable.""" + if plan is None: + return + run_ids = { + run_id + for pair in (plan.profile, plan.playbook) + if pair is not None + for run_id in pair[1].extraction_run_ids + } + for run_id in run_ids: + storage.update_agent_run_status( + run_id, + AgentRunStatus.FAILED, + last_error=f"durable learning attempt abandoned: {exc}", + ) + def _release_user_lock( self, gen: GenerationService | None, job: LearningJob ) -> None: diff --git a/reflexio/server/services/playbook/service.py b/reflexio/server/services/playbook/service.py index 4f79508fb..4ff6dd7ca 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -9,7 +9,10 @@ if TYPE_CHECKING: 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.deferred_learning_plan import ( + ExtractorBookmarkAdvance, + GenerationComputePlan, + ) from reflexio.server.services.storage.storage_base import ( AgentRunRecord, BaseStorage, @@ -663,6 +666,11 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: superseded one) — the phantom-billing gate. """ super().emit_generation_side_effects(plan) + if ( + plan.finalization_result is not None + and not plan.finalization_result.won_receipt + ): + return write_plan = plan.write_plan if write_plan is not None: self._dispatch_playbook_schedulers(write_plan) @@ -700,11 +708,10 @@ def _finalize_extracted_items_with_outcome( finalization_run_id: str | None = None, ) -> FinalizationResult: """Finalize playbooks and expose the atomic receipt winner internally.""" - entity_type = "user_playbook" if finalization_run_id is not None: receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] run_id=finalization_run_id, - entity_type=entity_type, + entity_type="user_playbook", ) if receipt is not None: # Receipts make persistence/billing idempotent. Derived schedulers @@ -722,7 +729,6 @@ def _finalize_extracted_items_with_outcome( finally: self._review_run = previous_review_run self._review_window_cache = previous_review_window - learning_ids: list[str] = [] if finalization_run_id is None: if plan is not None: with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] @@ -741,6 +747,32 @@ def _finalize_extracted_items_with_outcome( won_receipt=False, ) + result = self._finalize_write_plan_with_outcome( + plan, + finalization_run_id=finalization_run_id, + bookmark_advance=None, + ) + if plan is not None and result.won_receipt: + self._dispatch_playbook_schedulers(plan) + return result + + def _finalize_write_plan_with_outcome( + self, + write_plan: PlaybookWritePlan | None, + *, + finalization_run_id: str, + bookmark_advance: ExtractorBookmarkAdvance | None, + ) -> FinalizationResult: + """Commit a resolved playbook plan, bookmark, and receipt atomically.""" + entity_type = "user_playbook" + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + learning_ids: list[str] = [] + try: with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] @@ -749,13 +781,14 @@ def _finalize_extracted_items_with_outcome( ) if receipt is not None: return FinalizationResult(receipt, won_receipt=False) - if plan is not None: - self._persist_write_plan(plan) + if write_plan is not None: + self._persist_write_plan(write_plan) learning_ids = [ str(playbook.user_playbook_id) - for playbook in plan.new_playbooks + for playbook in write_plan.new_playbooks if playbook.user_playbook_id ] + self._apply_bookmark_advance(bookmark_advance) inserted = self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] run_id=finalization_run_id, entity_type=entity_type, @@ -773,8 +806,6 @@ def _finalize_extracted_items_with_outcome( "finalization receipt disappeared after insert conflict" ) from None return FinalizationResult(receipt, won_receipt=False) - if plan is not None: - self._dispatch_playbook_schedulers(plan) return FinalizationResult(learning_ids, won_receipt=True) def _apply_consolidation_lineage( diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index dc847f68e..b75671d95 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -31,6 +31,7 @@ StatusChangeOperation, ) from reflexio.server.services.deferred_learning_plan import ( + ExtractorBookmarkAdvance, FinalizationResult, ProfileWritePlan, _FinalizationReceiptAlreadyExistsError, @@ -365,11 +366,10 @@ def _finalize_extracted_items_with_outcome( finalization_run_id: str | None = None, ) -> FinalizationResult: """Finalize profiles and expose the atomic receipt winner internally.""" - entity_type = "profile" if finalization_run_id is not None: receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] run_id=finalization_run_id, - entity_type=entity_type, + entity_type="profile", ) if receipt is not None: return FinalizationResult(receipt, won_receipt=False) @@ -392,6 +392,37 @@ def _finalize_extracted_items_with_outcome( self._persist_write_plan(plan) return FinalizationResult(learning_ids, won_receipt=False) + return self._finalize_write_plan_with_outcome( + plan, + finalization_run_id=finalization_run_id, + bookmark_advance=None, + ) + + def _finalize_write_plan_with_outcome( + self, + write_plan: ProfileWritePlan | None, + *, + finalization_run_id: str, + bookmark_advance: ExtractorBookmarkAdvance | None, + ) -> FinalizationResult: + """Commit a resolved profile plan, bookmark, and receipt atomically.""" + entity_type = "profile" + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + learning_ids = ( + [ + str(profile.profile_id) + for profile in write_plan.new_profiles + if profile.profile_id + ] + if write_plan is not None + else [] + ) + try: with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] @@ -400,8 +431,9 @@ def _finalize_extracted_items_with_outcome( ) if receipt is not None: return FinalizationResult(receipt, won_receipt=False) - if plan is not None: - self._persist_write_plan(plan) + if write_plan is not None: + self._persist_write_plan(write_plan) + self._apply_bookmark_advance(bookmark_advance) inserted = self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] run_id=finalization_run_id, entity_type=entity_type, diff --git a/tests/server/services/durable_learning/test_worker.py b/tests/server/services/durable_learning/test_worker.py index fdc766a8b..212cbe252 100644 --- a/tests/server/services/durable_learning/test_worker.py +++ b/tests/server/services/durable_learning/test_worker.py @@ -28,6 +28,8 @@ import tempfile import threading import time +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace from unittest import mock import pytest @@ -37,6 +39,11 @@ from reflexio.server.services.deferred_learning_plan import DeferredLearningPlan from reflexio.server.services.durable_learning.worker import DurableLearningWorker from reflexio.server.services.generation_service import GenerationService +from reflexio.server.services.storage.storage_base import ( + AgentBinding, + AgentRunRecord, + AgentRunStatus, +) # --------------------------------------------------------------------------- # Module-level fixture: disable the local ONNX embedder for all tests. @@ -124,6 +131,62 @@ def _setup_job( ) +def _seed_completed_agent_run( + storage, + *, + org_id: str, + request_id: str, + run_id: str, + extractor_kind: str, +) -> None: + storage.create_agent_run( + AgentRunRecord( + id=run_id, + binding=AgentBinding( + org_id=org_id, + extractor_kind=extractor_kind, + user_id="test_user", + request_id=request_id, + agent_version="v1", + source="test_src", + ), + status=AgentRunStatus.AGENT_COMPLETED, + generation_request_snapshot={"request_id": request_id}, + committed_output={f"{extractor_kind}s": []}, + ) + ) + + +def _deferred_plan_with_runs( + *, + request_id: str, + profile_run_id: str, + playbook_run_id: str, +) -> DeferredLearningPlan: + return DeferredLearningPlan( + request_id=request_id, + user_id="test_user", + agent_version="v1", + lock_acquired=True, + profile=(mock.Mock(), SimpleNamespace(extraction_run_ids=[profile_run_id])), + playbook=(mock.Mock(), SimpleNamespace(extraction_run_ids=[playbook_run_id])), + ) + + +def _assert_runs_abandoned(storage, *run_ids: str) -> None: + for run_id in run_ids: + run = storage.get_agent_run(run_id) + assert run is not None and run.status == AgentRunStatus.FAILED + assert ( + storage.claim_finalization_failed_agent_run( + org_id=storage.org_id, + worker_id="resume-worker", + now=datetime.now(UTC) + timedelta(days=1), + ) + is None + ) + + # --------------------------------------------------------------------------- # Test 1: Exactly-once under a claim race (the headline) # --------------------------------------------------------------------------- @@ -556,6 +619,233 @@ def test_superseded_job_does_not_commit_outputs(): ) +def test_persist_failure_abandons_computed_agent_runs(): + """A rolled-back durable persist cannot leave resumable agent runs behind.""" + with tempfile.TemporaryDirectory() as tmp_dir: + factory = _factory(tmp_dir) + ctx = factory("org_persist_failure") + assert ctx.storage is not None + _setup_job( + ctx.storage, + org_id="org_persist_failure", + user_id="test_user", + request_id="req_persist_failure", + ) + run_ids = ("profile-persist-failure", "playbook-persist-failure") + _seed_completed_agent_run( + ctx.storage, + org_id="org_persist_failure", + request_id="req_persist_failure", + run_id=run_ids[0], + extractor_kind="profile", + ) + _seed_completed_agent_run( + ctx.storage, + org_id="org_persist_failure", + request_id="req_persist_failure", + run_id=run_ids[1], + extractor_kind="playbook", + ) + [job] = ctx.storage.claim_learning_jobs( + claimed_by="persist-failure-worker", limit=1, lease_seconds=300 + ) + plan = _deferred_plan_with_runs( + request_id="req_persist_failure", + profile_run_id=run_ids[0], + playbook_run_id=run_ids[1], + ) + + with ( + mock.patch.object( + GenerationService, "compute_deferred_learning", return_value=plan + ), + mock.patch.object( + GenerationService, + "persist_deferred_learning", + side_effect=RuntimeError("persist failed"), + ), + ): + processed = DurableLearningWorker(factory)._process_job( + factory("org_persist_failure"), job + ) + + assert processed is False + _assert_runs_abandoned(ctx.storage, *run_ids) + + +def test_superseded_job_abandons_computed_agent_runs(): + """Fence loss abandons every run computed by the superseded attempt.""" + with tempfile.TemporaryDirectory() as tmp_dir: + factory = _factory(tmp_dir) + ctx = factory("org_run_supersede") + assert ctx.storage is not None + _setup_job( + ctx.storage, + org_id="org_run_supersede", + user_id="test_user", + request_id="req_run_supersede", + ) + [stale_job] = ctx.storage.claim_learning_jobs( + claimed_by="stale-worker", limit=1, lease_seconds=-1 + ) + [live_job] = ctx.storage.claim_learning_jobs( + claimed_by="live-worker", limit=1, lease_seconds=300 + ) + assert stale_job.claim_token != live_job.claim_token + + run_ids = ("profile-run-supersede", "playbook-run-supersede") + _seed_completed_agent_run( + ctx.storage, + org_id="org_run_supersede", + request_id="req_run_supersede", + run_id=run_ids[0], + extractor_kind="profile", + ) + _seed_completed_agent_run( + ctx.storage, + org_id="org_run_supersede", + request_id="req_run_supersede", + run_id=run_ids[1], + extractor_kind="playbook", + ) + plan = _deferred_plan_with_runs( + request_id="req_run_supersede", + profile_run_id=run_ids[0], + playbook_run_id=run_ids[1], + ) + + with mock.patch.object( + GenerationService, "compute_deferred_learning", return_value=plan + ): + processed = DurableLearningWorker(factory)._process_job( + factory("org_run_supersede"), stale_job + ) + + assert processed is False + _assert_runs_abandoned(ctx.storage, *run_ids) + + +def test_post_commit_emit_failure_preserves_completed_job_and_agent_runs(): + """A post-commit side-effect failure cannot undo the durable winner.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserProfile + from reflexio.server.services.deferred_learning_plan import ( + GenerationComputePlan, + ProfileWritePlan, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + with tempfile.TemporaryDirectory() as tmp_dir: + factory = _factory(tmp_dir) + ctx = factory("org_post_commit_emit_failure") + assert ctx.storage is not None + _setup_job( + ctx.storage, + org_id="org_post_commit_emit_failure", + user_id="test_user", + request_id="req_post_commit_emit_failure", + force_extraction=True, + skip_aggregation=True, + ) + [job] = ctx.storage.claim_learning_jobs( + claimed_by="post-commit-worker", limit=1, lease_seconds=300 + ) + + computed_plans: list[DeferredLearningPlan] = [] + run_id = "profile-post-commit-emit-failure" + _seed_completed_agent_run( + ctx.storage, + org_id="org_post_commit_emit_failure", + request_id="req_post_commit_emit_failure", + run_id=run_id, + extractor_kind="profile", + ) + + def build_plan(self, **_kwargs): + profile = UserProfile( + profile_id="profile-post-commit", + user_id="test_user", + content="This durable profile survived a post-commit emit failure.", + last_modified_timestamp=1_000, + generated_from_request_id="req_post_commit_emit_failure", + ) + generation_plan = GenerationComputePlan( + prepared=SimpleNamespace( + identifier="profile_generation", + extractor_name="profile", + ), + generated_count=1, + billable_count=1, + write_plan=ProfileWritePlan( + user_id="test_user", + request_id="req_post_commit_emit_failure", + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ), + bookmark_advance=None, + generation_start=0.0, + extraction_run_ids=[run_id], + token_totals=None, + ) + plan = DeferredLearningPlan( + request_id="req_post_commit_emit_failure", + user_id="test_user", + agent_version="v1", + lock_acquired=True, + profile=( + ProfileGenerationService(self.client, self.request_context), + generation_plan, + ), + playbook=None, + ) + computed_plans.append(plan) + return plan + + def emit_then_raise(self, plan): + assert plan.profile is not None + for extraction_run_id in plan.profile[1].extraction_run_ids: + self.request_context.storage.update_agent_run_status( + extraction_run_id, + AgentRunStatus.FINALIZED, + pending_tool_call_ids=[], + ) + raise RuntimeError("post-commit emit failed") + + with ( + mock.patch.object( + GenerationService, + "compute_deferred_learning", + build_plan, + ), + mock.patch.object( + GenerationService, + "emit_deferred_learning_side_effects", + emit_then_raise, + ), + ): + processed = DurableLearningWorker(factory)._process_job( + factory("org_post_commit_emit_failure"), job + ) + + assert len(computed_plans) == 1 + plan = computed_plans[0] + assert plan.profile is not None + generation_plan = plan.profile[1] + assert generation_plan.finalization_result is not None + assert generation_plan.finalization_result.won_receipt is True + run = ctx.storage.get_agent_run(run_id) + assert run is not None and run.status == AgentRunStatus.FINALIZED + assert ctx.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) == ["profile-post-commit"] + + assert processed is True + assert ctx.storage.count_learning_jobs_by_status("done") == 1 + assert ctx.storage.count_all_profiles() == 1 + + # --------------------------------------------------------------------------- # Test 6: Failed job is re-claimable without manual status reset # --------------------------------------------------------------------------- diff --git a/tests/server/services/test_base_generation_service.py b/tests/server/services/test_base_generation_service.py index 2359f6af3..aac56574d 100644 --- a/tests/server/services/test_base_generation_service.py +++ b/tests/server/services/test_base_generation_service.py @@ -30,7 +30,11 @@ _weighted_content_length, ) from reflexio.server.services.extraction.outcome import ExtractionOutcome -from reflexio.server.services.storage.storage_base import AgentRunStatus +from reflexio.server.services.storage.storage_base import ( + AgentBinding, + AgentRunRecord, + AgentRunStatus, +) # =============================== # Test Data Classes @@ -2052,6 +2056,427 @@ def _create_extractor(self, extractor_config, service_config): pending_tool_call_ids=[], ) + @staticmethod + def _seed_inline_agent_run( + request_context, + *, + run_id: str, + request_id: str, + extractor_kind: str, + pending_tool_call_ids: list[str] | None = None, + ) -> None: + request_context.storage.create_agent_run( + AgentRunRecord( + id=run_id, + binding=AgentBinding( + org_id=request_context.org_id, + extractor_kind=extractor_kind, + user_id="test_user", + request_id=request_id, + agent_version="v1", + source="api", + ), + status=AgentRunStatus.AGENT_COMPLETED, + generation_request_snapshot={"request_id": request_id}, + committed_output={f"{extractor_kind}s": []}, + pending_tool_call_ids=pending_tool_call_ids or [], + ) + ) + + @staticmethod + def _build_receipt_required_plan( + entity_kind: str, + llm_client, + request_context, + *, + run_ids: list[str], + ): + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserPlaybook, UserProfile + from reflexio.server.services.deferred_learning_plan import ( + GenerationComputePlan, + PlaybookWritePlan, + ProfileWritePlan, + ) + from reflexio.server.services.playbook.service import PlaybookGenerationService + from reflexio.server.services.profile.service import ProfileGenerationService + + request_id = f"request-{entity_kind}-receipt-boundary" + if entity_kind == "profile": + profile = UserProfile( + profile_id="profile-receipt-boundary", + user_id="test_user", + content="This row requires a finalization receipt.", + last_modified_timestamp=1_000, + generated_from_request_id=request_id, + ) + service = ProfileGenerationService(llm_client, request_context) + write_plan = ProfileWritePlan( + user_id="test_user", + request_id=request_id, + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + else: + playbook = UserPlaybook( + user_id="test_user", + agent_version="v1", + request_id=request_id, + content="This row requires a finalization receipt.", + trigger="when validating persistence", + ) + service = PlaybookGenerationService( + llm_client, request_context, skip_aggregation=True + ) + write_plan = PlaybookWritePlan( + request_id=request_id, + output_pending_status=False, + skip_aggregation=True, + new_playbooks=[playbook], + superseded_ids=[], + merge_groups=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + return service, GenerationComputePlan( + prepared=SimpleNamespace(), + generated_count=1, + billable_count=1, + write_plan=write_plan, + bookmark_advance=None, + generation_start=0.0, + extraction_run_ids=run_ids, + token_totals=None, + ) + + @staticmethod + def _stored_inline_learning_count(request_context, entity_kind: str) -> int: + if entity_kind == "profile": + return len(request_context.storage.get_user_profile("test_user")) + return len( + request_context.storage.get_user_playbooks( + user_id="test_user", agent_version="v1" + ) + ) + + @pytest.mark.parametrize("entity_kind", ["profile", "playbook"]) + def test_receipt_required_plan_rejects_missing_run( + self, entity_kind, llm_client, request_context + ): + service, plan = self._build_receipt_required_plan( + entity_kind, + llm_client, + request_context, + run_ids=[f"missing-{entity_kind}-run"], + ) + + with pytest.raises(RuntimeError): + service.persist_generation(plan) + + assert self._stored_inline_learning_count(request_context, entity_kind) == 0 + + @pytest.mark.parametrize("entity_kind", ["profile", "playbook"]) + def test_receipt_required_plan_rejects_multiple_runs( + self, entity_kind, llm_client, request_context + ): + run_ids = [f"{entity_kind}-run-one", f"{entity_kind}-run-two"] + for run_id in run_ids: + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=f"request-{entity_kind}-receipt-boundary", + extractor_kind=entity_kind, + ) + service, plan = self._build_receipt_required_plan( + entity_kind, + llm_client, + request_context, + run_ids=run_ids, + ) + + with pytest.raises(RuntimeError): + service.persist_generation(plan) + + assert self._stored_inline_learning_count(request_context, entity_kind) == 0 + + def test_inline_profile_outcome_persists_receipt_before_terminal_status( + self, llm_client, request_context + ): + """A no-tool inline profile run retains an immutable receipt.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserProfile + from reflexio.server.services.deferred_learning_plan import ProfileWritePlan + from reflexio.server.services.profile.profile_generation_service_utils import ( + ProfileGenerationRequest, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + run_id = "run_inline_profile" + request_id = "request_inline_profile" + profile = UserProfile( + profile_id="profile-inline", + user_id="test_user", + content="The user deploys services to AWS ECS.", + last_modified_timestamp=1_000, + generated_from_request_id=request_id, + ) + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="profile", + ) + + class InlineProfileService(ProfileGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="profile") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([profile], run_id=run_id) + ) + + def _resolve_write_plan(self, results): + return ProfileWritePlan( + user_id="test_user", + request_id=request_id, + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + + service = InlineProfileService(llm_client, request_context) + service.run( + ProfileGenerationRequest( + user_id="test_user", + request_id=request_id, + source="api", + auto_run=False, + ) + ) + + stored = request_context.storage.get_user_profile("test_user") + run = request_context.storage.get_agent_run(run_id) + receipt = request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) + assert [item.profile_id for item in stored] == ["profile-inline"] + assert run is not None and run.status == AgentRunStatus.FINALIZED + assert receipt == ["profile-inline"] + + def test_inline_playbook_outcome_persists_receipt_before_terminal_status( + self, llm_client, request_context + ): + """A no-tool inline playbook run retains its database-assigned ids.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserPlaybook + from reflexio.server.services.deferred_learning_plan import PlaybookWritePlan + from reflexio.server.services.playbook.playbook_service_utils import ( + PlaybookGenerationRequest, + ) + from reflexio.server.services.playbook.service import PlaybookGenerationService + + run_id = "run_inline_playbook" + request_id = "request_inline_playbook" + playbook = UserPlaybook( + user_id="test_user", + agent_version="v1", + request_id=request_id, + content="Prefer AWS ECS for production deployments.", + trigger="when selecting a deployment target", + ) + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="playbook", + ) + + class InlinePlaybookService(PlaybookGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="playbook") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([playbook], run_id=run_id) + ) + + def _resolve_write_plan(self, results): + return PlaybookWritePlan( + request_id=request_id, + output_pending_status=False, + skip_aggregation=True, + new_playbooks=[playbook], + superseded_ids=[], + merge_groups=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + + def _dispatch_playbook_schedulers(self, plan): + return None + + service = InlinePlaybookService( + llm_client, + request_context, + skip_aggregation=True, + ) + service.run( + PlaybookGenerationRequest( + request_id=request_id, + agent_version="v1", + user_id="test_user", + source="api", + auto_run=False, + ) + ) + + stored = request_context.storage.get_user_playbooks( + user_id="test_user", + agent_version="v1", + ) + run = request_context.storage.get_agent_run(run_id) + receipt = request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="user_playbook", + ) + assert len(stored) == 1 + assert run is not None and run.status == AgentRunStatus.FINALIZED + assert receipt == [str(stored[0].user_playbook_id)] + + def test_inline_pending_tool_outcome_does_not_create_finalization_receipt( + self, llm_client, request_context + ): + """A pending-tool run remains available for resume finalization.""" + from reflexio.server.services.profile.profile_generation_service_utils import ( + ProfileGenerationRequest, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + run_id = "run_inline_pending_tool" + request_id = "request_inline_pending_tool" + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="profile", + pending_tool_call_ids=["pending-tool-call"], + ) + + class InlinePendingToolService(ProfileGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="profile") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([], run_id=run_id) + ) + + service = InlinePendingToolService(llm_client, request_context) + service.run( + ProfileGenerationRequest( + user_id="test_user", + request_id=request_id, + source="api", + auto_run=False, + ) + ) + + run = request_context.storage.get_agent_run(run_id) + assert run is not None + assert run.status == AgentRunStatus.FINALIZED_PENDING_TOOL + assert ( + request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) + is None + ) + + def test_inline_profile_persistence_failure_leaves_run_retryable( + self, llm_client, request_context + ): + """A failed inline receipt transaction never exposes terminal state.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserProfile + from reflexio.server.services.deferred_learning_plan import ProfileWritePlan + from reflexio.server.services.profile.profile_generation_service_utils import ( + ProfileGenerationRequest, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + run_id = "run_inline_profile_failure" + request_id = "request_inline_profile_failure" + profile = UserProfile( + profile_id="profile-inline-failure", + user_id="test_user", + content="This profile must roll back.", + last_modified_timestamp=1_000, + generated_from_request_id=request_id, + ) + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="profile", + ) + + class FailingInlineProfileService(ProfileGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="profile") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([profile], run_id=run_id) + ) + + def _resolve_write_plan(self, results): + return ProfileWritePlan( + user_id="test_user", + request_id=request_id, + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + + def _persist_write_plan(self, plan): + super()._persist_write_plan(plan) + raise RuntimeError("receipt persistence failed") + + service = FailingInlineProfileService(llm_client, request_context) + service.run( + ProfileGenerationRequest( + user_id="test_user", + request_id=request_id, + source="api", + auto_run=False, + ) + ) + + run = request_context.storage.get_agent_run(run_id) + assert request_context.storage.get_user_profile("test_user") == [] + assert ( + request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) + is None + ) + assert run is not None and run.status == AgentRunStatus.FINALIZATION_FAILED + assert run.finalization_attempts == 1 + def test_extraction_outcome_finalization_failure_marks_run_retryable( self, llm_client, request_context ): @@ -2086,6 +2511,41 @@ def _process_results(self, results): assert kwargs["last_error"] == "persist failed" assert kwargs["increment_finalization_attempts"] is True + def test_extraction_outcome_terminal_status_failure_marks_run_retryable( + self, llm_client, request_context + ): + """A failed terminal status write remains eligible for finalization retry.""" + + class FailingTerminalStatusService(ConcreteGenerationService): + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([{"name": "x"}], run_id="run_1") + ) + + service = FailingTerminalStatusService( + llm_client, + request_context, + extractor_configs=[MockExtractorConfig(extractor_name="ext1")], + ) + service.storage = MagicMock() + service.storage.get_agent_run.return_value = SimpleNamespace( + pending_tool_call_ids=[], + committed_output={"items": []}, + finalization_attempts=0, + ) + service.storage.update_agent_run_status.side_effect = [ + RuntimeError("terminal status failed"), + None, + ] + + service.run(MockServiceConfig(user_id="test_user", request_id="test_request")) + + _, status = service.storage.update_agent_run_status.call_args.args[:2] + kwargs = service.storage.update_agent_run_status.call_args.kwargs + assert status == AgentRunStatus.FINALIZATION_FAILED + assert kwargs["last_error"] == "terminal status failed" + assert kwargs["increment_finalization_attempts"] is True + def test_configured_extractor_timeout_fails_generation( self, llm_client, request_context, monkeypatch ): From 25ac8dd4c29ae0169fb60ddb65e7257291b659d1 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Sat, 15 Aug 2026 04:43:23 +0000 Subject: [PATCH 59/63] fix: keep tuner evidence production-only --- reflexio/server/routes/search.py | 21 ++++++------ .../routes/test_search_exposure_boundary.py | 33 +++++++++++++++++-- .../services/durable_learning/test_worker.py | 31 +++++++++++++---- .../playbook/test_playbook_reviewer.py | 4 ++- 4 files changed, 70 insertions(+), 19 deletions(-) diff --git a/reflexio/server/routes/search.py b/reflexio/server/routes/search.py index 1adb23301..7253a7572 100644 --- a/reflexio/server/routes/search.py +++ b/reflexio/server/routes/search.py @@ -251,7 +251,7 @@ def search_user_playbooks_endpoint( msg=response.msg, experiment=assignment, ) - if response.user_playbooks: + if caller_type == "production_agent" and response.user_playbooks: record_search_exposures( SearchExposureBatch( org_id=org_id, @@ -427,16 +427,17 @@ def run_search() -> Any: rehydrated_text=response.rehydrated_text, experiment=assignment, ) - record_search_exposures( - SearchExposureBatch( - org_id=org_id, - request_id=payload.request_id, - session_id=payload.session_id, - interaction_id=payload.interaction_id, - user_id=payload.user_id, - user_playbooks=tuple(response.user_playbooks), + if caller_type == "production_agent": + record_search_exposures( + SearchExposureBatch( + org_id=org_id, + request_id=payload.request_id, + session_id=payload.session_id, + interaction_id=payload.interaction_id, + user_id=payload.user_id, + user_playbooks=tuple(response.user_playbooks), + ) ) - ) enqueue_search_metering( org_id=org_id, caller_type=caller_type, diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py index 1530aef7f..bb0dc661f 100644 --- a/tests/server/routes/test_search_exposure_boundary.py +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -57,11 +57,11 @@ def _search_results(playbooks: list[UserPlaybook]) -> Iterator[MagicMock]: yield reflexio -def _client() -> TestClient: +def _client(caller_type: str = "production_agent") -> TestClient: return TestClient( create_app( get_org_id=lambda: "org-1", - get_caller_type=lambda: "production_agent", + get_caller_type=lambda: caller_type, ), raise_server_exceptions=False, ) @@ -80,6 +80,35 @@ def record(self, batch: Any) -> None: self.completed = True +@pytest.mark.parametrize("path", ["/api/search", "/api/search_user_playbooks"]) +def test_only_production_searches_record_exposures(path: str) -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([_playbook(11, "First")]): + dashboard_response = _client("dashboard").post( + path, + json={ + "query": "inspect", + "user_id": "user-1", + "request_id": "dashboard-request", + }, + ) + production_response = _client("production_agent").post( + path, + json={ + "query": "answer", + "user_id": "user-1", + "request_id": "production-request", + }, + ) + + assert dashboard_response.status_code == 200, dashboard_response.text + assert production_response.status_code == 200, production_response.text + assert len(recorder.batches) == 1 + assert recorder.batches[0].request_id == "production-request" + + def test_unified_search_records_the_final_user_playbook_set_before_return() -> None: playbooks = [_playbook(11, "First"), _playbook(12, "Second")] recorder = _Recorder() diff --git a/tests/server/services/durable_learning/test_worker.py b/tests/server/services/durable_learning/test_worker.py index 212cbe252..4e606fa06 100644 --- a/tests/server/services/durable_learning/test_worker.py +++ b/tests/server/services/durable_learning/test_worker.py @@ -29,14 +29,17 @@ import threading import time from datetime import UTC, datetime, timedelta -from types import SimpleNamespace from unittest import mock import pytest from reflexio.models.api_schema.domain.entities import Interaction, Request from reflexio.server.api_endpoints.request_context import RequestContext -from reflexio.server.services.deferred_learning_plan import DeferredLearningPlan +from reflexio.server.services.base_generation_service import PreparedGenerationRun +from reflexio.server.services.deferred_learning_plan import ( + DeferredLearningPlan, + GenerationComputePlan, +) from reflexio.server.services.durable_learning.worker import DurableLearningWorker from reflexio.server.services.generation_service import GenerationService from reflexio.server.services.storage.storage_base import ( @@ -163,13 +166,29 @@ def _deferred_plan_with_runs( profile_run_id: str, playbook_run_id: str, ) -> DeferredLearningPlan: + def generation_plan(run_id: str, extractor_name: str) -> GenerationComputePlan: + return GenerationComputePlan( + prepared=PreparedGenerationRun( + extractor_config=mock.Mock(), + extractor_name=extractor_name, + identifier=f"{extractor_name}_generation", + ), + generated_count=0, + billable_count=0, + write_plan=None, + bookmark_advance=None, + generation_start=0.0, + extraction_run_ids=[run_id], + token_totals=None, + ) + return DeferredLearningPlan( request_id=request_id, user_id="test_user", agent_version="v1", lock_acquired=True, - profile=(mock.Mock(), SimpleNamespace(extraction_run_ids=[profile_run_id])), - playbook=(mock.Mock(), SimpleNamespace(extraction_run_ids=[playbook_run_id])), + profile=(mock.Mock(), generation_plan(profile_run_id, "profile")), + playbook=(mock.Mock(), generation_plan(playbook_run_id, "playbook")), ) @@ -730,7 +749,6 @@ def test_post_commit_emit_failure_preserves_completed_job_and_agent_runs(): from reflexio.models.api_schema.domain.entities import LineageContext from reflexio.models.api_schema.service_schemas import UserProfile from reflexio.server.services.deferred_learning_plan import ( - GenerationComputePlan, ProfileWritePlan, ) from reflexio.server.services.profile.service import ProfileGenerationService @@ -770,7 +788,8 @@ def build_plan(self, **_kwargs): generated_from_request_id="req_post_commit_emit_failure", ) generation_plan = GenerationComputePlan( - prepared=SimpleNamespace( + prepared=PreparedGenerationRun( + extractor_config=mock.Mock(), identifier="profile_generation", extractor_name="profile", ), diff --git a/tests/server/services/playbook/test_playbook_reviewer.py b/tests/server/services/playbook/test_playbook_reviewer.py index cde7141f7..e0d8a193c 100644 --- a/tests/server/services/playbook/test_playbook_reviewer.py +++ b/tests/server/services/playbook/test_playbook_reviewer.py @@ -465,7 +465,9 @@ def test_reviewer_prompt_preserves_grounded_procedures_and_forbids_substitutes() # fixes is the reviewer going down the evidence axis and never asking. "First name, in your own words, what the entry is ABOUT", "before asking\nwhether any of its clauses are supported".replace("\n", " "), - "is not a core to preserve, it is\nthe same subject in gentler words".replace("\n", " "), + "is not a core to preserve, it is\nthe same subject in gentler words".replace( + "\n", " " + ), ) for invariant in required_invariants: assert invariant in normalized From e2e7334cdcc72bbfb45eb64b8d6c52603aac3ee9 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Sat, 15 Aug 2026 06:19:06 +0000 Subject: [PATCH 60/63] fix: expose playbook search correlation ids --- docs/lib/methods/user-playbooks.ts | 12 ++++++++++++ reflexio/client/client.py | 8 ++++++++ tests/client/test_search.py | 31 +++++++++++++++++++++++++----- 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/docs/lib/methods/user-playbooks.ts b/docs/lib/methods/user-playbooks.ts index ef1d55e40..d7d88f0a9 100644 --- a/docs/lib/methods/user-playbooks.ts +++ b/docs/lib/methods/user-playbooks.ts @@ -70,6 +70,18 @@ export const userPlaybookMethods: MethodDef[] = [ required: false, description: "Filter by user (via request_id linkage to requests table)", }, + { + name: "request_id", + type: "string", + required: false, + description: "Caller correlation ID, at most 255 characters", + }, + { + name: "session_id", + type: "string", + required: false, + description: "Caller session ID, at most 255 characters", + }, { name: "agent_version", type: "string", diff --git a/reflexio/client/client.py b/reflexio/client/client.py index 272efe918..b16acb1b3 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -895,6 +895,8 @@ def search_user_playbooks( threshold: float | None = None, enable_reformulation: bool | None = None, search_mode: SearchMode | None = None, + request_id: str | None = None, + session_id: str | None = None, ) -> SearchUserPlaybooksViewResponse: """Search for user playbooks with semantic/text search and filtering. @@ -912,6 +914,10 @@ def search_user_playbooks( threshold (Optional[float]): Similarity threshold for vector search. When omitted, the embedding model's default is used. enable_reformulation (Optional[bool]): Enable LLM query reformulation (default: False) + request_id (Optional[str]): Caller correlation ID for the search turn, + at most 255 characters. + session_id (Optional[str]): Caller session ID for the search turn, + at most 255 characters. Returns: SearchUserPlaybooksViewResponse: Response containing matching user playbooks @@ -931,6 +937,8 @@ def search_user_playbooks( threshold=threshold, enable_reformulation=enable_reformulation, search_mode=search_mode, + request_id=request_id, + session_id=session_id, ) response = self._make_request( "POST", "/api/search_user_playbooks", json=req.model_dump(mode="json") diff --git a/tests/client/test_search.py b/tests/client/test_search.py index a1e644b98..ee9dc43c6 100644 --- a/tests/client/test_search.py +++ b/tests/client/test_search.py @@ -59,6 +59,24 @@ def fake_make_request(method: str, path: str, **kwargs: Any) -> dict[str, Any]: assert captured["json"]["user_id"] == "user-1" +def test_user_playbook_search_serializes_correlation_ids(monkeypatch) -> None: + client = ReflexioClient(api_key="test-key", url_endpoint="http://localhost:8000") + captured: dict[str, Any] = {} + + def fake_make_request(method: str, path: str, **kwargs: Any) -> dict[str, Any]: + captured.update(method=method, path=path, **kwargs) + return {"success": True, "user_playbooks": []} + + monkeypatch.setattr(client, "_make_request", fake_make_request) + + client.search_user_playbooks( + query="billing", request_id="request-1", session_id="session-1" + ) + + assert captured["json"]["request_id"] == "request-1" + assert captured["json"]["session_id"] == "session-1" + + @pytest.mark.asyncio async def test_unified_search_async_uses_native_transport(monkeypatch) -> None: client = ReflexioClient(api_key="test-key", url_endpoint="http://localhost:8000") @@ -110,13 +128,16 @@ def test_unified_search_docs_track_schema_contract() -> None: def test_user_playbook_search_docs_track_top_k_schema_contract() -> None: - top_k_limit = _non_null_schema("top_k", request_model=SearchUserPlaybookRequest)[ - "maximum" - ] + top_k_schema = _non_null_schema("top_k", request_model=SearchUserPlaybookRequest) + identifier_limit = _non_null_schema( + "request_id", request_model=SearchUserPlaybookRequest + )["maxLength"] client_docs = inspect.getdoc(ReflexioClient.search_user_playbooks) or "" registry_docs = ( Path(__file__).parents[2] / "docs/lib/methods/user-playbooks.ts" ).read_text(encoding="utf-8") - assert f"1 to {top_k_limit}" in client_docs - assert f"1 to {top_k_limit}" in registry_docs + assert f"1 to {top_k_schema['maximum']}" in client_docs + assert client_docs.count(f"at most {identifier_limit} characters") >= 2 + assert f"1 to {top_k_schema['maximum']}" in registry_docs + assert registry_docs.count(f"at most {identifier_limit} characters") >= 2 From 19864acd9c965b62f89ade534421d9f053152905 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Sat, 15 Aug 2026 06:19:06 +0000 Subject: [PATCH 61/63] test: mark readiness match as regex --- tests/cli/test_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index 300d46653..98b0f1f5a 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -346,7 +346,7 @@ def fake_popen(command, **_kwargs) -> _FakeProcess: monkeypatch.setattr(utils.subprocess, "Popen", fake_popen) monkeypatch.setattr(utils, "_wait_for_all_ready", lambda *_args: False) - with pytest.raises(RuntimeError, match="embedding.*ready"): + with pytest.raises(RuntimeError, match=r"embedding.*ready"): utils.run_services( [ utils.ServiceConfig(name="embedding", command=["embedding"]), From f614717414c2705e3709bdca962fb96b2966efa1 Mon Sep 17 00:00:00 2001 From: Guangyu Date: Sat, 15 Aug 2026 06:19:06 +0000 Subject: [PATCH 62/63] docs: make SQLite check work without uv --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 5289a8459..31d820b02 100644 --- a/README.md +++ b/README.md @@ -111,7 +111,7 @@ Publish conversations from your agent, and Reflexio closes the self-improvement Check the SQLite runtime linked to Python with: ```shell -uv run python -c "import sqlite3; print(sqlite3.sqlite_version_info)" +python -c "import sqlite3; print(sqlite3.sqlite_version_info)" ```

From 3a2462919d9e50cff29c6352db809396fd0459dc Mon Sep 17 00:00:00 2001 From: Guangyu Date: Sat, 15 Aug 2026 18:44:47 +0000 Subject: [PATCH 63/63] test: model missing finalization receipts --- .../playbook/test_playbook_generation_service_integration.py | 1 + 1 file changed, 1 insertion(+) 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 ffd8e77f9..599db9e79 100644 --- a/tests/server/services/playbook/test_playbook_generation_service_integration.py +++ b/tests/server/services/playbook/test_playbook_generation_service_integration.py @@ -80,6 +80,7 @@ def mock_request_context(): context = MagicMock(spec=RequestContext) context.org_id = "test_org_123" context.storage = MagicMock() + context.storage.get_agent_run_finalization_receipt.return_value = None agent_runs = {} def create_agent_run(record):