From ab2c7603b2272ef074bb38f031b0b6a5cf4aa93c Mon Sep 17 00:00:00 2001 From: Yi Lu Date: Tue, 4 Aug 2026 00:49:56 -0700 Subject: [PATCH 1/4] revert: remove open-world evidence foundation Revert squash commit 85a4b225 (#407) and its dependent follow-ups 785a9e05 (#408) and eb88f44f (#409). Preserve the independent aggregation work that landed afterward. --- 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 | 32 +- .../server/services/deferred_learning_plan.py | 35 +- reflexio/server/services/extraction/README.md | 5 +- .../services/extraction/resume_worker.py | 88 +- .../server/services/governance/service.py | 332 +--- .../playbook/components/aggregator.py | 68 +- reflexio/server/services/playbook/service.py | 126 +- reflexio/server/services/profile/service.py | 86 +- reflexio/server/services/search_exposure.py | 166 -- .../services/storage/governance_claims.py | 30 - .../services/storage/governance_validation.py | 8 - reflexio/server/services/storage/retention.py | 15 - .../services/storage/retention_mixin.py | 36 +- .../storage/session_outcome_identity.py | 188 --- .../services/storage/sqlite_storage/_base.py | 245 +-- .../storage/sqlite_storage/_governance.py | 32 +- .../sqlite_storage/_session_outcomes.py | 184 +-- .../agent_run/_agent_run_store.py | 90 -- .../storage/sqlite_storage/base/_deletion.py | 9 +- .../governance/_erase_execution.py | 48 +- .../sqlite_storage/governance/_purge.py | 509 +----- .../governance/_subject_barrier.py | 106 +- .../services/storage/storage_base/__init__.py | 9 +- .../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 | 1006 +----------- .../governance/test_governance_local_e2e.py | 703 +-------- .../test_subject_write_barrier_sqlite.py | 358 +---- .../sqlite_storage/test_agent_run_storage.py | 209 --- .../test_governance_retrieved_learning.py | 21 +- .../sqlite_storage/test_governance_storage.py | 1401 +---------------- .../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_retention.py | 46 - .../test_storage_contract_session_outcomes.py | 338 ---- .../test_generation_billing_emission.py | 149 +- .../test_non_extraction_learning_metering.py | 212 +-- tests/server/services/test_search_exposure.py | 292 ---- tests/server/test_billing_meter_events.py | 11 +- 57 files changed, 678 insertions(+), 8037 deletions(-) delete mode 100644 reflexio/server/services/search_exposure.py delete mode 100644 reflexio/server/services/storage/governance_claims.py delete mode 100644 reflexio/server/services/storage/session_outcome_identity.py delete mode 100644 tests/models/test_session_outcome_identity.py delete mode 100644 tests/server/routes/test_search_exposure_boundary.py delete mode 100644 tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py delete mode 100644 tests/server/services/test_search_exposure.py diff --git a/README.md b/README.md index 6ef9efbb7..f6ad8eab2 100644 --- a/README.md +++ b/README.md @@ -104,15 +104,8 @@ Publish conversations from your agent, and Reflexio closes the self-improvement | Tool | Description | | --- | --- | | [uv](https://docs.astral.sh/uv/getting-started/installation/) | Python package manager | -| Python's linked SQLite runtime | >= 3.35.0 (required for local SQLite storage; not the standalone `sqlite3` CLI) | | [Node.js](https://nodejs.org/) >= 18 | Frontend runtime | -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 9b7be37e7..4efbe928e 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 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.", + "Record the lifetime-durable first success or failure marker for an existing session. 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: success, failure, or unknown", - enumValues: ["success", "failure", "unknown"], + description: "Terminal outcome", + enumValues: ["success", "failure"], }, { 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: success, failure, or unknown", - enumValues: ["success", "failure", "unknown"], + description: "Exact outcome filter", + enumValues: ["success", "failure"], }, { 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 7283e7ae9..af2361b63 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -1112,20 +1112,13 @@ def mark_session_outcome( value: float | None = None, metadata: dict[str, Any] | None = None, ) -> SetSessionOutcomeResponse: - """Record the immutable first outcome for a published session. + """Record the first terminal 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)``. 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. + 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. """ request = SetSessionOutcomeRequest( session_id=session_id, diff --git a/reflexio/lib/_session_outcome.py b/reflexio/lib/_session_outcome.py index 7a39ba518..d49cbcead 100644 --- a/reflexio/lib/_session_outcome.py +++ b/reflexio/lib/_session_outcome.py @@ -53,44 +53,51 @@ def mark_session_outcome( try: for _attempt in range(3): context = storage.get_session_outcome_context(request.session_id) - 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: + 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: return SetSessionOutcomeResponse( success=False, - reason=SessionOutcomeFailureReason.OCCURRED_BEFORE_SESSION, - message="Outcome occurred before the session began", + reason=reason, + message="Outcome was not accepted", 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, @@ -105,10 +112,6 @@ 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, @@ -120,10 +123,6 @@ 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 799c20e72..46540eebf 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -823,8 +823,6 @@ 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,35 +831,8 @@ 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 @@ -907,35 +878,6 @@ 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): @@ -1012,8 +954,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 (session outcomes, interactions, -# requests, user playbooks, profiles). Used by paired-protocol harnesses to +# Clear all data scoped to a single user_id (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 a66a1759d..98866fc6c 100644 --- a/reflexio/models/api_schema/domain/enums.py +++ b/reflexio/models/api_schema/domain/enums.py @@ -18,11 +18,9 @@ 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 46c7e764f..76818fe81 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -90,10 +90,11 @@ def record_learnings_generated( ) -> None: """Emit the Learning value facet — number of profiles/playbooks 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 - durable ids. No-op when ``count <= 0``. + 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``. Emits a single event carrying a synthesized ``event_key=f"learn-batch:{uuid4()}"`` (distinct per call) so this aggregate event still has a dedup key, even @@ -177,7 +178,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``). + (e.g. ``profile_id`` / ``user_playbook_id`` / ``agent_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"``). @@ -230,20 +231,20 @@ def emit_learnings_generated( ) -> None: """Resolve ``platform_llm`` from config and emit the Learning value facet. - 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``. + 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``. 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. ``"online_extraction"``). + source: Metering source/path label (e.g. ``"offline_optimizer"``). pipeline: Optional pipeline tag (e.g. ``"playbook"``). user_id: Optional user ID tied to the generated learning. request_id: Optional request correlation ID. @@ -298,20 +299,22 @@ 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 + 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 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. ``"resumable_extraction"``). + source: Metering source/path label (e.g. ``"aggregation"``). 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 356626f78..e7b7bbe53 100644 --- a/reflexio/server/routes/search.py +++ b/reflexio/server/routes/search.py @@ -52,10 +52,6 @@ from reflexio.server.services.retrieval_experiment import ( active_retrieval_experiment_assignment, ) -from reflexio.server.services.search_exposure import ( - SearchExposureBatch, - record_search_exposures, -) from reflexio.server.tracing import profile_step logger = logging.getLogger(__name__) @@ -423,16 +419,6 @@ 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), - ) - ) background_tasks.add_task( _meter_search_request, org_id=org_id, diff --git a/reflexio/server/services/base_generation/_usage_billing.py b/reflexio/server/services/base_generation/_usage_billing.py index 5aee8a7d9..1b1304f90 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 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. + 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. Args: prepared: The prepared generation run (used for input-text computation). - generated_count: Number of retained write-plan learnings eligible for billing. + generated_count: Number of learnings produced by this extraction run. """ if not self.EMITS_LEARNING_BILLING: return @@ -260,16 +260,3 @@ 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 9fd3f9fe6..ebf78b7b5 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 online profile/playbook extraction services emit extraction-run billing - # here. Resumable-extraction finalization emits separately; derived mutation - # paths emit no additional learnings_generated events. + # 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. # Default is False so any future subclass is safe by default (opt-IN). EMITS_LEARNING_BILLING: bool = False """ @@ -342,20 +342,10 @@ 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, - *, - finalization_run_id: str | None = None, - ) -> list[str] | None: + def _finalize_extracted_items(self, items: list) -> 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 None @abstractmethod def _should_track_in_progress(self) -> bool: @@ -643,6 +633,7 @@ 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 @@ -651,17 +642,9 @@ 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, @@ -692,8 +675,7 @@ 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`` / ``billable_count`` / ``prepared`` / - ``generation_start``) so a + snapshot (``generated_count`` / ``prepared`` / ``generation_start``) so a fence-lost job never emits. Billing purity note (round-2 finding): ``_record_billing_learning_events`` @@ -725,7 +707,7 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: }, ) self._record_billing_learning_events( - prepared=plan.prepared, generated_count=plan.billable_count + prepared=plan.prepared, generated_count=plan.generated_count ) @abstractmethod diff --git a/reflexio/server/services/deferred_learning_plan.py b/reflexio/server/services/deferred_learning_plan.py index 7e00c2ad9..eb0a87cd8 100644 --- a/reflexio/server/services/deferred_learning_plan.py +++ b/reflexio/server/services/deferred_learning_plan.py @@ -26,24 +26,6 @@ ) -@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). @@ -152,20 +134,16 @@ class GenerationComputePlan: ``emit_generation_side_effects`` fires the post-commit telemetry + billing. The billing inputs (``extraction_run_ids`` / ``token_totals`` / - ``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. + ``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. Attributes: prepared: The prepared generation run (identifier / extractor_name / extractor_config), reused by emit for telemetry + billing input. - 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. + generated_count: Learnings produced by this extraction run. 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. @@ -181,7 +159,6 @@ 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 dea2d7584..3acb5a553 100644 --- a/reflexio/server/services/extraction/README.md +++ b/reflexio/server/services/extraction/README.md @@ -16,16 +16,13 @@ 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` | Schedules due paused extraction runs in a background singleton. | -| `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. | +| `resume_worker.py` | Resumes paused runs, rebuilds request context, and records retry state. | | `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 392081f1d..b6f78bc47 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -19,7 +19,6 @@ 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, @@ -318,9 +317,8 @@ def run_once(self) -> AgentRunRecord | None: try: self.storage.update_agent_run_status(run.id, AgentRunStatus.FINALIZING) - result = self._finalize_items(run, items, model_provenance=model_provenance) - if result.won_receipt: - self._schedule_finalized_tagging(run) + self._finalize_items(run, items, model_provenance=model_provenance) + self._schedule_finalized_tagging(run) self.storage.consume_run_tool_dependencies(run.id) finalized_status = ( AgentRunStatus.FINALIZED_PENDING_TOOL @@ -365,9 +363,8 @@ def _retry_finalization(self, run: AgentRunRecord) -> AgentRunRecord | None: items, pending_tool_call_ids, model_provenance = ( self._items_from_committed_output(run) ) - result = self._finalize_items(run, items, model_provenance=model_provenance) - if result.won_receipt: - self._schedule_finalized_tagging(run) + self._finalize_items(run, items, model_provenance=model_provenance) + self._schedule_finalized_tagging(run) self.storage.consume_run_tool_dependencies(run.id) finalized_status = ( AgentRunStatus.FINALIZED_PENDING_TOOL @@ -874,7 +871,7 @@ def _finalize_items( items: list[Any], *, model_provenance: ModelProvenance | None = None, - ) -> FinalizationResult: + ) -> None: if run.binding.extractor_kind == "profile": service = ProfileGenerationService( llm_client=self.client, @@ -887,16 +884,9 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - result = service._finalize_extracted_items_with_outcome( - 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" - ) - return result + service._finalize_extracted_items(items, model_provenance=model_provenance) + self._record_finalized_learnings(run, items, entity_type="profile") + return if run.binding.extractor_kind == "playbook": service = PlaybookGenerationService( llm_client=self.client, @@ -910,40 +900,68 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - result = service._finalize_extracted_items_with_outcome( - items, - model_provenance=model_provenance, - finalization_run_id=run.id, - ) - if result.won_receipt: - self._record_finalized_learnings( - run, result.learning_ids, entity_type="user_playbook" - ) - return result + service._finalize_extracted_items(items, model_provenance=model_provenance) + self._record_finalized_learnings(run, items, entity_type="user_playbook") + return raise ResumeWorkerError( f"Unsupported extractor kind {run.binding.extractor_kind!r}" ) def _record_finalized_learnings( - self, run: AgentRunRecord, learning_ids: list[str], *, entity_type: str + self, run: AgentRunRecord, items: list[Any], *, entity_type: str ) -> None: """Emit ``learnings_generated`` for a finalized resumable-extraction batch. - Emits one event per durable learning id returned by finalization. - Per-record keys make finalization retries idempotent downstream. + 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. """ - from reflexio.server.billing_meter import emit_learnings_generated_records + from reflexio.server.billing_meter import ( + emit_learnings_generated, + emit_learnings_generated_records, + ) - if not learning_ids: + if not items: 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, } - emit_learnings_generated_records( + 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( org_id=self.request_context.org_id, configurator=self.request_context.configurator, - learning_ids=learning_ids, + count=len(items), source="resumable_extraction", pipeline=run.binding.extractor_kind, user_id=run.binding.user_id, diff --git a/reflexio/server/services/governance/service.py b/reflexio/server/services/governance/service.py index de3004e7a..fca64cdbf 100644 --- a/reflexio/server/services/governance/service.py +++ b/reflexio/server/services/governance/service.py @@ -1,8 +1,5 @@ from __future__ import annotations -import threading -import time -import uuid from contextlib import suppress from typing import Any, Literal, Protocol, TypedDict @@ -19,7 +16,6 @@ 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", @@ -33,18 +29,12 @@ "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): @@ -53,7 +43,7 @@ class GovernanceActorContext(TypedDict): class SubjectErasureLifecycle(Protocol): - """External erasure work invoked only after a synchronous live-claim check.""" + """Deployment-specific erasure work that must precede barrier completion.""" def erase_subject( self, @@ -61,78 +51,9 @@ 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, @@ -208,145 +129,84 @@ def erase_user( f"{self.org_id}:user_erasure:{subref}:{reqref}", ) purge_id = stable_id("purge", idempotency_key) - 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 + 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, + ) if purge.status == "complete": - return self._completed_erase_result_for_retry( + barrier = self._completed_barrier_for_retry( subject_ref=subref, purge_id=purge_id ) - 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, - ) - 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" + if barrier.status != "erased": + raise ValueError( + "Completed purge retry requires an erased subject barrier" ) - self._sleep(min(poll_seconds, remaining_seconds)) - poll_seconds = min( - poll_seconds * 2, - _DUPLICATE_ERASE_MAX_POLL_SECONDS, + 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) + ), ) try: - with _PurgeExecutionHeartbeat( - storage=self.storage, - purge_id=purge_id, - execution_claim=execution_claim, - ) as heartbeat: - self.storage.begin_subject_erasure_barrier( - subref, + self.storage.begin_subject_erasure_barrier(subref, purge_id) + if not self.storage.purge_targets_prepared(purge_id): + self.storage.prepare_governance_erase_targets( purge_id, - execution_claim=heartbeat.claim(), + user_id, ) - 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, - 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, - }, - ), - authoritative_user_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) + + 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, + }, + ), + ) 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( @@ -357,28 +217,6 @@ 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: @@ -397,57 +235,6 @@ 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]]]: @@ -523,11 +310,7 @@ 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, - execution_claim: PurgeExecutionClaim, - ) -> None: + def _record_subject_erasure_lifecycle_complete(self, purge_id: str) -> None: snapshot = self._prepared_target_snapshot(purge_id) if snapshot is None or snapshot.status != "complete": raise ValueError( @@ -544,7 +327,6 @@ def _record_subject_erasure_lifecycle_complete( 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 72efee5f7..0a827797f 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, Sequence +from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractContextManager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Protocol, cast @@ -1961,6 +1961,17 @@ 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: @@ -2003,6 +2014,61 @@ 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 40bb856cf..2c67f66ed 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -30,11 +30,7 @@ BaseGenerationService, StatusChangeOperation, ) -from reflexio.server.services.deferred_learning_plan import ( - FinalizationResult, - PlaybookWritePlan, - _FinalizationReceiptAlreadyExistsError, -) +from reflexio.server.services.deferred_learning_plan import PlaybookWritePlan from reflexio.server.services.playbook.aggregation_trigger import ( maybe_trigger_user_playbook_aggregation, ) @@ -600,21 +596,15 @@ 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 ``_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. + 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. """ - try: - self._enqueue_user_playbook_optimization(plan.new_playbooks) - except Exception: - logger.exception("Failed to schedule user playbook optimization") + self._enqueue_user_playbook_optimization(plan.new_playbooks) if not plan.output_pending_status and not plan.skip_aggregation: - try: - logger.info("Trigger playbook aggregation") - self._trigger_playbook_aggregation() - except Exception: - logger.exception("Failed to schedule user playbook aggregation") + logger.info("Trigger playbook aggregation") + self._trigger_playbook_aggregation() def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: """Post-commit side-effects — base telemetry/billing + playbook schedulers. @@ -634,97 +624,23 @@ def _finalize_extracted_items( all_playbooks: list[UserPlaybook], *, model_provenance: ModelProvenance | None = None, - finalization_run_id: str | None = None, - ) -> list[str]: - """Finalize extracted playbooks for synchronous resume/manual callers. - - 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. + ) -> None: + """Permanent V3 wrapper: compute→persist→schedulers together (no 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`` — then dispatches the same off-thread schedulers, so the + result is identical to the pre-split monolith. """ - 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] - 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 FinalizationResult(receipt, won_receipt=False) if model_provenance is not None: self._last_model_provenance = model_provenance plan = self._resolve_write_plan([all_playbooks]) - learning_ids: list[str] = [] - if finalization_run_id is None: - if plan is not None: - self._persist_write_plan(plan) - self._dispatch_playbook_schedulers(plan) - 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, - ) - - 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 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 FinalizationResult(learning_ids, won_receipt=True) + if plan is None: + return + self._persist_write_plan(plan) + self._dispatch_playbook_schedulers(plan) def _apply_consolidation_lineage( self, diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index 40cbd5fab..3e253ddaa 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -30,11 +30,7 @@ BaseGenerationService, StatusChangeOperation, ) -from reflexio.server.services.deferred_learning_plan import ( - FinalizationResult, - ProfileWritePlan, - _FinalizationReceiptAlreadyExistsError, -) +from reflexio.server.services.deferred_learning_plan import ProfileWritePlan from reflexio.server.services.profile.components.extractor import ProfileExtractor from reflexio.server.services.profile.profile_generation_service_utils import ( ProfileGenerationRequest, @@ -341,84 +337,20 @@ def _finalize_extracted_items( all_new_profiles: list[UserProfile], *, model_provenance: ModelProvenance | None = None, - finalization_run_id: str | None = None, - ) -> list[str]: + ) -> None: """Permanent V3 wrapper: compute-then-persist together (no external fence). - 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. + 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. """ - 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] - run_id=finalization_run_id, - entity_type=entity_type, - ) - if receipt is not None: - 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]) - # 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: - self._persist_write_plan(plan) - return FinalizationResult(learning_ids, won_receipt=False) - - 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 None: - raise RuntimeError( - "finalization receipt disappeared after insert conflict" - ) from None - return FinalizationResult(receipt, won_receipt=False) - return FinalizationResult(learning_ids, won_receipt=True) + if plan is not None: + self._persist_write_plan(plan) 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 deleted file mode 100644 index 63e887fff..000000000 --- a/reflexio/server/services/search_exposure.py +++ /dev/null @@ -1,166 +0,0 @@ -"""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) - ) - 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) - - -@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 - 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) -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, - 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, - "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=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, - 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, - playbook_owner_governance_subject_ref=(playbook_owner_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 deleted file mode 100644 index 147d3386b..000000000 --- a/reflexio/server/services/storage/governance_claims.py +++ /dev/null @@ -1,30 +0,0 @@ -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 3358604ec..a89fa1def 100644 --- a/reflexio/server/services/storage/governance_validation.py +++ b/reflexio/server/services/storage/governance_validation.py @@ -90,7 +90,6 @@ { "affected_agent_playbook_ids", "agent_playbook_id", - "authoritative_user_digest", "count", "deleted_counts", "deleted_count", @@ -167,7 +166,6 @@ "user_playbooks", "profiles", "requests", - "session_outcomes", "agent_success_evaluation_results", "retrieved_learning_evaluation_results", "evaluation_operation_states", @@ -478,12 +476,6 @@ 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 9dcb86292..d48035689 100644 --- a/reflexio/server/services/storage/retention.py +++ b/reflexio/server/services/storage/retention.py @@ -8,7 +8,6 @@ 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") @@ -21,7 +20,6 @@ class RetentionTarget: order_column: str id_columns: tuple[str, ...] priority_statuses: tuple[str, ...] = () - minimum_age_seconds: int = 0 @dataclass(frozen=True, slots=True) @@ -50,12 +48,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", @@ -138,13 +130,6 @@ class OptimizationRetentionClass: "created_at", ("retrieval_log_id",), ), - RetentionTarget( - "user_playbook_exposure_events", - "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 8a0e4db29..969cab1e2 100644 --- a/reflexio/server/services/storage/retention_mixin.py +++ b/reflexio/server/services/storage/retention_mixin.py @@ -9,7 +9,6 @@ from __future__ import annotations -import time from abc import ABC, abstractmethod from collections.abc import Iterator, Sequence from typing import Any @@ -99,16 +98,7 @@ 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 - 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, - ) + keys = self._retention_select_keys(target, count) if not keys: return 0 self._retention_perform_delete(target, keys) @@ -147,11 +137,7 @@ def _retention_gc_retired_optimization_jobs( return 0 def _retention_select_keys( - self, - target: RetentionTarget, - count: int, - *, - older_than_epoch: int | None, + self, target: RetentionTarget, count: int ) -> list[tuple[Any, ...]]: """Select tombstones first, then oldest rows when a target opts in. @@ -161,27 +147,16 @@ def _retention_select_keys( table holds that many rows. """ if not target.priority_statuses: - return self._retention_select_oldest_keys( - target, - count, - older_than_epoch=older_than_epoch, - ) + return self._retention_select_oldest_keys(target, count) keys = self._retention_select_oldest_keys( - target, - count, - statuses=target.priority_statuses, - older_than_epoch=older_than_epoch, + target, count, statuses=target.priority_statuses ) if len(keys) >= count: return keys seen = set(keys) - for key in self._retention_select_oldest_keys( - target, - count, - older_than_epoch=older_than_epoch, - ): + for key in self._retention_select_oldest_keys(target, count): if key not in seen: keys.append(key) seen.add(key) @@ -219,7 +194,6 @@ 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/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py deleted file mode 100644 index 273b0195a..000000000 --- a/reflexio/server/services/storage/session_outcome_identity.py +++ /dev/null @@ -1,188 +0,0 @@ -"""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 a94f0a158..9fec58e38 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -17,12 +17,10 @@ import threading 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,13 +57,6 @@ 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 @@ -74,9 +65,6 @@ logger = logging.getLogger(__name__) -_OUTCOME_ALLOWED_VALUES = tuple(kind.value for kind in SessionOutcomeKind) -_MINIMUM_SQLITE_VERSION = (3, 35, 0) - # --------------------------------------------------------------------------- # Module-level helpers @@ -97,47 +85,6 @@ 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"[a-zA-Z0-9_]+") @@ -716,12 +663,6 @@ def __init__( logger.info("SQLite Storage for org %s using db_path: %s", org_id, db_path) - if sqlite3.sqlite_version_info < _MINIMUM_SQLITE_VERSION: - detected_version = ".".join(map(str, sqlite3.sqlite_version_info)) - raise RuntimeError( - f"SQLite 3.35.0 or newer is required; detected {detected_version}" - ) - # Ensure parent directory exists Path(db_path).parent.mkdir(parents=True, exist_ok=True) @@ -830,26 +771,11 @@ def migrate(self) -> bool: # before executescript so index creation cannot fail on missing columns. self._migrate_request_retrieval_experiment() 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() @@ -984,139 +910,6 @@ 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: @@ -2149,9 +1942,6 @@ 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. @@ -2166,8 +1956,7 @@ 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, - governance_subject_ref TEXT + retrieval_experiment_arm TEXT ); INSERT INTO requests_new ( @@ -2179,8 +1968,7 @@ def _migrate_request_session_id_required(self) -> None: session_id, evaluation_only, retrieval_experiment_id, - retrieval_experiment_arm, - governance_subject_ref + retrieval_experiment_arm ) SELECT request_id, @@ -2195,8 +1983,7 @@ def _migrate_request_session_id_required(self) -> None: END, {evaluation_only_expr}, {retrieval_experiment_id_expr}, - {retrieval_experiment_arm_expr}, - {governance_subject_ref_expr} + {retrieval_experiment_arm_expr} FROM requests; DROP TABLE requests; ALTER TABLE requests_new RENAME TO requests; @@ -2205,8 +1992,6 @@ 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() @@ -2474,7 +2259,7 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: Returns: dict[str, int]: Per-entity counts with keys ``interactions``, - ``session_outcomes``, ``user_playbooks``, ``profiles``, ``requests``, + ``user_playbooks``, ``profiles``, ``requests``, ``purged_profiles``, and ``purged_user_playbooks``. ``profiles`` and ``user_playbooks`` reflect hard-deleted counts; purged rows are counted separately. @@ -2501,6 +2286,7 @@ 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] = { @@ -2550,9 +2336,9 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: interactions_cur = self.conn.execute( "DELETE FROM interactions WHERE user_id = ?", (user_id,) ) - session_outcomes_cur = self.conn.execute( - "DELETE FROM session_outcomes WHERE user_id = ?", - (user_id,), + self.conn.execute( + "DELETE FROM session_outcomes WHERE governance_subject_ref = ?", + (subject_ref,), ) requests_cur = self.conn.execute( "DELETE FROM requests WHERE user_id = ?", (user_id,) @@ -2589,7 +2375,6 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: self.purge_content(entity_type="user_playbook", entity_id=str(upid)) return { - "session_outcomes": session_outcomes_cur.rowcount, "interactions": interactions_cur.rowcount, "user_playbooks": upb_deleted_count, "profiles": profile_deleted_count, @@ -2681,18 +2466,14 @@ 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', 'unknown')), + 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) @@ -3071,14 +2852,6 @@ 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 922691fc7..8689b0787 100644 --- a/reflexio/server/services/storage/sqlite_storage/_governance.py +++ b/reflexio/server/services/storage/sqlite_storage/_governance.py @@ -72,16 +72,12 @@ 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 @@ -113,29 +109,6 @@ 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: @@ -431,9 +404,10 @@ 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 = ?", - (user_id,), + "SELECT COUNT(*) AS cnt FROM session_outcomes WHERE governance_subject_ref = ?", + (subject_ref,), ).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 d0657f4eb..c3f95d4a0 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -3,7 +3,6 @@ import json import sqlite3 from typing import Any -from uuid import uuid4 from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, @@ -12,51 +11,12 @@ 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 ( - _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 +from ._base import SQLiteStorageBase, _iso_to_epoch class SessionOutcomeStoreMixin: @@ -107,94 +67,15 @@ def record_session_outcome( try: self.conn.execute("BEGIN IMMEDIATE") existing = self.conn.execute( - "SELECT * FROM session_outcomes WHERE session_id = ?", + "SELECT user_id, source 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 @@ -241,44 +122,31 @@ 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 - (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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (user_id, session_id, outcome, occurred_at, source, label, value, + metadata, governance_subject_ref, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( - outcome_id, - 1, user_id, request.session_id, - str(request.outcome), + request.outcome.value, request.occurred_at, source, request.label, request.value, - self._metadata_json(request), - contract_digest, - snapshot_digest, + json.dumps( + request.metadata, sort_keys=True, separators=(",", ":") + ) + if request.metadata is not None + else None, subject_ref, created_at, ), ) self.conn.commit() return SessionOutcomeWriteResult( - 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, + recorded=True, user_id=user_id, source=source ) except Exception: self.conn.rollback() @@ -311,17 +179,13 @@ def get_session_outcomes( params.append(request.end_time) where = f" WHERE {' AND '.join(clauses)}" if clauses else "" rows = self.conn.execute( - 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} + f"""SELECT user_id, session_id, outcome, occurred_at, source, label, value, + metadata, 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"], @@ -330,8 +194,6 @@ 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 @@ -339,25 +201,13 @@ 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 = ?", - (user_id,), + "DELETE FROM session_outcomes WHERE governance_subject_ref = ?", + (subject_ref,), ) self.conn.commit() 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 4ecc0fc05..3858b29b3 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,7 +15,6 @@ from __future__ import annotations -import json import sqlite3 from datetime import UTC, datetime, timedelta from typing import Any @@ -30,13 +29,6 @@ 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.""" @@ -45,7 +37,6 @@ 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: @@ -222,87 +213,6 @@ 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], - ) -> bool: - 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" - ) - 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), - ).rowcount - == 1 - ) - 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: - 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( self, diff --git a/reflexio/server/services/storage/sqlite_storage/base/_deletion.py b/reflexio/server/services/storage/sqlite_storage/base/_deletion.py index 3f3394170..a7227d222 100644 --- a/reflexio/server/services/storage/sqlite_storage/base/_deletion.py +++ b/reflexio/server/services/storage/sqlite_storage/base/_deletion.py @@ -50,21 +50,16 @@ 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) - predicates: list[str] = [] + where_sql = "" params: list[Any] = [] if statuses: placeholders = ", ".join("?" for _ in statuses) - predicates.append(f"status IN ({placeholders})") + where_sql = f"WHERE 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/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py index 4333632a3..9ed6f5b9b 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py @@ -39,7 +39,6 @@ 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, @@ -89,11 +88,6 @@ 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, @@ -149,9 +143,10 @@ 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 = ?", - (user_id,), + "DELETE FROM session_outcomes WHERE governance_subject_ref = ?", + (subject_ref,), ) interaction_ids = [ int(row["interaction_id"]) @@ -274,7 +269,11 @@ def _clear_user_data_for_governance_locked( ) return { - "session_outcomes": session_outcomes_cur.rowcount, + **( + {"session_outcomes": session_outcomes_cur.rowcount} + if session_outcomes_cur.rowcount + else {} + ), "interactions": interactions_cur.rowcount, "user_playbooks": len(delete_upb_ids), "profiles": len(delete_profile_ids), @@ -350,11 +349,7 @@ def _delete_evaluation_operation_states_locked( return deleted def apply_governance_user_data_delete( - self, - purge_id: str, - user_id: str, - *, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, user_id: str ) -> dict[str, int]: purge_id = _validate_governance_purge_id("purge_id", purge_id) name_map = { @@ -377,11 +372,7 @@ 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.conn.execute("BEGIN") self._validate_prepared_delete_target_matrix_locked(purge_id) self._validate_hide_for_rebuild_targets_locked(purge_id) expected_user_playbook_ids = ( @@ -416,12 +407,7 @@ def apply_governance_user_data_delete( return counts def complete_purge_operation_with_audit( - self, - purge_id: str, - audit_event: AuditEvent, - *, - authoritative_user_id: str, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, audit_event: AuditEvent ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) if audit_event.org_id != self.org_id: @@ -437,9 +423,6 @@ 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), @@ -472,11 +455,6 @@ 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' @@ -556,9 +534,7 @@ def complete_purge_operation_with_audit( error_code = NULL, error_detail = NULL, updated_at = ?, - completed_at = ?, - execution_claim_owner = NULL, - execution_claim_expires_at = NULL + completed_at = ? 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 2087ee2c6..386cbdd95 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py @@ -18,8 +18,6 @@ from __future__ import annotations -import hashlib -import hmac import sqlite3 import threading from collections.abc import Callable @@ -29,11 +27,6 @@ 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, @@ -55,12 +48,7 @@ _validate_governance_target_ref, ) -from .._governance import ( - _json_dumps, - _json_loads, - _row_to_purge_operation, - _row_to_purge_target, -) +from .._governance import _json_dumps, _row_to_purge_operation, _row_to_purge_target if TYPE_CHECKING: from .._governance import _SQLiteGovernanceDeps @@ -77,114 +65,6 @@ 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, @@ -285,7 +165,6 @@ 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", @@ -308,254 +187,55 @@ 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: - self.conn.execute("BEGIN IMMEDIATE") - existing = self.conn.execute( - """SELECT * FROM purge_operations - WHERE org_id = ? AND idempotency_key = ?""", - (self.org_id, validated_idempotency_key), - ).fetchone() - if existing is not None: - existing_operation = _row_to_purge_operation(existing) - expected_identity = { - "purge_id": validated_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(existing_operation, field_name) != expected_value: - raise ValueError( - "Existing purge operation for idempotency_key has " - f"mismatched {field_name}" - ) - 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: + existing = self.conn.execute( + """SELECT * FROM purge_operations + WHERE org_id = ? AND idempotency_key = ?""", + (self.org_id, validated_idempotency_key), + ).fetchone() + if existing is not None: + existing_operation = _row_to_purge_operation(existing) + expected_identity = { + "purge_id": validated_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(existing_operation, field_name) != expected_value: raise ValueError( - "Existing purge operation has mismatched authoritative user identity" + "Existing purge operation for idempotency_key has " + f"mismatched {field_name}" ) - 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, authoritative_user_digest, - status, created_at, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)""", - ( - validated_purge_id, - self.org_id, - operation_type, - scope_type, - subject_ref, - request_ref, - validated_idempotency_key, - authoritative_user_digest, - now, - now, - ), - ) - self.conn.commit() - except Exception: - self.conn.rollback() - raise + 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', ?, ?)""", + ( + validated_purge_id, + self.org_id, + operation_type, + scope_type, + subject_ref, + request_ref, + validated_idempotency_key, + now, + now, + ), + ) + self.conn.commit() 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, @@ -578,25 +258,17 @@ def record_purge_target( allowed=_ALLOWED_PURGE_TARGET_STATUSES, ) 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, - target_ref=target_ref, - phase=phase, - status=status, - detail=detail, - deleted_count=deleted_count, - error_detail=error_detail, - ) - self.conn.commit() - except Exception: - self.conn.rollback() - raise + self._record_purge_target_locked( + purge_id=purge_id, + target_name=target_name, + target_ref=target_ref, + phase=phase, + status=status, + detail=detail, + deleted_count=deleted_count, + error_detail=error_detail, + ) + self.conn.commit() def list_purge_targets( self, purge_id: str, phase: str | None = None @@ -626,29 +298,14 @@ 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: + if self.purge_targets_prepared(purge_id): + return try: self.conn.execute("BEGIN IMMEDIATE") - 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) if owned_user_playbook_ids is not None @@ -676,7 +333,6 @@ 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, @@ -688,51 +344,36 @@ def prepare_governance_erase_targets( raise def fail_purge_operation( - self, - purge_id: str, - error_code: str, - error_detail: str, - *, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, error_code: str, error_detail: str ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) validated_error_code = _validate_governance_error_code(error_code) validated_error_detail = _validate_governance_error_detail(error_detail) now = _epoch_now() 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 = ?, - execution_claim_owner = NULL, - execution_claim_expires_at = NULL - WHERE purge_id = ? AND org_id = ? AND status != 'complete'""", - ( - validated_error_code, - validated_error_detail, - now, - now, - purge_id, - self.org_id, - ), - ) - if cur.rowcount == 0: - existing = self.conn.execute( - "SELECT status FROM purge_operations WHERE purge_id = ? AND org_id = ?", - (purge_id, self.org_id), - ).fetchone() - if existing is not None and str(existing["status"]) == "complete": - raise ValueError("Purge operation is already complete") - raise ValueError(f"Purge operation {purge_id!r} not found") - self.conn.commit() - except Exception: - self.conn.rollback() - raise + cur = self.conn.execute( + """UPDATE purge_operations + SET status = 'failed', error_code = ?, error_detail = ?, + updated_at = ?, completed_at = ? + WHERE purge_id = ? AND org_id = ? AND status != 'complete'""", + ( + validated_error_code, + validated_error_detail, + now, + now, + purge_id, + self.org_id, + ), + ) + if cur.rowcount == 0: + existing = self.conn.execute( + "SELECT status FROM purge_operations WHERE purge_id = ? AND org_id = ?", + (purge_id, self.org_id), + ).fetchone() + if existing is not None and str(existing["status"]) == "complete": + raise ValueError("Purge operation is already complete") + raise ValueError(f"Purge operation {purge_id!r} not found") + self.conn.commit() return self.get_purge_operation(purge_id) def get_purge_operation(self, purge_id: str) -> PurgeOperation: 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 da3506a99..3390533f0 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py @@ -38,7 +38,6 @@ 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, @@ -54,7 +53,6 @@ ) from .._governance import ( - _json_loads, _row_to_audit_event, _row_to_purge_operation, _row_to_subject_write_barrier, @@ -77,10 +75,6 @@ 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, @@ -164,59 +158,7 @@ def _legacy_user_id_rows_remain_locked( return True return False - 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: + def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: legacy_request_ids = self._legacy_request_ids_for_subject_locked(subject_ref) for table in ( "requests", @@ -225,6 +167,7 @@ def _same_subject_rows_remain_locked( "user_playbooks", "agent_success_evaluation_result", "retrieved_learning_evaluation", + "session_outcomes", ): row = self.conn.execute( f"""SELECT 1 FROM {table} @@ -236,10 +179,6 @@ def _same_subject_rows_remain_locked( 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, @@ -272,11 +211,7 @@ def _same_subject_rows_remain_locked( ) def begin_subject_erasure_barrier( - self, - subject_ref: str, - purge_id: str, - *, - execution_claim: PurgeExecutionClaim, + self, subject_ref: str, purge_id: str ) -> SubjectWriteBarrier: _validate_governance_prefixed_ref( "subject_ref", subject_ref, prefix="subref_v1_" @@ -286,9 +221,6 @@ 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 = ?""", @@ -361,12 +293,7 @@ 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, - *, - authoritative_user_id: str, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, audit_event: AuditEvent ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) if audit_event.org_id != self.org_id: @@ -382,9 +309,6 @@ 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), @@ -410,14 +334,7 @@ def complete_subject_erasure_barrier_after_empty_check( 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, - ) - if self._same_subject_rows_remain_locked( - audit_event.subject_ref or "", authoritative_user_id - ): + if self._same_subject_rows_remain_locked(audit_event.subject_ref or ""): raise ValueError("same-subject rows remain") delete_rows = self.conn.execute( """SELECT target_name, status FROM purge_operation_targets @@ -507,9 +424,7 @@ def complete_subject_erasure_barrier_after_empty_check( error_code = NULL, error_detail = NULL, updated_at = ?, - completed_at = ?, - execution_claim_owner = NULL, - execution_claim_expires_at = NULL + completed_at = ? WHERE purge_id = ? AND org_id = ?""", (now, now, purge_id, self.org_id), ) @@ -525,8 +440,6 @@ 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_" @@ -538,9 +451,6 @@ 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', @@ -570,9 +480,7 @@ def fail_subject_erasure_barrier( self.conn.execute( """UPDATE purge_operations SET status = 'failed', error_code = ?, error_detail = ?, - updated_at = ?, completed_at = ?, - execution_claim_owner = NULL, - execution_claim_expires_at = NULL + updated_at = ?, completed_at = ? 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 d45f5ff75..1dae46dd3 100644 --- a/reflexio/server/services/storage/storage_base/__init__.py +++ b/reflexio/server/services/storage/storage_base/__init__.py @@ -147,8 +147,8 @@ 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 session outcomes, interactions, user playbooks, - profiles, and requests. Intentionally does NOT touch ``agent_playbooks`` — those + Removes the user's 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 by paired protocols (e.g. SWE-bench) that share a single backend across @@ -177,12 +177,12 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: Returns: dict[str, int]: Per-entity counts with keys ``interactions``, - ``session_outcomes``, ``user_playbooks``, ``profiles``, ``requests``, + ``user_playbooks``, ``profiles``, ``requests``, ``purged_profiles``, and ``purged_user_playbooks``. ``profiles`` and ``user_playbooks`` reflect hard-deleted counts; purged rows are counted separately. """ - session_outcome_counts = self.clear_session_outcomes_for_user(user_id) + self.clear_session_outcomes_for_user(user_id) interaction_count = len(self.get_user_interaction(user_id)) # All statuses a user's row can have — including tombstones (SUPERSEDED, @@ -264,7 +264,6 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: self.purge_content(entity_type="user_playbook", entity_id=upid) return { - "session_outcomes": session_outcome_counts.get("session_outcomes", 0), "interactions": interaction_count, "user_playbooks": deleted_user_playbooks, "profiles": deleted_profiles, diff --git a/reflexio/server/services/storage/storage_base/_session_outcomes.py b/reflexio/server/services/storage/storage_base/_session_outcomes.py index 14695d9a3..d417e6065 100644 --- a/reflexio/server/services/storage/storage_base/_session_outcomes.py +++ b/reflexio/server/services/storage/storage_base/_session_outcomes.py @@ -18,24 +18,6 @@ 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 3acd3c9fd..b8bba0371 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,25 +24,6 @@ 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], - ) -> 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( 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 4650f726a..cc4016f72 100644 --- a/reflexio/server/services/storage/storage_base/governance/_erase_execution.py +++ b/reflexio/server/services/storage/storage_base/governance/_erase_execution.py @@ -6,7 +6,6 @@ AuditEvent, PurgeOperation, ) -from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class GovernanceEraseExecutionMixin(ABC): @@ -20,21 +19,12 @@ class GovernanceEraseExecutionMixin(ABC): @abstractmethod def apply_governance_user_data_delete( - self, - purge_id: str, - user_id: str, - *, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, user_id: str ) -> dict[str, int]: raise NotImplementedError @abstractmethod def complete_purge_operation_with_audit( - self, - purge_id: str, - audit_event: AuditEvent, - *, - authoritative_user_id: str, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, audit_event: AuditEvent ) -> 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 f3f1b4620..ef63fdc2c 100644 --- a/reflexio/server/services/storage/storage_base/governance/_purge.py +++ b/reflexio/server/services/storage/storage_base/governance/_purge.py @@ -7,7 +7,6 @@ PurgeOperation, PurgeOperationTarget, ) -from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class PurgeOperationStoreMixin(ABC): @@ -28,39 +27,9 @@ 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, @@ -68,8 +37,6 @@ 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, @@ -92,20 +59,13 @@ 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, - *, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, error_code: str, error_detail: str ) -> 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 2de2f5a26..7752ec5c9 100644 --- a/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py +++ b/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py @@ -7,7 +7,6 @@ PurgeOperation, SubjectWriteBarrier, ) -from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class SubjectBarrierMixin(ABC): @@ -21,11 +20,7 @@ class SubjectBarrierMixin(ABC): @abstractmethod def begin_subject_erasure_barrier( - self, - subject_ref: str, - purge_id: str, - *, - execution_claim: PurgeExecutionClaim, + self, subject_ref: str, purge_id: str ) -> SubjectWriteBarrier: raise NotImplementedError @@ -35,12 +30,7 @@ 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, - *, - authoritative_user_id: str, - execution_claim: PurgeExecutionClaim, + self, purge_id: str, audit_event: AuditEvent ) -> PurgeOperation: raise NotImplementedError @@ -51,8 +41,6 @@ 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 9b72405e1..662a7436b 100644 --- a/tests/client/test_session_outcomes_client.py +++ b/tests/client/test_session_outcomes_client.py @@ -3,9 +3,6 @@ import inspect from typing import Any -import pytest -from pydantic import ValidationError - from reflexio import ReflexioClient @@ -56,23 +53,3 @@ 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 deleted file mode 100644 index 41e55df8f..000000000 --- a/tests/models/test_session_outcome_identity.py +++ /dev/null @@ -1,336 +0,0 @@ -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 1b3d408a9..5471dd743 100644 --- a/tests/server/api_endpoints/test_session_outcomes_integration.py +++ b/tests/server/api_endpoints/test_session_outcomes_integration.py @@ -54,18 +54,13 @@ def test_source_is_derived_from_tiebroken_first_request( ) assert response.status_code == 200 - body = response.json() - assert { + assert response.json() == { "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 @@ -110,169 +105,6 @@ 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 deleted file mode 100644 index 94cef8c68..000000000 --- a/tests/server/routes/test_search_exposure_boundary.py +++ /dev/null @@ -1,153 +0,0 @@ -"""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 51355e51f..c362b0e45 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -2,45 +2,24 @@ 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, - UserPlaybook, - UserProfile, -) +from reflexio.models.api_schema.service_schemas import Interaction, Request from reflexio.models.config_schema import ( Config, PendingToolCallConfig, - PlaybookConfig, ProfileExtractorConfig, 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, _run_uses_strict_playbook_evidence, ) -from reflexio.server.services.playbook.components.consolidator import ( - PlaybookConsolidationOutput, - UnifyDecision, -) -from reflexio.server.services.playbook.service import ( - PlaybookGenerationService, - PlaybookGenerationServiceConfig, -) -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, @@ -53,10 +32,6 @@ build_scope_hash, human_feedback_scope, ) -from reflexio.server.usage_metrics import ( - UsageEvent, - configure_usage_event_recorder, -) @pytest.fixture @@ -90,131 +65,6 @@ 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"), [ @@ -446,7 +296,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_with_outcome", + "ProfileGenerationService._finalize_extracted_items", side_effect=RuntimeError("storage write failed"), ), ): @@ -473,8 +323,8 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( ), patch( "reflexio.server.services.profile.service." - "ProfileGenerationService._finalize_extracted_items_with_outcome", - return_value=FinalizationResult([], won_receipt=False), + "ProfileGenerationService._finalize_extracted_items", + return_value=None, ) as finalize, ): resumed = worker.drain(max_runs=1) @@ -512,854 +362,6 @@ 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": 1, - "profile_distinct_keys": 1, - "playbook_persisted": 1, - "playbook_lineage_events": 1, - "playbook_events": 1, - "playbook_distinct_keys": 1, - } - assert profile_receipt_ids == profile_survivor_ids - 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 - assert playbook_keys == [f"learn:user_playbook:{playbook_survivor_ids[0]}"] - - -@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}", - ) - 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, - ), - 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_empty_playbook", - agent_version="v1", - source="api", - ), - 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" - ) - == [] - ) - assert [ - event for event in events if event.event_name == "learnings_generated" - ] == [] - finally: - configure_usage_event_recorder(None) - - -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) - - 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"] - - -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) - ] - 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: - 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( - "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 == [] - 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.id, entity_type="profile" - ) - 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 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( 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 576419903..d8ee06b06 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -1,7 +1,5 @@ from __future__ import annotations -import json -import threading from collections.abc import Generator from datetime import UTC, datetime from pathlib import Path @@ -20,14 +18,12 @@ 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 @@ -121,27 +117,6 @@ 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, @@ -331,7 +306,6 @@ 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") == [] @@ -463,55 +437,6 @@ 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: @@ -592,54 +517,6 @@ 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: @@ -752,21 +629,10 @@ 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", ) - 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, - ) + storage.begin_subject_erasure_barrier(subject_ref, first_purge.purge_id) service = GovernanceService( storage=storage, org_id=storage.org_id, @@ -855,7 +721,6 @@ def erase_subject( storage: SQLiteStorage, subject_ref: str, purge_id: str, - execution_claim: PurgeExecutionClaim, ) -> None: del subject_ref self.calls += 1 @@ -875,7 +740,6 @@ def erase_subject( status="complete", detail={"count": 2}, deleted_count=2, - execution_claim=execution_claim, ) lifecycle = RetrySafeLifecycle() @@ -921,571 +785,6 @@ 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 9632af24a..1e60813ea 100644 --- a/tests/server/services/governance/test_subject_write_barrier_sqlite.py +++ b/tests/server/services/governance/test_subject_write_barrier_sqlite.py @@ -1,7 +1,5 @@ from __future__ import annotations -import hashlib -import hmac from datetime import UTC, datetime from pathlib import Path @@ -16,13 +14,12 @@ UserPlaybook, UserProfile, ) -from reflexio.models.api_schema.domain.governance import AuditEvent, SubjectWriteBarrier +from reflexio.models.api_schema.domain.governance import AuditEvent 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, ) @@ -39,78 +36,14 @@ 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", - execution_claim=claim, - detail={ - "prepared": True, - "authoritative_user_digest": _authoritative_user_digest(storage, purge_id), - }, + detail={"prepared": True}, ) # Single source of truth — a stale local copy of the canonical tuple is # exactly how this suite went red when new delete targets landed. @@ -121,7 +54,6 @@ def _mark_all_completion_targets(storage: SQLiteStorage, purge_id: str) -> None: phase="delete", status="complete", target_ref="all", - execution_claim=claim, detail={"count": 0}, ) @@ -145,190 +77,8 @@ 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, @@ -341,14 +91,11 @@ 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 = _begin_claimed_subject_erasure_barrier( - storage, subject_ref, purge.purge_id - ) + barrier = storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) assert barrier.status == "erasing" with pytest.raises(SubjectWriteBarrierError): @@ -398,7 +145,6 @@ 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", ) @@ -428,7 +174,7 @@ def test_barrier_blocks_playbook_eval_and_source_window_writes( ] )[0] - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.save_user_playbooks( @@ -482,21 +228,16 @@ 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"): - _begin_claimed_subject_erasure_barrier(storage, bob_subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(bob_subject_ref, purge.purge_id) with pytest.raises(ValueError, match="not found"): storage.begin_subject_erasure_barrier( - alice_subject_ref, - "purge_barrier_missing", - execution_claim=_typed_test_claim_for_unvalidated_purge_id( - "purge_barrier_missing" - ), + alice_subject_ref, "purge_barrier_missing" ) @@ -511,7 +252,6 @@ 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", ) @@ -520,12 +260,11 @@ 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", ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, first_purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, first_purge.purge_id) with pytest.raises(ValueError, match="matching barrier"): storage.fail_subject_erasure_barrier( @@ -533,7 +272,6 @@ 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) @@ -553,7 +291,7 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( ) -def test_begin_subject_erasure_barrier_rejects_inactive_claim_after_completion( +def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -565,11 +303,10 @@ def test_begin_subject_erasure_barrier_rejects_inactive_claim_after_completion( 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, ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -577,14 +314,7 @@ def test_begin_subject_erasure_barrier_rejects_inactive_claim_after_completion( request_ref=request_ref, ) - 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 + barrier = storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) stored_barrier = storage.get_subject_write_barrier(subject_ref) stored_purge = storage.get_purge_operation(purge.purge_id) @@ -595,7 +325,7 @@ def test_begin_subject_erasure_barrier_rejects_inactive_claim_after_completion( assert stored_purge.status == "complete" -def test_fail_subject_erasure_barrier_rejects_inactive_claim_after_completion( +def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -607,11 +337,10 @@ def test_fail_subject_erasure_barrier_rejects_inactive_claim_after_completion( 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, ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -619,13 +348,12 @@ def test_fail_subject_erasure_barrier_rejects_inactive_claim_after_completion( request_ref=request_ref, ) - with pytest.raises(ValueError, match="purge execution claim"): + with pytest.raises(ValueError, match="matching barrier"): 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) @@ -638,7 +366,7 @@ def test_fail_subject_erasure_barrier_rejects_inactive_claim_after_completion( assert purge_after_failure.error_code is None -def test_fail_purge_operation_rejects_inactive_claim_after_completion( +def test_fail_purge_operation_rejects_terminal_complete_state( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -650,11 +378,10 @@ def test_fail_purge_operation_rejects_inactive_claim_after_completion( 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, ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -662,12 +389,11 @@ def test_fail_purge_operation_rejects_inactive_claim_after_completion( request_ref=request_ref, ) - with pytest.raises(ValueError, match="purge execution claim"): + with pytest.raises(ValueError, match="already complete"): 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) @@ -709,11 +435,10 @@ 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", ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) assert storage.purge_content(entity_type="profile", entity_id=profile.profile_id) assert storage.purge_content( @@ -775,7 +500,6 @@ 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", ) @@ -789,20 +513,14 @@ def test_guarded_completion_requires_empty_subject_rows( created_at=_now(), ) ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) storage.record_purge_target( purge.purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - execution_claim=_claim_purge(storage, purge.purge_id), - detail={ - "prepared": True, - "authoritative_user_digest": _authoritative_user_digest( - storage, purge.purge_id - ), - }, + detail={"prepared": True}, ) with pytest.raises(ValueError, match="same-subject rows remain"): @@ -817,8 +535,6 @@ 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), ) @@ -833,7 +549,6 @@ 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", ) @@ -855,20 +570,14 @@ def test_guarded_completion_requires_empty_legacy_null_subject_rows( ) storage.conn.commit() - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) storage.record_purge_target( purge.purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - execution_claim=_claim_purge(storage, purge.purge_id), - detail={ - "prepared": True, - "authoritative_user_digest": _authoritative_user_digest( - storage, purge.purge_id - ), - }, + detail={"prepared": True}, ) with pytest.raises(ValueError, match="same-subject rows remain"): @@ -883,8 +592,6 @@ 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), ) @@ -899,7 +606,6 @@ 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", ) @@ -917,8 +623,6 @@ 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), ) @@ -933,17 +637,15 @@ 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", ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(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) @@ -959,8 +661,6 @@ 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), ) @@ -983,11 +683,10 @@ 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", ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.update_user_profile_tags("alice", profile.profile_id, ["blocked"]) @@ -1024,11 +723,10 @@ 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", ) - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.archive_user_playbook_by_id("alice", playbook.user_playbook_id) @@ -1056,7 +754,6 @@ 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, ) @@ -1065,9 +762,7 @@ def test_assert_subject_writable_blocks_only_barriered_subject( storage.assert_subject_writable(barriered_subject_ref) storage.assert_subject_writable(other_subject_ref) - _begin_claimed_subject_erasure_barrier( - storage, barriered_subject_ref, purge.purge_id - ) + storage.begin_subject_erasure_barrier(barriered_subject_ref, purge.purge_id) # The 'erasing' barrier blocks only its own subject. with pytest.raises(SubjectWriteBarrierError, match="blocked by erasure barrier"): @@ -1101,7 +796,6 @@ 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", ) @@ -1138,7 +832,7 @@ def test_source_window_write_blocks_legacy_null_subject_ref_user_playbook( ] )[0] - _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) + storage.begin_subject_erasure_barrier(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 cb369c7d4..78300ee39 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,8 +5,6 @@ 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 ( @@ -99,213 +97,6 @@ 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)) - - inserted = storage.save_agent_run_finalization_receipt( - run_id="run_empty", entity_type="profile", learning_ids=[] - ) - 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" - ) - == [] - ) - 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_reports_existing_immutable_value(storage): - storage.create_agent_run(_agent_run("run_immutable", AgentRunStatus.FINALIZING)) - inserted = storage.save_agent_run_finalization_receipt( - run_id="run_immutable", - entity_type="profile", - learning_ids=["profile-1"], - ) - - 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"] - - -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 17d1c385d..4ff0df379 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,26 +117,11 @@ 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", - authoritative_user_id=user_id, - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, ) - 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, - ) + storage.prepare_governance_erase_targets(purge_id, user_id) + return storage.apply_governance_user_data_delete(purge_id, user_id) 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 0e0496f75..89783fe86 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -1,13 +1,7 @@ from __future__ import annotations -import ast -import hashlib -import inspect import json import sqlite3 -import threading -from collections.abc import Callable -from pathlib import Path from typing import Any, Literal, cast from unittest.mock import patch @@ -27,8 +21,6 @@ ) 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, ) @@ -42,50 +34,17 @@ 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 = governance_subject_ref("org1", "alice", "test-governance-secret") -OTHER_SUBJECT_REF = governance_subject_ref("org1", "bob", "test-governance-secret") +SUBJECT_REF = "subref_v1_" + "a" * 32 +OTHER_SUBJECT_REF = "subref_v1_" + "c" * 32 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 @@ -106,155 +65,29 @@ def _make_storage(org_id: str) -> SQLiteStorage: yield _make_storage -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, +def _begin_purge(storage: SQLiteStorage, purge_id: str) -> str: + purge = storage.begin_purge_operation( 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 _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: +def _add_complete_delete_target_matrix(storage: SQLiteStorage, purge_id: str) -> None: for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -262,34 +95,13 @@ def _add_complete_delete_target_matrix( target_ref="all", phase="delete", status="complete", - execution_claim=execution_claim, ) -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, - ) +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) return purge_id @@ -298,13 +110,12 @@ 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, @@ -532,7 +343,6 @@ 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 [ @@ -674,8 +484,7 @@ def test_list_audit_events_rejects_cross_org_override(storage_factory): def test_purge_targets_require_snapshot_marker(storage): - purge = _begin_test_purge_operation( - storage, + purge = storage.begin_purge_operation( purge_id="purge_snapshot_marker", idempotency_key="idem_snapshot_marker", operation_type="user_erasure", @@ -690,15 +499,10 @@ 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, - execution_claim=_claim_purge(storage, purge.purge_id), - ) + storage.begin_subject_erasure_barrier(SUBJECT_REF, purge.purge_id) with pytest.raises(ValueError, match="target snapshot"): storage.complete_purge_operation_with_audit( purge.purge_id, @@ -710,293 +514,25 @@ 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"] - 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( + same = 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( @@ -1009,8 +545,6 @@ 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) @@ -1052,8 +586,6 @@ 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" @@ -1083,12 +615,7 @@ 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, - authoritative_user_id="alice", - execution_claim=_claim_purge(storage, purge_id), - ) + storage.complete_purge_operation_with_audit(purge_id, event) assert storage.get_purge_operation(purge_id).status == "running" assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] @@ -1111,8 +638,7 @@ 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 ): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_begin_retry", idempotency_key="idem_begin_retry", operation_type="user_erasure", @@ -1122,8 +648,7 @@ def test_begin_purge_operation_rejects_mismatched_idempotent_retry( ) with pytest.raises(ValueError, match=match): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=retry_kwargs.get("purge_id", "purge_begin_retry"), idempotency_key="idem_begin_retry", operation_type=retry_kwargs.get("operation_type", "user_erasure"), @@ -1140,8 +665,7 @@ 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"): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_numeric_idem", idempotency_key="12345", operation_type="user_erasure", @@ -1157,8 +681,7 @@ def test_begin_purge_operation_rejects_numeric_idempotency_key(storage): def test_begin_purge_operation_accepts_code_shaped_idempotency_key_with_content( storage, ): - purge = _begin_test_purge_operation( - storage, + purge = storage.begin_purge_operation( purge_id="purge_content_retry", idempotency_key="content_purge_retry_1", operation_type="user_erasure", @@ -1173,8 +696,7 @@ 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"): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key=f"idem_{purge_id}", operation_type="user_erasure", @@ -1232,12 +754,7 @@ 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, - authoritative_user_id="alice", - execution_claim=_claim_purge(storage, purge_id), - ) + storage.complete_purge_operation_with_audit(purge_id, event) assert storage.get_purge_operation(purge_id).status == "running" assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] @@ -1266,10 +783,7 @@ 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), - authoritative_user_id="alice", - execution_claim=_claim_purge(storage, purge_id), + purge_id, _erase_event(purge_id=purge_id) ) assert storage.get_purge_operation(purge_id).status == "running" @@ -1326,10 +840,7 @@ 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), - authoritative_user_id="alice", - execution_claim=_claim_purge(storage, purge_id), + purge_id, _erase_event(purge_id=purge_id) ) assert storage.get_purge_operation(purge_id).status == "running" @@ -1360,18 +871,12 @@ 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, - execution_claim=_claim_purge(storage, purge_id), - ) + storage.begin_subject_erasure_barrier(SUBJECT_REF, 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" @@ -1385,7 +890,6 @@ 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 @@ -1393,641 +897,25 @@ 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): - user_id = "user_123@example.com" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_detail", idempotency_key="idem_purge_detail", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, - authoritative_user_id=user_id, ) storage.prepare_governance_erase_targets( purge_id="purge_detail", - user_id=user_id, + user_id="user_123@example.com", owned_user_playbook_ids={7}, - execution_claim=_claim_purge(storage, "purge_detail"), ) snapshot = next( @@ -2037,33 +925,23 @@ def test_prepare_governance_erase_targets_sanitizes_snapshot_detail(storage): ) if target.target_name == "target_snapshot" ) - 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], - } + assert snapshot.detail == {"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) - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_snapshot_drift", idempotency_key="idem_purge_snapshot_drift", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, 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 ( @@ -2080,11 +958,7 @@ 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, - execution_claim=_claim_purge(storage, "purge_snapshot_drift"), - ) + storage.apply_governance_user_data_delete("purge_snapshot_drift", user_id) remaining_ids = { int(row["user_playbook_id"]) @@ -2099,16 +973,13 @@ 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, ): - user_id = "user-rebuild-windows" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_rebuild_windows", idempotency_key="idem_purge_rebuild_windows", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, - authoritative_user_id=user_id, ) agent_playbook_id = _seed_agent_playbook( storage, @@ -2122,9 +993,8 @@ 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_id, + user_id="user-rebuild-windows", owned_user_playbook_ids={7}, - execution_claim=_claim_purge(storage, "purge_rebuild_windows"), ) assert ( @@ -2146,22 +1016,19 @@ def test_prepare_governance_erase_targets_records_full_delete_matrix_counts(stor session_id="session_seed", evaluation_name="governance_prepare_counts", ) - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_prepare_counts", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, 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 = { @@ -2185,7 +1052,6 @@ def test_prepare_governance_erase_targets_records_full_delete_matrix_counts(stor counts = storage.clear_user_data(user_id) assert counts == { - "session_outcomes": 0, "interactions": 1, "user_playbooks": 1, "profiles": 1, @@ -2199,15 +1065,13 @@ def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_an storage, ): purge_id = "purge_hide_rebuild" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( 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, @@ -2263,8 +1127,7 @@ 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" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_complete", operation_type="user_erasure", @@ -2280,14 +1143,12 @@ 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]}, @@ -2305,7 +1166,6 @@ 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": [ @@ -2367,8 +1227,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_ad_hoc_rebuild_without_ storage, ): purge_id = "purge_rebuild_requires_target" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_requires_target", operation_type="user_erasure", @@ -2431,8 +1290,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_rebuild_before_hide_pha storage, ): purge_id = "purge_rebuild_requires_hide" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_requires_hide", operation_type="user_erasure", @@ -2448,14 +1306,12 @@ 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]}, @@ -2529,8 +1385,7 @@ def test_apply_governance_agent_playbook_rebuild_succeeds_after_prepare_and_hide storage, ): purge_id = "purge_rebuild_prepare_hide" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_prepare_hide", operation_type="user_erasure", @@ -2596,8 +1451,7 @@ def test_apply_governance_agent_playbook_rebuild_does_not_complete_target_when_s storage, monkeypatch ): purge_id = "purge_rebuild_search_refresh_failure" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_search_refresh_failure", operation_type="user_erasure", @@ -2613,14 +1467,12 @@ 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]}, @@ -2638,7 +1490,6 @@ 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 @@ -2713,8 +1564,7 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when storage, ): purge_id = "purge_rebuild_remove_orphan" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_remove_orphan", operation_type="user_erasure", @@ -2729,14 +1579,12 @@ 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]}, @@ -2751,7 +1599,6 @@ 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( @@ -2818,8 +1665,7 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta storage, ): purge_id = "purge_rebuild_restore_archived" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_restore_archived", operation_type="user_erasure", @@ -2835,14 +1681,12 @@ 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]}, @@ -2860,7 +1704,6 @@ 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( @@ -2888,8 +1731,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl storage, ): purge_id = "purge_rebuild_second_call" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_second_call", operation_type="user_erasure", @@ -3001,8 +1843,7 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet storage, ): purge_id = "purge_hide_after_complete" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_hide_after_complete", operation_type="user_erasure", @@ -3094,8 +1935,7 @@ def test_hide_governance_agent_playbooks_for_rebuild_does_not_reopen_complete_ta storage, monkeypatch ): purge_id = "purge_hide_stale_prelock" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_hide_stale_prelock", operation_type="user_erasure", @@ -3206,22 +2046,19 @@ 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" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_prepare_idempotent_after_snapshot", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, 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 = [ @@ -3241,7 +2078,6 @@ 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 = [ @@ -3274,9 +2110,8 @@ 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=storage_instance._subject_ref_for_user_id("alice"), + subject_ref=SUBJECT_REF, request_ref=request_ref, - authoritative_user_id="alice", ) storage_org1.record_purge_target( @@ -3285,7 +2120,6 @@ 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( @@ -3294,7 +2128,6 @@ 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( @@ -3303,7 +2136,6 @@ 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, ) @@ -3331,7 +2163,6 @@ 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}, ) @@ -3428,7 +2259,6 @@ 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) @@ -3467,7 +2297,6 @@ def test_record_purge_target_rejects_invalid_deleted_count( phase="delete", status="complete", deleted_count=deleted_count, - execution_claim=_claim_purge(storage, purge_id), ) @@ -3485,7 +2314,6 @@ 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}, ) @@ -3508,7 +2336,6 @@ 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), ) @@ -3567,16 +2394,6 @@ 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" ), @@ -3613,7 +2430,6 @@ 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": [ @@ -3631,7 +2447,6 @@ 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, ) @@ -3683,7 +2498,6 @@ 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 @@ -3697,7 +2511,6 @@ 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 @@ -3710,7 +2523,6 @@ 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" @@ -3729,7 +2541,6 @@ 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" @@ -3744,7 +2555,6 @@ 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 @@ -3767,7 +2577,6 @@ 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 @@ -3778,7 +2587,6 @@ 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 @@ -3917,8 +2725,7 @@ def test_begin_purge_operation_validates_top_level_refs( storage, subject_ref, request_ref, match ): with pytest.raises(ValueError, match=match): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_top_level_refs", idempotency_key="idem_purge_top_level_refs", operation_type="user_erasure", @@ -3943,8 +2750,7 @@ def test_begin_purge_operation_rejects_invalid_enum_values( storage, operation_type, scope_type, match ): with pytest.raises(ValueError, match=match): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_invalid_enum", idempotency_key="idem_purge_invalid_enum", operation_type=operation_type, @@ -3965,8 +2771,7 @@ 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"): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_invalid_id", operation_type="user_erasure", @@ -4007,7 +2812,6 @@ 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"}]}, ) @@ -4043,7 +2847,6 @@ 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]}]}, ) @@ -4065,7 +2868,6 @@ 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]} @@ -4144,7 +2946,6 @@ 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, ) @@ -4172,7 +2973,6 @@ 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 @@ -4183,7 +2983,6 @@ 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), ) @@ -4218,7 +3017,6 @@ 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"): @@ -4232,8 +3030,6 @@ 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"): @@ -4251,9 +3047,6 @@ 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) @@ -4268,11 +3061,7 @@ 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", - authoritative_user_id="user-delete-seed", - ) + purge_id = _begin_purge(storage, "purge_internal_target_name") for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -4280,7 +3069,6 @@ 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}, ) @@ -4303,7 +3091,6 @@ 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") @@ -4313,12 +3100,8 @@ 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) @@ -4328,7 +3111,6 @@ 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( @@ -4337,7 +3119,6 @@ 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, ) @@ -4363,7 +3144,6 @@ 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 @@ -4380,15 +3160,13 @@ def test_apply_governance_user_data_delete_preserves_org_agent_playbooks_without ): purge_id = "purge_delete_requires_hide" user_id = "user-delete-hide-required" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_delete_requires_hide", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, 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( @@ -4412,13 +3190,11 @@ 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 @@ -4446,15 +3222,13 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( """ purge_id = "purge_delete_after_hide" user_id = "user-delete-hide-complete" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_delete_after_hide", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, 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( @@ -4488,18 +3262,15 @@ 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 == { - "session_outcomes": 0, "interactions": 1, "user_playbooks": 1, "profiles": 1, @@ -4576,15 +3347,13 @@ 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" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_delete_atomic", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id(user_id), + subject_ref=SUBJECT_REF, 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) @@ -4602,7 +3371,6 @@ 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) @@ -4658,7 +3426,6 @@ 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 @@ -4730,7 +3497,6 @@ 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) @@ -4743,8 +3509,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so storage, ): purge_id = "purge_rebuild_windows_mismatch" - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_rebuild_windows_mismatch", operation_type="user_erasure", @@ -4766,7 +3531,6 @@ 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]}, @@ -4784,7 +3548,6 @@ 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 @@ -4925,7 +3688,6 @@ def test_record_purge_target_rejects_invalid_enum_values( target_ref="all", phase=phase, status=status, - execution_claim=_claim_purge(storage, purge_id), ) @@ -4978,8 +3740,7 @@ def test_governance_persistence_rejects_unsafe_idempotency_keys( storage.append_audit_event(event) with pytest.raises(ValueError, match="idempotency_key"): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id="purge_unsafe_idem", idempotency_key=idempotency_key, operation_type="user_erasure", @@ -5084,7 +3845,6 @@ 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, ) @@ -5095,8 +3855,7 @@ 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"): - _begin_test_purge_operation( - storage, + storage.begin_purge_operation( purge_id=purge_id, idempotency_key="idem_purge_identifier_suffix", operation_type="user_erasure", @@ -5132,7 +3891,6 @@ 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, ) @@ -5168,7 +3926,6 @@ 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, ) @@ -5296,7 +4053,6 @@ 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 @@ -5307,7 +4063,6 @@ 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), ) @@ -5368,15 +4123,13 @@ 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 = _begin_test_purge_operation( - storage, + purge_id = storage.begin_purge_operation( purge_id="purge_legacy_upgrade", idempotency_key="idem_legacy_upgrade", operation_type="user_erasure", scope_type="user", - subject_ref=storage._subject_ref_for_user_id("alice"), + subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, - authoritative_user_id="alice", ).purge_id storage.record_purge_target( purge_id=purge_id, @@ -5384,7 +4137,6 @@ 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 @@ -5591,23 +4343,18 @@ 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), - authoritative_user_id="alice", - execution_claim=_claim_purge(storage, purge_id), + purge_id, _erase_event(purge_id=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) - 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), - ) + # 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) + ) 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 2a6512833..fb40107bd 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,7 +20,6 @@ source TEXT NOT NULL DEFAULT '', agent_version TEXT NOT NULL DEFAULT '', session_id TEXT, - governance_subject_ref TEXT, metadata TEXT NOT NULL DEFAULT '{}' ); """ @@ -30,35 +29,12 @@ 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, - governance_subject_ref - ) VALUES (?, ?, ?, ?, ?, ?)""", + "INSERT INTO requests (request_id, user_id, created_at, source, session_id) " + "VALUES (?, ?, ?, ?, ?)", [ - ( - "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", - ), + ("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"), ], ) conn.commit() @@ -113,28 +89,6 @@ 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 deleted file mode 100644 index 73439ffc7..000000000 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py +++ /dev/null @@ -1,310 +0,0 @@ -"""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 dd7ca1800..07ee32ce5 100644 --- a/tests/server/services/storage/test_sqlite_storage.py +++ b/tests/server/services/storage/test_sqlite_storage.py @@ -29,7 +29,6 @@ _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, @@ -56,42 +55,6 @@ 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 e7def4662..fea348a9c 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 @@ -20,8 +20,6 @@ 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 @@ -168,69 +166,6 @@ 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_retention.py b/tests/server/services/storage/test_storage_contract_retention.py index 2612dc1c4..6cdd17083 100644 --- a/tests/server/services/storage/test_storage_contract_retention.py +++ b/tests/server/services/storage/test_storage_contract_retention.py @@ -1,8 +1,6 @@ """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 @@ -190,50 +188,6 @@ 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/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index ba8be714b..b544500c9 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -1,18 +1,11 @@ """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 @@ -46,345 +39,14 @@ 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 cdc597054..6138f19f4 100644 --- a/tests/server/services/test_generation_billing_emission.py +++ b/tests/server/services/test_generation_billing_emission.py @@ -15,14 +15,7 @@ from typing import Any from unittest.mock import MagicMock, patch -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.api_schema.service_schemas import Interaction, Request from reflexio.models.config_schema import ( Config, ProfileExtractorConfig, @@ -35,14 +28,6 @@ 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, ) @@ -180,123 +165,12 @@ 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 - # This fixture retains every generated output, so billing and success telemetry - # have equal counts here. + # learnings_generated.count_value must equal the existing generation_succeeded count. 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. @@ -502,25 +376,6 @@ 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 b5313d181..87f9875de 100644 --- a/tests/server/services/test_non_extraction_learning_metering.py +++ b/tests/server/services/test_non_extraction_learning_metering.py @@ -1,18 +1,12 @@ from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock 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.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 ( - PlaybookAggregatorRequest, -) from reflexio.server.services.storage.storage_base import ( AgentBinding, AgentRunRecord, @@ -53,8 +47,11 @@ def _agent_run(*, extractor_kind: str) -> AgentRunRecord: ) -def test_resumable_profile_bills_only_ids_returned_by_finalization() -> None: - """A preassigned candidate ID is not billed when finalization drops it.""" +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. + """ events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -63,22 +60,28 @@ def test_resumable_profile_bills_only_ids_returned_by_finalization() -> None: ) run = _agent_run(extractor_kind="profile") - dropped_candidate = MagicMock(profile_id="dropped-before-persist") - with patch( - "reflexio.server.services.extraction.resume_worker.ProfileGenerationService" - ) as service_class: - 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 + worker._record_finalized_learnings( + run, + [object(), object()], + entity_type="profile", ) - assert events == [] + + 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:") def test_resumable_finalization_emits_one_event_per_profile_id() -> None: - """Finalization survivor IDs emit one entity-backed event per profile.""" + """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. + """ events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -87,9 +90,13 @@ 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, - ["prof-1", "prof-2"], + [_FakeProfile("prof-1"), _FakeProfile("prof-2")], entity_type="profile", ) @@ -109,7 +116,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.""" + """Same as above for the playbook (``user_playbook_id``) kind.""" events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -118,9 +125,13 @@ 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, - ["11", "12", "13"], + [_FakePlaybook(11), _FakePlaybook(12), _FakePlaybook(13)], entity_type="user_playbook", ) @@ -137,8 +148,11 @@ def test_resumable_finalization_emits_one_event_per_playbook_id() -> None: assert event.event_key == f"learn:{event.entity_type}:{event.entity_id}" -def test_resumable_playbook_bills_consolidation_replacement_id() -> None: - """Billing follows the persisted replacement, not its input candidate.""" +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. + """ events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) worker = ExtractionResumeWorker( @@ -147,91 +161,93 @@ def test_resumable_playbook_bills_consolidation_replacement_id() -> None: ) run = _agent_run(extractor_kind="playbook") - original_candidate = MagicMock(user_playbook_id=21) - with patch( - "reflexio.server.services.extraction.resume_worker.PlaybookGenerationService" - ) as service_class: - finalizer = service_class.return_value._finalize_extracted_items_with_outcome - finalizer.return_value = FinalizationResult(["88"], won_receipt=True) - worker._finalize_items(run, [original_candidate]) + class _FakePlaybook: + def __init__(self, user_playbook_id: int) -> None: + self.user_playbook_id = user_playbook_id - finalizer.assert_called_once_with( - [original_candidate], model_provenance=None, finalization_run_id=run.id + worker._record_finalized_learnings( + run, + [_FakePlaybook(21), _FakePlaybook(0)], + entity_type="user_playbook", ) + assert len(events) == 1 - assert events[0].count_value == 1 - assert events[0].event_key == "learn:user_playbook:88" - assert events[0].entity_id == "88" + 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:" + ) -def test_aggregation_emits_no_learnings_generated( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A completed aggregation remains observable but adds no billable learning.""" +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. + """ 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", ) - 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, - ), + + aggregator._record_learnings_generated( + learning_ids=["101", "102", "103"], + playbook_name="agent_rules", + request_id="agg-run-1", + metadata={"playbooks_generated": 3}, ) - 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", + + 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(), agent_version="v1", - content="Deploy changes and verify results.", - ) - 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}, - ) - monkeypatch.setattr( - aggregator, - "_generate_playbooks_with_source_clusters", - lambda *_args, **_kwargs: [(generated, user_playbooks, None)], - ) - monkeypatch.setattr( - aggregator, "_enqueue_playbook_optimization", lambda _items: None ) - stats = aggregator.run(PlaybookAggregatorRequest(agent_version="v1", rerun=True)) + 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, + ) - 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) + 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].entity_type == "agent_playbook" + assert events[0].source == "aggregation" diff --git a/tests/server/services/test_search_exposure.py b/tests/server/services/test_search_exposure.py deleted file mode 100644 index c8e330a80..000000000 --- a/tests/server/services/test_search_exposure.py +++ /dev/null @@ -1,292 +0,0 @@ -"""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", - playbook_owner_governance_subject_ref="owner: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_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 -): - 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 ab933db1a..dee35f07e 100644 --- a/tests/server/test_billing_meter_events.py +++ b/tests/server/test_billing_meter_events.py @@ -13,11 +13,12 @@ mint the same ``event_key`` and collapse into one event downstream. The existing count-based ``record_learnings_generated`` / ``emit_learnings_generated`` -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. +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. 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 f83d4fed640ad5ab479eb9b54d8e017e386434fa Mon Sep 17 00:00:00 2001 From: Yi Lu Date: Tue, 4 Aug 2026 00:50:01 -0700 Subject: [PATCH 2/4] fix: emit first callback drop anomaly Use an explicit unset sentinel for the last anomaly time so hosts with less than one hour of monotonic uptime do not suppress the first threshold crossing. --- reflexio/server/callback_executor.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/reflexio/server/callback_executor.py b/reflexio/server/callback_executor.py index ea1d30253..4978d003c 100644 --- a/reflexio/server/callback_executor.py +++ b/reflexio/server/callback_executor.py @@ -58,7 +58,7 @@ def __init__( self._cond = threading.Condition() self._active = 0 self._drop_times: deque[float] = deque() - self._last_drop_anomaly = 0.0 + self._last_drop_anomaly: float | None = None for i in range(workers): threading.Thread( target=self._worker_loop, @@ -113,9 +113,9 @@ def _record_drop_locked(self, dropped_name: str) -> _DropFacts: while self._drop_times and now - self._drop_times[0] > 60.0: self._drop_times.popleft() drops_last_minute = len(self._drop_times) - fire_anomaly = ( - drops_last_minute > _DROP_RATE_PER_MINUTE - and now - self._last_drop_anomaly > _DROP_ANOMALY_THROTTLE_SECONDS + fire_anomaly = drops_last_minute > _DROP_RATE_PER_MINUTE and ( + self._last_drop_anomaly is None + or now - self._last_drop_anomaly > _DROP_ANOMALY_THROTTLE_SECONDS ) if fire_anomaly: self._last_drop_anomaly = now From 3c054d5d3a0adf00f274bf877b174b565e3d4f84 Mon Sep 17 00:00:00 2001 From: Yi Lu Date: Tue, 4 Aug 2026 09:16:37 -0700 Subject: [PATCH 3/4] fix: address revert review findings Preserve session-outcome compatibility across the #407 downgrade, make governance erasure and transactions robust, and keep learning metering retry-safe and aligned with persisted survivors. --- reflexio/models/api_schema/domain/entities.py | 4 +- reflexio/server/billing_meter.py | 14 +- reflexio/server/routes/system.py | 4 +- .../services/base_generation_service.py | 3 +- .../services/extraction/resume_worker.py | 17 +- .../playbook/components/aggregator.py | 11 ++ reflexio/server/services/playbook/service.py | 24 ++- reflexio/server/services/profile/service.py | 8 +- .../services/storage/sqlite_storage/_base.py | 130 ++++++++++++- .../storage/sqlite_storage/_governance.py | 5 +- .../sqlite_storage/_session_outcomes.py | 5 +- .../governance/_erase_execution.py | 13 +- .../sqlite_storage/governance/_purge.py | 51 ++--- .../services/storage/storage_base/__init__.py | 21 ++- .../services/extraction/test_resume_worker.py | 46 +++++ .../test_playbook_generation_service.py | 25 +++ .../sqlite_storage/test_governance_storage.py | 44 +++++ ...est_session_outcome_downgrade_migration.py | 174 ++++++++++++++++++ ..._playbook_aggregation_state_integration.py | 12 ++ .../test_storage_contract_clear_user_data.py | 15 ++ .../test_storage_contract_session_outcomes.py | 31 ++++ .../test_non_extraction_learning_metering.py | 18 ++ tests/server/test_billing_meter_events.py | 13 ++ 23 files changed, 617 insertions(+), 71 deletions(-) create mode 100644 tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 46540eebf..9165b746f 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -954,8 +954,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, user -# playbooks, profiles). Used by paired-protocol harnesses (e.g. SWE-bench) to +# 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 # 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/server/billing_meter.py b/reflexio/server/billing_meter.py index 76818fe81..4e0380d5e 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -86,6 +86,7 @@ def record_learnings_generated( agent_version: str | None = None, playbook_name: str | None = None, entity_type: str | None = None, + event_key: str | None = None, metadata: Mapping[str, Any] | None = None, ) -> None: """Emit the Learning value facet — number of profiles/playbooks generated. @@ -96,9 +97,10 @@ def record_learnings_generated( :func:`record_learnings_generated_records` whenever the caller has the durable learning ids in scope. No-op when ``count <= 0``. - Emits a single event carrying a synthesized ``event_key=f"learn-batch:{uuid4()}"`` - (distinct per call) so this aggregate event still has a dedup key, even - though it is not entity-backed. + Emits a single event carrying ``event_key`` when the caller has a durable + retry identity, otherwise synthesizes ``f"learn-batch:{uuid4()}"``. This + gives retryable callers an idempotent aggregate event without forcing an + unstable key on event-moment callers. Args: org_id: Organisation identifier. @@ -113,6 +115,7 @@ def record_learnings_generated( agent_version: Optional agent version tied to the generated learning. playbook_name: Optional playbook name for playbook learnings. entity_type: Optional entity type (e.g. ``"profile"``). + event_key: Optional caller-supplied, retry-stable event key. metadata: Optional path-specific usage metadata. """ if count <= 0: @@ -129,7 +132,7 @@ def record_learnings_generated( agent_version=agent_version, playbook_name=playbook_name, entity_type=entity_type, - event_key=f"learn-batch:{uuid.uuid4()}", + event_key=event_key or f"learn-batch:{uuid.uuid4()}", count_value=count, platform_llm=platform_llm, platform_storage=platform_storage, @@ -227,6 +230,7 @@ def emit_learnings_generated( agent_version: str | None = None, playbook_name: str | None = None, entity_type: str | None = None, + event_key: str | None = None, metadata: Mapping[str, Any] | None = None, ) -> None: """Resolve ``platform_llm`` from config and emit the Learning value facet. @@ -251,6 +255,7 @@ def emit_learnings_generated( agent_version: Optional agent version tied to the generated learning. playbook_name: Optional playbook name for playbook learnings. entity_type: Optional entity type (e.g. ``"profile"``). + event_key: Optional caller-supplied, retry-stable event key. metadata: Optional path-specific usage metadata. """ if count <= 0: @@ -271,6 +276,7 @@ def emit_learnings_generated( agent_version=agent_version, playbook_name=playbook_name, entity_type=entity_type, + event_key=event_key, metadata=metadata, ) except Exception: diff --git a/reflexio/server/routes/system.py b/reflexio/server/routes/system.py index d97b01e8f..03f9cf60b 100644 --- a/reflexio/server/routes/system.py +++ b/reflexio/server/routes/system.py @@ -140,8 +140,8 @@ def clear_user_data( ) -> ClearUserDataResponse: """Delete all rows scoped to a single ``user_id``. - Removes the user's interactions, user playbooks, profiles, and - requests. Does NOT touch ``agent_playbooks`` — they are + Removes the user's interactions, session outcomes, user playbooks, + profiles, and requests. Does NOT touch ``agent_playbooks`` — they are intentionally shared cross-project. Used by paired-protocol harnesses (e.g. SWE-bench) to isolate per-task data on a shared backend without one task's clear-all nuking another in-flight diff --git a/reflexio/server/services/base_generation_service.py b/reflexio/server/services/base_generation_service.py index ebf78b7b5..2334d3741 100644 --- a/reflexio/server/services/base_generation_service.py +++ b/reflexio/server/services/base_generation_service.py @@ -342,10 +342,11 @@ 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) -> None: + def _finalize_extracted_items(self, items: list) -> list: """Persist already-flattened extracted items through the service path.""" if items: self._process_results([items]) + return items @abstractmethod def _should_track_in_progress(self) -> bool: diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index b6f78bc47..c66b9daa4 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -884,8 +884,12 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - service._finalize_extracted_items(items, model_provenance=model_provenance) - self._record_finalized_learnings(run, items, entity_type="profile") + persisted_items = service._finalize_extracted_items( + items, model_provenance=model_provenance + ) + self._record_finalized_learnings( + run, persisted_items or [], entity_type="profile" + ) return if run.binding.extractor_kind == "playbook": service = PlaybookGenerationService( @@ -900,8 +904,12 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - service._finalize_extracted_items(items, model_provenance=model_provenance) - self._record_finalized_learnings(run, items, entity_type="user_playbook") + persisted_items = service._finalize_extracted_items( + items, model_provenance=model_provenance + ) + self._record_finalized_learnings( + run, persisted_items or [], entity_type="user_playbook" + ) return raise ResumeWorkerError( f"Unsupported extractor kind {run.binding.extractor_kind!r}" @@ -968,6 +976,7 @@ def _record_finalized_learnings( 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/playbook/components/aggregator.py b/reflexio/server/services/playbook/components/aggregator.py index 0a827797f..305a72c57 100644 --- a/reflexio/server/services/playbook/components/aggregator.py +++ b/reflexio/server/services/playbook/components/aggregator.py @@ -1073,6 +1073,17 @@ 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 2c67f66ed..c96b28eac 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -12,6 +12,7 @@ from reflexio.server.services.deferred_learning_plan import GenerationComputePlan from reflexio.server.services.storage.storage_base import BaseStorage +from reflexio.models.api_schema.common import sanitise_for_log from reflexio.models.api_schema.domain.entities import LineageContext from reflexio.models.api_schema.internal_schema import RequestInteractionDataModel from reflexio.models.api_schema.service_schemas import ( @@ -601,10 +602,22 @@ def _dispatch_playbook_schedulers(self, plan: PlaybookWritePlan) -> None: keeping that path identical to the pre-split monolith. The two callers are mutually exclusive, so the schedulers fire exactly once per run. """ - self._enqueue_user_playbook_optimization(plan.new_playbooks) + try: + self._enqueue_user_playbook_optimization(plan.new_playbooks) + except Exception: + logger.exception( + "Failed to schedule post-persist playbook optimization for request %s", + sanitise_for_log(plan.request_id), + ) if not plan.output_pending_status and not plan.skip_aggregation: - logger.info("Trigger playbook aggregation") - self._trigger_playbook_aggregation() + try: + logger.info("Trigger playbook aggregation") + self._trigger_playbook_aggregation() + except Exception: + logger.exception( + "Failed to schedule post-persist playbook aggregation for request %s", + sanitise_for_log(plan.request_id), + ) def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: """Post-commit side-effects — base telemetry/billing + playbook schedulers. @@ -624,7 +637,7 @@ def _finalize_extracted_items( all_playbooks: list[UserPlaybook], *, model_provenance: ModelProvenance | None = None, - ) -> None: + ) -> list[UserPlaybook]: """Permanent V3 wrapper: compute→persist→schedulers together (no fence). Kept for the synchronous resume/manual callers @@ -638,9 +651,10 @@ def _finalize_extracted_items( self._last_model_provenance = model_provenance plan = self._resolve_write_plan([all_playbooks]) if plan is None: - return + return [] self._persist_write_plan(plan) self._dispatch_playbook_schedulers(plan) + return plan.new_playbooks def _apply_consolidation_lineage( self, diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index 3e253ddaa..3118d9305 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -337,7 +337,7 @@ def _finalize_extracted_items( all_new_profiles: list[UserProfile], *, model_provenance: ModelProvenance | None = None, - ) -> None: + ) -> list[UserProfile]: """Permanent V3 wrapper: compute-then-persist together (no external fence). Kept for the synchronous resume/manual callers @@ -349,8 +349,10 @@ def _finalize_extracted_items( if model_provenance is not None: self._last_model_provenance = model_provenance plan = self._resolve_write_plan([all_new_profiles]) - if plan is not None: - self._persist_write_plan(plan) + if plan is None: + return [] + self._persist_write_plan(plan) + return plan.new_profiles 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/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 9fec58e38..688265711 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -65,6 +65,8 @@ logger = logging.getLogger(__name__) +_MINIMUM_SQLITE_VERSION = (3, 35, 0) + # --------------------------------------------------------------------------- # Module-level helpers @@ -663,6 +665,12 @@ def __init__( logger.info("SQLite Storage for org %s using db_path: %s", org_id, db_path) + if sqlite3.sqlite_version_info < _MINIMUM_SQLITE_VERSION: + detected_version = ".".join(map(str, sqlite3.sqlite_version_info)) + raise RuntimeError( + f"SQLite 3.35.0 or newer is required; detected {detected_version}" + ) + # Ensure parent directory exists Path(db_path).parent.mkdir(parents=True, exist_ok=True) @@ -770,6 +778,7 @@ 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: cur = self.conn.cursor() cur.executescript(_DDL) @@ -807,6 +816,115 @@ def migrate(self) -> bool: self._migrate_learning_jobs() return True + 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") + 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 "'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. @@ -2258,8 +2376,8 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: user_id (str): The user id whose rows should be deleted. Returns: - dict[str, int]: Per-entity counts with keys ``interactions``, - ``user_playbooks``, ``profiles``, ``requests``, + dict[str, int]: Per-entity counts with keys ``session_outcomes``, + ``interactions``, ``user_playbooks``, ``profiles``, ``requests``, ``purged_profiles``, and ``purged_user_playbooks``. ``profiles`` and ``user_playbooks`` reflect hard-deleted counts; purged rows are counted separately. @@ -2336,9 +2454,10 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: interactions_cur = self.conn.execute( "DELETE FROM interactions WHERE user_id = ?", (user_id,) ) - self.conn.execute( - "DELETE FROM session_outcomes WHERE governance_subject_ref = ?", - (subject_ref,), + session_outcomes_cur = self.conn.execute( + """DELETE FROM session_outcomes + WHERE user_id = ? OR governance_subject_ref = ?""", + (user_id, subject_ref), ) requests_cur = self.conn.execute( "DELETE FROM requests WHERE user_id = ?", (user_id,) @@ -2375,6 +2494,7 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: self.purge_content(entity_type="user_playbook", entity_id=str(upid)) return { + "session_outcomes": session_outcomes_cur.rowcount, "interactions": interactions_cur.rowcount, "user_playbooks": upb_deleted_count, "profiles": profile_deleted_count, diff --git a/reflexio/server/services/storage/sqlite_storage/_governance.py b/reflexio/server/services/storage/sqlite_storage/_governance.py index 8689b0787..5e4abad09 100644 --- a/reflexio/server/services/storage/sqlite_storage/_governance.py +++ b/reflexio/server/services/storage/sqlite_storage/_governance.py @@ -406,8 +406,9 @@ def _planned_governance_delete_counts( ).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 governance_subject_ref = ?", - (subject_ref,), + """SELECT COUNT(*) AS cnt FROM session_outcomes + WHERE user_id = ? OR governance_subject_ref = ?""", + (user_id, subject_ref), ).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 c3f95d4a0..aaec5a1b2 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -204,8 +204,9 @@ 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 governance_subject_ref = ?", - (subject_ref,), + """DELETE FROM session_outcomes + WHERE user_id = ? OR governance_subject_ref = ?""", + (user_id, subject_ref), ) 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 9ed6f5b9b..954eae6c1 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py @@ -145,8 +145,9 @@ def _clear_user_data_for_governance_locked( deps = self._deps() subject_ref = deps._subject_ref_for_user_id(user_id) session_outcomes_cur = self.conn.execute( - "DELETE FROM session_outcomes WHERE governance_subject_ref = ?", - (subject_ref,), + """DELETE FROM session_outcomes + WHERE user_id = ? OR governance_subject_ref = ?""", + (user_id, subject_ref), ) interaction_ids = [ int(row["interaction_id"]) @@ -269,11 +270,7 @@ def _clear_user_data_for_governance_locked( ) return { - **( - {"session_outcomes": session_outcomes_cur.rowcount} - if session_outcomes_cur.rowcount - else {} - ), + "session_outcomes": session_outcomes_cur.rowcount, "interactions": interactions_cur.rowcount, "user_playbooks": len(delete_upb_ids), "profiles": len(delete_profile_ids), @@ -372,7 +369,7 @@ def apply_governance_user_data_delete( } with self._lock: try: - self.conn.execute("BEGIN") + self.conn.execute("BEGIN IMMEDIATE") self._validate_prepared_delete_target_matrix_locked(purge_id) self._validate_hide_for_rebuild_targets_locked(purge_id) expected_user_playbook_ids = ( diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py index 386cbdd95..0a62e91ea 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py @@ -351,29 +351,34 @@ def fail_purge_operation( validated_error_detail = _validate_governance_error_detail(error_detail) now = _epoch_now() with self._lock: - cur = self.conn.execute( - """UPDATE purge_operations - SET status = 'failed', error_code = ?, error_detail = ?, - updated_at = ?, completed_at = ? - WHERE purge_id = ? AND org_id = ? AND status != 'complete'""", - ( - validated_error_code, - validated_error_detail, - now, - now, - purge_id, - self.org_id, - ), - ) - if cur.rowcount == 0: - existing = self.conn.execute( - "SELECT status FROM purge_operations WHERE purge_id = ? AND org_id = ?", - (purge_id, self.org_id), - ).fetchone() - if existing is not None and str(existing["status"]) == "complete": - raise ValueError("Purge operation is already complete") - raise ValueError(f"Purge operation {purge_id!r} not found") - self.conn.commit() + try: + self.conn.execute("BEGIN IMMEDIATE") + cur = self.conn.execute( + """UPDATE purge_operations + SET status = 'failed', error_code = ?, error_detail = ?, + updated_at = ?, completed_at = ? + WHERE purge_id = ? AND org_id = ? AND status != 'complete'""", + ( + validated_error_code, + validated_error_detail, + now, + now, + purge_id, + self.org_id, + ), + ) + if cur.rowcount == 0: + existing = self.conn.execute( + "SELECT status FROM purge_operations WHERE purge_id = ? AND org_id = ?", + (purge_id, self.org_id), + ).fetchone() + if existing is not None and str(existing["status"]) == "complete": + raise ValueError("Purge operation is already complete") + raise ValueError(f"Purge operation {purge_id!r} not found") + self.conn.commit() + except Exception: + self.conn.rollback() + raise return self.get_purge_operation(purge_id) def get_purge_operation(self, purge_id: str) -> PurgeOperation: diff --git a/reflexio/server/services/storage/storage_base/__init__.py b/reflexio/server/services/storage/storage_base/__init__.py index 1dae46dd3..4a790ee8d 100644 --- a/reflexio/server/services/storage/storage_base/__init__.py +++ b/reflexio/server/services/storage/storage_base/__init__.py @@ -147,13 +147,13 @@ 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, 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 by paired - protocols (e.g. SWE-bench) that share a single backend across - parallel tasks without one task's clear-all nuking another - in-flight task's rows. + Removes the user's interactions, session outcomes, 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 + by paired protocols (e.g. SWE-bench) that share a single backend across + parallel tasks without one task's clear-all nuking another in-flight + task's rows. **Lineage-aware erasure:** rows that are tombstones (``merged_into`` or ``superseded_by`` is set) *or* are pointed to by another row @@ -176,13 +176,13 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: user_id (str): The user id whose rows should be deleted. Returns: - dict[str, int]: Per-entity counts with keys ``interactions``, - ``user_playbooks``, ``profiles``, ``requests``, + dict[str, int]: Per-entity counts with keys ``session_outcomes``, + ``interactions``, ``user_playbooks``, ``profiles``, ``requests``, ``purged_profiles``, and ``purged_user_playbooks``. ``profiles`` and ``user_playbooks`` reflect hard-deleted counts; purged rows are counted separately. """ - self.clear_session_outcomes_for_user(user_id) + session_outcome_counts = self.clear_session_outcomes_for_user(user_id) interaction_count = len(self.get_user_interaction(user_id)) # All statuses a user's row can have — including tombstones (SUPERSEDED, @@ -264,6 +264,7 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: self.purge_content(entity_type="user_playbook", entity_id=upid) return { + "session_outcomes": session_outcome_counts.get("session_outcomes", 0), "interactions": interaction_count, "user_playbooks": deleted_user_playbooks, "profiles": deleted_profiles, diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index c362b0e45..be2880686 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -337,6 +337,52 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( assert storage.list_run_tool_dependencies("run_1")[0].consumed_at is not None +@pytest.mark.parametrize( + ("extractor_kind", "finalize_path", "entity_type"), + [ + ( + "profile", + "reflexio.server.services.profile.service." + "ProfileGenerationService._finalize_extracted_items", + "profile", + ), + ( + "playbook", + "reflexio.server.services.playbook.service." + "PlaybookGenerationService._finalize_extracted_items", + "user_playbook", + ), + ], +) +def test_resume_bills_only_items_that_survive_finalization( + request_context, extractor_kind, finalize_path, entity_type +): + run = AgentRunRecord( + id="survivor-run", + binding=AgentBinding( + org_id="org_1", + extractor_kind=extractor_kind, + user_id="user_1", + request_id="request_1", + agent_version="v1", + source="api", + ), + status=AgentRunStatus.FINALIZING, + generation_request_snapshot={}, + ) + dropped = object() + survivor = object() + worker = ExtractionResumeWorker(request_context=request_context) + + with ( + patch(finalize_path, return_value=[survivor]), + 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) + + def test_resume_worker_tagging_schedule_failure_is_best_effort( request_context, ): diff --git a/tests/server/services/playbook/test_playbook_generation_service.py b/tests/server/services/playbook/test_playbook_generation_service.py index f8c9fa12f..fcce32a2d 100644 --- a/tests/server/services/playbook/test_playbook_generation_service.py +++ b/tests/server/services/playbook/test_playbook_generation_service.py @@ -92,6 +92,31 @@ def test_inline_aggregation_default_path_schedules_durably(): _storage(service).schedule_playbook_aggregation.assert_called_once_with("v1") +def test_post_persist_scheduler_failures_are_best_effort(caplog): + configurator = MagicMock() + configurator.get_config.return_value = _aggregation_enabled_config() + service = _service_for_inline_aggregation(configurator) + plan = MagicMock() + plan.new_playbooks = [MagicMock()] + plan.output_pending_status = False + plan.skip_aggregation = False + service._enqueue_user_playbook_optimization = MagicMock( + side_effect=RuntimeError("optimization unavailable") + ) + service._trigger_playbook_aggregation = MagicMock( + side_effect=RuntimeError("aggregation unavailable") + ) + + service._dispatch_playbook_schedulers(plan) + + service._enqueue_user_playbook_optimization.assert_called_once_with( + plan.new_playbooks + ) + service._trigger_playbook_aggregation.assert_called_once_with() + assert "optimization unavailable" in caplog.text + assert "aggregation unavailable" in caplog.text + + def test_maybe_trigger_user_playbook_aggregation_durably_schedules(): from reflexio.server.services.playbook.aggregation_trigger import ( maybe_trigger_user_playbook_aggregation, 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 89783fe86..045cf394b 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -560,6 +560,36 @@ def test_complete_purge_operation_with_audit_begins_immediate_transaction_before assert begin_index < first_validation_read_index +def test_apply_governance_delete_begins_immediate_transaction_before_reads(storage): + purge_id = _begin_purge(storage, "purge_delete_begin_immediate") + for target_name in CANONICAL_DELETE_TARGET_NAMES: + storage.record_purge_target( + purge_id=purge_id, + target_name=target_name, + target_ref="all", + phase="delete", + status="pending", + 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") + finally: + storage.conn.set_trace_callback(None) + + begin_index = next( + i for i, statement in enumerate(statements) if statement == "BEGIN IMMEDIATE" + ) + first_validation_read_index = next( + i + for i, statement in enumerate(statements) + if statement.lstrip().upper().startswith("SELECT") + ) + assert begin_index < first_validation_read_index + + def test_complete_purge_operation_with_audit_accepts_planned_success_detail(storage): purge_id = _begin_completeable_purge(storage, "purge_success_detail") deleted_counts = { @@ -1052,6 +1082,7 @@ def test_prepare_governance_erase_targets_records_full_delete_matrix_counts(stor counts = storage.clear_user_data(user_id) assert counts == { + "session_outcomes": 0, "interactions": 1, "user_playbooks": 1, "profiles": 1, @@ -2529,6 +2560,18 @@ def test_fail_purge_operation_persists_code_shaped_error_detail(storage): assert failed.error_detail == "target_delete_failed" +def test_fail_missing_purge_rolls_back_implicit_transaction(storage): + with pytest.raises(ValueError, match="not found"): + storage.fail_purge_operation( + "purge_missing", + error_code="PURGE_TARGET_FAILED", + error_detail="target_delete_failed", + ) + + assert storage.conn.in_transaction is False + _begin_purge(storage, "purge_after_missing_failure") + + @pytest.mark.parametrize( "error_code", ["content_purge_failed", "prompt_redaction_route"] ) @@ -3271,6 +3314,7 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( ) assert counts == { + "session_outcomes": 0, "interactions": 1, "user_playbooks": 1, "profiles": 1, 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 new file mode 100644 index 000000000..24220d270 --- /dev/null +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py @@ -0,0 +1,174 @@ +"""Regression coverage for reverting the session-outcome identity schema.""" + +from unittest.mock import patch + +import pytest + +from reflexio.models.api_schema.domain import ( + GetSessionOutcomesRequest, + Request, + SessionOutcomeKind, + SetSessionOutcomeRequest, +) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage import _base as sqlite_base + +pytestmark = pytest.mark.integration + + +_IDENTITY_SCHEMA = """ +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) +); +""" + + +def _storage(db_path: str) -> SQLiteStorage: + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + return SQLiteStorage(org_id="downgrade-org", db_path=db_path) + + +def test_identity_schema_is_downgraded_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"): + storage.add_request( + Request( + request_id=f"request-{session_id}", + user_id="user-1", + session_id=session_id, + source="test", + created_at=100, + ) + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_IDENTITY_SCHEMA) + storage.conn.executemany( + """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, 'user-1', ?, ?, 101, 'test', ?, ?, ?, 102)""", + [ + ("outcome-kept", "kept", "success", "a" * 64, "b" * 64, "old-ref"), + ( + "outcome-unknown", + "unknown", + "unknown", + "c" * 64, + "d" * 64, + "old-ref", + ), + ], + ) + storage.conn.commit() + storage.conn.close() + + migrated = _storage(db_path) + + assert { + row["name"] + for row in migrated.conn.execute("PRAGMA table_info(session_outcomes)") + } == { + "user_id", + "session_id", + "outcome", + "occurred_at", + "source", + "label", + "value", + "metadata", + "governance_subject_ref", + "created_at", + } + records = migrated.get_session_outcomes(GetSessionOutcomesRequest()) + assert [(record.session_id, record.outcome) for record in records] == [ + ("kept", SessionOutcomeKind.SUCCESS) + ] + + request = SetSessionOutcomeRequest( + session_id="new", outcome=SessionOutcomeKind.FAILURE, occurred_at=101 + ) + result = migrated.record_session_outcome( + request, + created_at=102, + expected_context=migrated.get_session_outcome_context("new"), + ) + assert result.recorded is True + + +@pytest.mark.parametrize("with_subject_column", [False, True]) +def test_legacy_schema_backfills_required_subject_ref( + tmp_path, with_subject_column: bool +) -> None: + db_path = str(tmp_path / f"legacy-subject-{with_subject_column}.db") + storage = _storage(db_path) + storage.conn.execute("DROP TABLE session_outcomes") + subject_column = "governance_subject_ref TEXT," if with_subject_column else "" + storage.conn.executescript( + 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) + );""" + ) + columns = "user_id, session_id, outcome, occurred_at, source, created_at" + values: tuple[object, ...] = ( + "legacy-user", + "legacy-session", + "success", + 101, + "legacy-source", + 102, + ) + if with_subject_column: + columns += ", governance_subject_ref" + values += (None,) + placeholders = ", ".join("?" for _ in values) + storage.conn.execute( + f"INSERT INTO session_outcomes ({columns}) VALUES ({placeholders})", # noqa: S608 + values, + ) + storage.conn.commit() + storage.conn.close() + + migrated = _storage(db_path) + + row = migrated.conn.execute( + "SELECT governance_subject_ref FROM session_outcomes" + ).fetchone() + assert row is not None + assert row["governance_subject_ref"] == migrated._subject_ref_for_user_id( + "legacy-user" + ) + + +def test_sqlite_versions_without_returning_and_drop_column_are_rejected( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(sqlite_base.sqlite3, "sqlite_version_info", (3, 34, 1)) + + with pytest.raises(RuntimeError, match="SQLite 3.35.0 or newer"): + _storage(str(tmp_path / "old-sqlite.db")) 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 3ebcd1e3e..85be9337b 100644 --- a/tests/server/services/storage/test_playbook_aggregation_state_integration.py +++ b/tests/server/services/storage/test_playbook_aggregation_state_integration.py @@ -1032,9 +1032,21 @@ 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")) 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=learning_meter.call_args.kwargs["request_id"], + 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/storage/test_storage_contract_clear_user_data.py b/tests/server/services/storage/test_storage_contract_clear_user_data.py index fea348a9c..cf550ae22 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 @@ -11,6 +11,10 @@ import pytest +from reflexio.models.api_schema.domain import ( + SessionOutcomeKind, + SetSessionOutcomeRequest, +) from reflexio.models.api_schema.service_schemas import ( AgentPlaybook, Interaction, @@ -186,9 +190,20 @@ def test_returned_counts_match_seeded_rows(self, storage: BaseStorage) -> None: _make_profile("userA", "pa2"), ], ) + outcome = SetSessionOutcomeRequest( + session_id="sess_ra1", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=int(datetime.now(UTC).timestamp()), + ) + storage.record_session_outcome( + outcome, + created_at=int(datetime.now(UTC).timestamp()), + expected_context=storage.get_session_outcome_context("sess_ra1"), + ) counts = storage.clear_user_data("userA") + assert counts["session_outcomes"] == 1 assert counts["interactions"] == 2 assert counts["user_playbooks"] == 2 assert counts["profiles"] == 2 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 b544500c9..b07a2223b 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -125,3 +125,34 @@ def test_empty_request_source_is_preserved(storage: BaseStorage) -> None: GetSessionOutcomesRequest(session_ids=["empty-source"]) ) assert record.source == "" + + +def test_clear_outcomes_survives_governance_secret_rotation( + storage: BaseStorage, monkeypatch +) -> None: + monkeypatch.setenv("REFLEXIO_GOVERNANCE_REF_SECRET", "old-secret") + storage.add_request( + Request( + request_id="rotated-r1", + user_id="rotated-user", + session_id="rotated-session", + source="test", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="rotated-session", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("rotated-session"), + ) + + monkeypatch.setenv("REFLEXIO_GOVERNANCE_REF_SECRET", "new-secret") + counts = storage.clear_session_outcomes_for_user("rotated-user") + + assert counts == {"session_outcomes": 1} + assert storage.get_session_outcomes(GetSessionOutcomesRequest()) == [] diff --git a/tests/server/services/test_non_extraction_learning_metering.py b/tests/server/services/test_non_extraction_learning_metering.py index 87f9875de..3536b39cc 100644 --- a/tests/server/services/test_non_extraction_learning_metering.py +++ b/tests/server/services/test_non_extraction_learning_metering.py @@ -77,6 +77,24 @@ def test_resumable_finalization_falls_back_when_items_lack_ids() -> None: assert event.event_key is not None and event.event_key.startswith("learn-batch:") +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: """When every item carries a durable ``profile_id`` (the common case -- profile ids are assigned by the extractor before finalize runs), emit one diff --git a/tests/server/test_billing_meter_events.py b/tests/server/test_billing_meter_events.py index dee35f07e..b9ab61401 100644 --- a/tests/server/test_billing_meter_events.py +++ b/tests/server/test_billing_meter_events.py @@ -126,6 +126,19 @@ def test_records_emits_distinct_keys(): assert len(keys) == len(set(keys)) == 3 +def test_batch_record_accepts_a_retry_stable_event_key(): + with patch(HOOK) as hook: + record_learnings_generated( + org_id="org1", + count=2, + platform_llm=True, + platform_storage=None, + event_key="learn-batch:resumable:run-1:profile", + ) + + assert hook.call_args.kwargs["event_key"] == ("learn-batch:resumable:run-1:profile") + + def test_records_totals_preserved(): with patch(HOOK) as hook: record_learnings_generated_records( From eebc0e5876ea887d9e4fcb91430df4dccc6a45d9 Mon Sep 17 00:00:00 2001 From: Yi Lu Date: Tue, 4 Aug 2026 09:50:19 -0700 Subject: [PATCH 4/4] fix: address follow-up revert review findings Make synchronous persistence atomic, serialize purge state transitions across SQLite connections, repair legacy outcome schemas with empty governance defaults, and strengthen metering coverage. --- reflexio/server/services/playbook/service.py | 3 +- reflexio/server/services/profile/service.py | 3 +- .../services/storage/sqlite_storage/_base.py | 5 + .../sqlite_storage/governance/_purge.py | 115 ++++++++------- .../test_playbook_generation_service.py | 49 +++++++ .../test_dedup_always_soft_integration.py | 42 ++++++ .../sqlite_storage/test_governance_storage.py | 135 ++++++++++++++++++ ...est_session_outcome_downgrade_migration.py | 38 +++++ ..._playbook_aggregation_state_integration.py | 6 +- 9 files changed, 340 insertions(+), 56 deletions(-) diff --git a/reflexio/server/services/playbook/service.py b/reflexio/server/services/playbook/service.py index c96b28eac..f3ba0cc67 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -652,7 +652,8 @@ def _finalize_extracted_items( plan = self._resolve_write_plan([all_playbooks]) if plan is None: return [] - self._persist_write_plan(plan) + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + self._persist_write_plan(plan) self._dispatch_playbook_schedulers(plan) return plan.new_playbooks diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index 3118d9305..4b2bd1053 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -351,7 +351,8 @@ def _finalize_extracted_items( plan = self._resolve_write_plan([all_new_profiles]) if plan is None: return [] - self._persist_write_plan(plan) + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + self._persist_write_plan(plan) return plan.new_profiles def check_and_update_profiles(self, profiles: list[UserProfile]) -> None: diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 688265711..ea2f1b72e 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -843,6 +843,10 @@ def _migrate_session_outcomes_schema(self) -> None: ).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", @@ -854,6 +858,7 @@ def _migrate_session_outcomes_schema(self) -> None: 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 diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py index 0a62e91ea..1f2a7f0e9 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py @@ -189,45 +189,51 @@ def begin_purge_operation( ) now = _epoch_now() with self._lock: - existing = self.conn.execute( - """SELECT * FROM purge_operations - WHERE org_id = ? AND idempotency_key = ?""", - (self.org_id, validated_idempotency_key), - ).fetchone() - if existing is not None: - existing_operation = _row_to_purge_operation(existing) - expected_identity = { - "purge_id": validated_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(existing_operation, field_name) != expected_value: - raise ValueError( - "Existing purge operation for idempotency_key has " - f"mismatched {field_name}" - ) - 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', ?, ?)""", - ( - validated_purge_id, - self.org_id, - operation_type, - scope_type, - subject_ref, - request_ref, - validated_idempotency_key, - now, - now, - ), - ) - self.conn.commit() + try: + self.conn.execute("BEGIN IMMEDIATE") + existing = self.conn.execute( + """SELECT * FROM purge_operations + WHERE org_id = ? AND idempotency_key = ?""", + (self.org_id, validated_idempotency_key), + ).fetchone() + if existing is not None: + existing_operation = _row_to_purge_operation(existing) + expected_identity = { + "purge_id": validated_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(existing_operation, field_name) != expected_value: + raise ValueError( + "Existing purge operation for idempotency_key has " + f"mismatched {field_name}" + ) + self.conn.rollback() + return existing_operation + 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', ?, ?)""", + ( + validated_purge_id, + self.org_id, + operation_type, + scope_type, + subject_ref, + request_ref, + validated_idempotency_key, + now, + now, + ), + ) + self.conn.commit() + except Exception: + self.conn.rollback() + raise return self.get_purge_operation(validated_purge_id) def record_purge_target( @@ -258,17 +264,21 @@ def record_purge_target( allowed=_ALLOWED_PURGE_TARGET_STATUSES, ) with self._lock: - self._record_purge_target_locked( - purge_id=purge_id, - target_name=target_name, - target_ref=target_ref, - phase=phase, - status=status, - detail=detail, - deleted_count=deleted_count, - error_detail=error_detail, - ) - self.conn.commit() + try: + self._record_purge_target_locked( + purge_id=purge_id, + target_name=target_name, + target_ref=target_ref, + phase=phase, + status=status, + detail=detail, + deleted_count=deleted_count, + error_detail=error_detail, + ) + self.conn.commit() + except Exception: + self.conn.rollback() + raise def list_purge_targets( self, purge_id: str, phase: str | None = None @@ -302,10 +312,11 @@ def prepare_governance_erase_targets( ) -> None: purge_id = _validate_governance_purge_id("purge_id", purge_id) with self._lock: - if self.purge_targets_prepared(purge_id): - return try: self.conn.execute("BEGIN IMMEDIATE") + if self.purge_targets_prepared(purge_id): + self.conn.rollback() + return owned_user_playbook_ids = ( set(owned_user_playbook_ids) if owned_user_playbook_ids is not None diff --git a/tests/server/services/playbook/test_playbook_generation_service.py b/tests/server/services/playbook/test_playbook_generation_service.py index fcce32a2d..20833ea4f 100644 --- a/tests/server/services/playbook/test_playbook_generation_service.py +++ b/tests/server/services/playbook/test_playbook_generation_service.py @@ -560,6 +560,55 @@ def test_finalize_without_provenance_emits_create_with_null_model_fields(): ) +def test_synchronous_finalize_rolls_back_creation_when_lineage_fails(): + with tempfile.TemporaryDirectory() as temp_dir: + service = PlaybookGenerationService( + llm_client=LiteLLMClient(LiteLLMConfig(model="gpt-4o-mini")), + request_context=RequestContext(org_id="0", storage_base_dir=temp_dir), + ) + service.service_config = PlaybookGenerationServiceConfig( + request_id="atomic-request", + agent_version="1.0", + user_id="atomic-user", + source="test", + ) + playbook = UserPlaybook( + agent_version="1.0", + request_id="atomic-request", + content="Keep creation and lineage atomic.", + trigger="When synchronous finalization retries", + ) + + with ( + patch( + "reflexio.server.services.playbook.components.consolidator.PlaybookConsolidator", + ) as consolidator_class, + patch.object( + service, + "_apply_consolidation_lineage", + side_effect=RuntimeError("lineage failed"), + ), + ): + consolidator_class.return_value.deduplicate.return_value = ( + [playbook], + [], + [], + ) + consolidator_class.return_value.model_provenance = None + consolidator_class.return_value.consolidated_output_indices = set() + with pytest.raises(RuntimeError, match="lineage failed"): + service._finalize_extracted_items([playbook]) + + assert ( + _storage(service) + .conn.execute( + "SELECT count(*) FROM user_playbooks WHERE user_id = 'atomic-user'" + ) + .fetchone()[0] + == 0 + ) + + def test_resolve_write_plan_honors_reject_all_consolidation(): """An intentionally empty consolidation result must not restore candidates.""" with tempfile.TemporaryDirectory() as temp_dir: diff --git a/tests/server/services/profile/test_dedup_always_soft_integration.py b/tests/server/services/profile/test_dedup_always_soft_integration.py index e7afe98cf..5139fea4c 100644 --- a/tests/server/services/profile/test_dedup_always_soft_integration.py +++ b/tests/server/services/profile/test_dedup_always_soft_integration.py @@ -186,6 +186,48 @@ def test_dedup_removal_soft_supersedes_and_reconstructs(tmp_path) -> None: assert [p.profile_id for p in recon_log.added_profiles] == ["p_new_A"] +def test_synchronous_finalize_rolls_back_profile_creation_when_supersede_fails( + tmp_path, +) -> None: + org_id = "always-soft-org-atomic" + user_id = "u-atomic" + request_id = "manual-atomic-1" + storage = SQLiteStorage(org_id=org_id, db_path=str(tmp_path / "atomic.db")) + old_profile = _make_profile(user_id, "old-atomic", generated_from_request_id="seed") + storage.add_user_profile(user_id, [old_profile]) + new_profile = _make_profile( + user_id, + "new-atomic", + generated_from_request_id=request_id, + ) + service = _build_service( + storage, + org_id=org_id, + user_id=user_id, + request_id=request_id, + ) + mock_dedup, mock_dedup_cls = _patch_dedup( + all_new=[new_profile], + existing_ids=[old_profile.profile_id], + superseded=[old_profile], + ) + + with ( + mock_dedup_cls as cls, + patch.object( + storage, + "supersede_profiles_by_ids", + side_effect=RuntimeError("supersede failed"), + ), + ): + cls.return_value = mock_dedup + with pytest.raises(RuntimeError, match="supersede failed"): + service._finalize_extracted_items([new_profile]) + + assert storage.get_profile_by_id("new-atomic", include_tombstones=True) is None + assert storage.get_profile_by_id("old-atomic") is not None + + # =========================================================================== # B. Failure-path (mock storage): atomicity guard — no phantom removal # =========================================================================== 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 045cf394b..880cbad09 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -2,6 +2,9 @@ import json import sqlite3 +import threading +import time +from concurrent.futures import ThreadPoolExecutor from typing import Any, Literal, cast from unittest.mock import patch @@ -2572,6 +2575,138 @@ def test_fail_missing_purge_rolls_back_implicit_transaction(storage): _begin_purge(storage, "purge_after_missing_failure") +def test_record_purge_target_rolls_back_after_write_failure(storage, monkeypatch): + purge_id = _begin_purge(storage, "purge_target_write_failure") + + def _write_then_raise(**_kwargs: object) -> None: + storage.conn.execute( + "UPDATE purge_operations SET status = 'running' WHERE purge_id = ?", + (purge_id,), + ) + raise RuntimeError("target write failed") + + monkeypatch.setattr(storage, "_record_purge_target_locked", _write_then_raise) + + with pytest.raises(RuntimeError, match="target write failed"): + storage.record_purge_target( + purge_id=purge_id, + target_name="request", + target_ref="all", + phase="delete", + status="running", + ) + + assert storage.conn.in_transaction is False + + +def test_begin_purge_operation_serializes_idempotent_two_connection_retry( + storage_factory, +) -> None: + first = storage_factory("org1") + second = storage_factory("org1") + purge_id = "purge_two_connection_retry" + idempotency_key = "idem_two_connection_retry" + first.conn.execute("BEGIN IMMEDIATE") + first.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 (?, 'org1', 'user_erasure', 'user', ?, ?, ?, 'pending', 1, 1)""", + (purge_id, SUBJECT_REF, REQUEST_REF, idempotency_key), + ) + entered = threading.Event() + + def _trace(statement: str) -> None: + if ( + statement.startswith("BEGIN IMMEDIATE") + or "FROM purge_operations" in statement + ): + entered.set() + + second.conn.set_trace_callback(_trace) + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + second.begin_purge_operation, + purge_id, + idempotency_key, + "user_erasure", + "user", + SUBJECT_REF, + REQUEST_REF, + ) + assert entered.wait(timeout=1) + time.sleep(0.05) + first.conn.commit() + operation = future.result(timeout=2) + + assert operation.purge_id == purge_id + + +def test_prepare_targets_rechecks_snapshot_after_two_connection_write_lock( + storage_factory, +) -> None: + first = storage_factory("org1") + second = storage_factory("org1") + purge_id = "purge_two_connection_prepare" + first.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_two_connection_prepare", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + ) + first.conn.execute("BEGIN IMMEDIATE") + first._record_purge_target_locked( + purge_id=purge_id, + target_name="request", + target_ref="all", + phase="delete", + status="running", + detail={"count": 1}, + deleted_count=0, + error_detail=None, + ) + first._record_purge_target_locked( + purge_id=purge_id, + target_name="target_snapshot", + target_ref="all", + phase="prepare_targets", + status="complete", + detail={"owned_user_playbook_ids": []}, + deleted_count=0, + error_detail=None, + ) + entered = threading.Event() + + def _trace(statement: str) -> None: + if ( + statement.startswith("BEGIN IMMEDIATE") + or "purge_operation_targets" in statement + ): + entered.set() + + second.conn.set_trace_callback(_trace) + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + second.prepare_governance_erase_targets, + purge_id, + "two-connection-user", + set(), + ) + assert entered.wait(timeout=1) + time.sleep(0.05) + first.conn.commit() + future.result(timeout=2) + + request_target = next( + target + for target in second.list_purge_targets(purge_id) + if target.target_name == "request" and target.phase == "delete" + ) + assert request_target.status == "running" + + @pytest.mark.parametrize( "error_code", ["content_purge_failed", "prompt_redaction_route"] ) 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 24220d270..54a8d5d13 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 @@ -165,6 +165,44 @@ def test_legacy_schema_backfills_required_subject_ref( ) +def test_legacy_empty_subject_default_is_rebuilt_and_backfilled(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) + );""" + ) + 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)""" + ) + storage.conn.commit() + storage.conn.close() + + migrated = _storage(db_path) + + row = migrated.conn.execute( + "SELECT governance_subject_ref FROM session_outcomes" + ).fetchone() + assert row is not None + assert row["governance_subject_ref"] == migrated._subject_ref_for_user_id( + "legacy-user" + ) + + def test_sqlite_versions_without_returning_and_drop_column_are_rejected( tmp_path, monkeypatch: pytest.MonkeyPatch ) -> None: 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 85be9337b..b0a2de76f 100644 --- a/tests/server/services/storage/test_playbook_aggregation_state_integration.py +++ b/tests/server/services/storage/test_playbook_aggregation_state_integration.py @@ -1035,7 +1035,9 @@ def test_incremental_run_refreshes_agent_and_centroid_after_match( learning_meter = MagicMock() monkeypatch.setattr(aggregator, "_record_learnings_generated", learning_meter) - first = aggregator.run(PlaybookAggregatorRequest(agent_version="v1")) + first = aggregator.run( + PlaybookAggregatorRequest(agent_version="v1", operation_key="test-run-1") + ) assert first["playbooks_generated"] == 1 first_saved_id = store.conn.execute( "SELECT agent_playbook_id FROM agent_playbooks" @@ -1043,7 +1045,7 @@ def test_incremental_run_refreshes_agent_and_centroid_after_match( learning_meter.assert_called_once_with( learning_ids=[str(first_saved_id)], playbook_name="playbook", - request_id=learning_meter.call_args.kwargs["request_id"], + request_id="test-run-1", metadata=first, total_count=1, )