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..9165b746f 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, 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/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..4e0380d5e 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -86,18 +86,21 @@ 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. - 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 - 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. @@ -112,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: @@ -128,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, @@ -177,7 +181,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"``). @@ -226,30 +230,32 @@ 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. - 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. 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: @@ -270,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: @@ -298,20 +305,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/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 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/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/_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..2334d3741 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,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, - *, - finalization_run_id: str | None = None, - ) -> list[str] | None: + def _finalize_extracted_items(self, items: list) -> list: """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 + return items @abstractmethod def _should_track_in_progress(self) -> bool: @@ -643,6 +634,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 +643,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 +676,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 +708,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..c66b9daa4 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,13 @@ 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, + persisted_items = service._finalize_extracted_items( + items, model_provenance=model_provenance ) - if result.won_receipt: - self._record_finalized_learnings( - run, result.learning_ids, entity_type="profile" - ) - return result + self._record_finalized_learnings( + run, persisted_items or [], entity_type="profile" + ) + return if run.binding.extractor_kind == "playbook": service = PlaybookGenerationService( llm_client=self.client, @@ -910,46 +904,79 @@ 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, + persisted_items = service._finalize_extracted_items( + items, model_provenance=model_provenance ) - if result.won_receipt: - self._record_finalized_learnings( - run, result.learning_ids, entity_type="user_playbook" - ) - return result + self._record_finalized_learnings( + run, persisted_items or [], 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, request_id=run.binding.request_id, agent_version=run.binding.agent_version, entity_type=entity_type, + event_key=f"learn-batch:resumable:{run.id}:{entity_type}", metadata=metadata, ) diff --git a/reflexio/server/services/governance/service.py b/reflexio/server/services/governance/service.py index 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..305a72c57 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 @@ -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: @@ -1961,6 +1972,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 +2025,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..f3ba0cc67 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 ( @@ -30,11 +31,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 +597,27 @@ 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") + 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: try: logger.info("Trigger playbook aggregation") self._trigger_playbook_aggregation() except Exception: - logger.exception("Failed to schedule user playbook aggregation") + 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. @@ -634,97 +637,25 @@ 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. + ) -> list[UserPlaybook]: + """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 [] + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + 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 40cbd5fab..4b2bd1053 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,23 @@ def _finalize_extracted_items( all_new_profiles: list[UserProfile], *, model_provenance: ModelProvenance | None = None, - finalization_run_id: str | None = None, - ) -> list[str]: + ) -> list[UserProfile]: """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 None: + return [] + 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: """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..ea2f1b72e 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,7 +65,6 @@ logger = logging.getLogger(__name__) -_OUTCOME_ALLOWED_VALUES = tuple(kind.value for kind in SessionOutcomeKind) _MINIMUM_SQLITE_VERSION = (3, 35, 0) @@ -97,47 +87,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_]+") @@ -829,27 +778,13 @@ def migrate(self) -> bool: # _DDL creates an index over these columns. Upgrade legacy request tables # before executescript so index creation cannot fail on missing columns. self._migrate_request_retrieval_experiment() + self._migrate_session_outcomes_schema() with self._lock: - session_outcome_columns = { - row["name"] - for row in self.conn.execute( - "PRAGMA table_info(session_outcomes)" - ).fetchall() - } - if ( - session_outcome_columns - and "governance_subject_ref" not in session_outcome_columns - ): - self.conn.execute( - "ALTER TABLE session_outcomes " - "ADD COLUMN governance_subject_ref TEXT" - ) cur = self.conn.cursor() cur.executescript(_DDL) init_governance_tables(self.conn) init_playbook_aggregation_tables(self.conn) self.conn.commit() - self._migrate_session_outcomes_schema() if self._has_sqlite_vec: self._create_vec_tables() self._migrate_vec_tables() @@ -881,6 +816,120 @@ 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") + has_empty_subject_default = ( + governance_column is not None + and governance_column["dflt_value"] in ("''", '""') + ) + identity_columns = { + "outcome_id", + "outcome_revision", + "outcome_contract_digest", + "finalized_trajectory_digest", + } + if ( + expected_columns.issubset(columns) + and not identity_columns.intersection(columns) + and governance_column is not None + and int(governance_column["notnull"]) == 1 + and not has_empty_subject_default + and "'unknown'" not in table_sql + ): + return + + self.conn.execute("BEGIN IMMEDIATE") + try: + legacy_rows = self.conn.execute( + "SELECT * FROM session_outcomes" + ).fetchall() + self.conn.execute( + "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" + ) + self.conn.execute( + """CREATE TABLE session_outcomes ( + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), + occurred_at INTEGER NOT NULL, + source TEXT NOT NULL, + label TEXT, + value REAL, + metadata TEXT, + governance_subject_ref TEXT NOT NULL, + created_at INTEGER NOT NULL, + PRIMARY KEY (user_id, session_id) + )""" + ) + compatible_rows = [ + row + for row in legacy_rows + if row["outcome"] in ("success", "failure") + ] + for row in compatible_rows: + subject_ref = ( + row["governance_subject_ref"] + if "governance_subject_ref" in columns + else None + ) + if subject_ref is None or not str(subject_ref).strip(): + subject_ref = self._subject_ref_for_user_id(str(row["user_id"])) + self.conn.execute( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, + label, value, metadata, governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + row["user_id"], + row["session_id"], + row["outcome"], + row["occurred_at"], + row["source"], + row["label"] if "label" in columns else None, + row["value"] if "value" in columns else None, + row["metadata"] if "metadata" in columns else None, + subject_ref, + row["created_at"], + ), + ) + dropped_unknown = len(legacy_rows) - len(compatible_rows) + if dropped_unknown: + logger.warning( + "Dropped %d unrepresentable 'unknown' session outcomes while " + "restoring the pre-#407 schema", + dropped_unknown, + ) + self.conn.execute("DROP TABLE session_outcomes_legacy") + self.conn.commit() + except Exception: + self.conn.rollback() + raise + def _try_load_sqlite_vec(self) -> bool: """Attempt to load the sqlite-vec extension for native KNN search. @@ -984,139 +1033,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 +2065,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 +2079,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 +2091,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 +2106,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 +2115,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() @@ -2473,8 +2381,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``, - ``session_outcomes``, ``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. @@ -2501,6 +2409,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] = { @@ -2551,8 +2460,9 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: "DELETE FROM interactions WHERE user_id = ?", (user_id,) ) session_outcomes_cur = self.conn.execute( - "DELETE FROM session_outcomes WHERE user_id = ?", - (user_id,), + """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,) @@ -2681,18 +2591,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 +2977,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..5e4abad09 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,11 @@ 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 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 d0657f4eb..aaec5a1b2 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,14 @@ 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 user_id = ? OR governance_subject_ref = ?""", + (user_id, 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..954eae6c1 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,11 @@ 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 user_id = ? OR governance_subject_ref = ?""", + (user_id, subject_ref), ) interaction_ids = [ int(row["interaction_id"]) @@ -350,11 +346,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 = { @@ -378,10 +370,6 @@ def apply_governance_user_data_delete( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") - self._assert_purge_operation_execution_claim_locked( - purge_id, execution_claim - ) - self._assert_authoritative_user_identity_locked(purge_id, user_id) self._validate_prepared_delete_target_matrix_locked(purge_id) self._validate_hide_for_rebuild_targets_locked(purge_id) expected_user_playbook_ids = ( @@ -416,12 +404,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 +420,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 +452,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 +531,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..1f2a7f0e9 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,20 +187,6 @@ 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: @@ -346,26 +211,13 @@ def begin_purge_operation( "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: - raise ValueError( - "Existing purge operation has mismatched authoritative user identity" - ) - self.conn.commit() - return _row_to_purge_operation(existing) + 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, authoritative_user_digest, - status, created_at, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)""", + request_ref, idempotency_key, status, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)""", ( validated_purge_id, self.org_id, @@ -374,7 +226,6 @@ def begin_purge_operation( subject_ref, request_ref, validated_idempotency_key, - authoritative_user_digest, now, now, ), @@ -385,177 +236,12 @@ def begin_purge_operation( raise return self.get_purge_operation(validated_purge_id) - def claim_purge_operation_execution( - self, - purge_id: str, - *, - lease_owner: str, - lease_ttl_seconds: int, - ) -> PurgeExecutionClaim | None: - validated_purge_id = _validate_governance_purge_id("purge_id", purge_id) - if not lease_owner.strip(): - raise ValueError("lease_owner is required") - if lease_ttl_seconds <= 0: - raise ValueError("lease_ttl_seconds must be positive") - now = _epoch_now() - expires_at = now + lease_ttl_seconds - with self._lock: - try: - self.conn.execute("BEGIN IMMEDIATE") - cursor = self.conn.execute( - """UPDATE purge_operations - SET status = 'running', error_code = NULL, error_detail = NULL, - completed_at = NULL, updated_at = ?, - execution_claim_owner = ?, - execution_claim_fence = execution_claim_fence + 1, - execution_claim_expires_at = ? - WHERE purge_id = ? AND org_id = ? - AND ( - status IN ('pending', 'failed') - OR ( - status = 'running' - AND ( - execution_claim_expires_at IS NULL - OR execution_claim_expires_at <= ? - ) - ) - ) - RETURNING execution_claim_owner, - execution_claim_fence, - execution_claim_expires_at""", - ( - now, - lease_owner, - expires_at, - validated_purge_id, - self.org_id, - now, - ), - ) - row = cursor.fetchone() - self.conn.commit() - if row is None: - return None - return PurgeExecutionClaim( - purge_id=validated_purge_id, - owner=str(row["execution_claim_owner"]), - fence=int(row["execution_claim_fence"]), - expires_at=int(row["execution_claim_expires_at"]), - ) - except Exception: - self.conn.rollback() - raise - - def assert_purge_operation_execution_claim( - self, purge_id: str, execution_claim: PurgeExecutionClaim - ) -> None: - purge_id = _validate_governance_purge_id("purge_id", purge_id) - claim = validate_purge_execution_claim(purge_id, execution_claim) - now = _epoch_now() - row = self._deps()._fetchone( - """SELECT status, execution_claim_owner, execution_claim_fence, - execution_claim_expires_at - FROM purge_operations - WHERE purge_id = ? AND org_id = ?""", - (purge_id, self.org_id), - ) - if row is None: - raise ValueError(f"Purge operation {purge_id!r} not found") - if ( - row["status"] != "running" - or row["execution_claim_owner"] != claim.owner - or int(row["execution_claim_fence"]) != claim.fence - or row["execution_claim_expires_at"] is None - or int(row["execution_claim_expires_at"]) <= now - ): - raise ValueError("purge execution claim is no longer active") - - def _assert_purge_operation_execution_claim_locked( - self, - purge_id: str, - execution_claim: PurgeExecutionClaim, - ) -> None: - purge_id = _validate_governance_purge_id("purge_id", purge_id) - claim = validate_purge_execution_claim(purge_id, execution_claim) - now = _epoch_now() - row = self.conn.execute( - """SELECT status, execution_claim_owner, execution_claim_fence, - execution_claim_expires_at - FROM purge_operations - WHERE purge_id = ? AND org_id = ?""", - (purge_id, self.org_id), - ).fetchone() - if row is None: - raise ValueError(f"Purge operation {purge_id!r} not found") - if ( - row["status"] != "running" - or row["execution_claim_owner"] != claim.owner - or int(row["execution_claim_fence"]) != claim.fence - or row["execution_claim_expires_at"] is None - or int(row["execution_claim_expires_at"]) <= now - ): - raise ValueError("purge execution claim is no longer active") - - def renew_purge_operation_execution_claim( - self, - purge_id: str, - execution_claim: PurgeExecutionClaim, - *, - lease_ttl_seconds: int, - ) -> PurgeExecutionClaim: - purge_id = _validate_governance_purge_id("purge_id", purge_id) - claim = validate_purge_execution_claim(purge_id, execution_claim) - if lease_ttl_seconds <= 0: - raise ValueError("lease_ttl_seconds must be positive") - now = _epoch_now() - expires_at = now + lease_ttl_seconds - with self._lock: - try: - self.conn.execute("BEGIN IMMEDIATE") - cursor = self.conn.execute( - """UPDATE purge_operations - SET execution_claim_expires_at = ?, updated_at = ? - WHERE purge_id = ? AND org_id = ? - AND status = 'running' - AND execution_claim_owner = ? - AND execution_claim_fence = ? - AND execution_claim_expires_at IS NOT NULL - AND execution_claim_expires_at > ? - RETURNING execution_claim_owner, - execution_claim_fence, - execution_claim_expires_at""", - ( - expires_at, - now, - purge_id, - self.org_id, - claim.owner, - claim.fence, - now, - ), - ) - row = cursor.fetchone() - self.conn.commit() - except Exception: - self.conn.rollback() - raise - if row is None: - raise ValueError("purge execution claim is no longer active") - return PurgeExecutionClaim( - purge_id=purge_id, - owner=str(row["execution_claim_owner"]), - fence=int(row["execution_claim_fence"]), - expires_at=int(row["execution_claim_expires_at"]), - ) - def record_purge_target( self, purge_id: str, target_name: str, phase: str, status: Literal["pending", "running", "failed", "complete"], - *, - execution_claim: PurgeExecutionClaim, target_ref: str = "", detail: dict[str, object] | None = None, deleted_count: int = 0, @@ -579,10 +265,6 @@ def record_purge_target( ) with self._lock: try: - self.conn.execute("BEGIN IMMEDIATE") - self._assert_purge_operation_execution_claim_locked( - purge_id, execution_claim - ) self._record_purge_target_locked( purge_id=purge_id, target_name=target_name, @@ -626,28 +308,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: 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() + if self.purge_targets_prepared(purge_id): + self.conn.rollback() return owned_user_playbook_ids = ( set(owned_user_playbook_ids) @@ -676,7 +344,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,12 +355,7 @@ 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) @@ -702,15 +364,10 @@ def fail_purge_operation( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") - self._assert_purge_operation_execution_claim_locked( - purge_id, execution_claim - ) cur = self.conn.execute( """UPDATE purge_operations SET status = 'failed', error_code = ?, error_detail = ?, - updated_at = ?, completed_at = ?, - execution_claim_owner = NULL, - execution_claim_expires_at = NULL + updated_at = ?, completed_at = ? WHERE purge_id = ? AND org_id = ? AND status != 'complete'""", ( validated_error_code, diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py b/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py index 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..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 session outcomes, interactions, user playbooks, - profiles, and requests. Intentionally does NOT touch ``agent_playbooks`` — those - are the cross-project rollup of skills and have no ``user_id`` - column. This is the data-isolation primitive used 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,8 +176,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``, - ``session_outcomes``, ``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. 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..be2880686 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) @@ -487,37 +337,57 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( assert storage.list_run_tool_dependencies("run_1")[0].consumed_at is not None -def test_resume_worker_tagging_schedule_failure_is_best_effort( - request_context, +@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="run_tagging", + id="survivor-run", binding=AgentBinding( org_id="org_1", - extractor_kind="profile", + extractor_kind=extractor_kind, user_id="user_1", request_id="request_1", agent_version="v1", source="api", ), - status=AgentRunStatus.FINALIZED_PENDING_TOOL, - generation_request_snapshot={"request_id": "request_1"}, + status=AgentRunStatus.FINALIZING, + generation_request_snapshot={}, ) + dropped = object() + survivor = object() worker = ExtractionResumeWorker(request_context=request_context) - with patch( - "reflexio.server.services.extraction.resume_worker.schedule_tagging", - side_effect=RuntimeError("scheduler unavailable"), + with ( + patch(finalize_path, return_value=[survivor]), + patch.object(worker, "_record_finalized_learnings") as record, ): - worker._schedule_finalized_tagging(run) + worker._finalize_items(run, [dropped, survivor]) + + record.assert_called_once_with(run, [survivor], entity_type=entity_type) -def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( +def test_resume_worker_tagging_schedule_failure_is_best_effort( request_context, ): - """A mixed batch charges its persisted profile once across finalization retries.""" run = AgentRunRecord( - id="run_mixed_billing", + id="run_tagging", binding=AgentBinding( org_id="org_1", extractor_kind="profile", @@ -526,838 +396,16 @@ def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( agent_version="v1", source="api", ), - status=AgentRunStatus.FINALIZATION_FAILED, + status=AgentRunStatus.FINALIZED_PENDING_TOOL, 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) + with patch( + "reflexio.server.services.extraction.resume_worker.schedule_tagging", + side_effect=RuntimeError("scheduler unavailable"), + ): + worker._schedule_finalized_tagging(run) def test_resume_worker_fails_run_when_step_budget_exhausted( 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/playbook/test_playbook_generation_service.py b/tests/server/services/playbook/test_playbook_generation_service.py index f8c9fa12f..20833ea4f 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, @@ -535,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_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..880cbad09 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,10 @@ 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 +import time +from concurrent.futures import ThreadPoolExecutor from typing import Any, Literal, cast from unittest.mock import patch @@ -27,8 +24,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 +37,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 +68,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 +98,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 +113,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 +346,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 +487,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 +502,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 +517,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 +548,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) @@ -1026,6 +563,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,8 +619,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 +648,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 +671,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 +681,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 +698,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 +714,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 +729,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 +787,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) == [] @@ -1262,772 +812,143 @@ def test_complete_purge_operation_requires_matching_existing_erase_row( storage, seed_event, match ): purge_id = _begin_completeable_purge(storage, "purge_seeded") - assert storage.append_audit_event(seed_event) is True - - 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), - ) - - assert storage.get_purge_operation(purge_id).status == "running" - rows = storage.list_audit_events(subject_ref=SUBJECT_REF) - assert len(rows) == 1 - assert rows[0].operation == seed_event.operation - assert rows[0].status == seed_event.status - - -@pytest.mark.parametrize( - ("field_name", "seed_kwargs"), - [ - pytest.param("entity_type", {"entity_type": "session"}, id="entity-type"), - pytest.param( - "subject_ref", {"subject_ref": OTHER_SUBJECT_REF}, id="subject-ref" - ), - pytest.param( - "request_ref", {"request_ref": OTHER_REQUEST_REF}, id="request-ref" - ), - pytest.param("actor_type", {"actor_type": "jwt"}, id="actor-type"), - pytest.param("actor_ref", {"actor_ref": ACTOR_REF}, id="actor-ref"), - pytest.param("entity_id", {"entity_id": "17"}, id="entity-id"), - pytest.param("detail", {"detail": {"count": 2}}, id="detail"), - ], -) -def test_complete_purge_operation_rejects_mismatched_existing_erase_row( - storage, field_name, seed_kwargs -): - purge_id = _begin_completeable_purge(storage, "purge_seeded_mismatch") - seeded_event = _erase_event(purge_id=purge_id).model_copy(update=seed_kwargs) - storage.conn.execute( - """INSERT INTO audit_events ( - org_id, actor_type, actor_ref, operation, entity_type, entity_id, - subject_ref, request_ref, idempotency_key, status, detail, created_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - seeded_event.org_id, - seeded_event.actor_type, - seeded_event.actor_ref, - seeded_event.operation, - seeded_event.entity_type, - seeded_event.entity_id, - seeded_event.subject_ref, - seeded_event.request_ref, - seeded_event.idempotency_key, - seeded_event.status, - json.dumps(seeded_event.detail) - if seeded_event.detail is not None - else None, - seeded_event.created_at, - ), - ) - storage.conn.commit() - - 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), - ) - - assert storage.get_purge_operation(purge_id).status == "running" - rows = storage.list_audit_events(subject_ref=seeded_event.subject_ref) - assert len(rows) == 1 - assert getattr(rows[0], field_name) == getattr(seeded_event, field_name) - - -def test_append_audit_event_rejects_successful_erase(storage): - with pytest.raises(ValueError, match="Successful ERASE audit rows"): - storage.append_audit_event(_erase_event(purge_id="purge_append")) - - -def test_append_audit_event_rejects_successful_erase_without_idempotency_key(storage): - with pytest.raises(ValueError, match="Successful ERASE audit rows"): - storage.append_audit_event( - AuditEvent( - org_id="org1", - operation="ERASE", - entity_type="request", - subject_ref=SUBJECT_REF, - request_ref=REQUEST_REF, - idempotency_key=None, - status="ok", - ) - ) - - -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), - ) - - 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" - assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] - - -def test_complete_retry_replaces_failed_completed_at(storage): - purge_id = _begin_completeable_purge(storage, "purge_retry_completion_time") - with patch.object(purge_module, "_epoch_now", return_value=111): - failed = storage.fail_purge_operation( - purge_id, - error_code="governance_erase_failed", - error_detail="RuntimeError", - execution_claim=_claim_purge(storage, purge_id), - ) - assert failed.completed_at == 111 - - with patch.object(erase_execution_module, "_epoch_now", return_value=222): - 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) + assert storage.append_audit_event(seed_event) is True - with pytest.raises(ValueError, match="purge execution claim"): - storage.begin_subject_erasure_barrier( - SUBJECT_REF, - purge_id, - execution_claim=first_claim, + with pytest.raises(ValueError, match=match): + storage.complete_purge_operation_with_audit( + purge_id, _erase_event(purge_id=purge_id) ) - 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" + assert storage.get_purge_operation(purge_id).status == "running" + rows = storage.list_audit_events(subject_ref=SUBJECT_REF) + assert len(rows) == 1 + assert rows[0].operation == seed_event.operation + assert rows[0].status == seed_event.status -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) +@pytest.mark.parametrize( + ("field_name", "seed_kwargs"), + [ + pytest.param("entity_type", {"entity_type": "session"}, id="entity-type"), + pytest.param( + "subject_ref", {"subject_ref": OTHER_SUBJECT_REF}, id="subject-ref" + ), + pytest.param( + "request_ref", {"request_ref": OTHER_REQUEST_REF}, id="request-ref" + ), + pytest.param("actor_type", {"actor_type": "jwt"}, id="actor-type"), + pytest.param("actor_ref", {"actor_ref": ACTOR_REF}, id="actor-ref"), + pytest.param("entity_id", {"entity_id": "17"}, id="entity-id"), + pytest.param("detail", {"detail": {"count": 2}}, id="detail"), + ], +) +def test_complete_purge_operation_rejects_mismatched_existing_erase_row( + storage, field_name, seed_kwargs +): + purge_id = _begin_completeable_purge(storage, "purge_seeded_mismatch") + seeded_event = _erase_event(purge_id=purge_id).model_copy(update=seed_kwargs) + storage.conn.execute( + """INSERT INTO audit_events ( + org_id, actor_type, actor_ref, operation, entity_type, entity_id, + subject_ref, request_ref, idempotency_key, status, detail, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + seeded_event.org_id, + seeded_event.actor_type, + seeded_event.actor_ref, + seeded_event.operation, + seeded_event.entity_type, + seeded_event.entity_id, + seeded_event.subject_ref, + seeded_event.request_ref, + seeded_event.idempotency_key, + seeded_event.status, + json.dumps(seeded_event.detail) + if seeded_event.detail is not None + else None, + seeded_event.created_at, + ), + ) + storage.conn.commit() - 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, + with pytest.raises(ValueError, match="matching successful ERASE"): + storage.complete_purge_operation_with_audit( + purge_id, _erase_event(purge_id=purge_id) ) - 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) + assert storage.get_purge_operation(purge_id).status == "running" + rows = storage.list_audit_events(subject_ref=seeded_event.subject_ref) + assert len(rows) == 1 + assert getattr(rows[0], field_name) == getattr(seeded_event, field_name) -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, - ) +def test_append_audit_event_rejects_successful_erase(storage): + with pytest.raises(ValueError, match="Successful ERASE audit rows"): + storage.append_audit_event(_erase_event(purge_id="purge_append")) - 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_append_audit_event_rejects_successful_erase_without_idempotency_key(storage): + with pytest.raises(ValueError, match="Successful ERASE audit rows"): + storage.append_audit_event( + AuditEvent( + org_id="org1", + operation="ERASE", + entity_type="request", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + idempotency_key=None, + status="ok", + ) + ) -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) +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) - with pytest.raises(ValueError, match="purge execution claim"): - storage.complete_subject_erasure_barrier_after_empty_check( + 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=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, +def test_complete_retry_replaces_failed_completed_at(storage): + purge_id = _begin_completeable_purge(storage, "purge_retry_completion_time") + with patch.object(purge_module, "_epoch_now", return_value=111): + failed = storage.fail_purge_operation( + 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" + assert failed.completed_at == 111 - 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( + with patch.object(erase_execution_module, "_epoch_now", return_value=222): + completed = storage.complete_purge_operation_with_audit( purge_id, - error_code="governance_erase_failed", - error_detail="RuntimeError", - execution_claim=first_claim, + _erase_event(purge_id=purge_id), ) - 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" + assert completed.status == "complete" + assert completed.completed_at == 222 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 +958,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 +991,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 +1006,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 +1026,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 +1049,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 = { @@ -2199,15 +1099,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 +1161,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 +1177,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 +1200,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 +1261,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 +1324,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 +1340,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 +1419,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 +1485,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 +1501,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 +1524,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 +1598,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 +1613,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 +1633,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 +1699,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 +1715,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 +1738,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 +1765,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 +1877,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 +1969,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 +2080,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 +2112,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 +2144,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 +2154,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 +2162,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 +2170,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 +2197,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 +2293,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 +2331,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 +2348,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 +2370,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 +2428,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 +2464,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 +2481,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 +2532,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 +2545,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,13 +2557,156 @@ 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" 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") + + +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"] ) @@ -3729,7 +2719,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 +2733,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 +2755,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 +2765,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 +2903,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 +2928,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 +2949,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 +2990,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 +3025,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 +3046,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 +3124,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 +3151,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 +3161,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 +3195,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 +3208,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 +3225,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 +3239,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 +3247,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 +3269,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 +3278,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 +3289,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 +3297,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 +3322,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 +3338,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 +3368,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 +3400,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,14 +3440,12 @@ 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 == { @@ -4576,15 +3526,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 +3550,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 +3605,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 +3676,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 +3688,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 +3710,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 +3727,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 +3867,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 +3919,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 +4024,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 +4034,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 +4070,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 +4105,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 +4232,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 +4242,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 +4302,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 +4316,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 +4522,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_downgrade_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py new file mode 100644 index 000000000..54a8d5d13 --- /dev/null +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py @@ -0,0 +1,212 @@ +"""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_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: + 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/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_playbook_aggregation_state_integration.py b/tests/server/services/storage/test_playbook_aggregation_state_integration.py index 3ebcd1e3e..b0a2de76f 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,23 @@ 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")) + 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" + ).fetchone()[0] + learning_meter.assert_called_once_with( + learning_ids=[str(first_saved_id)], + playbook_name="playbook", + request_id="test-run-1", + metadata=first, + total_count=1, + ) assert store.conn.execute("SELECT count(*) FROM agent_playbooks").fetchone()[0] == 1 _insert_current(store, 3, embedding=encoded, trigger="same trigger") diff --git a/tests/server/services/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..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, @@ -20,8 +24,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 +170,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. @@ -251,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_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..b07a2223b 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( @@ -463,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_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..3536b39cc 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,46 @@ 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]) + worker._record_finalized_learnings( + run, + [object(), object()], + entity_type="profile", + ) + + 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:") + - finalizer.assert_called_once_with( - [dropped_candidate], model_provenance=None, finalization_run_id=run.id +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(), ) - assert events == [] + run = _agent_run(extractor_kind="profile") + + worker._record_finalized_learnings(run, [object()], entity_type="profile") + worker._record_finalized_learnings(run, [object()], entity_type="profile") + + assert [event.event_key for event in events] == [ + "learn-batch:resumable:run-1:profile", + "learn-batch:resumable:run-1:profile", + ] def test_resumable_finalization_emits_one_event_per_profile_id() -> None: - """Finalization survivor IDs emit one entity-backed event per profile.""" + """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 +108,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 +134,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 +143,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 +166,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 +179,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_records_attributed_learnings_generated() -> None: + """Aggregation emits one entity-backed event per generated playbook. -def test_aggregation_emits_no_learnings_generated( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A completed aggregation remains observable but adds no billable learning.""" + ``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..b9ab61401 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 @@ -125,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(