diff --git a/AI_AGENT_INTEGRATION.md b/AI_AGENT_INTEGRATION.md index 4807eb2aa..5952c0a85 100644 --- a/AI_AGENT_INTEGRATION.md +++ b/AI_AGENT_INTEGRATION.md @@ -95,7 +95,7 @@ learns, what stays private to a user scope, and what transfers to other users. | `user_id` | Scope for profiles and user playbooks | Use the human user, tenant, workspace, repo, or project whose preferences should be isolated. For example, use a project id when repo-specific rules should not leak into unrelated repos. | | `agent_version` | Scope for shared agent playbooks | Use a stable agent name plus major behavior version, for example `my-agent-v1`. Keep it stable if learnings should transfer across users/projects. If you omit it, the SDK uses `DEFAULT_AGENT_VERSION` (`"agent-v0"`) — fine for a single agent, but set an explicit value before you run more than one. | | `session_id` | Group turns for one conversation | Use the host session/conversation id. Generate a UUID if the host does not provide one. | -| `source` | Audit label | Use the integration name, for example `my-agent-plugin`. | +| `source` | Producer/workflow label | Use a non-sensitive machine label such as `support-agent:v2`. Where the API permits an empty value, empty means the source is absent. Every non-empty value must match `^[a-z0-9][a-z0-9._:-]{0,127}$` and is limited to 128 ASCII characters. Do not include user identifiers, email addresses, or other PII. | `user_id` and `agent_version` work together: @@ -234,7 +234,7 @@ Publish request fields: | --- | --- | --- | --- | | `user_id` | Yes | The user, tenant, workspace, repo, or project scope whose private profiles and user playbooks should be isolated. | `"alice"`, `"tenant-acme"`, `"repo-reflexio"` | | `interactions` / `interaction_data_list` | Yes | Ordered conversation turns to publish. Include at least one turn; multi-turn correction examples are best for learning. | `[{"role": "User", "content": "Use pnpm here."}, {"role": "Assistant", "content": "Got it, I will use pnpm."}]` | -| `source` | No | Integration label for debugging and filtering. Use a stable name for the plugin, framework, or adapter. | `"my-agent-plugin"`, `"vscode-assistant"`, `"support-chatbot"` | +| `source` | No | Non-sensitive producer/workflow machine label for debugging and filtering. Empty means the source is absent where the API permits it. Every non-empty value must match `^[a-z0-9][a-z0-9._:-]{0,127}$` and is limited to 128 ASCII characters. Do not include user identifiers, email addresses, or other PII. | `"support-agent:v2"` | | `agent_version` | Strongly recommended | The shared-agent learning boundary. Use the same value when playbooks should transfer across users. Change it when old playbooks should not transfer. | `"support-agent-v1"`, `"coding-agent-2026-05"` | | `session_id` | Yes | Host conversation/session id. Generate a UUID if the host has no session id, and reuse it for all turns in that conversation. | `"sess_01HX8Y..."`, `"3f02b7f8-..."` | | `skip_aggregation` | No | `False` when user playbooks should be eligible to roll up into shared agent playbooks. `True` when you want user-level extraction only. | `false` | @@ -315,7 +315,7 @@ def publish_turns( response = client.publish_interaction( user_id=user_id, interactions=interactions, - source="my-agent-plugin", + source="support-agent:v2", agent_version=agent_version, session_id=session_id, wait_for_response=False, diff --git a/README.md b/README.md index a7b60af7b..097e52898 100644 --- a/README.md +++ b/README.md @@ -104,8 +104,15 @@ 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 4efbe928e..9b7be37e7 100644 --- a/docs/lib/methods/requests-sessions.ts +++ b/docs/lib/methods/requests-sessions.ts @@ -7,7 +7,7 @@ export const requestSessionMethods: MethodDef[] = [ displayName: "Mark Session Outcome", group: "requests-sessions", description: - "Record the lifetime-durable first success or failure marker for an existing session. User and source are derived from the first request.", + "Record the lifetime-durable first terminal outcome for an existing session. Idempotency requires unchanged payload, governance context, and finalized trajectory/session context; changed values conflict. User and source are derived from the first request.", httpMethod: "POST", endpoint: "/api/session_outcome", requestStyle: "json_body", @@ -17,8 +17,8 @@ export const requestSessionMethods: MethodDef[] = [ name: "outcome", type: "enum", required: true, - description: "Terminal outcome", - enumValues: ["success", "failure"], + description: "Terminal outcome: success, failure, or unknown", + enumValues: ["success", "failure", "unknown"], }, { name: "occurred_at", type: "number", required: true, description: "Unix epoch seconds when the outcome occurred" }, { name: "label", type: "string", required: false, description: "Optional label, at most 128 characters" }, @@ -43,8 +43,8 @@ export const requestSessionMethods: MethodDef[] = [ name: "outcome", type: "enum", required: false, - description: "Exact outcome filter", - enumValues: ["success", "failure"], + description: "Exact outcome filter: success, failure, or unknown", + enumValues: ["success", "failure", "unknown"], }, { name: "label", type: "string", required: false, description: "Exact label filter" }, { name: "start_time", type: "number", required: false, description: "Inclusive minimum event time" }, diff --git a/docs/lib/methods/unified-search.ts b/docs/lib/methods/unified-search.ts index 76e2004b9..b63470725 100644 --- a/docs/lib/methods/unified-search.ts +++ b/docs/lib/methods/unified-search.ts @@ -23,7 +23,7 @@ export const unifiedSearchMethods: MethodDef[] = [ type: "number", required: false, default: 5, - description: "Maximum results per entity type", + description: "Maximum results per entity type, from 1 to 100", }, { name: "threshold", @@ -50,7 +50,8 @@ export const unifiedSearchMethods: MethodDef[] = [ name: "user_id", type: "string", required: false, - description: "Filter by user ID (profiles, user_playbooks)", + description: + "Filter by user ID (profiles, user_playbooks), at most 255 characters", }, { name: "entity_types", @@ -90,12 +91,26 @@ export const unifiedSearchMethods: MethodDef[] = [ "Search mode: vector (embedding similarity), fts (full-text search), or hybrid (combined with RRF)", enumValues: ["vector", "fts", "hybrid"], }, + { + name: "request_id", + type: "string", + required: false, + description: + "Caller correlation ID for the search turn, at most 255 characters", + }, { name: "session_id", type: "string", required: false, description: - "Agent session this search serves. When set, results already returned to the same session are skipped and next-best matches backfilled; searches without it neither read nor record session dedup state", + "Agent session this search serves, at most 255 characters. When set, results already returned to the same session are skipped and next-best matches backfilled; searches without it neither read nor record session dedup state", + }, + { + name: "interaction_id", + type: "number", + required: false, + description: + "Caller interaction ID for the search turn; must be a positive integer (minimum 1)", }, ], }, diff --git a/docs/lib/methods/user-playbooks.ts b/docs/lib/methods/user-playbooks.ts index 36333c5f5..ef1d55e40 100644 --- a/docs/lib/methods/user-playbooks.ts +++ b/docs/lib/methods/user-playbooks.ts @@ -105,7 +105,7 @@ export const userPlaybookMethods: MethodDef[] = [ type: "number", required: false, default: 10, - description: "Maximum number of results to return", + description: "Maximum number of results to return, from 1 to 100", }, { name: "threshold", diff --git a/reflexio/cli/README.md b/reflexio/cli/README.md index c0680c50a..eb420b8b8 100644 --- a/reflexio/cli/README.md +++ b/reflexio/cli/README.md @@ -144,7 +144,7 @@ Apply to all three modes: | `--wait` | Block until server-side extraction finishes (returns real counts). | | `--session-id` | Required unless each payload includes `session_id`; group publishes into one session. | | `--agent-version` | Tag the interaction with an agent version (used by playbook filtering). | -| `--source` | Free-form source tag (defaults to `cli`). | +| `--source` | Non-sensitive producer/workflow label matching `^[a-z0-9][a-z0-9._:-]{0,127}$`; never use user identifiers or PII (defaults to `cli`). | | `--skip-aggregation`| Extract profiles/playbooks but skip playbook aggregation. | | `--force-extraction`| Bypass all extraction gates (`stride_size`, cheap pre-filter, LLM `should_run`) and always run extractors. | | `--evaluation-only` | Store the request for session-level evaluation only; requires `--session-id` and skips profile/playbook extraction. | diff --git a/reflexio/client/client.py b/reflexio/client/client.py index af2361b63..a6b0d3920 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -449,7 +449,9 @@ def publish_interaction( Args: user_id: The user ID. interactions: List of interaction data. - source: The source of the interaction. + source: Non-sensitive producer/workflow label. A non-empty value + must match ``^[a-z0-9][a-z0-9._:-]{0,127}$`` and must not + contain user identifiers or PII. agent_version: The agent version. session_id: Required non-empty session ID for grouping requests. wait_for_response: If True, the **server** waits for @@ -808,7 +810,7 @@ def search_user_playbooks( end_time (Optional[datetime]): End time for created_at filter status_filter (Optional[list[Optional[Status]]]): Filter by status (None for CURRENT, PENDING, ARCHIVED) tags (Optional[list[str]]): Match playbooks having any of these tags. - top_k (Optional[int]): Maximum number of results to return (default: 10) + top_k (Optional[int]): Maximum results to return, from 1 to 100 (default: 10) threshold (Optional[float]): Similarity threshold for vector search. When omitted, the embedding model's default is used. enable_reformulation (Optional[bool]): Enable LLM query reformulation (default: False) @@ -1112,13 +1114,20 @@ def mark_session_outcome( value: float | None = None, metadata: dict[str, Any] | None = None, ) -> SetSessionOutcomeResponse: - """Record the first terminal outcome for a published session. + """Record the immutable first outcome for a published session. The session must already contain at least one published request. Reflexio derives both ``user_id`` and ``source`` from the earliest request ordered - by ``(created_at, request_id)``. Only the first outcome is recorded; - retries return ``success=True`` and ``recorded=False``. Sessions are not - required to report an outcome. + by ``(created_at, request_id)``. New canonical rows bind the outcome to + the server-owned outcome contract and canonical finalized trajectory. An + exact canonical retry must match the payload, contract, and trajectory; + otherwise it is rejected with ``reason="conflicting_finalization"``. + Rolling-upgrade rows with all four identity fields null compare the + caller payload and any available server-derived session context, but + cannot compare absent contract or trajectory digests. An accepted retry + preserves all four null identity fields and returns ``success=True`` and + ``recorded=False``. Sessions may report ``success``, ``failure``, or + ``unknown`` and are not required to report an outcome. """ request = SetSessionOutcomeRequest( session_id=session_id, @@ -2724,12 +2733,14 @@ def search( Args: request (Optional[UnifiedSearchRequest]): The search request object (alternative to kwargs) query (str): Search query text - top_k (Optional[int]): Maximum results per entity type (default: 5) + top_k (Optional[int]): Maximum results per entity type, from 1 to 100 + (default: 5). threshold (Optional[float]): Similarity threshold for vector search. When omitted, the embedding model's default is used. agent_version (Optional[str]): Filter by agent version (agent_playbooks, user_playbooks) playbook_name (Optional[str]): Filter by playbook name (agent_playbooks, user_playbooks) - user_id (Optional[str]): Filter by user ID (profiles, user_playbooks) + user_id (Optional[str]): Filter by user ID (profiles, user_playbooks), + at most 255 characters. tags (Optional[list[str]]): Match entities having any requested tag. entity_types (Optional[list[str]]): Entity types to search. Valid values: "profiles", "user_playbooks", "agent_playbooks". @@ -2745,9 +2756,12 @@ def search( the configured search backend supports it (default: False). conversation_history (Optional[list[ConversationTurn] | list[dict]]): Prior conversation turns for context-aware query reformulation. Accepts ConversationTurn objects or dicts with "role" and "content" keys. search_mode (Optional[SearchMode | str]): Search mode to use. Accepts SearchMode enum or string value ("vector", "fts", "hybrid"). - request_id (Optional[str]): Caller correlation id for the search turn. - session_id (Optional[str]): Caller session id for the search turn. - interaction_id (Optional[int]): Caller interaction id for the search turn. + request_id (Optional[str]): Caller correlation ID for the search turn, + at most 255 characters. + session_id (Optional[str]): Caller session ID for the search turn, + at most 255 characters. Also enables session-scoped result deduplication. + interaction_id (Optional[int]): Caller interaction ID for the search + turn; must be a positive integer (minimum 1). Returns: UnifiedSearchViewResponse: Combined search results from all entity types diff --git a/reflexio/lib/_session_outcome.py b/reflexio/lib/_session_outcome.py index d49cbcead..7a39ba518 100644 --- a/reflexio/lib/_session_outcome.py +++ b/reflexio/lib/_session_outcome.py @@ -53,51 +53,44 @@ def mark_session_outcome( try: for _attempt in range(3): context = storage.get_session_outcome_context(request.session_id) - if context.existing: - return SetSessionOutcomeResponse( - success=True, - recorded=False, - user_id=context.user_id, - source=context.source, - message="Outcome already exists", - ) - if context.user_id is None or context.first_request_at is None: - return SetSessionOutcomeResponse( - success=False, - reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, - message="Session has no published requests", - ) - if context.user_contract_violation: - logger.warning( - "Session outcome contract violation: multiple users for session %s", - sanitise_for_log(request.session_id), - ) - if context.source_contract_violation: - logger.warning( - "Session outcome contract violation: multiple sources for session %s", - sanitise_for_log(request.session_id), - ) - if request.occurred_at < context.first_request_at: - return SetSessionOutcomeResponse( - success=False, - reason=SessionOutcomeFailureReason.OCCURRED_BEFORE_SESSION, - message="Outcome occurred before the session began", - user_id=context.user_id, - source=context.source, - ) - provider = get_service(SESSION_OUTCOME_ACCEPTANCE) - if provider is not None: - reason = provider( - self.org_id, request, received_at, context.user_id - ) - if reason is not None: + if not context.existing: + if context.user_id is None or context.first_request_at is None: + return SetSessionOutcomeResponse( + success=False, + reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, + message="Session has no published requests", + ) + if context.user_contract_violation: + logger.warning( + "Session outcome contract violation: multiple users for session %s", + sanitise_for_log(request.session_id), + ) + if context.source_contract_violation: + logger.warning( + "Session outcome contract violation: multiple sources for session %s", + sanitise_for_log(request.session_id), + ) + if request.occurred_at < context.first_request_at: return SetSessionOutcomeResponse( success=False, - reason=reason, - message="Outcome was not accepted", + reason=SessionOutcomeFailureReason.OCCURRED_BEFORE_SESSION, + message="Outcome occurred before the session began", user_id=context.user_id, source=context.source, ) + provider = get_service(SESSION_OUTCOME_ACCEPTANCE) + if provider is not None: + reason = provider( + self.org_id, request, received_at, context.user_id + ) + if reason is not None: + return SetSessionOutcomeResponse( + success=False, + reason=reason, + message="Outcome was not accepted", + user_id=context.user_id, + source=context.source, + ) result = storage.record_session_outcome( request, created_at=received_at, @@ -112,6 +105,10 @@ def mark_session_outcome( message="Outcome was not recorded", user_id=result.user_id, source=result.source, + outcome_id=result.outcome_id, + outcome_revision=result.outcome_revision, + outcome_contract_digest=result.outcome_contract_digest, + finalized_trajectory_digest=result.finalized_trajectory_digest, ) return SetSessionOutcomeResponse( success=True, @@ -123,6 +120,10 @@ def mark_session_outcome( if result.recorded else "Outcome already exists" ), + outcome_id=result.outcome_id, + outcome_revision=result.outcome_revision, + outcome_contract_digest=result.outcome_contract_digest, + finalized_trajectory_digest=result.finalized_trajectory_digest, ) except Exception: logger.exception("Failed to record session outcome") diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 295f0ac2c..59bc3f27a 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -26,6 +26,7 @@ from ..validators import ( EmbeddingVector, NonEmptyStr, + SessionOutcomeSource, TimeRangeValidatorMixin, _validate_image_url, ) @@ -278,7 +279,8 @@ class Request(BaseModel): user_id (str): Owner of the request. created_at (int): Unix epoch seconds at request creation. Defaults to the current UTC time. - source (str): Free-form origin tag (integration name, etc.). + source (str): Non-sensitive producer/workflow label. Non-empty values + use the session outcome source contract. agent_version (str): The agent version that handled this request. session_id (str): Non-empty session this request belongs to. evaluation_only (bool): Whether this request is stored for @@ -293,7 +295,7 @@ class Request(BaseModel): request_id: str user_id: str created_at: int = Field(default_factory=lambda: int(datetime.now(UTC).timestamp())) - source: str = "" + source: SessionOutcomeSource = "" agent_version: str = "" session_id: NonEmptyStr evaluation_only: bool = False @@ -364,6 +366,8 @@ class UserPlaybook(BaseModel): reader_angle: str | None = None merged_into: int | None = None superseded_by: int | None = None + governance_subject_ref: str | None = Field(default=None, exclude=True) + retired_at: int | None = Field(default=None, exclude=True) class ProfileChangeLog(BaseModel): @@ -399,6 +403,7 @@ class AgentPlaybook(BaseModel): OptimizerKind = Literal[ "gepa", "offline_tuner_replay", + "offline_tuner_open_world", "offline_tuner_legacy", "optimizer_legacy_unknown", ] @@ -439,6 +444,7 @@ class AgentPlaybook(BaseModel): "replay_manifest", "candidate", "candidate_search_projection", + "open_world_evidence_bundle", ] Sha256Digest = str @@ -825,16 +831,45 @@ 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 occurred_at: int = Field(ge=0) - source: str + source: SessionOutcomeSource label: str | None = Field(default=None, max_length=128) value: float | None = Field(default=None, allow_inf_nan=False) metadata: dict[str, Any] | None = None + 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 @@ -879,13 +914,42 @@ class SetSessionOutcomeResponse(BaseModel): reason: SessionOutcomeFailureReason | None = None message: str = "" user_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None + outcome_id: NonEmptyStr | None = None + outcome_revision: int | None = Field(default=None, ge=1) + outcome_contract_digest: Sha256Digest | None = None + 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): session_ids: list[NonEmptyStr] | None = Field(default=None, max_length=100) user_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None outcome: SessionOutcomeKind | None = None label: str | None = None start_time: int | None = Field(default=None, ge=0) @@ -956,8 +1020,8 @@ class DeleteUserPlaybooksByIdsRequest(BaseModel): user_playbook_ids: list[int] = Field(min_length=1, max_length=10_000) -# Clear all data scoped to a single user_id (interactions, requests, session -# outcomes, user playbooks, profiles). Used by paired-protocol harnesses (e.g. SWE-bench) to +# Clear all data scoped to a single user_id (session outcomes, interactions, +# requests, user playbooks, profiles). Used by paired-protocol harnesses (e.g. SWE-bench) to # isolate per-task data on a shared storage backend without nuking sibling # tasks' rows. Intentionally does NOT touch agent_playbooks — they are the # cross-project rollup of skills and have no user_id column. @@ -1182,7 +1246,7 @@ class PublishUserInteractionRequest(CapturesUnknownFields): request_id: NonEmptyStr | None = None user_id: NonEmptyStr interaction_data_list: list[InteractionData] = Field(min_length=1, max_length=1_000) - source: str = Field(default="", max_length=1_000) + source: SessionOutcomeSource = "" # this is used for aggregating interactions for generating agent playbooks agent_version: str = Field(default="", max_length=1_000) session_id: NonEmptyStr # used for grouping requests together @@ -1670,7 +1734,7 @@ class RerunProfileGenerationRequest(BaseModel): user_id: str | None = None start_time: datetime | None = None end_time: datetime | None = None - source: str | None = None + source: SessionOutcomeSource | None = None extractor_names: list[str] | None = ( None # Deprecated compatibility field; ignored for selection. ) @@ -1697,7 +1761,7 @@ class ManualProfileGenerationRequest(BaseModel): """ user_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None extractor_names: list[str] | None = None @@ -1717,7 +1781,7 @@ class ManualPlaybookGenerationRequest(BaseModel): """ agent_version: str = DEFAULT_AGENT_VERSION - source: str | None = None + source: SessionOutcomeSource | None = None playbook_name: str | None = ( None # Deprecated compatibility field; ignored for selection. ) @@ -1743,7 +1807,7 @@ class RerunPlaybookGenerationRequest(BaseModel): playbook_name: str | None = ( None # Deprecated compatibility field; ignored for selection. ) - source: str | None = None + source: SessionOutcomeSource | None = None @field_validator("agent_version") @classmethod diff --git a/reflexio/models/api_schema/domain/enums.py b/reflexio/models/api_schema/domain/enums.py index 41c99be37..6ea0d9fa8 100644 --- a/reflexio/models/api_schema/domain/enums.py +++ b/reflexio/models/api_schema/domain/enums.py @@ -52,9 +52,11 @@ class SessionOutcomeKind(StrEnum): SUCCESS = "success" FAILURE = "failure" + UNKNOWN = "unknown" class SessionOutcomeFailureReason(StrEnum): + CONFLICTING_FINALIZATION = "conflicting_finalization" UNKNOWN_SESSION = "unknown_session" OCCURRED_BEFORE_SESSION = "occurred_before_session" OCCURRED_IN_FUTURE = "occurred_in_future" diff --git a/reflexio/models/api_schema/retriever_schema.py b/reflexio/models/api_schema/retriever_schema.py index eab732023..b34be97a6 100644 --- a/reflexio/models/api_schema/retriever_schema.py +++ b/reflexio/models/api_schema/retriever_schema.py @@ -29,6 +29,7 @@ ) from .validators import ( NonEmptyStr, + SessionOutcomeSource, TimeRangeValidatorMixin, ) @@ -63,7 +64,7 @@ class SearchUserProfileRequest(BaseModel): start_time: datetime | None = None end_time: datetime | None = None top_k: int | None = Field(default=10, gt=0) - source: str | None = None + source: SessionOutcomeSource | None = None custom_feature: str | None = None extractor_name: str | None = ( None # Deprecated compatibility field; accepted but ignored. @@ -196,7 +197,7 @@ class GetUserProfilesRequest(BaseModel): start_time: datetime | None = None end_time: datetime | None = None top_k: int | None = Field(default=30, gt=0) - source: str | None = None + source: SessionOutcomeSource | None = None profile_time_to_live: str | None = None status_filter: list[Status | None] | None = None tags: list[str] | None = None @@ -318,27 +319,27 @@ class SearchUserPlaybookRequest(BaseModel): start_time (datetime, optional): Start time for created_at filter end_time (datetime, optional): End time for created_at filter status_filter (list[Optional[Status]], optional): Filter by status (None for CURRENT, PENDING, ARCHIVED) - top_k (int, optional): Maximum number of results to return. Defaults to 10 + top_k (int, optional): Maximum results to return, up to 100. Defaults to 10 threshold (float, optional): Similarity threshold for vector search. When omitted, the embedding model's default is used. """ query: str | None = None - user_id: str | None = None + user_id: str | None = Field(default=None, max_length=255) agent_version: str | None = None playbook_name: str | None = None start_time: datetime | None = None end_time: datetime | None = None status_filter: list[Status | None] | None = None tags: list[str] | None = None - top_k: int | None = Field(default=10, gt=0) + top_k: int | None = Field(default=10, gt=0, le=100) threshold: float | None = Field(default=None, ge=0.0, le=1.0) enable_reformulation: bool | None = False search_mode: SearchMode = SearchMode.HYBRID # Caller correlation IDs for billing attribution on the Application line. # Optional; consumed by _meter_applied_learnings in server/api.py. - request_id: str | None = None - session_id: str | None = None + request_id: str | None = Field(default=None, max_length=255) + session_id: str | None = Field(default=None, max_length=255) @model_validator(mode="after") def check_time_range(self) -> Self: @@ -478,7 +479,7 @@ class GetRequestsRequest(BaseModel): user_id: str | None = None request_id: str | None = None session_id: str | None = None - source: str | None = None + source: SessionOutcomeSource | None = None start_time: datetime | None = None end_time: datetime | None = None top_k: int | None = Field( @@ -794,7 +795,8 @@ class UnifiedSearchRequest(BaseModel): Args: query (str): Search query text - top_k (int, optional): Maximum results per entity type. Defaults to 5 + top_k (int, optional): Maximum results per entity type, up to 100. + Defaults to 5. threshold (float, optional): Similarity threshold for vector search. When omitted, the embedding model's default is used. agent_version (str, optional): Filter by agent version (agent_playbooks, user_playbooks) @@ -812,11 +814,11 @@ class UnifiedSearchRequest(BaseModel): """ query: NonEmptyStr - top_k: int | None = Field(default=5, gt=0) + top_k: int | None = Field(default=5, gt=0, le=100) threshold: float | None = Field(default=None, ge=0.0, le=1.0) agent_version: str | None = None playbook_name: str | None = None - user_id: str | None = None + user_id: str | None = Field(default=None, max_length=255) tags: list[str] | None = None entity_types: list[UnifiedSearchEntityType] | None = None agent_playbook_status_filter: list[PlaybookStatus] | None = None @@ -829,8 +831,8 @@ class UnifiedSearchRequest(BaseModel): # ``session_id`` additionally enables session-scoped result dedup: items # already served to the same (org, session) are skipped and the next-best # matches backfilled (see server/services/retrieval/session_dedup.py). - request_id: str | None = None - session_id: str | None = None + request_id: str | None = Field(default=None, max_length=255) + session_id: str | None = Field(default=None, max_length=255) interaction_id: int | None = Field(default=None, gt=0) diff --git a/reflexio/models/api_schema/validators.py b/reflexio/models/api_schema/validators.py index 50bf9d9d6..f4a875f53 100644 --- a/reflexio/models/api_schema/validators.py +++ b/reflexio/models/api_schema/validators.py @@ -18,10 +18,10 @@ import os import re import socket -from typing import Annotated, Any +from typing import Annotated, Any, Literal from urllib.parse import urlparse -from pydantic import AfterValidator, HttpUrl +from pydantic import AfterValidator, HttpUrl, StringConstraints # Embedding vector dimensions — must match config_schema.EMBEDDING_DIMENSIONS. # Duplicated here to avoid circular imports (config_schema imports from this module). @@ -104,6 +104,18 @@ def _check_embedding_dimensions(v: list[float]) -> list[float]: """Embedding vector that must be either empty or exactly EMBEDDING_DIMENSIONS (512) floats.""" +SESSION_OUTCOME_SOURCE_PATTERN = r"^[a-z0-9][a-z0-9._:-]{0,127}$" + +SessionOutcomeSource = ( + Literal[""] + | Annotated[ + str, + StringConstraints(max_length=128, pattern=SESSION_OUTCOME_SOURCE_PATTERN), + ] +) +"""Outcome producer/workflow label; empty preserves the existing absent-source value.""" + + # ============================================================================= # Security Validators — SSRF Prevention # ============================================================================= diff --git a/reflexio/server/README.md b/reflexio/server/README.md index 35707f65c..b4b8d9572 100644 --- a/reflexio/server/README.md +++ b/reflexio/server/README.md @@ -100,6 +100,8 @@ Description: FastAPI backend server that processes user interactions to generate **Authentication Pattern**: The open-source app uses `default_get_org_id` and `DEFAULT_ORG_ID` for local/no-auth starts. Enterprise deployments wrap `create_app()` with authenticated org resolution, additional account routers, admin checks, observability hooks, and usage metrics. +**User-playbook exposure boundary**: Authenticated enterprise deployments install `ExposureLedgerRecorder` through their capability registry. In those deployments, `POST /api/search` and `POST /api/search_user_playbooks` synchronously record the final served user-playbook set through `services/search_exposure.py` before metering or releasing success, and recorder failures fail closed. Shared `create_app()` intentionally installs no default recorder; local/no-auth OSS and custom shared app constructions persist exposures only when their deployment registers `SEARCH_EXPOSURE_RECORDER`. The direct route skips recording when its final result set is empty. + **Pattern**: Core route handlers call `Reflexio` through `get_reflexio(org_id)`; endpoint helper files should not instantiate `Reflexio` directly. ## Extension Registry diff --git a/reflexio/server/api.py b/reflexio/server/api.py index 5117b2aac..ed860fe15 100644 --- a/reflexio/server/api.py +++ b/reflexio/server/api.py @@ -474,64 +474,73 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # noqa: ARG001 durable_learning_scheduler = None aggregation_scheduler = None started_caps: list = [] - if mounts_data_plane: - _log_multi_worker_daemons() - log_publish_hardware_capacity() - validate_llm_availability() - # The scheduler discovers every org with resumable work each tick and - # drives a per-org worker with org-scoped claims, so it is not limited - # to the bootstrap org. The bootstrap org is only used to read config - # and to seed cross-org discovery. - bootstrap_org_id = _resolve_lifespan_org_id(get_org_id) - scheduler = maybe_start_resume_scheduler( - lambda org_id: RequestContext(org_id=org_id), - bootstrap_org_id=bootstrap_org_id, - org_id_provider=resume_org_ids_provider, + if capabilities is None: + from reflexio.server.usage_metrics import ( + configure_usage_event_recorder, + exempt_usage_event_recorder, ) - # Register the missing-vector backfill sweep (opt-in via - # REFLEXIO_MISSING_VECTOR_BACKFILL_ENABLED) BEFORE starting the GC - # scheduler, so its per-org hook is visible when maybe_start_lineage_gc - # evaluates its start conditions. No-op when the flag is off. - install_missing_vector_backfill_sweep() - gc_scheduler = maybe_start_lineage_gc( - lambda org_id: RequestContext(org_id=org_id), - bootstrap_org_id=bootstrap_org_id, - ) - # Durable learning drains the learning_jobs queue per org. Gated on - # REFLEXIO_DURABLE_LEARNING_QUEUE; the default provider discovers - # orgs-with-work via the bootstrap storage (single-ref). A deployment - # may inject its own discovery via ``durable_org_ids_provider`` (e.g. - # enterprise cross-ref fan-out); when None the single-ref default runs. - durable_learning_scheduler = maybe_start_durable_learning( - lambda org_id: RequestContext(org_id=org_id), - bootstrap_org_id=bootstrap_org_id, - org_ids_provider=durable_org_ids_provider, - ) - # Enterprise owns cross-org aggregation through its capability. - # OSS still needs a startup-owned scheduler so pending SQLite work - # drains after a restart even when no new publish arrives. - if capabilities is None: - from reflexio.server.services.playbook.aggregation_scheduler import ( - ensure_local_playbook_aggregation_scheduler, - ) - try: - aggregation_context = RequestContext(org_id=bootstrap_org_id) - aggregation_scheduler = ensure_local_playbook_aggregation_scheduler( - aggregation_context - ) - except Exception: # noqa: BLE001 - # A composition root may intentionally install an enterprise - # configurator while exercising the OSS app shape without - # enterprise deployment settings. The post-persist trigger - # still starts the same scheduler lazily once a real context - # exists. - logger.warning( - "playbook aggregation scheduler startup deferred: " - "bootstrap context is unavailable", - exc_info=True, - ) + configure_usage_event_recorder(exempt_usage_event_recorder) try: + if mounts_data_plane: + _log_multi_worker_daemons() + log_publish_hardware_capacity() + validate_llm_availability() + # The scheduler discovers every org with resumable work each tick and + # drives a per-org worker with org-scoped claims, so it is not limited + # to the bootstrap org. The bootstrap org is only used to read config + # and to seed cross-org discovery. + bootstrap_org_id = _resolve_lifespan_org_id(get_org_id) + scheduler = maybe_start_resume_scheduler( + lambda org_id: RequestContext(org_id=org_id), + bootstrap_org_id=bootstrap_org_id, + org_id_provider=resume_org_ids_provider, + ) + # Register the missing-vector backfill sweep (opt-in via + # REFLEXIO_MISSING_VECTOR_BACKFILL_ENABLED) BEFORE starting the GC + # scheduler, so its per-org hook is visible when maybe_start_lineage_gc + # evaluates its start conditions. No-op when the flag is off. + install_missing_vector_backfill_sweep() + gc_scheduler = maybe_start_lineage_gc( + lambda org_id: RequestContext(org_id=org_id), + bootstrap_org_id=bootstrap_org_id, + ) + # Durable learning drains the learning_jobs queue per org. Gated on + # REFLEXIO_DURABLE_LEARNING_QUEUE; the default provider discovers + # orgs-with-work via the bootstrap storage (single-ref). A deployment + # may inject its own discovery via ``durable_org_ids_provider`` (e.g. + # enterprise cross-ref fan-out); when None the single-ref default runs. + durable_learning_scheduler = maybe_start_durable_learning( + lambda org_id: RequestContext(org_id=org_id), + bootstrap_org_id=bootstrap_org_id, + org_ids_provider=durable_org_ids_provider, + ) + # Enterprise owns cross-org aggregation through its capability. + # OSS still needs a startup-owned scheduler so pending SQLite work + # drains after a restart even when no new publish arrives. + if capabilities is None: + from reflexio.server.services.playbook.aggregation_scheduler import ( + ensure_local_playbook_aggregation_scheduler, + ) + + try: + aggregation_context = RequestContext(org_id=bootstrap_org_id) + aggregation_scheduler = ( + ensure_local_playbook_aggregation_scheduler( + aggregation_context + ) + ) + except Exception: # noqa: BLE001 + # A composition root may intentionally install an enterprise + # configurator while exercising the OSS app shape without + # enterprise deployment settings. The post-persist trigger + # still starts the same scheduler lazily once a real context + # exists. + logger.warning( + "playbook aggregation scheduler startup deferred: " + "bootstrap context is unavailable", + exc_info=True, + ) if capabilities is not None: ctx = ( app_context_factory() @@ -555,6 +564,12 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # noqa: ARG001 ) stop_search_metering_worker(timeout=5.0) + if capabilities is None: + from reflexio.server.usage_metrics import ( + configure_usage_event_recorder, + ) + + configure_usage_event_recorder(None) for cap in reversed(started_caps): try: await cap.on_shutdown() diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py index 4e0380d5e..4a5cdf6a4 100644 --- a/reflexio/server/billing_meter.py +++ b/reflexio/server/billing_meter.py @@ -13,7 +13,12 @@ from collections.abc import Mapping from typing import Any -from reflexio.server.usage_metrics import record_usage_event +from reflexio.server.usage_metrics import ( + UsageEventDeliveryError, + UsageEventDeliveryStatus, + record_usage_event, + record_usage_event_strict, +) logger = logging.getLogger(__name__) @@ -22,6 +27,18 @@ ) +class ReceiptBillingDeliveryError(RuntimeError): + """A durable finalization receipt still has an undelivered billing event.""" + + def __init__( + self, + status: UsageEventDeliveryStatus, + message: str = "receipt-backed learning billing delivery failed", + ) -> None: + super().__init__(message) + self.status = status + + def record_extraction_tokens( *, org_id: str, @@ -91,11 +108,10 @@ def record_learnings_generated( ) -> None: """Emit the Learning value facet — number of profiles/playbooks generated. - Documented FALLBACK for callers that genuinely lack a per-record id list - (e.g. dedup/consolidation can reduce the persisted count below the raw - extracted count, so there is no safe 1:1 id per unit of ``count``). Prefer - :func:`record_learnings_generated_records` whenever the caller has the - durable learning ids in scope. No-op when ``count <= 0``. + Intended for online extraction paths that have a known billable count but + do not retain a complete per-record id list. Resumable finalization must + use :func:`record_learnings_generated_records_strict` and skip items without + durable ids. No-op when ``count <= 0``. Emits a single event carrying ``event_key`` when the caller has a durable retry identity, otherwise synthesizes ``f"learn-batch:{uuid4()}"``. This @@ -181,7 +197,7 @@ def record_learnings_generated_records( Args: org_id: Organisation identifier. learning_ids: Ids of the learnings durably generated in this run - (e.g. ``profile_id`` / ``user_playbook_id`` / ``agent_playbook_id``). + (e.g. ``profile_id`` / ``user_playbook_id``). platform_llm: True iff the platform supplies the LLM for this org. platform_storage: True iff the platform supplies storage; None defers to rollup. pipeline: Optional pipeline tag (e.g. ``"playbook"``). @@ -194,9 +210,82 @@ def record_learnings_generated_records( entity_type: Optional entity type (e.g. ``"profile"``). metadata: Optional path-specific usage metadata (shared across events). """ + _record_learnings_generated_records( + record_event=record_usage_event, + org_id=org_id, + learning_ids=learning_ids, + platform_llm=platform_llm, + platform_storage=platform_storage, + pipeline=pipeline, + user_id=user_id, + request_id=request_id, + session_id=session_id, + source=source, + agent_version=agent_version, + playbook_name=playbook_name, + entity_type=entity_type, + metadata=metadata, + ) + + +def record_learnings_generated_records_strict( + *, + org_id: str, + learning_ids: list[str], + platform_llm: bool | None, + platform_storage: bool | None, + pipeline: str | None = None, + user_id: str | None = None, + request_id: str | None = None, + session_id: str | None = None, + source: str | None = None, + agent_version: str | None = None, + playbook_name: str | None = None, + entity_type: str | None = None, + metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, +) -> None: + """Strict per-record emission for receipt-backed finalization only.""" + _record_learnings_generated_records( + record_event=record_usage_event_strict, + org_id=org_id, + learning_ids=learning_ids, + platform_llm=platform_llm, + platform_storage=platform_storage, + pipeline=pipeline, + user_id=user_id, + request_id=request_id, + session_id=session_id, + source=source, + agent_version=agent_version, + playbook_name=playbook_name, + entity_type=entity_type, + metadata=metadata, + created_at=created_at, + ) + + +def _record_learnings_generated_records( + *, + record_event: Any, + org_id: str, + learning_ids: list[str], + platform_llm: bool | None, + platform_storage: bool | None, + pipeline: str | None, + user_id: str | None, + request_id: str | None, + session_id: str | None, + source: str | None, + agent_version: str | None, + playbook_name: str | None, + entity_type: str | None, + metadata: Mapping[str, Any] | None, + created_at: float | None = None, +) -> None: key_entity_type = entity_type or "_" for learning_id in learning_ids: - record_usage_event( + record_event( org_id=org_id, event_name="learnings_generated", event_category="learning", @@ -215,6 +304,7 @@ def record_learnings_generated_records( platform_storage=platform_storage, caller_type=_INTERNAL, metadata=metadata, + created_at=created_at, ) @@ -235,20 +325,21 @@ def emit_learnings_generated( ) -> None: """Resolve ``platform_llm`` from config and emit the Learning value facet. - Convenience wrapper for non-extraction learning-mutation paths such as - resumable-extraction finalization, aggregation, and offline-tuner auto-apply. It - owns the ``configurator.get_config()`` + ``platform_llm_from_config`` lookup so - each call site stays a thin one-liner, and — critically — is **guarded**: the - product path must never fail because metering failed, so config resolution and - emission are wrapped and any exception is logged and swallowed (mirroring the - extraction path's ``_record_billing_learning_events``). No-op when - ``count <= 0``. + Convenience wrapper for count-based online extraction callers. It owns the + ``configurator.get_config()`` + ``platform_llm_from_config`` lookup so the + call site stays a thin one-liner, and — critically — is **guarded**: the + product path must never fail because metering failed, so config resolution + and emission are wrapped and any exception is logged and swallowed + (mirroring the extraction path's ``_record_billing_learning_events``). + Resumable finalization must use + :func:`emit_learnings_generated_records_strict`. + No-op when ``count <= 0``. Args: org_id: Organisation identifier. configurator: Object exposing ``get_config()`` for platform-LLM resolution. count: Number of learnings durably produced by this path. - source: Metering source/path label (e.g. ``"offline_optimizer"``). + source: Metering source/path label (e.g. ``"online_extraction"``). pipeline: Optional pipeline tag (e.g. ``"playbook"``). user_id: Optional user ID tied to the generated learning. request_id: Optional request correlation ID. @@ -305,22 +396,20 @@ def emit_learnings_generated_records( ) -> None: """Resolve ``platform_llm`` from config and emit one event per learning id. - Entity-backed counterpart to :func:`emit_learnings_generated`, currently - adopted by two of the non-extraction learning-mutation paths — - resumable-extraction finalization and aggregation — the callers with - durable per-record ids in scope. Extraction and offline-tuner auto-apply do - not have a safe 1:1 id per unit of count (see - :func:`record_learnings_generated_records`) and use the count-based - :func:`emit_learnings_generated` fallback instead. Same guard semantics: - config resolution and emission are wrapped and any exception is logged - and swallowed — the product path must never fail because metering - failed. No-op when ``learning_ids`` is empty. + Ordinary fail-open per-record counterpart to + :func:`emit_learnings_generated`. Receipt-backed resumable finalization uses + :func:`emit_learnings_generated_records_strict`; items without durable ids + are not billable on that path. Online extraction uses the count-based + :func:`record_learnings_generated` helper because it does not retain a safe + 1:1 id per generated unit. Config resolution and emission are wrapped and + any exception is logged and swallowed — the product path must never fail + because metering failed. No-op when ``learning_ids`` is empty. Args: org_id: Organisation identifier. configurator: Object exposing ``get_config()`` for platform-LLM resolution. learning_ids: Ids of the learnings durably produced by this path. - source: Metering source/path label (e.g. ``"aggregation"``). + source: Metering source/path label (e.g. ``"resumable_extraction"``). pipeline: Optional pipeline tag (e.g. ``"playbook"``). user_id: Optional user ID tied to the generated learning. request_id: Optional request correlation ID. @@ -359,6 +448,49 @@ def emit_learnings_generated_records( ) +def emit_learnings_generated_records_strict( + *, + org_id: str, + configurator: Any, + learning_ids: list[str], + source: str, + pipeline: str | None = None, + user_id: str | None = None, + request_id: str | None = None, + agent_version: str | None = None, + playbook_name: str | None = None, + entity_type: str | None = None, + metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, +) -> None: + """Strict receipt-backed counterpart to the ordinary fail-open emitter.""" + if not learning_ids: + return + from reflexio.server.billing_signals import platform_llm_from_config + + try: + config = configurator.get_config() + record_learnings_generated_records_strict( + org_id=org_id, + learning_ids=learning_ids, + platform_llm=platform_llm_from_config(config), + platform_storage=None, + pipeline=pipeline, + user_id=user_id, + request_id=request_id, + source=source, + agent_version=agent_version, + playbook_name=playbook_name, + entity_type=entity_type, + metadata=metadata, + created_at=created_at, + ) + except UsageEventDeliveryError as exc: + raise ReceiptBillingDeliveryError(exc.status) from exc + except Exception as exc: + raise ReceiptBillingDeliveryError(UsageEventDeliveryStatus.UNKNOWN) from exc + + def record_applied_learnings( *, org_id: str, diff --git a/reflexio/server/routes/search.py b/reflexio/server/routes/search.py index cb83bf6b1..1adb23301 100644 --- a/reflexio/server/routes/search.py +++ b/reflexio/server/routes/search.py @@ -48,6 +48,10 @@ 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.services.search_metering_worker import enqueue_search_metering from reflexio.server.tracing import profile_step @@ -247,6 +251,17 @@ def search_user_playbooks_endpoint( msg=response.msg, experiment=assignment, ) + if response.user_playbooks: + record_search_exposures( + SearchExposureBatch( + org_id=org_id, + request_id=payload.request_id, + session_id=payload.session_id, + interaction_id=None, + user_id=payload.user_id, + user_playbooks=tuple(response.user_playbooks), + ) + ) enqueue_search_metering( org_id=org_id, caller_type=caller_type, @@ -412,6 +427,16 @@ def run_search() -> Any: rehydrated_text=response.rehydrated_text, experiment=assignment, ) + record_search_exposures( + SearchExposureBatch( + org_id=org_id, + request_id=payload.request_id, + session_id=payload.session_id, + interaction_id=payload.interaction_id, + user_id=payload.user_id, + user_playbooks=tuple(response.user_playbooks), + ) + ) enqueue_search_metering( org_id=org_id, caller_type=caller_type, diff --git a/reflexio/server/services/README.md b/reflexio/server/services/README.md index 2e2062aa0..cce700102 100644 --- a/reflexio/server/services/README.md +++ b/reflexio/server/services/README.md @@ -58,6 +58,7 @@ strings before deleting old import paths in the same PR. | `pre_retrieval/` | `QueryReformulator` (`_query_reformulator.py`) + `DocumentExpander` (`_document_expander.py`) - query rewrite and doc expansion for recall. Compact by design; see [README](pre_retrieval/README.md). | | `tagging/` | `TaggingService` (`service.py`) + deferred `tagging_scheduler.py` - post-generation profile/playbook tagging. Compact by design; see [README](tagging/README.md). | | `unified_search_service.py` | `run_unified_search()` — two-phase parallel search across profiles / agent playbooks / user playbooks. | +| `search_exposure.py` | Optional synchronous recorder contract for final user-playbook result sets. Enterprise capability registration installs the recorder and makes authenticated unified/direct search fail closed before metering/response. Shared `create_app()` has no default recorder; other constructions persist exposures only when they register one. Direct search omits empty batches. | | `retrieval/` | `relevance_floor.py` — result relevance thresholding. `temporal.py` — temporal post-processing driven by reformulation signals: query time windows → per-arm SQL filters, near-duplicate freshness collapse for current-value questions, timestamp ordering for latest-value questions. `user_context_guard.py` — high-precision detection of explicit personalization opt-outs, including Simplified and Traditional Chinese, before user-context retrieval. (Superseded/expired rows are already excluded by storage search SQL.) | ## Persistence & Config diff --git a/reflexio/server/services/base_generation/_usage_billing.py b/reflexio/server/services/base_generation/_usage_billing.py index 1b1304f90..5aee8a7d9 100644 --- a/reflexio/server/services/base_generation/_usage_billing.py +++ b/reflexio/server/services/base_generation/_usage_billing.py @@ -188,14 +188,14 @@ def _record_billing_learning_events( (cost facet) via the OSS emission helpers. ``platform_storage`` is left ``None`` here and resolved enterprise-side at rollup (Phase 1). - Gated by ``EMITS_LEARNING_BILLING`` — only profile/playbook generation - services opt in here. Non-extraction learning mutation paths emit their - own ``learnings_generated`` value-facet events when they durably apply - revisions/successors. + Gated by ``EMITS_LEARNING_BILLING`` — only online profile/playbook + extraction services opt in here. Resumable-extraction finalization emits + the same value facet separately. Derived mutation paths emit no additional + ``learnings_generated`` events. Args: prepared: The prepared generation run (used for input-text computation). - generated_count: Number of learnings produced by this extraction run. + generated_count: Number of retained write-plan learnings eligible for billing. """ if not self.EMITS_LEARNING_BILLING: return @@ -260,3 +260,16 @@ def _count_generated_results(result: Any) -> int: if isinstance(result, list): return len(result) return 1 if result else 0 + + @staticmethod + def _count_retained_online_learnings(write_plan: Any) -> int: + from reflexio.server.services.deferred_learning_plan import ( + PlaybookWritePlan, + ProfileWritePlan, + ) + + if isinstance(write_plan, ProfileWritePlan): + return len(write_plan.new_profiles) + if isinstance(write_plan, PlaybookWritePlan): + return len(write_plan.new_playbooks) + return 0 diff --git a/reflexio/server/services/base_generation_service.py b/reflexio/server/services/base_generation_service.py index 2334d3741..9fd3f9fe6 100644 --- a/reflexio/server/services/base_generation_service.py +++ b/reflexio/server/services/base_generation_service.py @@ -210,9 +210,9 @@ class BaseGenerationService( ABC, Generic[TExtractorConfig, TExtractor, TGenerationServiceConfig, TRequest], # noqa: UP046 ): - # Only profile/playbook GENERATION services emit extraction-run billing here. - # Non-extraction learning mutation paths emit their value facet at their own - # durable-success point. + # Only online profile/playbook extraction services emit extraction-run billing + # here. Resumable-extraction finalization emits separately; derived mutation + # paths emit no additional learnings_generated events. # Default is False so any future subclass is safe by default (opt-IN). EMITS_LEARNING_BILLING: bool = False """ @@ -342,11 +342,20 @@ def _process_results(self, results: list) -> None: results: List of all results from extractors (one per successful extractor) """ - def _finalize_extracted_items(self, items: list) -> list: + def _finalize_extracted_items( + self, + items: list, + *, + finalization_run_id: str | None = None, + ) -> list[str] | None: """Persist already-flattened extracted items through the service path.""" + if finalization_run_id is not None: + raise NotImplementedError( + "Receipt-aware finalization must be implemented by resumable services" + ) if items: self._process_results([items]) - return items + return None @abstractmethod def _should_track_in_progress(self) -> bool: @@ -634,7 +643,6 @@ def compute_generation(self, request: TRequest) -> GenerationComputePlan | None: self._last_bookmark_advance = None self._last_model_provenance = None result = self._execute_extractor(prepared.extractor_config, prepared.identifier) - generated_count = self._count_generated_results(result) try: write_plan = self._resolve_write_plan([result]) if result else None @@ -643,9 +651,17 @@ def compute_generation(self, request: TRequest) -> GenerationComputePlan | None: self._mark_extraction_runs_finalization_failed(exc) raise + generated_count = self._count_generated_results(result) + billable_count = ( + self._count_retained_online_learnings(write_plan) + if self.EMITS_LEARNING_BILLING + else 0 + ) + return GenerationComputePlan( prepared=prepared, generated_count=generated_count, + billable_count=billable_count, write_plan=write_plan, bookmark_advance=self._last_bookmark_advance, generation_start=generation_start, @@ -676,7 +692,8 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: Runs only for a fence-winning job (the durable worker calls it after the scope commits; ``.run()`` calls it inline). Reads the plan's compute-time - snapshot (``generated_count`` / ``prepared`` / ``generation_start``) so a + snapshot (``generated_count`` / ``billable_count`` / ``prepared`` / + ``generation_start``) so a fence-lost job never emits. Billing purity note (round-2 finding): ``_record_billing_learning_events`` @@ -708,7 +725,7 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: }, ) self._record_billing_learning_events( - prepared=plan.prepared, generated_count=plan.generated_count + prepared=plan.prepared, generated_count=plan.billable_count ) @abstractmethod diff --git a/reflexio/server/services/deferred_learning_plan.py b/reflexio/server/services/deferred_learning_plan.py index eb0a87cd8..e540357e6 100644 --- a/reflexio/server/services/deferred_learning_plan.py +++ b/reflexio/server/services/deferred_learning_plan.py @@ -26,6 +26,26 @@ ) +@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 and may replay retry-safe durable + obligations, such as billing records keyed by those ids. They must not + replay winner-only derived work such as optimization, aggregation, or + tagging. + """ + + learning_ids: list[str] + 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). @@ -134,16 +154,20 @@ class GenerationComputePlan: ``emit_generation_side_effects`` fires the post-commit telemetry + billing. The billing inputs (``extraction_run_ids`` / ``token_totals`` / - ``generated_count`` / ``prepared``) are **snapshotted at compute time** so - the fence-crossing emit reads this plan rather than the reused service - instance's mutable ``_last_*`` accumulators (purity contract, plan §File - Structure). See ``emit_generation_side_effects`` for the single-use-instance - invariant that also keeps the money helper's ``self._last_*`` reads safe. + ``billable_count`` / ``prepared``) and telemetry's ``generated_count`` are + **snapshotted at compute time** so the fence-crossing emit reads this plan + rather than the reused service instance's mutable ``_last_*`` accumulators + (purity contract, plan §File Structure). See + ``emit_generation_side_effects`` for the single-use-instance invariant that + also keeps the money helper's ``self._last_*`` reads safe. Attributes: prepared: The prepared generation run (identifier / extractor_name / extractor_config), reused by emit for telemetry + billing input. - generated_count: Learnings produced by this extraction run. + generated_count: Raw learnings produced by the extractor, used for + generation-success telemetry. + billable_count: Retained write-plan learnings eligible for online billing; + zero for services that do not emit learning billing. write_plan: The resolved write-plan (``ProfileWritePlan`` / ``PlaybookWritePlan`` in Tasks 6-7, a ``_LegacyItems`` shim marker until then) or ``None`` when the extractor produced nothing. @@ -159,6 +183,7 @@ class GenerationComputePlan: prepared: PreparedGenerationRun[Any] generated_count: int + billable_count: int write_plan: Any bookmark_advance: ExtractorBookmarkAdvance | None generation_start: float diff --git a/reflexio/server/services/extraction/README.md b/reflexio/server/services/extraction/README.md index 668455202..a0dfe2b5f 100644 --- a/reflexio/server/services/extraction/README.md +++ b/reflexio/server/services/extraction/README.md @@ -16,13 +16,16 @@ information, and resumes outside the request path. | `prior_answer_search.py` | Finds and formats previous human answers for async extraction context. | | `agent_run_records.py` | Builds durable extraction-agent run records and source interaction identity. | | `resume_scheduler.py` | Discovers and schedules due paused/finalization work in a background singleton. | -| `resume_worker.py` | Resumes paused runs, rebuilds request context, and records retry state. | +| `resume_worker.py` | Resumes paused runs, rebuilds request context, and records retry state. Finalization uses an immutable run-keyed receipt so learning writes and retry billing reuse the same persisted IDs. | | `outcome.py` | Provides the generic extraction outcome wrapper used by callers. | ## Boundary Rules - Keep profile-specific and playbook-specific extraction behavior in their own modules; call this package only for shared async runtime concerns. +- Commit resumable learning writes, lineage changes, and the agent-run + finalization receipt in one storage transaction. Retries must return the + receipt's persisted IDs instead of repeating finalization. - Add a new file here when the behavior is shared by more than one extraction caller or is part of the resumable runtime itself. - Split into subpackages only when a responsibility grows large enough that a diff --git a/reflexio/server/services/extraction/resume_worker.py b/reflexio/server/services/extraction/resume_worker.py index ff1732027..fc68e8de7 100644 --- a/reflexio/server/services/extraction/resume_worker.py +++ b/reflexio/server/services/extraction/resume_worker.py @@ -16,10 +16,15 @@ from reflexio.models.api_schema.service_schemas import Interaction, Request from reflexio.models.config_schema import PlaybookConfig, ProfileExtractorConfig from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.billing_meter import ( + ReceiptBillingDeliveryError, + emit_learnings_generated_records_strict, +) from reflexio.server.error_reporting import error_tags from reflexio.server.llm._litellm_types import ModelProvenance from reflexio.server.llm.litellm_client import LiteLLMClient, LiteLLMConfig 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, @@ -78,6 +83,7 @@ ) from reflexio.server.services.tagging.tagging_scheduler import schedule_tagging from reflexio.server.site_var.site_var_manager import SiteVarManager +from reflexio.server.usage_metrics import UsageEventDeliveryStatus logger = logging.getLogger(__name__) @@ -91,6 +97,28 @@ def _next_retry_at(attempt_count: int) -> datetime: return datetime.now(UTC) + timedelta(seconds=delay_seconds) +def _finalization_failure_status( + exc: Exception, + *, + next_attempt_count: int, + max_finalization_attempts: int, +) -> AgentRunStatus: + """Classify finalization failures under the receipt-delivery contract. + + Transient receipt delivery failures remain retryable beyond the ordinary + finalization-attempt ceiling because their committed billing obligation must + eventually be delivered. Permanent receipt rejection is terminal, as are + ordinary finalization failures at or above the configured ceiling. + """ + if isinstance(exc, ReceiptBillingDeliveryError): + if exc.status is UsageEventDeliveryStatus.REJECTED: + return AgentRunStatus.FAILED + return AgentRunStatus.FINALIZATION_FAILED + if next_attempt_count >= max_finalization_attempts: + return AgentRunStatus.FAILED + return AgentRunStatus.FINALIZATION_FAILED + + def _create_llm_client(request_context: RequestContext) -> LiteLLMClient: # The tool loop re-resolves the model per call via ModelRole.EXTRACTION_AGENT # (see run_tool_loop); this client's model name is only a fallback. We resolve @@ -335,8 +363,9 @@ def run_once(self) -> AgentRunRecord | None: try: self.storage.update_agent_run_status(run.id, AgentRunStatus.FINALIZING) - self._finalize_items(run, items, model_provenance=model_provenance) - self._schedule_finalized_tagging(run) + result = self._finalize_items(run, items, model_provenance=model_provenance) + if result.won_receipt: + self._schedule_finalized_tagging(run) self.storage.consume_run_tool_dependencies(run.id) finalized_status = ( AgentRunStatus.FINALIZED_PENDING_TOOL @@ -361,10 +390,10 @@ def run_once(self) -> AgentRunRecord | None: run.id, ) next_attempt_count = run.finalization_attempts + 1 - failed_status = ( - AgentRunStatus.FAILED - if next_attempt_count >= pending_config.max_finalization_attempts - else AgentRunStatus.FINALIZATION_FAILED + failed_status = _finalization_failure_status( + exc, + next_attempt_count=next_attempt_count, + max_finalization_attempts=pending_config.max_finalization_attempts, ) return self.storage.update_agent_run_status( run.id, @@ -382,8 +411,9 @@ def _retry_finalization(self, run: AgentRunRecord) -> AgentRunRecord | None: items, pending_tool_call_ids, model_provenance = ( self._items_from_committed_output(run) ) - self._finalize_items(run, items, model_provenance=model_provenance) - self._schedule_finalized_tagging(run) + result = self._finalize_items(run, items, model_provenance=model_provenance) + if result.won_receipt: + self._schedule_finalized_tagging(run) self.storage.consume_run_tool_dependencies(run.id) finalized_status = ( AgentRunStatus.FINALIZED_PENDING_TOOL @@ -408,10 +438,10 @@ def _retry_finalization(self, run: AgentRunRecord) -> AgentRunRecord | None: run.id, ) next_attempt_count = run.finalization_attempts + 1 - failed_status = ( - AgentRunStatus.FAILED - if next_attempt_count >= pending_config.max_finalization_attempts - else AgentRunStatus.FINALIZATION_FAILED + failed_status = _finalization_failure_status( + exc, + next_attempt_count=next_attempt_count, + max_finalization_attempts=pending_config.max_finalization_attempts, ) return self.storage.update_agent_run_status( run.id, @@ -896,7 +926,7 @@ def _finalize_items( items: list[Any], *, model_provenance: ModelProvenance | None = None, - ) -> None: + ) -> FinalizationResult: if run.binding.extractor_kind == "profile": service = ProfileGenerationService( llm_client=self.client, @@ -909,13 +939,15 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - persisted_items = service._finalize_extracted_items( - items, model_provenance=model_provenance + result = service._finalize_extracted_items_with_outcome( + items, + model_provenance=model_provenance, + finalization_run_id=run.id, ) self._record_finalized_learnings( - run, persisted_items or [], entity_type="profile" + run, result.learning_ids, entity_type="profile" ) - return + return result if run.binding.extractor_kind == "playbook": user_id = run.binding.user_id if not user_id: @@ -932,82 +964,52 @@ def _finalize_items( auto_run=False, force_extraction=True, ) - persisted_items = service._finalize_extracted_items( + result = service._finalize_extracted_items_with_outcome( items, model_provenance=model_provenance, extraction_run=run, + finalization_run_id=run.id, ) self._record_finalized_learnings( - run, persisted_items or [], entity_type="user_playbook" + run, result.learning_ids, entity_type="user_playbook" ) - return + return result raise ResumeWorkerError( f"Unsupported extractor kind {run.binding.extractor_kind!r}" ) def _record_finalized_learnings( - self, run: AgentRunRecord, items: list[Any], *, entity_type: str + self, run: AgentRunRecord, learning_ids: list[str], *, entity_type: str ) -> None: """Emit ``learnings_generated`` for a finalized resumable-extraction batch. - Prefers one event per learning id (``entity_id``/``profile_id`` for - profiles, ``user_playbook_id`` for playbooks) when every item in - ``items`` carries a durable id — the common case, since these ids are - assigned by the extractor (profile) or by ``save_user_playbooks`` - in-place during ``_finalize_extracted_items`` (playbook), which has - already run by the time this is called. Falls back to the - count-based aggregate event when any item lacks one (e.g. dropped by - within-batch/consolidator dedup before persist, leaving a default - ``user_playbook_id=0``) — this avoids both fabricating an id for a row - that never persisted and colliding on the shared default-id key. - Totals are preserved either way: ``len(items)`` learnings are counted - whether via ``len(learning_ids)`` per-record events or one aggregate - ``count=len(items)`` event. + Emits one event per durable learning id returned by finalization. + Per-record keys make finalization retries idempotent downstream. """ - from reflexio.server.billing_meter import ( - emit_learnings_generated, - emit_learnings_generated_records, - ) - - if not items: + if not learning_ids: return - - id_attr = "profile_id" if entity_type == "profile" else "user_playbook_id" - learning_ids = [ - str(getattr(item, id_attr)) - for item in items - if getattr(item, id_attr, None) - ] + billing_timestamp = run.created_at or run.agent_completed_at + if billing_timestamp is None: + raise ReceiptBillingDeliveryError( + UsageEventDeliveryStatus.UNKNOWN, + "receipt-backed learning billing timestamp is not durable", + ) metadata = { "run_id": run.id, "extractor_kind": run.binding.extractor_kind, } - if len(learning_ids) == len(items): - emit_learnings_generated_records( - org_id=self.request_context.org_id, - configurator=self.request_context.configurator, - learning_ids=learning_ids, - source="resumable_extraction", - pipeline=run.binding.extractor_kind, - user_id=run.binding.user_id, - request_id=run.binding.request_id, - agent_version=run.binding.agent_version, - entity_type=entity_type, - metadata=metadata, - ) - return - emit_learnings_generated( + emit_learnings_generated_records_strict( org_id=self.request_context.org_id, configurator=self.request_context.configurator, - count=len(items), + learning_ids=learning_ids, source="resumable_extraction", pipeline=run.binding.extractor_kind, user_id=run.binding.user_id, request_id=run.binding.request_id, agent_version=run.binding.agent_version, entity_type=entity_type, - event_key=f"learn-batch:resumable:{run.id}:{entity_type}", metadata=metadata, + created_at=billing_timestamp.timestamp(), ) def _schedule_finalized_tagging(self, run: AgentRunRecord) -> None: diff --git a/reflexio/server/services/governance/service.py b/reflexio/server/services/governance/service.py index fca64cdbf..5b60a9be4 100644 --- a/reflexio/server/services/governance/service.py +++ b/reflexio/server/services/governance/service.py @@ -1,5 +1,8 @@ from __future__ import annotations +import threading +import time +import uuid from contextlib import suppress from typing import Any, Literal, Protocol, TypedDict @@ -16,6 +19,7 @@ governance_subject_ref, ) from reflexio.server.services.governance.subject_refs import stable_id +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim _DELETE_TARGET_NAME_TO_RESULT_KEY = { "interaction": "interactions", @@ -29,12 +33,18 @@ "offline_tuner_reward_label_target_by_target_owner": ( "offline_tuner_reward_label_targets_by_target_owner" ), + "session_outcome": "session_outcomes", "profile_purge": "purged_profiles", "user_playbook_purge": "purged_user_playbooks", } _REQUIRED_DELETE_TARGET_NAMES = tuple(_DELETE_TARGET_NAME_TO_RESULT_KEY) _USER_PLAYBOOK_PAGE_SIZE = 1000 _LIFECYCLE_COMPLETION_STATUS = "complete" +_DUPLICATE_ERASE_POLL_SECONDS = 0.05 +_DUPLICATE_ERASE_MAX_POLL_SECONDS = 1.0 +_DUPLICATE_ERASE_WAIT_SECONDS = 5.0 +_PURGE_EXECUTION_LEASE_SECONDS = 300 +_PURGE_EXECUTION_HEARTBEAT_SECONDS = 30 class GovernanceActorContext(TypedDict): @@ -43,7 +53,7 @@ class GovernanceActorContext(TypedDict): class SubjectErasureLifecycle(Protocol): - """Deployment-specific erasure work that must precede barrier completion.""" + """External erasure work invoked only after a synchronous live-claim check.""" def erase_subject( self, @@ -51,9 +61,80 @@ 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_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: + with self._renewal_lock: + try: + renewed = self._storage.renew_purge_operation_execution_claim( + self._purge_id, + self.claim(), + lease_ttl_seconds=_PURGE_EXECUTION_LEASE_SECONDS, + ) + except _PurgeExecutionHeartbeatLostError: + raise + except Exception as exc: + with self._lock: + self._renewal_error = exc + raise _PurgeExecutionHeartbeatLostError( + "purge execution heartbeat renewal was lost" + ) from exc + with self._lock: + self._claim = renewed + return renewed + + def _run(self) -> None: + while not self._stop.wait(_PURGE_EXECUTION_HEARTBEAT_SECONDS): + try: + self.renew_now() + except _PurgeExecutionHeartbeatLostError: + return + + class GovernanceService: def __init__( self, @@ -129,84 +210,145 @@ def erase_user( f"{self.org_id}:user_erasure:{subref}:{reqref}", ) purge_id = stable_id("purge", idempotency_key) - purge = self.storage.begin_purge_operation( - purge_id=purge_id, - idempotency_key=idempotency_key, - operation_type="user_erasure", - scope_type="user", - subject_ref=subref, - request_ref=reqref, - ) + try: + purge = self.storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key=idempotency_key, + operation_type="user_erasure", + scope_type="user", + subject_ref=subref, + request_ref=reqref, + authoritative_user_id=user_id, + ) + except Exception as begin_exc: + try: + purge = self._matching_user_erasure_purge_for_retry( + purge_id=purge_id, + operation_type="user_erasure", + scope_type="user", + subject_ref=subref, + request_ref=reqref, + authoritative_user_id=user_id, + ) + except Exception: + raise begin_exc from None + if purge.status == "complete": + raise begin_exc from None if purge.status == "complete": - barrier = self._completed_barrier_for_retry( + return self._completed_erase_result_for_retry( subject_ref=subref, purge_id=purge_id ) - if barrier.status != "erased": - raise ValueError( - "Completed purge retry requires an erased subject barrier" + lease_owner = f"governance-erase-{uuid.uuid4().hex}" + execution_claim: PurgeExecutionClaim | None = None + claim_deadline = self._monotonic() + _DUPLICATE_ERASE_WAIT_SECONDS + poll_seconds = _DUPLICATE_ERASE_POLL_SECONDS + while execution_claim is None: + execution_claim = self.storage.claim_purge_operation_execution( + purge_id, + lease_owner=lease_owner, + lease_ttl_seconds=_PURGE_EXECUTION_LEASE_SECONDS, + ) + if execution_claim is not None: + break + purge = self.storage.get_purge_operation(purge_id) + if purge.status == "complete": + return self._completed_erase_result_for_retry( + subject_ref=subref, + purge_id=purge_id, ) - return UserEraseResult( - subject_ref=subref, - purge_id=purge_id, - status="complete", - deleted_counts=self._deleted_counts_from_targets(purge_id), - rebuilt_agent_playbook_ids=( - self._rebuilt_agent_playbook_ids_from_targets(purge_id) - ), + if purge.status not in {"pending", "running", "failed"}: + raise ValueError(f"Unsupported purge operation status: {purge.status}") + remaining_seconds = claim_deadline - self._monotonic() + if remaining_seconds <= 0: + raise GovernanceEraseRetryLaterError( + "Another erase request still owns the execution claim; retry later" + ) + self._sleep(min(poll_seconds, remaining_seconds)) + poll_seconds = min( + poll_seconds * 2, + _DUPLICATE_ERASE_MAX_POLL_SECONDS, ) try: - self.storage.begin_subject_erasure_barrier(subref, purge_id) - if not self.storage.purge_targets_prepared(purge_id): - self.storage.prepare_governance_erase_targets( + with _PurgeExecutionHeartbeat( + storage=self.storage, + purge_id=purge_id, + execution_claim=execution_claim, + ) as heartbeat: + self.storage.begin_subject_erasure_barrier( + subref, purge_id, - user_id, + execution_claim=heartbeat.claim(), ) + if not self.storage.purge_targets_prepared(purge_id): + self.storage.prepare_governance_erase_targets( + purge_id, + user_id, + execution_claim=heartbeat.claim(), + ) - if not self._delete_targets_complete(purge_id): - self.storage.apply_governance_user_data_delete(purge_id, user_id) - if ( - self.subject_erasure_lifecycle is not None - and not self._subject_erasure_lifecycle_complete(purge_id) - ): - self.subject_erasure_lifecycle.erase_subject( - storage=self.storage, - subject_ref=subref, - purge_id=purge_id, - ) - self._record_subject_erasure_lifecycle_complete(purge_id) - deleted_counts = self._deleted_counts_from_targets(purge_id) + if not self._delete_targets_complete(purge_id): + self.storage.apply_governance_user_data_delete( + purge_id, + user_id, + execution_claim=heartbeat.claim(), + ) + if ( + self.subject_erasure_lifecycle is not None + and not self._subject_erasure_lifecycle_complete(purge_id) + ): + heartbeat.renew_now() + self._assert_execution_claim(purge_id, heartbeat.claim()) + self.subject_erasure_lifecycle.erase_subject( + storage=self.storage, + subject_ref=subref, + purge_id=purge_id, + execution_claim=heartbeat.claim(), + ) + self._record_subject_erasure_lifecycle_complete( + purge_id, + execution_claim=heartbeat.claim(), + ) + deleted_counts = self._deleted_counts_from_targets(purge_id) - rebuilt_agent_playbook_ids: list[int] = [] - completed = self.storage.complete_subject_erasure_barrier_after_empty_check( - purge_id, - AuditEvent( - org_id=self.org_id, - actor_type=actor_type, - actor_ref=actor_ref, - operation="ERASE", - entity_type="request", - subject_ref=subref, - request_ref=reqref, - idempotency_key=purge_id, - detail={ - "deleted_counts": deleted_counts, - "rebuilt_agent_playbook_ids": rebuilt_agent_playbook_ids, - }, - ), - ) + rebuilt_agent_playbook_ids: list[int] = [] + completed = self.storage.complete_subject_erasure_barrier_after_empty_check( + purge_id, + AuditEvent( + org_id=self.org_id, + actor_type=actor_type, + actor_ref=actor_ref, + operation="ERASE", + entity_type="request", + subject_ref=subref, + request_ref=reqref, + idempotency_key=purge_id, + detail={ + "deleted_counts": deleted_counts, + "rebuilt_agent_playbook_ids": rebuilt_agent_playbook_ids, + }, + ), + authoritative_user_id=user_id, + execution_claim=heartbeat.claim(), + ) except Exception as exc: + if isinstance(exc, _PurgeExecutionHeartbeatLostError): + raise + if not self._execution_claim_is_current(purge_id, execution_claim): + raise with suppress(Exception): self.storage.fail_subject_erasure_barrier( subref, purge_id, error_code="governance_erase_failed", error_detail=type(exc).__name__, + execution_claim=execution_claim, ) with suppress(Exception): self.storage.fail_purge_operation( purge_id, error_code="governance_erase_failed", error_detail=type(exc).__name__, + execution_claim=execution_claim, ) raise return UserEraseResult( @@ -217,6 +359,28 @@ def erase_user( rebuilt_agent_playbook_ids=rebuilt_agent_playbook_ids, ) + @staticmethod + def _monotonic() -> float: + return time.monotonic() + + @staticmethod + def _sleep(seconds: float) -> None: + time.sleep(seconds) + + def _assert_execution_claim( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> None: + self.storage.assert_purge_operation_execution_claim(purge_id, execution_claim) + + def _execution_claim_is_current( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> bool: + try: + self._assert_execution_claim(purge_id, execution_claim) + except Exception: + return False + return True + def _assert_storage_ref_secret_matches(self) -> None: storage_secret = get_governance_ref_secret() if storage_secret != self.ref_secret: @@ -235,6 +399,57 @@ def _completed_barrier_for_retry( ) return barrier + def _completed_erase_result_for_retry( + self, *, subject_ref: str, purge_id: str + ) -> UserEraseResult: + barrier = self._completed_barrier_for_retry( + subject_ref=subject_ref, purge_id=purge_id + ) + if barrier.status != "erased": + raise ValueError("Completed purge retry requires an erased subject barrier") + return UserEraseResult( + subject_ref=subject_ref, + purge_id=purge_id, + status="complete", + deleted_counts=self._deleted_counts_from_targets(purge_id), + rebuilt_agent_playbook_ids=( + self._rebuilt_agent_playbook_ids_from_targets(purge_id) + ), + ) + + def _matching_user_erasure_purge_for_retry( + self, + *, + purge_id: str, + operation_type: str, + scope_type: str, + subject_ref: str, + request_ref: str, + authoritative_user_id: str, + ) -> Any: + purge = self.storage.get_purge_operation(purge_id) + expected_identity = { + "purge_id": purge_id, + "operation_type": operation_type, + "scope_type": scope_type, + "subject_ref": subject_ref, + "request_ref": request_ref, + } + for field_name, expected_value in expected_identity.items(): + if getattr(purge, field_name) != expected_value: + raise ValueError( + "Existing purge operation for idempotency_key has " + f"mismatched {field_name}" + ) + if ( + governance_subject_ref(self.org_id, authoritative_user_id, self.ref_secret) + != purge.subject_ref + ): + raise ValueError( + "Existing purge operation has mismatched authoritative user" + ) + return purge + def _load_user_requests_and_sessions( self, user_id: str ) -> tuple[list[Any], list[dict[str, Any]]]: @@ -310,7 +525,11 @@ def _subject_erasure_lifecycle_complete(self, purge_id: str) -> bool: and (snapshot.detail or {}).get("status") == _LIFECYCLE_COMPLETION_STATUS ) - def _record_subject_erasure_lifecycle_complete(self, purge_id: str) -> None: + def _record_subject_erasure_lifecycle_complete( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + ) -> None: snapshot = self._prepared_target_snapshot(purge_id) if snapshot is None or snapshot.status != "complete": raise ValueError( @@ -327,6 +546,7 @@ def _record_subject_erasure_lifecycle_complete(self, purge_id: str) -> None: detail=detail, deleted_count=snapshot.deleted_count, error_detail=snapshot.error_detail, + execution_claim=execution_claim, ) def _prepared_target_snapshot(self, purge_id: str) -> PurgeOperationTarget | None: @@ -356,7 +576,12 @@ def _rebuilt_agent_playbook_ids_from_targets(self, purge_id: str) -> list[int]: ) ] - def _rebuild_agent_playbooks(self, purge_id: str) -> list[int]: + def _rebuild_agent_playbooks( + self, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, + ) -> list[int]: rebuilt_ids: list[int] = [] for target in self.storage.list_purge_targets( purge_id, @@ -382,6 +607,7 @@ def _rebuild_agent_playbooks(self, purge_id: str) -> list[int]: blocking_issue=rebuild_fields["blocking_issue"], expanded_terms=rebuild_fields["expanded_terms"], tags=rebuild_fields["tags"], + execution_claim=execution_claim, ) rebuilt_ids.append(agent_playbook_id) return rebuilt_ids diff --git a/reflexio/server/services/lineage/gc_scheduler.py b/reflexio/server/services/lineage/gc_scheduler.py index 37d6e9dc4..55cd405cb 100644 --- a/reflexio/server/services/lineage/gc_scheduler.py +++ b/reflexio/server/services/lineage/gc_scheduler.py @@ -126,6 +126,30 @@ def clear_global_sweeps() -> None: _global_sweep_hooks.clear() +# Always-global sweeps run once per elected tick and own their applicability +# checks. Enterprise uses this for global maintenance that must not inherit the +# unrelated expiry-reclamation feature gate. +_always_global_sweep_hooks: list[Callable[[int], int]] = [] + + +def register_always_global_sweep(fn: Callable[[int], int]) -> None: + """Register an always-evaluated global sweep. + + The scheduler's leader gate still controls the tick. The sweep itself must + decide whether its backing service is configured for this deployment. + + Args: + fn (Callable[[int], int]): Called with the current unix epoch; returns + the number of rows it processed. + """ + _always_global_sweep_hooks.append(fn) + + +def clear_always_global_sweeps() -> None: + """Clear all always-evaluated global sweeps (tests restore defaults).""" + _always_global_sweep_hooks.clear() + + # Per-org sweeps run once per org per tick (for per-org reclamation concerns). # Each fn takes (org_id, now) and returns a deleted-row count. Enterprise # registers its closure here at startup so governance retention folds into the @@ -447,7 +471,21 @@ def _run_global_sweeps(self, cfg: object) -> None: capture_anomaly("lineage.global_sweep.failed", sweep=sweep_id) logger.exception("event=global_sweep_failed sweep=%s", sweep_id) + def _run_always_global_sweeps(self) -> None: + """Invoke each applicability-owning global sweep once per elected tick.""" + now = int(time.time()) + for sweep in _always_global_sweep_hooks: + try: + processed = sweep(now) + if processed: + logger.info("event=always_global_sweep processed=%d", processed) + except Exception: + sweep_id = getattr(sweep, "__qualname__", repr(sweep)) + capture_anomaly("lineage.always_global_sweep.failed", sweep=sweep_id) + logger.exception("event=always_global_sweep_failed sweep=%s", sweep_id) + def _run_once(self) -> float: + self._run_always_global_sweeps() poll_interval = _DEFAULT_POLL_INTERVAL_SECONDS try: if ( @@ -537,7 +575,9 @@ def maybe_start_lineage_gc( # must not silence sweeps that were explicitly registered. Production ran # with 8 registered sweeps (metering force-seal/retention/derive, governance # retention, missing-vector backfill) never firing because this read raised. - has_registered_sweeps = bool(_per_org_sweep_hooks) or bool(_global_sweep_hooks) + has_registered_sweeps = bool( + _per_org_sweep_hooks or _global_sweep_hooks or _always_global_sweep_hooks + ) try: ctx = request_context_factory(bootstrap_org_id) diff --git a/reflexio/server/services/playbook/components/aggregator.py b/reflexio/server/services/playbook/components/aggregator.py index b0887d0e7..0b53055cd 100644 --- a/reflexio/server/services/playbook/components/aggregator.py +++ b/reflexio/server/services/playbook/components/aggregator.py @@ -5,7 +5,7 @@ import os import time import uuid -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Sequence from contextlib import AbstractContextManager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Protocol, cast @@ -1079,17 +1079,6 @@ def save_generated_outcome( duration_ms=int((time.perf_counter() - aggregation_start) * 1000), metadata=stats, ) - self._record_learnings_generated( - learning_ids=[ - str(saved.agent_playbook_id) - for saved in saved_playbooks - if getattr(saved, "agent_playbook_id", None) - ], - playbook_name=SINGLETON_USER_PLAYBOOK_NAME, - request_id=run_id, - metadata=stats, - total_count=len(saved_playbooks), - ) return stats def _stable_aggregation_cluster_id(self, fingerprint: str) -> str: @@ -1989,17 +1978,6 @@ def run(self, playbook_aggregator_request: PlaybookAggregatorRequest) -> dict: duration_ms=int((time.perf_counter() - aggregation_start) * 1000), metadata=stats, ) - self._record_learnings_generated( - learning_ids=[ - str(saved.agent_playbook_id) - for saved in saved_playbook_list - if getattr(saved, "agent_playbook_id", None) - ], - playbook_name=playbook_name, - request_id=_run_id, - metadata=stats, - total_count=len(saved_playbook_list), - ) return stats except Exception as e: @@ -2042,61 +2020,6 @@ def run(self, playbook_aggregator_request: PlaybookAggregatorRequest) -> dict: # Re-raise the exception after restoring raise - def _record_learnings_generated( - self, - *, - learning_ids: list[str], - playbook_name: str, - request_id: str, - metadata: Mapping[str, Any], - total_count: int | None = None, - ) -> None: - """Emit ``learnings_generated`` for a completed aggregation run. - - Prefers one event per learning id (entity-backed) when every saved - playbook in this run carries a durable ``agent_playbook_id`` — the - common case, since ``save_agent_playbooks`` - raises rather than returning a partial row. Falls back to the - count-based aggregate event when ``learning_ids`` is short of - ``total_count`` (a falsy/unset id slipped through), mirroring - ``ExtractionResumeWorker._record_finalized_learnings`` — this avoids - emitting a colliding ``learn:agent_playbook:0`` key. ``total_count`` - defaults to ``len(learning_ids)`` so callers that already guarantee a - complete id list (e.g. existing tests) are unaffected. - """ - from reflexio.server.billing_meter import ( - emit_learnings_generated, - emit_learnings_generated_records, - ) - - total = len(learning_ids) if total_count is None else total_count - if len(learning_ids) == total: - emit_learnings_generated_records( - org_id=self.request_context.org_id, - configurator=self.configurator, - learning_ids=learning_ids, - source="aggregation", - pipeline="playbook", - request_id=request_id, - agent_version=self.agent_version, - playbook_name=playbook_name, - entity_type="agent_playbook", - metadata=metadata, - ) - return - emit_learnings_generated( - org_id=self.request_context.org_id, - configurator=self.configurator, - count=total, - source="aggregation", - pipeline="playbook", - request_id=request_id, - agent_version=self.agent_version, - playbook_name=playbook_name, - entity_type="agent_playbook", - metadata=metadata, - ) - def get_clusters( self, user_playbooks: list[UserPlaybook], diff --git a/reflexio/server/services/playbook/service.py b/reflexio/server/services/playbook/service.py index ce12ec491..4f79508fb 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -34,7 +34,11 @@ BaseGenerationService, StatusChangeOperation, ) -from reflexio.server.services.deferred_learning_plan import PlaybookWritePlan +from reflexio.server.services.deferred_learning_plan import ( + FinalizationResult, + PlaybookWritePlan, + _FinalizationReceiptAlreadyExistsError, +) from reflexio.server.services.playbook.aggregation_trigger import ( maybe_trigger_user_playbook_aggregation, ) @@ -628,10 +632,10 @@ def _dispatch_playbook_schedulers(self, plan: PlaybookWritePlan) -> None: Phantom-billing gate: on the durable / ``.run()`` path this is invoked from ``emit_generation_side_effects`` (post-commit), so a fence-lost (superseded) job never enqueues optimization or triggers aggregation. On - the synchronous resume/manual path the permanent - ``_finalize_extracted_items`` wrapper invokes it right after persist, - keeping that path identical to the pre-split monolith. The two callers - are mutually exclusive, so the schedulers fire exactly once per run. + the synchronous resume/manual path ``_finalize_extracted_items`` invokes + it after persistence. Dispatch is best-effort and at most once per + committed finalization attempt; derived scheduler work has no durable + replay idempotency. """ try: self._enqueue_user_playbook_optimization(plan.new_playbooks) @@ -669,16 +673,43 @@ def _finalize_extracted_items( *, model_provenance: ModelProvenance | None = None, extraction_run: AgentRunRecord | None = None, - ) -> list[UserPlaybook]: - """Permanent V3 wrapper: compute→persist→schedulers together (no fence). - - 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. + 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. """ + return self._finalize_extracted_items_with_outcome( + all_playbooks, + model_provenance=model_provenance, + extraction_run=extraction_run, + finalization_run_id=finalization_run_id, + ).learning_ids + + def _finalize_extracted_items_with_outcome( + self, + all_playbooks: list[UserPlaybook], + *, + model_provenance: ModelProvenance | None = None, + extraction_run: AgentRunRecord | None = None, + finalization_run_id: str | None = None, + ) -> FinalizationResult: + """Finalize playbooks and expose the atomic receipt winner internally.""" + 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 previous_review_run = self._review_run @@ -691,12 +722,60 @@ def _finalize_extracted_items( finally: self._review_run = previous_review_run self._review_window_cache = previous_review_window - if plan is None: - return [] - with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] - self._persist_write_plan(plan) - self._dispatch_playbook_schedulers(plan) - return plan.new_playbooks + learning_ids: list[str] = [] + if finalization_run_id is None: + if plan is not None: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + self._persist_write_plan(plan) + self._dispatch_playbook_schedulers(plan) + return 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) def _apply_consolidation_lineage( self, diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index 4b2bd1053..dc847f68e 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -30,7 +30,11 @@ BaseGenerationService, StatusChangeOperation, ) -from reflexio.server.services.deferred_learning_plan import ProfileWritePlan +from reflexio.server.services.deferred_learning_plan import ( + FinalizationResult, + ProfileWritePlan, + _FinalizationReceiptAlreadyExistsError, +) from reflexio.server.services.profile.components.extractor import ProfileExtractor from reflexio.server.services.profile.profile_generation_service_utils import ( ProfileGenerationRequest, @@ -337,23 +341,85 @@ def _finalize_extracted_items( all_new_profiles: list[UserProfile], *, model_provenance: ModelProvenance | None = None, - ) -> list[UserProfile]: + finalization_run_id: str | None = None, + ) -> list[str]: """Permanent V3 wrapper: compute-then-persist together (no external fence). - Kept for the synchronous resume/manual callers - (``ExtractionResumeWorker`` calls this directly). Routes them through the - same ``_resolve_write_plan`` (compute) + ``_persist_write_plan`` - (persist) split the durable worker uses — with no external - ``commit_scope`` — so the result is identical to the pre-split monolith. + Compatibility surface for synchronous callers that expect ordered + learning ids. Routes them through the same ``_resolve_write_plan`` + (compute) + ``_persist_write_plan`` (persist) split the durable worker + uses, with no external ``commit_scope``, so the result is identical to + the pre-split monolith. """ + return self._finalize_extracted_items_with_outcome( + all_new_profiles, + model_provenance=model_provenance, + finalization_run_id=finalization_run_id, + ).learning_ids + + def _finalize_extracted_items_with_outcome( + self, + all_new_profiles: list[UserProfile], + *, + model_provenance: ModelProvenance | None = None, + finalization_run_id: str | None = None, + ) -> FinalizationResult: + """Finalize profiles and expose the atomic receipt winner internally.""" + entity_type = "profile" + if finalization_run_id is not None: + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + 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]) - if plan is None: - return [] - with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] - self._persist_write_plan(plan) - return plan.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: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + 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) def check_and_update_profiles(self, profiles: list[UserProfile]) -> None: """check if the profiles are expired and update them if they are""" diff --git a/reflexio/server/services/search_exposure.py b/reflexio/server/services/search_exposure.py new file mode 100644 index 000000000..e2ab47a4c --- /dev/null +++ b/reflexio/server/services/search_exposure.py @@ -0,0 +1,183 @@ +"""Optional synchronous recording boundary for served user playbooks.""" + +from __future__ import annotations + +from collections.abc import Sized +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, +) + +MAX_EXPOSURE_EVENTS_PER_BATCH = 100 + + +@dataclass(frozen=True) +class SearchExposureBatch: + """The final user-playbook set returned by an authenticated 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 validate_exposure_batch_size(events: Sized) -> None: + """Reject exposure batches that exceed the fixed storage safety bound.""" + if len(events) > MAX_EXPOSURE_EVENTS_PER_BATCH: + raise ValueError( + f"exposure batch must contain at most {MAX_EXPOSURE_EVENTS_PER_BATCH} events" + ) + + +def user_playbook_full_version_fingerprint(playbook: UserPlaybook) -> str: + """Bind every persisted playbook field except its derived embedding vector. + + 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"}) + | { + "governance_subject_ref": playbook.governance_subject_ref, + "retired_at": playbook.retired_at, + }, + } + 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__ = [ + "MAX_EXPOSURE_EVENTS_PER_BATCH", + "SEARCH_EXPOSURE_RECORDER", + "ExposureEventWriteResult", + "SearchExposureBatch", + "SearchExposureRecorder", + "UserPlaybookExposureEvent", + "build_user_playbook_exposure_event", + "record_search_exposures", + "user_playbook_full_version_fingerprint", + "validate_exposure_batch_size", +] diff --git a/reflexio/server/services/storage/error.py b/reflexio/server/services/storage/error.py index cf696ed79..fe14734bd 100644 --- a/reflexio/server/services/storage/error.py +++ b/reflexio/server/services/storage/error.py @@ -22,6 +22,14 @@ class OptimizationJobLeaseLiveError(StorageError): """Raised when an optimizer recovery attempt finds a non-expired lease.""" +class OptimizationJobIdentityConflictError(StorageError): + """Raised when one durable optimizer identity resolves to conflicting jobs.""" + + +class OptimizationArtifactIntegrityError(StorageError): + """Raised when a durable optimizer artifact is malformed or conflicts.""" + + def require_non_empty_session_id(value: Any) -> str: """Return a stripped, non-empty request ``session_id`` or raise ``StorageError``. diff --git a/reflexio/server/services/storage/governance_claims.py b/reflexio/server/services/storage/governance_claims.py new file mode 100644 index 000000000..147d3386b --- /dev/null +++ b/reflexio/server/services/storage/governance_claims.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class PurgeExecutionClaim: + purge_id: str + owner: str + fence: int + expires_at: int + + +def validate_purge_execution_claim( + purge_id: str, + execution_claim: PurgeExecutionClaim | None, +) -> PurgeExecutionClaim: + if execution_claim is None: + raise ValueError("purge execution claim is required") + if type(execution_claim) is not PurgeExecutionClaim: + raise ValueError("purge execution claim must be typed") + if execution_claim.purge_id != purge_id: + raise ValueError("purge execution claim purge_id mismatch") + if not execution_claim.owner.strip(): + raise ValueError("purge execution claim owner is required") + if execution_claim.fence <= 0: + raise ValueError("purge execution claim fence is invalid") + if execution_claim.expires_at <= 0: + raise ValueError("purge execution claim expiry is invalid") + return execution_claim diff --git a/reflexio/server/services/storage/governance_validation.py b/reflexio/server/services/storage/governance_validation.py index a89fa1def..3358604ec 100644 --- a/reflexio/server/services/storage/governance_validation.py +++ b/reflexio/server/services/storage/governance_validation.py @@ -90,6 +90,7 @@ { "affected_agent_playbook_ids", "agent_playbook_id", + "authoritative_user_digest", "count", "deleted_counts", "deleted_count", @@ -166,6 +167,7 @@ "user_playbooks", "profiles", "requests", + "session_outcomes", "agent_success_evaluation_results", "retrieved_learning_evaluation_results", "evaluation_operation_states", @@ -476,6 +478,12 @@ def _validate_governance_detail_entry( if key in {"agent_playbook_id", "user_playbook_id"}: _validate_governance_int(field_name, value) return cast(int, value) + if key == "authoritative_user_digest": + if not isinstance(value, str) or re.fullmatch(r"[0-9a-f]{64}", value) is None: + _raise_governance_validation_error( + field_name, "expected 64 lowercase hex chars" + ) + return value if key == "deleted_counts": return _validate_governance_deleted_counts(field_name, value) if key in { diff --git a/reflexio/server/services/storage/retention.py b/reflexio/server/services/storage/retention.py index d48035689..9dcb86292 100644 --- a/reflexio/server/services/storage/retention.py +++ b/reflexio/server/services/storage/retention.py @@ -8,6 +8,7 @@ DEFAULT_ROW_RETENTION_LIMIT = 250_000 ROW_RETENTION_DELETE_FRACTION = 0.20 +OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS = 14 * 24 * 60 * 60 TOMBSTONE_STATUSES = ("archived", "merged", "superseded", "expired") @@ -20,6 +21,7 @@ class RetentionTarget: order_column: str id_columns: tuple[str, ...] priority_statuses: tuple[str, ...] = () + minimum_age_seconds: int = 0 @dataclass(frozen=True, slots=True) @@ -48,6 +50,12 @@ class OptimizationRetentionClass: ), RetentionTarget("interactions", "interactions", "created_at", ("interaction_id",)), RetentionTarget("requests", "requests", "created_at", ("request_id",)), + RetentionTarget( + "session_outcomes", + "session_outcomes", + "created_at", + ("user_id", "session_id"), + ), RetentionTarget( "user_playbooks", "user_playbooks", @@ -130,6 +138,13 @@ 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 969cab1e2..8a0e4db29 100644 --- a/reflexio/server/services/storage/retention_mixin.py +++ b/reflexio/server/services/storage/retention_mixin.py @@ -9,6 +9,7 @@ from __future__ import annotations +import time from abc import ABC, abstractmethod from collections.abc import Iterator, Sequence from typing import Any @@ -98,7 +99,16 @@ def delete_oldest_retention_target_rows(self, target_name: str, count: int) -> i target = get_retention_target(target_name) if not self._retention_table_exists(target.table_name): return 0 - keys = self._retention_select_keys(target, count) + older_than_epoch = ( + int(time.time()) - target.minimum_age_seconds + if target.minimum_age_seconds > 0 + else None + ) + keys = self._retention_select_keys( + target, + count, + older_than_epoch=older_than_epoch, + ) if not keys: return 0 self._retention_perform_delete(target, keys) @@ -137,7 +147,11 @@ def _retention_gc_retired_optimization_jobs( return 0 def _retention_select_keys( - self, target: RetentionTarget, count: int + self, + target: RetentionTarget, + count: int, + *, + older_than_epoch: int | None, ) -> list[tuple[Any, ...]]: """Select tombstones first, then oldest rows when a target opts in. @@ -147,16 +161,27 @@ def _retention_select_keys( table holds that many rows. """ if not target.priority_statuses: - return self._retention_select_oldest_keys(target, count) + return self._retention_select_oldest_keys( + target, + count, + older_than_epoch=older_than_epoch, + ) keys = self._retention_select_oldest_keys( - target, count, statuses=target.priority_statuses + target, + count, + statuses=target.priority_statuses, + older_than_epoch=older_than_epoch, ) if len(keys) >= count: return keys seen = set(keys) - for key in self._retention_select_oldest_keys(target, count): + for key in self._retention_select_oldest_keys( + target, + count, + older_than_epoch=older_than_epoch, + ): if key not in seen: keys.append(key) seen.add(key) @@ -194,6 +219,7 @@ def _retention_select_oldest_keys( target: RetentionTarget, count: int, statuses: tuple[str, ...] | None = None, + older_than_epoch: int | None = None, ) -> list[tuple[Any, ...]]: """Return up to ``count`` oldest key tuples for ``target``. diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py new file mode 100644 index 000000000..28e0e0064 --- /dev/null +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -0,0 +1,208 @@ +"""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_json_bytes", + "canonical_session_trajectory", + "canonical_trajectory_bytes", + "outcome_contract_digest", + "trajectory_digest", +] + +MAX_CANONICAL_TRAJECTORY_JSON_DEPTH = 100 + + +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, *, depth: int = 0) -> str: + """Encode trajectory JSON deterministically within a bounded nesting depth.""" + if ( + isinstance(value, tuple | list | Mapping) + and depth >= MAX_CANONICAL_TRAJECTORY_JSON_DEPTH + ): + raise ValueError("canonical trajectory JSON exceeds maximum depth") + if isinstance(value, float): + return json.dumps(value, allow_nan=False, separators=(",", ":")) + if isinstance(value, tuple | list): + return ( + "[" + + ",".join( + _canonical_trajectory_json(item, depth=depth + 1) for item in value + ) + + "]" + ) + if isinstance(value, Mapping): + if not all(isinstance(key, str) for key in value): + raise TypeError("canonical trajectory object keys must be strings") + keys = sorted(value, key=lambda key: key.encode("utf-16be")) + return ( + "{" + + ",".join( + f"{canonical_json_bytes(key).decode()}:" + f"{_canonical_trajectory_json(value[key], depth=depth + 1)}" + for key in keys + ) + + "}" + ) + return canonical_json_bytes(value).decode() + + +def canonical_trajectory_bytes(trajectory: object) -> bytes: + """Encode a finalized session trajectory to its canonical UTF-8 bytes.""" + return _canonical_trajectory_json(trajectory).encode("utf-8") + + +def trajectory_digest(trajectory: object) -> str: + """Hash the canonical finalized session trajectory.""" + return sha256(canonical_trajectory_bytes(trajectory)).hexdigest() diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index faba66279..dfb471f59 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -18,10 +18,12 @@ import unicodedata from collections.abc import Callable, Generator, Sequence from datetime import UTC, datetime +from hashlib import sha256 from pathlib import Path from typing import Any, ClassVar, Literal from reflexio.models.api_schema.common import BlockingIssue +from reflexio.models.api_schema.domain import SessionOutcomeKind from reflexio.models.api_schema.service_schemas import ( AgentPlaybook, AgentSuccessEvaluationResult, @@ -58,7 +60,14 @@ StorageError, require_non_empty_session_id, ) -from reflexio.server.services.storage.retention_mixin import RetentionMixin +from reflexio.server.services.storage.retention_mixin import RetentionMixin, chunked +from reflexio.server.services.storage.session_outcome_identity import ( + CanonicalSessionTrajectory, + canonical_json_bytes, + canonical_session_trajectory, + outcome_contract_digest, + trajectory_digest, +) from reflexio.server.services.storage.storage_base import BaseStorage from reflexio.server.site_var.site_var_manager import SiteVarManager @@ -67,6 +76,7 @@ logger = logging.getLogger(__name__) +_OUTCOME_ALLOWED_VALUES = tuple(kind.value for kind in SessionOutcomeKind) _MINIMUM_SQLITE_VERSION = (3, 35, 0) _SQLITE_INITIALIZATION_LOCK_STRIPES = 64 _sqlite_initialization_locks = tuple( @@ -101,6 +111,109 @@ 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 _prefetch_canonical_session_trajectory_digests( + conn: sqlite3.Connection, session_ids: Sequence[str] +) -> dict[str, str]: + """Derive trajectory digests while retaining only one payload chunk.""" + digests: dict[str, str] = {} + for session_id_chunk in chunked(list(dict.fromkeys(session_ids))): + placeholders = ",".join("?" for _ in session_id_chunk) + request_rows = conn.execute( + f"""SELECT request_id, user_id, created_at, source, agent_version, + session_id, evaluation_only, retrieval_experiment_id, + retrieval_experiment_arm + FROM requests WHERE session_id IN ({placeholders}) + ORDER BY session_id ASC, created_at ASC, request_id ASC""", # noqa: S608 + session_id_chunk, + ).fetchall() + interaction_rows = conn.execute( + f"""SELECT requests.session_id AS trajectory_session_id, + interactions.interaction_id, interactions.user_id, + interactions.request_id, interactions.created_at, + interactions.content, interactions.role, + interactions.token_count, interactions.user_action, + interactions.user_action_description, + interactions.interacted_image_url, + interactions.image_encoding, interactions.shadow_content, + interactions.expert_content, interactions.tools_used, + interactions.citations, interactions.retrieved_learnings + FROM interactions + JOIN requests ON requests.request_id = interactions.request_id + WHERE requests.session_id IN ({placeholders}) + ORDER BY requests.session_id ASC, interactions.request_id ASC, + interactions.created_at ASC, + interactions.interaction_id ASC""", # noqa: S608 + session_id_chunk, + ).fetchall() + + requests_by_session: dict[str, list[dict[str, object]]] = { + session_id: [] for session_id in session_id_chunk + } + interactions_by_session: dict[str, dict[str, list[dict[str, object]]]] = { + session_id: {} for session_id in session_id_chunk + } + for request in request_rows: + session_id = str(request["session_id"]) + requests_by_session[session_id].append(dict(request)) + interactions_by_session[session_id][str(request["request_id"])] = [] + for interaction in interaction_rows: + interaction_payload = dict(interaction) + session_id = str(interaction_payload.pop("trajectory_session_id")) + interactions_by_session[session_id][ + str(interaction_payload["request_id"]) + ].append(interaction_payload) + for session_id in session_id_chunk: + digests[session_id] = trajectory_digest( + canonical_session_trajectory( + session_id, + requests_by_session[session_id], + interactions_by_session[session_id], + ) + ) + return digests + + +def _legacy_session_outcome_id(user_id: str, session_id: str) -> str: + """Return a delimiter-safe immutable identity for a migrated legacy outcome.""" + return sha256(canonical_json_bytes([user_id, session_id])).hexdigest() + + _FTS5_OPERATORS = frozenset({"OR", "AND", "NOT"}) _FTS5_RESERVED = _FTS5_OPERATORS | {"NEAR"} _TOKEN_RE = re.compile(r"\w+", re.UNICODE) @@ -915,13 +1028,27 @@ def migrate(self) -> bool: # _DDL creates an index over these columns. Upgrade legacy request tables # before executescript so index creation cannot fail on missing columns. self._migrate_request_retrieval_experiment() - self._migrate_session_outcomes_schema() with self._lock: + session_outcome_columns = { + row["name"] + for row in self.conn.execute( + "PRAGMA table_info(session_outcomes)" + ).fetchall() + } + if ( + session_outcome_columns + and "governance_subject_ref" not in session_outcome_columns + ): + self.conn.execute( + "ALTER TABLE session_outcomes " + "ADD COLUMN governance_subject_ref TEXT" + ) cur = self.conn.cursor() cur.executescript(_DDL) init_governance_tables(self.conn) init_playbook_aggregation_tables(self.conn) self.conn.commit() + self._migrate_session_outcomes_schema() if self._has_sqlite_vec: self._create_vec_tables() self._migrate_vec_tables() @@ -948,6 +1075,7 @@ def migrate(self) -> bool: self._migrate_user_playbook_publication_staging_columns() self._classify_legacy_playbook_optimization_jobs() self._enforce_playbook_optimization_job_constraints() + self._enforce_playbook_optimization_artifact_constraints() self._migrate_retire_profile_change_logs() self._migrate_retire_playbook_aggregation_change_logs() init_stall_state_table(self.conn) @@ -1020,120 +1148,6 @@ def _migrate_unicode_lexical_indexes(self) -> None: self.conn.rollback() raise - def _migrate_session_outcomes_schema(self) -> None: - """Restore the pre-identity outcome schema after downgrading #407.""" - with self._lock: - table_info = self.conn.execute( - "PRAGMA table_info(session_outcomes)" - ).fetchall() - if not table_info: - return - expected_columns = { - "user_id", - "session_id", - "outcome", - "occurred_at", - "source", - "label", - "value", - "metadata", - "governance_subject_ref", - "created_at", - } - columns = {str(row["name"]): row for row in table_info} - table = self.conn.execute( - "SELECT sql FROM sqlite_master WHERE type = 'table' " - "AND name = 'session_outcomes'" - ).fetchone() - table_sql = str(table["sql"] or "") if table is not None else "" - governance_column = columns.get("governance_subject_ref") - has_empty_subject_default = ( - governance_column is not None - and governance_column["dflt_value"] in ("''", '""') - ) - identity_columns = { - "outcome_id", - "outcome_revision", - "outcome_contract_digest", - "finalized_trajectory_digest", - } - if ( - expected_columns.issubset(columns) - and not identity_columns.intersection(columns) - and governance_column is not None - and int(governance_column["notnull"]) == 1 - and not has_empty_subject_default - and "'unknown'" not in table_sql - ): - return - - self.conn.execute("BEGIN IMMEDIATE") - try: - legacy_rows = self.conn.execute( - "SELECT * FROM session_outcomes" - ).fetchall() - self.conn.execute( - "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" - ) - self.conn.execute( - """CREATE TABLE session_outcomes ( - user_id TEXT NOT NULL, - session_id TEXT NOT NULL, - outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), - occurred_at INTEGER NOT NULL, - source TEXT NOT NULL, - label TEXT, - value REAL, - metadata TEXT, - governance_subject_ref TEXT NOT NULL, - created_at INTEGER NOT NULL, - PRIMARY KEY (user_id, session_id) - )""" - ) - compatible_rows = [ - row - for row in legacy_rows - if row["outcome"] in ("success", "failure") - ] - for row in compatible_rows: - subject_ref = ( - row["governance_subject_ref"] - if "governance_subject_ref" in columns - else None - ) - if subject_ref is None or not str(subject_ref).strip(): - subject_ref = self._subject_ref_for_user_id(str(row["user_id"])) - self.conn.execute( - """INSERT INTO session_outcomes ( - user_id, session_id, outcome, occurred_at, source, - label, value, metadata, governance_subject_ref, created_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - row["user_id"], - row["session_id"], - row["outcome"], - row["occurred_at"], - row["source"], - row["label"] if "label" in columns else None, - row["value"] if "value" in columns else None, - row["metadata"] if "metadata" in columns else None, - subject_ref, - row["created_at"], - ), - ) - dropped_unknown = len(legacy_rows) - len(compatible_rows) - if dropped_unknown: - logger.warning( - "Dropped %d unrepresentable 'unknown' session outcomes while " - "restoring the pre-#407 schema", - dropped_unknown, - ) - self.conn.execute("DROP TABLE session_outcomes_legacy") - self.conn.commit() - except Exception: - self.conn.rollback() - raise - def _try_load_sqlite_vec(self) -> bool: """Attempt to load the sqlite-vec extension for native KNN search. @@ -1237,6 +1251,191 @@ 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: + table_info = self.conn.execute( + "PRAGMA table_info(session_outcomes)" + ).fetchall() + columns = {str(row["name"]): row for row in table_info} + table = self.conn.execute( + "SELECT sql FROM sqlite_master WHERE type = 'table' " + "AND name = 'session_outcomes'" + ).fetchone() + table_sql = (table["sql"] if table is not None else "") or "" + identity_columns = { + "outcome_id", + "outcome_revision", + "outcome_contract_digest", + "finalized_trajectory_digest", + } + expected_columns = identity_columns | { + "user_id", + "session_id", + "outcome", + "occurred_at", + "source", + "label", + "value", + "metadata", + "governance_subject_ref", + "created_at", + } + subject_column = columns.get("governance_subject_ref") + has_empty_subject_default = ( + subject_column is not None + and subject_column["dflt_value"] in ("''", '""') + ) + schema_is_canonical = ( + expected_columns.issubset(columns) + and "'unknown'" in table_sql + and subject_column is not None + and int(subject_column["notnull"]) == 1 + and not has_empty_subject_default + ) + if schema_is_canonical: + rows_missing_subject_ref = self.conn.execute( + """SELECT user_id, session_id FROM session_outcomes + WHERE governance_subject_ref IS NULL + 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() + trajectory_digests = _prefetch_canonical_session_trajectory_digests( + self.conn, + [ + str(row["session_id"]) + for row in legacy_rows + if "finalized_trajectory_digest" not in columns + or not row["finalized_trajectory_digest"] + ], + ) + self.conn.execute( + "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" + ) + 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 "governance_subject_ref" in columns + else None + ) + if subject_ref is None or not str(subject_ref).strip(): + subject_ref = self._subject_ref_for_user_id(str(row["user_id"])) + existing_outcome_id = ( + row["outcome_id"] if "outcome_id" in columns else None + ) + existing_contract_digest = ( + row["outcome_contract_digest"] + if "outcome_contract_digest" in columns + else None + ) + existing_trajectory_digest = ( + row["finalized_trajectory_digest"] + if "finalized_trajectory_digest" in columns + else None + ) + self.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + existing_outcome_id + or _legacy_session_outcome_id( + str(row["user_id"]), str(row["session_id"]) + ), + ( + row["outcome_revision"] + if "outcome_revision" in columns + and row["outcome_revision"] is not None + else 1 + ), + row["user_id"], + row["session_id"], + row["outcome"], + row["occurred_at"], + source, + row["label"] if "label" in columns else None, + row["value"] if "value" in columns else None, + row["metadata"] if "metadata" in columns else None, + existing_contract_digest + or outcome_contract_digest( + source=source, + schema_version=1, + allowed_values=_OUTCOME_ALLOWED_VALUES, + finalization_rule="first_write", + ), + existing_trajectory_digest + or trajectory_digests[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: @@ -1889,7 +2088,9 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: "CHECK (terminal_outcome IS NULL OR terminal_outcome IN", "'governance_erased'", ) - if all(check in table_sql for check in required_checks): + if all(check in table_sql for check in required_checks) and ( + "'offline_tuner_open_world'" not in table_sql + ): return foreign_keys_enabled = bool( self.conn.execute("PRAGMA foreign_keys").fetchone()[0] @@ -1899,6 +2100,10 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: self.conn.execute("PRAGMA foreign_keys=OFF") try: self.conn.execute("BEGIN IMMEDIATE") + sequence_high_water = self._autoincrement_high_water( + "playbook_optimization_jobs", + "job_id", + ) self.conn.execute("DROP TABLE IF EXISTS playbook_optimization_jobs_new") self.conn.execute( """ @@ -1990,6 +2195,10 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: "ALTER TABLE playbook_optimization_jobs_new " "RENAME TO playbook_optimization_jobs" ) + self._restore_autoincrement_high_water( + "playbook_optimization_jobs", + sequence_high_water, + ) self.conn.execute( "CREATE INDEX idx_poj_target " "ON playbook_optimization_jobs(target_kind, target_id)" @@ -2026,6 +2235,125 @@ def _enforce_playbook_optimization_job_constraints(self) -> None: f"PRAGMA foreign_keys={'ON' if foreign_keys_enabled else 'OFF'}" ) + def _enforce_playbook_optimization_artifact_constraints(self) -> None: + """Rebuild legacy artifact tables with the current kind allowlist.""" + table_sql_row = self.conn.execute( + """SELECT sql FROM sqlite_master + WHERE type = 'table' AND name = 'playbook_optimization_artifacts'""" + ).fetchone() + if table_sql_row is None: + return + table_sql = table_sql_row["sql"] + artifact_kinds = ( + "'expected_population_manifest'", + "'generation_selection'", + "'replay_manifest'", + "'candidate'", + "'candidate_search_projection'", + "'open_world_evidence_bundle'", + ) + if all(artifact_kind in table_sql for artifact_kind in artifact_kinds): + return + + foreign_keys_enabled = bool( + self.conn.execute("PRAGMA foreign_keys").fetchone()[0] + ) + self.conn.commit() + if foreign_keys_enabled: + self.conn.execute("PRAGMA foreign_keys=OFF") + try: + self.conn.execute("BEGIN IMMEDIATE") + sequence_high_water = self._autoincrement_high_water( + "playbook_optimization_artifacts", + "artifact_id", + ) + self.conn.execute( + "DROP TABLE IF EXISTS playbook_optimization_artifacts_new" + ) + self.conn.execute( + """ + CREATE TABLE playbook_optimization_artifacts_new ( + artifact_id INTEGER PRIMARY KEY AUTOINCREMENT, + job_id INTEGER NOT NULL, + artifact_kind TEXT NOT NULL CHECK (artifact_kind IN ( + 'expected_population_manifest', + 'generation_selection', + 'replay_manifest', + 'candidate', + 'candidate_search_projection', + 'open_world_evidence_bundle' + )), + content_json TEXT NOT NULL, + content_digest TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (job_id, artifact_kind), + FOREIGN KEY (job_id) + REFERENCES playbook_optimization_jobs(job_id) + ON DELETE CASCADE + ) + """ + ) + self.conn.execute( + """ + INSERT INTO playbook_optimization_artifacts_new ( + artifact_id, job_id, artifact_kind, content_json, + content_digest, created_at, updated_at + ) SELECT + artifact_id, job_id, artifact_kind, content_json, + content_digest, created_at, updated_at + FROM playbook_optimization_artifacts + """ + ) + self.conn.execute("DROP TABLE playbook_optimization_artifacts") + self.conn.execute( + "ALTER TABLE playbook_optimization_artifacts_new " + "RENAME TO playbook_optimization_artifacts" + ) + self._restore_autoincrement_high_water( + "playbook_optimization_artifacts", + sequence_high_water, + ) + self.conn.execute( + "CREATE INDEX idx_poa_job ON playbook_optimization_artifacts(job_id)" + ) + violations = self.conn.execute("PRAGMA foreign_key_check").fetchall() + if violations: + raise sqlite3.IntegrityError( + "foreign key check failed after optimizer artifact migration: " + f"{violations}" + ) + self.conn.commit() + except Exception: + self.conn.rollback() + raise + finally: + self.conn.execute( + f"PRAGMA foreign_keys={'ON' if foreign_keys_enabled else 'OFF'}" + ) + + def _autoincrement_high_water(self, table_name: str, id_column: str) -> int: + sequence_row = self.conn.execute( + "SELECT MAX(seq) FROM sqlite_sequence WHERE name = ?", + (table_name,), + ).fetchone() + maximum_row = self.conn.execute( + f"SELECT MAX({id_column}) FROM {table_name}" # noqa: S608 + ).fetchone() + return max(sequence_row[0] or 0, maximum_row[0] or 0) + + def _restore_autoincrement_high_water( + self, + table_name: str, + high_water: int, + ) -> None: + self.conn.execute("DELETE FROM sqlite_sequence WHERE name = ?", (table_name,)) + if high_water: + self.conn.execute( + "INSERT INTO sqlite_sequence(name, seq) VALUES (?, ?)", + (table_name, high_water), + ) + def _migrate_retire_profile_change_logs(self) -> None: """Retire the frozen ``profile_change_logs`` table via a reversible RENAME. @@ -2274,6 +2602,9 @@ def _migrate_request_session_id_required(self) -> None: retrieval_experiment_arm_expr = ( "retrieval_experiment_arm" if "retrieval_experiment_arm" in cols else "NULL" ) + governance_subject_ref_expr = ( + "governance_subject_ref" if "governance_subject_ref" in cols else "NULL" + ) # NOTE: this rebuild hardcodes the full `requests` column set. If a # future migration adds a column to `requests`, it MUST be added here # too (and to the SELECT below) or the rebuild will silently drop it. @@ -2288,7 +2619,8 @@ def _migrate_request_session_id_required(self) -> None: session_id TEXT NOT NULL CHECK (trim(session_id) != ''), evaluation_only INTEGER NOT NULL DEFAULT 0, retrieval_experiment_id TEXT, - retrieval_experiment_arm TEXT + retrieval_experiment_arm TEXT, + governance_subject_ref TEXT ); INSERT INTO requests_new ( @@ -2300,7 +2632,8 @@ def _migrate_request_session_id_required(self) -> None: session_id, evaluation_only, retrieval_experiment_id, - retrieval_experiment_arm + retrieval_experiment_arm, + governance_subject_ref ) SELECT request_id, @@ -2315,7 +2648,8 @@ def _migrate_request_session_id_required(self) -> None: END, {evaluation_only_expr}, {retrieval_experiment_id_expr}, - {retrieval_experiment_arm_expr} + {retrieval_experiment_arm_expr}, + {governance_subject_ref_expr} FROM requests; DROP TABLE requests; ALTER TABLE requests_new RENAME TO requests; @@ -2324,6 +2658,8 @@ def _migrate_request_session_id_required(self) -> None: CREATE INDEX IF NOT EXISTS idx_requests_created_at ON requests(created_at); CREATE INDEX IF NOT EXISTS idx_requests_retrieval_experiment ON requests(retrieval_experiment_id, user_id, session_id, created_at, request_id); + CREATE INDEX IF NOT EXISTS idx_requests_governance_subject_ref + ON requests(governance_subject_ref); """ ) self.conn.commit() @@ -2618,7 +2954,6 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: "SELECT rowid, profile_id FROM profiles WHERE user_id = ?", (user_id,), ).fetchall() - subject_ref = self._subject_ref_for_user_id(user_id) # Build a rowid lookup for FTS/vec cleanup (SQLite-specific need). profile_rowid_by_id: dict[str, int] = { @@ -2669,9 +3004,8 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: "DELETE FROM interactions WHERE user_id = ?", (user_id,) ) session_outcomes_cur = self.conn.execute( - """DELETE FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "DELETE FROM session_outcomes WHERE user_id = ?", + (user_id,), ) requests_cur = self.conn.execute( "DELETE FROM requests WHERE user_id = ?", (user_id,) @@ -2800,14 +3134,18 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: ON requests(retrieval_experiment_id, user_id, session_id, created_at, request_id); CREATE TABLE IF NOT EXISTS session_outcomes ( + outcome_id TEXT NOT NULL UNIQUE, + outcome_revision INTEGER NOT NULL CHECK (outcome_revision >= 1), user_id TEXT NOT NULL, session_id TEXT NOT NULL, - outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure', 'unknown')), occurred_at INTEGER NOT NULL, source TEXT NOT NULL, label TEXT, value REAL, metadata TEXT, + outcome_contract_digest TEXT NOT NULL, + finalized_trajectory_digest TEXT NOT NULL, governance_subject_ref TEXT NOT NULL, created_at INTEGER NOT NULL, PRIMARY KEY (user_id, session_id) @@ -3041,7 +3379,8 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: 'generation_selection', 'replay_manifest', 'candidate', - 'candidate_search_projection' + 'candidate_search_projection', + 'open_world_evidence_bundle' )), content_json TEXT NOT NULL, content_digest TEXT NOT NULL, @@ -3187,6 +3526,14 @@ def clear_user_data(self, user_id: str) -> dict[str, int]: CREATE INDEX IF NOT EXISTS idx_agent_runs_ready ON _agent_runs(status, next_resume_at, updated_at); CREATE INDEX IF NOT EXISTS idx_agent_runs_binding ON _agent_runs(org_id, extractor_kind, user_id); +CREATE TABLE IF NOT EXISTS _agent_run_finalization_receipts ( + run_id TEXT PRIMARY KEY, + entity_type TEXT NOT NULL CHECK (entity_type IN ('profile', 'user_playbook')), + learning_ids TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + FOREIGN KEY (run_id) REFERENCES _agent_runs(id) ON DELETE CASCADE +); + CREATE TABLE IF NOT EXISTS _pending_tool_calls ( id TEXT PRIMARY KEY, org_id TEXT NOT NULL, diff --git a/reflexio/server/services/storage/sqlite_storage/_governance.py b/reflexio/server/services/storage/sqlite_storage/_governance.py index 5e4abad09..922691fc7 100644 --- a/reflexio/server/services/storage/sqlite_storage/_governance.py +++ b/reflexio/server/services/storage/sqlite_storage/_governance.py @@ -72,12 +72,16 @@ subject_ref TEXT, request_ref TEXT NOT NULL, idempotency_key TEXT NOT NULL, + authoritative_user_digest TEXT, status TEXT NOT NULL DEFAULT 'pending', error_code TEXT, error_detail TEXT, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, completed_at INTEGER, + execution_claim_owner TEXT, + execution_claim_fence INTEGER NOT NULL DEFAULT 0, + execution_claim_expires_at INTEGER, PRIMARY KEY (org_id, purge_id) ); CREATE UNIQUE INDEX IF NOT EXISTS idx_purge_operations_org_idem @@ -109,6 +113,29 @@ def init_governance_tables(conn: sqlite3.Connection) -> None: conn.executescript(GOVERNANCE_DDL) _enforce_audit_request_ref_not_null(conn) _ensure_governance_subject_ref_columns(conn) + _ensure_purge_operation_execution_claim_columns(conn) + + +def _ensure_purge_operation_execution_claim_columns(conn: sqlite3.Connection) -> None: + columns = [row[1] for row in conn.execute("PRAGMA table_info(purge_operations)")] + if not columns: + return + if "execution_claim_owner" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN execution_claim_owner TEXT" + ) + if "execution_claim_fence" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN execution_claim_fence INTEGER NOT NULL DEFAULT 0" + ) + if "execution_claim_expires_at" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN execution_claim_expires_at INTEGER" + ) + if "authoritative_user_digest" not in columns: + conn.execute( + "ALTER TABLE purge_operations ADD COLUMN authoritative_user_digest TEXT" + ) def _ensure_governance_subject_ref_columns(conn: sqlite3.Connection) -> None: @@ -404,11 +431,9 @@ def _planned_governance_delete_counts( "SELECT COUNT(DISTINCT session_id) AS cnt FROM requests WHERE user_id = ?", (user_id,), ).fetchone() - subject_ref = self._deps()._subject_ref_for_user_id(user_id) session_outcome_row = self.conn.execute( - """SELECT COUNT(*) AS cnt FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "SELECT COUNT(*) AS cnt FROM session_outcomes WHERE user_id = ?", + (user_id,), ).fetchone() profile_rows = self.conn.execute( "SELECT profile_id FROM profiles WHERE user_id = ?", diff --git a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py index aaec5a1b2..d0657f4eb 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -3,6 +3,7 @@ import json import sqlite3 from typing import Any +from uuid import uuid4 from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, @@ -11,12 +12,51 @@ SetSessionOutcomeRequest, ) from reflexio.server.services.storage.error import SubjectWriteBarrierError +from reflexio.server.services.storage.session_outcome_identity import ( + outcome_contract_digest, + trajectory_digest, +) from reflexio.server.services.storage.storage_base._session_outcomes import ( SessionOutcomeContext, SessionOutcomeWriteResult, ) -from ._base import SQLiteStorageBase, _iso_to_epoch +from ._base import ( + _OUTCOME_ALLOWED_VALUES, + SQLiteStorageBase, + _canonical_session_snapshot, + _iso_to_epoch, +) + +_OUTCOME_SCHEMA_VERSION = 1 +_OUTCOME_FINALIZATION_RULE = "first_write" + + +def _canonical_metadata_json(metadata: object) -> str | None: + if metadata is None: + return None + return json.dumps( + metadata, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + + +def _metadata_matches(*, stored_metadata: str | None, request_metadata: object) -> bool: + if stored_metadata is None: + stored_value = None + else: + try: + stored_value = json.loads(stored_metadata) + except (RecursionError, TypeError, ValueError): + return False + try: + return _canonical_metadata_json(stored_value) == _canonical_metadata_json( + request_metadata + ) + except (RecursionError, TypeError, ValueError): + return False class SessionOutcomeStoreMixin: @@ -67,15 +107,94 @@ def record_session_outcome( try: self.conn.execute("BEGIN IMMEDIATE") existing = self.conn.execute( - "SELECT user_id, source FROM session_outcomes WHERE session_id = ?", + "SELECT * FROM session_outcomes WHERE session_id = ?", (request.session_id,), ).fetchone() if existing is not None: + first = self.conn.execute( + """SELECT user_id, source, governance_subject_ref + FROM requests WHERE session_id = ? + ORDER BY created_at ASC, request_id ASC LIMIT 1""", + (request.session_id,), + ).fetchone() + source = ( + str(first["source"]) + if first is not None + else str(existing["source"]) + ) + subject_ref = ( + str( + first["governance_subject_ref"] + or self._subject_ref_for_user_id(str(first["user_id"])) + ) + if first is not None + else str(existing["governance_subject_ref"]) + ) + contract_digest = self._outcome_contract_digest(source) + current_snapshot_digest = ( + trajectory_digest( + _canonical_session_snapshot(self.conn, request.session_id) + ) + if first is not None + else None + ) + stored_contract_digest = existing["outcome_contract_digest"] + stored_snapshot_digest = existing["finalized_trajectory_digest"] + server_context_matches = first is None or ( + str(existing["user_id"]) == str(first["user_id"]) + and str(existing["source"]) == str(first["source"]) + and str(existing["governance_subject_ref"]) == subject_ref + ) + exact_retry = ( + existing["outcome"] == str(request.outcome) + and int(existing["occurred_at"]) == request.occurred_at + and existing["label"] == request.label + and existing["value"] == request.value + and _metadata_matches( + stored_metadata=existing["metadata"], + request_metadata=request.metadata, + ) + and server_context_matches + and ( + stored_contract_digest is None + or stored_contract_digest == contract_digest + ) + and ( + stored_snapshot_digest is None + or current_snapshot_digest is None + or stored_snapshot_digest == current_snapshot_digest + ) + ) self.conn.rollback() return SessionOutcomeWriteResult( recorded=False, user_id=str(existing["user_id"]), source=str(existing["source"]), + reason=( + None + if exact_retry + else SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + ), + outcome_id=( + str(existing["outcome_id"]) + if existing["outcome_id"] is not None + else None + ), + outcome_revision=( + int(existing["outcome_revision"]) + if existing["outcome_revision"] is not None + else None + ), + outcome_contract_digest=( + str(stored_contract_digest) + if stored_contract_digest is not None + else None + ), + finalized_trajectory_digest=( + str(stored_snapshot_digest) + if stored_snapshot_digest is not None + else None + ), ) first = self.conn.execute( """SELECT user_id, source, created_at, request_id @@ -122,31 +241,44 @@ def record_session_outcome( user_id=user_id, reason=SessionOutcomeFailureReason.SUBJECT_NOT_WRITABLE, ) + contract_digest = self._outcome_contract_digest(source) + snapshot_digest = trajectory_digest( + _canonical_session_snapshot(self.conn, request.session_id) + ) + outcome_id = uuid4().hex self.conn.execute( """INSERT INTO session_outcomes - (user_id, session_id, outcome, occurred_at, source, label, value, - metadata, governance_subject_ref, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, + governance_subject_ref, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( + outcome_id, + 1, user_id, request.session_id, - request.outcome.value, + str(request.outcome), request.occurred_at, source, request.label, request.value, - json.dumps( - request.metadata, sort_keys=True, separators=(",", ":") - ) - if request.metadata is not None - else None, + self._metadata_json(request), + contract_digest, + snapshot_digest, subject_ref, created_at, ), ) self.conn.commit() return SessionOutcomeWriteResult( - recorded=True, user_id=user_id, source=source + recorded=True, + user_id=user_id, + source=source, + outcome_id=outcome_id, + outcome_revision=1, + outcome_contract_digest=contract_digest, + finalized_trajectory_digest=snapshot_digest, ) except Exception: self.conn.rollback() @@ -179,13 +311,17 @@ def get_session_outcomes( params.append(request.end_time) where = f" WHERE {' AND '.join(clauses)}" if clauses else "" rows = self.conn.execute( - f"""SELECT user_id, session_id, outcome, occurred_at, source, label, value, - metadata, created_at FROM session_outcomes{where} + f"""SELECT outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, created_at + FROM session_outcomes{where} ORDER BY occurred_at DESC, user_id ASC, session_id ASC LIMIT ? OFFSET ?""", [*params, request.top_k, request.offset], ).fetchall() return [ SessionOutcomeRecord( + outcome_id=row["outcome_id"], + outcome_revision=row["outcome_revision"], user_id=row["user_id"], session_id=row["session_id"], outcome=row["outcome"], @@ -194,6 +330,8 @@ def get_session_outcomes( label=row["label"], value=row["value"], metadata=json.loads(row["metadata"]) if row["metadata"] else None, + outcome_contract_digest=row["outcome_contract_digest"], + finalized_trajectory_digest=row["finalized_trajectory_digest"], created_at=row["created_at"], ) for row in rows @@ -201,14 +339,25 @@ def get_session_outcomes( @SQLiteStorageBase.handle_exceptions def clear_session_outcomes_for_user(self, user_id: str) -> dict[str, int]: - subject_ref = self._subject_ref_for_user_id(user_id) with self._lock: outcome_cursor = self.conn.execute( - """DELETE FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "DELETE FROM session_outcomes WHERE user_id = ?", + (user_id,), ) self.conn.commit() return { "session_outcomes": int(outcome_cursor.rowcount or 0), } + + @staticmethod + def _metadata_json(request: SetSessionOutcomeRequest) -> str | None: + return _canonical_metadata_json(request.metadata) + + @staticmethod + def _outcome_contract_digest(source: str) -> str: + return outcome_contract_digest( + source=source, + schema_version=_OUTCOME_SCHEMA_VERSION, + allowed_values=_OUTCOME_ALLOWED_VALUES, + finalization_rule=_OUTCOME_FINALIZATION_RULE, + ) diff --git a/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py b/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py index 3858b29b3..4ecc0fc05 100644 --- a/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py +++ b/reflexio/server/services/storage/sqlite_storage/agent_run/_agent_run_store.py @@ -15,6 +15,7 @@ from __future__ import annotations +import json import sqlite3 from datetime import UTC, datetime, timedelta from typing import Any @@ -29,6 +30,13 @@ from .._base import SQLiteStorageBase, _json_dumps +def _valid_finalized_learning_ids(value: object) -> bool: + return isinstance(value, list) and all( + isinstance(learning_id, str) and bool(learning_id.strip()) + for learning_id in value + ) + + class SQLiteAgentRunStoreMixin: """SQLite-backed resumable extraction run store primitives.""" @@ -37,6 +45,7 @@ class SQLiteAgentRunStoreMixin: _fetchone: Any _fetchall: Any _current_timestamp: Any + _own_transaction: Any org_id: str def _finalize_runs_without_pending_dependencies_unlocked(self, now_s: str) -> None: @@ -213,6 +222,87 @@ 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 a7227d222..3f3394170 100644 --- a/reflexio/server/services/storage/sqlite_storage/base/_deletion.py +++ b/reflexio/server/services/storage/sqlite_storage/base/_deletion.py @@ -50,16 +50,21 @@ def _retention_select_oldest_keys( target: RetentionTarget, count: int, statuses: tuple[str, ...] | None = None, + older_than_epoch: int | None = None, ) -> list[tuple[Any, ...]]: if statuses is not None and not statuses: return [] id_sql = ", ".join(target.id_columns) - where_sql = "" + predicates: list[str] = [] params: list[Any] = [] if statuses: placeholders = ", ".join("?" for _ in statuses) - where_sql = f"WHERE status IN ({placeholders}) " + predicates.append(f"status IN ({placeholders})") params.extend(statuses) + if older_than_epoch is not None: + predicates.append(f"{target.order_column} < ?") + params.append(older_than_epoch) + where_sql = f"WHERE {' AND '.join(predicates)} " if predicates else "" params.append(count) rows = self._fetchall( f"SELECT {id_sql} FROM {target.table_name} {where_sql}" # noqa: S608 diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py index 954eae6c1..4333632a3 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_erase_execution.py @@ -39,6 +39,7 @@ AuditEvent, PurgeOperation, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, _PREPARE_PHASE, @@ -88,6 +89,11 @@ class GovernanceEraseExecutionMixin: ] get_purge_operation: Callable[[str], PurgeOperation] _record_purge_target_locked: Callable[..., None] + _assert_purge_operation_execution_claim_locked: Callable[ + [str, PurgeExecutionClaim | None], None + ] + _assert_authoritative_user_identity_locked: Callable[[str, str], str] + _assert_bound_authoritative_user_identity_locked: Callable[[str, str, str], None] def _purge_governance_entity_content_locked( self, @@ -143,11 +149,9 @@ def _clear_user_data_for_governance_locked( expected_user_playbook_ids: set[int] | None = None, ) -> dict[str, int]: deps = self._deps() - subject_ref = deps._subject_ref_for_user_id(user_id) session_outcomes_cur = self.conn.execute( - """DELETE FROM session_outcomes - WHERE user_id = ? OR governance_subject_ref = ?""", - (user_id, subject_ref), + "DELETE FROM session_outcomes WHERE user_id = ?", + (user_id,), ) interaction_ids = [ int(row["interaction_id"]) @@ -346,7 +350,11 @@ def _delete_evaluation_operation_states_locked( return deleted def apply_governance_user_data_delete( - self, purge_id: str, user_id: str + self, + purge_id: str, + user_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> dict[str, int]: purge_id = _validate_governance_purge_id("purge_id", purge_id) name_map = { @@ -370,6 +378,10 @@ def apply_governance_user_data_delete( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) + self._assert_authoritative_user_identity_locked(purge_id, user_id) self._validate_prepared_delete_target_matrix_locked(purge_id) self._validate_hide_for_rebuild_targets_locked(purge_id) expected_user_playbook_ids = ( @@ -404,7 +416,12 @@ def apply_governance_user_data_delete( return counts def complete_purge_operation_with_audit( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) if audit_event.org_id != self.org_id: @@ -420,6 +437,9 @@ def complete_purge_operation_with_audit( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) row = self.conn.execute( "SELECT * FROM purge_operations WHERE purge_id = ? AND org_id = ?", (purge_id, self.org_id), @@ -452,6 +472,11 @@ def complete_purge_operation_with_audit( raise ValueError( "Cannot complete purge without target snapshot marker" ) + self._assert_bound_authoritative_user_identity_locked( + purge_id, + audit_event.subject_ref or "", + authoritative_user_id, + ) delete_rows = self.conn.execute( """SELECT target_name, status FROM purge_operation_targets WHERE org_id = ? AND purge_id = ? AND phase = 'delete' @@ -531,7 +556,9 @@ def complete_purge_operation_with_audit( error_code = NULL, error_detail = NULL, updated_at = ?, - completed_at = ? + completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ?""", (now, now, purge_id, self.org_id), ) diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py index 1f2a7f0e9..2087ee2c6 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_purge.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_purge.py @@ -18,6 +18,8 @@ from __future__ import annotations +import hashlib +import hmac import sqlite3 import threading from collections.abc import Callable @@ -27,6 +29,11 @@ PurgeOperation, PurgeOperationTarget, ) +from reflexio.server.services.governance.config import get_governance_ref_secret +from reflexio.server.services.storage.governance_claims import ( + PurgeExecutionClaim, + validate_purge_execution_claim, +) from reflexio.server.services.storage.governance_validation import ( _ALLOWED_PURGE_OPERATION_TYPES, _ALLOWED_PURGE_SCOPE_TYPES, @@ -48,7 +55,12 @@ _validate_governance_target_ref, ) -from .._governance import _json_dumps, _row_to_purge_operation, _row_to_purge_target +from .._governance import ( + _json_dumps, + _json_loads, + _row_to_purge_operation, + _row_to_purge_target, +) if TYPE_CHECKING: from .._governance import _SQLiteGovernanceDeps @@ -65,6 +77,114 @@ class PurgeOperationStoreMixin: _deps: Callable[[], _SQLiteGovernanceDeps] _owned_user_playbook_ids_locked: Callable[[str], set[int]] _planned_governance_delete_counts: Callable[[str, set[int]], dict[str, int]] + _subject_ref_for_user_id: Callable[[str], str] + + def _authoritative_user_digest(self, purge_id: str, user_id: str) -> str: + material = f"authoritative-user-v1\0{self.org_id}\0{purge_id}\0{user_id}" + return hmac.new( + get_governance_ref_secret().encode(), + material.encode(), + hashlib.sha256, + ).hexdigest() + + @staticmethod + def _legacy_authoritative_user_digest(purge_id: str, user_id: str) -> str: + return hashlib.sha256(f"{purge_id}\0{user_id}".encode()).hexdigest() + + def _assert_authoritative_user_identity_locked( + self, purge_id: str, user_id: str + ) -> str: + row = self.conn.execute( + """SELECT operation_type, scope_type, subject_ref, + authoritative_user_digest + FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (self.org_id, purge_id), + ).fetchone() + expected_digest = self._authoritative_user_digest(purge_id, user_id) + if ( + row is None + or row["operation_type"] != "user_erasure" + or row["scope_type"] != "user" + or row["subject_ref"] != self._subject_ref_for_user_id(user_id) + or row["authoritative_user_digest"] != expected_digest + ): + raise ValueError("Purge authoritative user identity does not match") + return expected_digest + + def _adopt_authoritative_user_digest_bindings_locked( + self, + *, + purge_id: str, + user_id: str, + existing_digest: object, + authoritative_user_digest: str, + now: int, + ) -> None: + legacy_digest = self._legacy_authoritative_user_digest(purge_id, user_id) + + def is_recognized(binding: object) -> bool: + return binding is None or ( + isinstance(binding, str) + and ( + hmac.compare_digest(binding, authoritative_user_digest) + or hmac.compare_digest(binding, legacy_digest) + ) + ) + + if not is_recognized(existing_digest): + raise ValueError( + "Existing purge operation has mismatched authoritative user identity" + ) + + snapshot_row = self.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = ? + AND target_ref = 'all' AND phase = ?""", + (self.org_id, purge_id, _SNAPSHOT_TARGET_NAME, _PREPARE_PHASE), + ).fetchone() + snapshot_detail = None + if snapshot_row is not None: + snapshot_detail = _json_loads(snapshot_row["detail"]) + if not isinstance(snapshot_detail, dict) or not is_recognized( + snapshot_detail.get("authoritative_user_digest") + ): + raise ValueError( + "Existing purge snapshot has mismatched authoritative user identity" + ) + + if existing_digest != authoritative_user_digest: + self.conn.execute( + """UPDATE purge_operations + SET authoritative_user_digest = ?, updated_at = ? + WHERE org_id = ? AND purge_id = ? + AND authoritative_user_digest IS ?""", + ( + authoritative_user_digest, + now, + self.org_id, + purge_id, + existing_digest, + ), + ) + if ( + snapshot_detail is not None + and snapshot_detail.get("authoritative_user_digest") + != authoritative_user_digest + ): + snapshot_detail["authoritative_user_digest"] = authoritative_user_digest + self.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = ? + AND target_ref = 'all' AND phase = ?""", + ( + _json_dumps(snapshot_detail), + self.org_id, + purge_id, + _SNAPSHOT_TARGET_NAME, + _PREPARE_PHASE, + ), + ) def _record_purge_target_locked( self, @@ -165,6 +285,7 @@ def begin_purge_operation( scope_type: Literal["user", "org"], subject_ref: str | None, request_ref: str, + authoritative_user_id: str | None = None, ) -> PurgeOperation: _validate_governance_enum( "operation_type", @@ -187,6 +308,20 @@ def begin_purge_operation( str, _validate_governance_idempotency_key("idempotency_key", idempotency_key), ) + if operation_type == "user_erasure" and scope_type == "user": + if not authoritative_user_id: + raise ValueError("authoritative user identity is required") + if subject_ref != self._subject_ref_for_user_id(authoritative_user_id): + raise ValueError("authoritative user identity must match subject_ref") + elif authoritative_user_id: + raise ValueError( + "authoritative user identity is only valid for user erasure" + ) + authoritative_user_digest = ( + self._authoritative_user_digest(validated_purge_id, authoritative_user_id) + if authoritative_user_id + else None + ) now = _epoch_now() with self._lock: try: @@ -211,13 +346,26 @@ def begin_purge_operation( "Existing purge operation for idempotency_key has " f"mismatched {field_name}" ) - self.conn.rollback() - return existing_operation + if authoritative_user_id and authoritative_user_digest: + self._adopt_authoritative_user_digest_bindings_locked( + purge_id=validated_purge_id, + user_id=authoritative_user_id, + existing_digest=existing["authoritative_user_digest"], + authoritative_user_digest=authoritative_user_digest, + now=now, + ) + elif existing["authoritative_user_digest"] is not None: + raise ValueError( + "Existing purge operation has mismatched authoritative user identity" + ) + self.conn.commit() + return _row_to_purge_operation(existing) self.conn.execute( """INSERT INTO purge_operations ( purge_id, org_id, operation_type, scope_type, subject_ref, - request_ref, idempotency_key, status, created_at, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)""", + request_ref, idempotency_key, authoritative_user_digest, + status, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)""", ( validated_purge_id, self.org_id, @@ -226,6 +374,7 @@ def begin_purge_operation( subject_ref, request_ref, validated_idempotency_key, + authoritative_user_digest, now, now, ), @@ -236,12 +385,177 @@ def begin_purge_operation( raise return self.get_purge_operation(validated_purge_id) + def claim_purge_operation_execution( + self, + purge_id: str, + *, + lease_owner: str, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim | None: + validated_purge_id = _validate_governance_purge_id("purge_id", purge_id) + if not lease_owner.strip(): + raise ValueError("lease_owner is required") + if lease_ttl_seconds <= 0: + raise ValueError("lease_ttl_seconds must be positive") + now = _epoch_now() + expires_at = now + lease_ttl_seconds + with self._lock: + try: + self.conn.execute("BEGIN IMMEDIATE") + cursor = self.conn.execute( + """UPDATE purge_operations + SET status = 'running', error_code = NULL, error_detail = NULL, + completed_at = NULL, updated_at = ?, + execution_claim_owner = ?, + execution_claim_fence = execution_claim_fence + 1, + execution_claim_expires_at = ? + WHERE purge_id = ? AND org_id = ? + AND ( + status IN ('pending', 'failed') + OR ( + status = 'running' + AND ( + execution_claim_expires_at IS NULL + OR execution_claim_expires_at <= ? + ) + ) + ) + RETURNING execution_claim_owner, + execution_claim_fence, + execution_claim_expires_at""", + ( + now, + lease_owner, + expires_at, + validated_purge_id, + self.org_id, + now, + ), + ) + row = cursor.fetchone() + self.conn.commit() + if row is None: + return None + return PurgeExecutionClaim( + purge_id=validated_purge_id, + owner=str(row["execution_claim_owner"]), + fence=int(row["execution_claim_fence"]), + expires_at=int(row["execution_claim_expires_at"]), + ) + except Exception: + self.conn.rollback() + raise + + def assert_purge_operation_execution_claim( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> None: + purge_id = _validate_governance_purge_id("purge_id", purge_id) + claim = validate_purge_execution_claim(purge_id, execution_claim) + now = _epoch_now() + row = self._deps()._fetchone( + """SELECT status, execution_claim_owner, execution_claim_fence, + execution_claim_expires_at + FROM purge_operations + WHERE purge_id = ? AND org_id = ?""", + (purge_id, self.org_id), + ) + if row is None: + raise ValueError(f"Purge operation {purge_id!r} not found") + if ( + row["status"] != "running" + or row["execution_claim_owner"] != claim.owner + or int(row["execution_claim_fence"]) != claim.fence + or row["execution_claim_expires_at"] is None + or int(row["execution_claim_expires_at"]) <= now + ): + raise ValueError("purge execution claim is no longer active") + + def _assert_purge_operation_execution_claim_locked( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + ) -> None: + purge_id = _validate_governance_purge_id("purge_id", purge_id) + claim = validate_purge_execution_claim(purge_id, execution_claim) + now = _epoch_now() + row = self.conn.execute( + """SELECT status, execution_claim_owner, execution_claim_fence, + execution_claim_expires_at + FROM purge_operations + WHERE purge_id = ? AND org_id = ?""", + (purge_id, self.org_id), + ).fetchone() + if row is None: + raise ValueError(f"Purge operation {purge_id!r} not found") + if ( + row["status"] != "running" + or row["execution_claim_owner"] != claim.owner + or int(row["execution_claim_fence"]) != claim.fence + or row["execution_claim_expires_at"] is None + or int(row["execution_claim_expires_at"]) <= now + ): + raise ValueError("purge execution claim is no longer active") + + def renew_purge_operation_execution_claim( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + *, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim: + purge_id = _validate_governance_purge_id("purge_id", purge_id) + claim = validate_purge_execution_claim(purge_id, execution_claim) + if lease_ttl_seconds <= 0: + raise ValueError("lease_ttl_seconds must be positive") + now = _epoch_now() + expires_at = now + lease_ttl_seconds + with self._lock: + try: + self.conn.execute("BEGIN IMMEDIATE") + cursor = self.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = ?, updated_at = ? + WHERE purge_id = ? AND org_id = ? + AND status = 'running' + AND execution_claim_owner = ? + AND execution_claim_fence = ? + AND execution_claim_expires_at IS NOT NULL + AND execution_claim_expires_at > ? + RETURNING execution_claim_owner, + execution_claim_fence, + execution_claim_expires_at""", + ( + expires_at, + now, + purge_id, + self.org_id, + claim.owner, + claim.fence, + now, + ), + ) + row = cursor.fetchone() + self.conn.commit() + except Exception: + self.conn.rollback() + raise + if row is None: + raise ValueError("purge execution claim is no longer active") + return PurgeExecutionClaim( + purge_id=purge_id, + owner=str(row["execution_claim_owner"]), + fence=int(row["execution_claim_fence"]), + expires_at=int(row["execution_claim_expires_at"]), + ) + def record_purge_target( self, purge_id: str, target_name: str, phase: str, status: Literal["pending", "running", "failed", "complete"], + *, + execution_claim: PurgeExecutionClaim, target_ref: str = "", detail: dict[str, object] | None = None, deleted_count: int = 0, @@ -265,6 +579,10 @@ def record_purge_target( ) with self._lock: try: + self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) self._record_purge_target_locked( purge_id=purge_id, target_name=target_name, @@ -308,14 +626,28 @@ def prepare_governance_erase_targets( self, purge_id: str, user_id: str, + *, + execution_claim: PurgeExecutionClaim, owned_user_playbook_ids: set[int] | None = None, ) -> None: purge_id = _validate_governance_purge_id("purge_id", purge_id) with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") - if self.purge_targets_prepared(purge_id): - self.conn.rollback() + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) + authoritative_user_digest = ( + self._assert_authoritative_user_identity_locked(purge_id, user_id) + ) + prepared = self.conn.execute( + """SELECT 1 FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = ? AND target_ref = 'all' + AND phase = ? AND status = 'complete'""", + (self.org_id, purge_id, _SNAPSHOT_TARGET_NAME, _PREPARE_PHASE), + ).fetchone() + if prepared is not None: + self.conn.commit() return owned_user_playbook_ids = ( set(owned_user_playbook_ids) @@ -344,6 +676,7 @@ def prepare_governance_erase_targets( phase=_PREPARE_PHASE, status="complete", detail={ + "authoritative_user_digest": authoritative_user_digest, "owned_user_playbook_ids": sorted(owned_user_playbook_ids), }, deleted_count=0, @@ -355,7 +688,12 @@ def prepare_governance_erase_targets( raise def fail_purge_operation( - self, purge_id: str, error_code: str, error_detail: str + self, + purge_id: str, + error_code: str, + error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) validated_error_code = _validate_governance_error_code(error_code) @@ -364,10 +702,15 @@ def fail_purge_operation( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) cur = self.conn.execute( """UPDATE purge_operations SET status = 'failed', error_code = ?, error_detail = ?, - updated_at = ?, completed_at = ? + updated_at = ?, completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ? AND status != 'complete'""", ( validated_error_code, diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py b/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py index eea1397b3..ead650ff6 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_rebuild_hide.py @@ -35,6 +35,7 @@ from reflexio.models.api_schema.domain import AgentPlaybookSourceWindow from reflexio.models.api_schema.domain.enums import Status from reflexio.server.services.embedding_text import playbook_trigger_embedding_text +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _canonicalize_governance_windows, _parse_governance_window_list, @@ -81,6 +82,9 @@ class RebuildHideMixin: # Provided via MRO by the co-composed PurgeOperationStoreMixin (purge bucket); # reached here by the cross-bucket rebuild-hide method. _record_purge_target_locked: Callable[..., None] + _assert_purge_operation_execution_claim_locked: Callable[ + [str, PurgeExecutionClaim | None], None + ] def _replace_agent_playbook_source_windows_locked( self, agent_playbook_id: int, windows: list[AgentPlaybookSourceWindow] @@ -137,11 +141,19 @@ def _upsert_agent_playbook_search_rows_locked( (agent_playbook_id, json.dumps(embedding)), ) - def hide_governance_agent_playbooks_for_rebuild(self, purge_id: str) -> list[int]: + def hide_governance_agent_playbooks_for_rebuild( + self, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, + ) -> list[int]: purge_id = _validate_governance_purge_id("purge_id", purge_id) with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) target_rows = self.conn.execute( """SELECT target_ref FROM purge_operation_targets @@ -202,6 +214,8 @@ def apply_governance_agent_playbook_rebuild( blocking_issue: dict[str, object] | None, expanded_terms: str | None, tags: list[str] | None, + *, + execution_claim: PurgeExecutionClaim, ) -> None: purge_id = _validate_governance_purge_id("purge_id", purge_id) windows = _parse_governance_window_list( @@ -210,13 +224,18 @@ def apply_governance_agent_playbook_rebuild( canonical_remaining_windows = [window.model_dump() for window in windows] content_value = content or "" trigger_value = trigger or None - embedding_text = playbook_trigger_embedding_text(trigger_value) - embedding = ( - self._deps()._get_embedding(embedding_text) if embedding_text else [] - ) + embedding: list[float] = [] + if windows: + embedding_text = playbook_trigger_embedding_text(trigger_value) + embedding = ( + self._deps()._get_embedding(embedding_text) if embedding_text else [] + ) with self._lock: try: - self.conn.execute("BEGIN") + self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) rebuild_target_row = self.conn.execute( """SELECT status, detail FROM purge_operation_targets diff --git a/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py b/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py index 3390533f0..da3506a99 100644 --- a/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py +++ b/reflexio/server/services/storage/sqlite_storage/governance/_subject_barrier.py @@ -38,6 +38,7 @@ governance_subject_ref, ) from reflexio.server.services.storage.error import SubjectWriteBarrierError +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, _PREPARE_PHASE, @@ -53,6 +54,7 @@ ) from .._governance import ( + _json_loads, _row_to_audit_event, _row_to_purge_operation, _row_to_subject_write_barrier, @@ -75,6 +77,10 @@ class SubjectBarrierMixin: [sqlite3.Connection | sqlite3.Cursor, AuditEvent], bool ] get_purge_operation: Callable[[str], PurgeOperation] + _assert_purge_operation_execution_claim_locked: Callable[ + [str, PurgeExecutionClaim | None], None + ] + _authoritative_user_digest: Callable[[str, str], str] def _barrier_from_purge( self, @@ -158,7 +164,59 @@ def _legacy_user_id_rows_remain_locked( return True return False - def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: + def _authoritative_user_session_outcome_remains_locked(self, user_id: str) -> bool: + return ( + self.conn.execute( + "SELECT 1 FROM session_outcomes WHERE user_id = ? LIMIT 1", + (user_id,), + ).fetchone() + is not None + ) + + def _assert_bound_authoritative_user_identity_locked( + self, purge_id: str, subject_ref: str, authoritative_user_id: str + ) -> None: + purge_row = self.conn.execute( + """SELECT operation_type, scope_type, subject_ref, + authoritative_user_digest + FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (self.org_id, purge_id), + ).fetchone() + if purge_row is None: + raise ValueError(f"Purge operation {purge_id!r} not found") + if ( + purge_row["operation_type"] != "user_erasure" + or purge_row["scope_type"] != "user" + ): + raise ValueError("Completion requires a user erasure purge") + authoritative_user_digest = purge_row["authoritative_user_digest"] + snapshot_row = self.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = ? + AND target_ref = 'all' AND phase = ? AND status = 'complete'""", + (self.org_id, purge_id, _SNAPSHOT_TARGET_NAME, _PREPARE_PHASE), + ).fetchone() + snapshot_detail = ( + _json_loads(snapshot_row["detail"]) if snapshot_row is not None else None + ) + expected_digest = self._authoritative_user_digest( + purge_id, authoritative_user_id + ) + if ( + purge_row["scope_type"] != "user" + or purge_row["subject_ref"] != subject_ref + or self._subject_ref_for_user_id(authoritative_user_id) != subject_ref + or not isinstance(authoritative_user_digest, str) + or authoritative_user_digest != expected_digest + or not isinstance(snapshot_detail, dict) + or snapshot_detail.get("authoritative_user_digest") != expected_digest + ): + raise ValueError("Purge authoritative user identity does not match") + + def _same_subject_rows_remain_locked( + self, subject_ref: str, authoritative_user_id: str + ) -> bool: legacy_request_ids = self._legacy_request_ids_for_subject_locked(subject_ref) for table in ( "requests", @@ -167,7 +225,6 @@ def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: "user_playbooks", "agent_success_evaluation_result", "retrieved_learning_evaluation", - "session_outcomes", ): row = self.conn.execute( f"""SELECT 1 FROM {table} @@ -179,6 +236,10 @@ def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: return True if legacy_request_ids: return True + if self._authoritative_user_session_outcome_remains_locked( + authoritative_user_id + ): + return True if self._legacy_user_id_rows_remain_locked( table="interactions", subject_ref=subject_ref, @@ -211,7 +272,11 @@ def _same_subject_rows_remain_locked(self, subject_ref: str) -> bool: ) def begin_subject_erasure_barrier( - self, subject_ref: str, purge_id: str + self, + subject_ref: str, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: _validate_governance_prefixed_ref( "subject_ref", subject_ref, prefix="subref_v1_" @@ -221,6 +286,9 @@ def begin_subject_erasure_barrier( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + validated_purge_id, execution_claim + ) purge_row = self.conn.execute( """SELECT * FROM purge_operations WHERE purge_id = ? AND org_id = ?""", @@ -293,7 +361,12 @@ def assert_subject_writable(self, subject_ref: str) -> None: raise def complete_subject_erasure_barrier_after_empty_check( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: purge_id = _validate_governance_purge_id("purge_id", purge_id) if audit_event.org_id != self.org_id: @@ -309,6 +382,9 @@ def complete_subject_erasure_barrier_after_empty_check( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + purge_id, execution_claim + ) row = self.conn.execute( "SELECT * FROM purge_operations WHERE purge_id = ? AND org_id = ?", (purge_id, self.org_id), @@ -334,7 +410,14 @@ def complete_subject_erasure_barrier_after_empty_check( raise ValueError( "Cannot complete purge without target snapshot marker" ) - if self._same_subject_rows_remain_locked(audit_event.subject_ref or ""): + self._assert_bound_authoritative_user_identity_locked( + purge_id, + audit_event.subject_ref or "", + authoritative_user_id, + ) + if self._same_subject_rows_remain_locked( + audit_event.subject_ref or "", authoritative_user_id + ): raise ValueError("same-subject rows remain") delete_rows = self.conn.execute( """SELECT target_name, status FROM purge_operation_targets @@ -424,7 +507,9 @@ def complete_subject_erasure_barrier_after_empty_check( error_code = NULL, error_detail = NULL, updated_at = ?, - completed_at = ? + completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ?""", (now, now, purge_id, self.org_id), ) @@ -440,6 +525,8 @@ def fail_subject_erasure_barrier( purge_id: str, error_code: str, error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: _validate_governance_prefixed_ref( "subject_ref", subject_ref, prefix="subref_v1_" @@ -451,6 +538,9 @@ def fail_subject_erasure_barrier( with self._lock: try: self.conn.execute("BEGIN IMMEDIATE") + self._assert_purge_operation_execution_claim_locked( + validated_purge_id, execution_claim + ) update_cursor = self.conn.execute( """UPDATE subject_write_barriers SET status = 'failed', @@ -480,7 +570,9 @@ def fail_subject_erasure_barrier( self.conn.execute( """UPDATE purge_operations SET status = 'failed', error_code = ?, error_detail = ?, - updated_at = ?, completed_at = ? + updated_at = ?, completed_at = ?, + execution_claim_owner = NULL, + execution_claim_expires_at = NULL WHERE purge_id = ? AND org_id = ?""", ( validated_error_code, diff --git a/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py b/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py index d97fd6f24..41b938d95 100644 --- a/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py +++ b/reflexio/server/services/storage/sqlite_storage/playbook/_optimization.py @@ -6,6 +6,8 @@ from hashlib import sha256 from typing import Any +from pydantic import ValidationError + from reflexio.models.api_schema.domain.entities import canonicalize_artifact_json from reflexio.models.api_schema.service_schemas import ( OptimizationArtifactKind, @@ -25,6 +27,8 @@ canonical_json_bytes, ) from reflexio.server.services.storage.error import ( + OptimizationArtifactIntegrityError, + OptimizationJobIdentityConflictError, OptimizationJobLeaseLiveError, StorageError, ) @@ -86,15 +90,20 @@ def _row_to_playbook_optimization_job(row: sqlite3.Row) -> PlaybookOptimizationJ def _row_to_playbook_optimization_artifact( row: sqlite3.Row, ) -> PlaybookOptimizationArtifact: - return PlaybookOptimizationArtifact( - artifact_id=row["artifact_id"], - job_id=row["job_id"], - artifact_kind=row["artifact_kind"], - content_json=row["content_json"], - content_digest=row["content_digest"], - created_at=row["created_at"], - updated_at=row["updated_at"], - ) + try: + return PlaybookOptimizationArtifact( + artifact_id=row["artifact_id"], + job_id=row["job_id"], + artifact_kind=row["artifact_kind"], + content_json=row["content_json"], + content_digest=row["content_digest"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + except (ValidationError, ValueError) as exc: + raise OptimizationArtifactIntegrityError( + "optimizer artifact row is malformed" + ) from exc def _job_insert_values(job: PlaybookOptimizationJob) -> tuple[Any, ...]: @@ -531,7 +540,9 @@ def create_or_get_playbook_optimization_job( and by_attempt is not None and by_discovery["job_id"] != by_attempt["job_id"] ): - raise ValueError("conflicting immutable optimizer job identity") + raise OptimizationJobIdentityConflictError( + "conflicting immutable optimizer job identity" + ) existing = by_discovery or by_attempt if existing is not None: if ( @@ -542,7 +553,9 @@ def create_or_get_playbook_optimization_job( and existing["attempt_key"] != job.attempt_key ) ): - raise ValueError("conflicting immutable optimizer job identity") + raise OptimizationJobIdentityConflictError( + "conflicting immutable optimizer job identity" + ) result = _row_to_playbook_optimization_job(existing) else: cur = self.conn.execute(_JOB_INSERT_SQL, _job_insert_values(job)) @@ -801,7 +814,9 @@ def upsert_playbook_optimization_artifact( sha256(artifact_content_json.encode()).hexdigest() != artifact.content_digest ): - raise ValueError("optimizer artifact digest does not match content") + raise OptimizationArtifactIntegrityError( + "optimizer artifact digest does not match content" + ) written_at = self._lease_now(now) with self._lock: owns_transaction = self._own_transaction() @@ -825,9 +840,13 @@ def upsert_playbook_optimization_artifact( ).fetchone() if existing is not None: if existing["content_digest"] != artifact.content_digest: - raise ValueError("optimizer artifact digest conflict") + raise OptimizationArtifactIntegrityError( + "optimizer artifact digest conflict" + ) if existing["content_json"] != artifact_content_json: - raise ValueError("optimizer artifact content conflict") + raise OptimizationArtifactIntegrityError( + "optimizer artifact content conflict" + ) result = _row_to_playbook_optimization_artifact(existing) else: cur = self.conn.execute( diff --git a/reflexio/server/services/storage/storage_base/__init__.py b/reflexio/server/services/storage/storage_base/__init__.py index 4a790ee8d..2ac8a0f6b 100644 --- a/reflexio/server/services/storage/storage_base/__init__.py +++ b/reflexio/server/services/storage/storage_base/__init__.py @@ -147,7 +147,7 @@ def _partition_purge_vs_delete( def clear_user_data(self, user_id: str) -> dict[str, int]: """Delete all rows scoped to a single ``user_id``. - Removes the user's interactions, session outcomes, user playbooks, + Removes the user's session outcomes, interactions, user playbooks, profiles, and requests. Intentionally does NOT touch ``agent_playbooks`` — those are the cross-project rollup of skills and have no ``user_id`` column. This is the data-isolation primitive used diff --git a/reflexio/server/services/storage/storage_base/_session_outcomes.py b/reflexio/server/services/storage/storage_base/_session_outcomes.py index d417e6065..14695d9a3 100644 --- a/reflexio/server/services/storage/storage_base/_session_outcomes.py +++ b/reflexio/server/services/storage/storage_base/_session_outcomes.py @@ -18,6 +18,24 @@ class SessionOutcomeWriteResult: source: str | None = None reason: SessionOutcomeFailureReason | None = None context_changed: bool = False + outcome_id: str | None = None + outcome_revision: int | None = None + outcome_contract_digest: str | None = None + finalized_trajectory_digest: str | None = None + + def __post_init__(self) -> None: + identity = ( + self.outcome_id, + self.outcome_revision, + self.outcome_contract_digest, + self.finalized_trajectory_digest, + ) + if any(value is None for value in identity) and not all( + value is None for value in identity + ): + raise ValueError( + "outcome identity fields must be all populated or all null" + ) @dataclass(frozen=True) diff --git a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py index b8bba0371..d5f2eec6c 100644 --- a/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py +++ b/reflexio/server/services/storage/storage_base/agent_run/_agent_run_store.py @@ -24,6 +24,29 @@ 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 the committed learning ids, or ``None`` when no receipt exists. + + A committed empty receipt returns ``[]`` and must remain distinguishable + from a missing receipt so finalization stays idempotent. + """ + raise NotImplementedError(f"{type(self).__name__} does not support agent runs") + + def save_agent_run_finalization_receipt( + 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 cc4016f72..4650f726a 100644 --- a/reflexio/server/services/storage/storage_base/governance/_erase_execution.py +++ b/reflexio/server/services/storage/storage_base/governance/_erase_execution.py @@ -6,6 +6,7 @@ AuditEvent, PurgeOperation, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class GovernanceEraseExecutionMixin(ABC): @@ -19,12 +20,21 @@ class GovernanceEraseExecutionMixin(ABC): @abstractmethod def apply_governance_user_data_delete( - self, purge_id: str, user_id: str + self, + purge_id: str, + user_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> dict[str, int]: raise NotImplementedError @abstractmethod def complete_purge_operation_with_audit( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: raise NotImplementedError diff --git a/reflexio/server/services/storage/storage_base/governance/_purge.py b/reflexio/server/services/storage/storage_base/governance/_purge.py index ef63fdc2c..f3f1b4620 100644 --- a/reflexio/server/services/storage/storage_base/governance/_purge.py +++ b/reflexio/server/services/storage/storage_base/governance/_purge.py @@ -7,6 +7,7 @@ PurgeOperation, PurgeOperationTarget, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class PurgeOperationStoreMixin(ABC): @@ -27,9 +28,39 @@ def begin_purge_operation( scope_type: Literal["user", "org"], subject_ref: str | None, request_ref: str, + authoritative_user_id: str | None = None, ) -> PurgeOperation: raise NotImplementedError + @abstractmethod + def claim_purge_operation_execution( + self, + purge_id: str, + *, + lease_owner: str, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim | None: + """Atomically claim or take over a stale purge execution.""" + raise NotImplementedError + + @abstractmethod + def assert_purge_operation_execution_claim( + self, purge_id: str, execution_claim: PurgeExecutionClaim + ) -> None: + """Raise when the purge execution claim no longer owns the live fence.""" + raise NotImplementedError + + @abstractmethod + def renew_purge_operation_execution_claim( + self, + purge_id: str, + execution_claim: PurgeExecutionClaim, + *, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim: + """Atomically renew an active purge execution claim.""" + raise NotImplementedError + @abstractmethod def record_purge_target( self, @@ -37,6 +68,8 @@ def record_purge_target( target_name: str, phase: str, status: Literal["pending", "running", "failed", "complete"], + *, + execution_claim: PurgeExecutionClaim, target_ref: str = "", detail: dict[str, object] | None = None, deleted_count: int = 0, @@ -59,13 +92,20 @@ def prepare_governance_erase_targets( self, purge_id: str, user_id: str, + *, + execution_claim: PurgeExecutionClaim, owned_user_playbook_ids: set[int] | None = None, ) -> None: raise NotImplementedError @abstractmethod def fail_purge_operation( - self, purge_id: str, error_code: str, error_detail: str + self, + purge_id: str, + error_code: str, + error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: raise NotImplementedError diff --git a/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py b/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py index 6020245a2..66f905a92 100644 --- a/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py +++ b/reflexio/server/services/storage/storage_base/governance/_rebuild_hide.py @@ -2,6 +2,8 @@ from abc import ABC, abstractmethod +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim + class RebuildHideMixin(ABC): """Backend-neutral governance rebuild-hide store contract. @@ -13,7 +15,12 @@ class RebuildHideMixin(ABC): """ @abstractmethod - def hide_governance_agent_playbooks_for_rebuild(self, purge_id: str) -> list[int]: + def hide_governance_agent_playbooks_for_rebuild( + self, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, + ) -> list[int]: raise NotImplementedError @abstractmethod @@ -28,5 +35,7 @@ def apply_governance_agent_playbook_rebuild( blocking_issue: dict[str, object] | None, expanded_terms: str | None, tags: list[str] | None, + *, + execution_claim: PurgeExecutionClaim, ) -> None: raise NotImplementedError diff --git a/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py b/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py index 7752ec5c9..2de2f5a26 100644 --- a/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py +++ b/reflexio/server/services/storage/storage_base/governance/_subject_barrier.py @@ -7,6 +7,7 @@ PurgeOperation, SubjectWriteBarrier, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim class SubjectBarrierMixin(ABC): @@ -20,7 +21,11 @@ class SubjectBarrierMixin(ABC): @abstractmethod def begin_subject_erasure_barrier( - self, subject_ref: str, purge_id: str + self, + subject_ref: str, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: raise NotImplementedError @@ -30,7 +35,12 @@ def assert_subject_writable(self, subject_ref: str) -> None: @abstractmethod def complete_subject_erasure_barrier_after_empty_check( - self, purge_id: str, audit_event: AuditEvent + self, + purge_id: str, + audit_event: AuditEvent, + *, + authoritative_user_id: str, + execution_claim: PurgeExecutionClaim, ) -> PurgeOperation: raise NotImplementedError @@ -41,6 +51,8 @@ def fail_subject_erasure_barrier( purge_id: str, error_code: str, error_detail: str, + *, + execution_claim: PurgeExecutionClaim, ) -> SubjectWriteBarrier: raise NotImplementedError diff --git a/reflexio/server/usage_metrics.py b/reflexio/server/usage_metrics.py index de7f6f1b0..d505f0916 100644 --- a/reflexio/server/usage_metrics.py +++ b/reflexio/server/usage_metrics.py @@ -11,6 +11,7 @@ import time from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from enum import Enum from typing import Any logger = logging.getLogger(__name__) @@ -47,7 +48,26 @@ class UsageEvent: created_at: float = field(default_factory=time.time) -UsageEventRecorder = Callable[[UsageEvent], None] +class UsageEventDeliveryStatus(Enum): + """Recorder-reported durability outcome for one usage event.""" + + APPENDED = "appended" + DUPLICATE = "duplicate" + EXEMPT = "exempt" + UNKNOWN = "unknown" + FAILED = "failed" + REJECTED = "rejected" + + +class UsageEventDeliveryError(RuntimeError): + """Raised when strict usage delivery was not durably accepted.""" + + def __init__(self, status: UsageEventDeliveryStatus) -> None: + self.status = status + super().__init__(f"usage event delivery {status.value}") + + +UsageEventRecorder = Callable[[UsageEvent], UsageEventDeliveryStatus | None] _recorder: UsageEventRecorder | None = None @@ -62,6 +82,11 @@ def configure_usage_event_recorder(recorder: UsageEventRecorder | None) -> None: _recorder = recorder +def exempt_usage_event_recorder(_event: UsageEvent) -> UsageEventDeliveryStatus: + """Explicitly exempt a deployment from durable usage-event delivery.""" + return UsageEventDeliveryStatus.EXEMPT + + def record_usage_event( *, org_id: str, @@ -90,45 +115,134 @@ def record_usage_event( duration_ms: int | None = None, error_kind: str | None = None, metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, ) -> None: """Record one usage event if a recorder is configured. - The product path must never fail because metrics failed, so this function - catches and logs all recorder errors. + The product path must never fail because metrics failed. Missing and legacy + recorders are silent on this ordinary fail-open path; explicit delivery + failures and recorder exceptions are logged. + """ + try: + record_usage_event_strict( + org_id=org_id, + event_name=event_name, + event_category=event_category, + user_id=user_id, + request_id=request_id, + session_id=session_id, + pipeline=pipeline, + entity_type=entity_type, + entity_id=entity_id, + event_key=event_key, + extractor_name=extractor_name, + playbook_name=playbook_name, + source=source, + agent_version=agent_version, + backend=backend, + outcome=outcome, + count_value=count_value, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + billing_input_tokens=billing_input_tokens, + platform_llm=platform_llm, + platform_storage=platform_storage, + caller_type=caller_type, + duration_ms=duration_ms, + error_kind=error_kind, + metadata=metadata, + created_at=created_at, + ) + except UsageEventDeliveryError as exc: + if exc.status is not UsageEventDeliveryStatus.UNKNOWN: + logger.warning("Usage metrics recorder failed: %s", exc) + except Exception: # noqa: BLE001 + logger.warning("Usage metrics recorder failed", exc_info=True) + + +def record_usage_event_strict( + *, + org_id: str, + event_name: str, + event_category: str, + user_id: str | None = None, + request_id: str | None = None, + session_id: str | None = None, + pipeline: str | None = None, + entity_type: str | None = None, + entity_id: str | None = None, + event_key: str | None = None, + extractor_name: str | None = None, + playbook_name: str | None = None, + source: str | None = None, + agent_version: str | None = None, + backend: str | None = None, + outcome: str | None = None, + count_value: int = 1, + prompt_tokens: int | None = None, + completion_tokens: int | None = None, + billing_input_tokens: int | None = None, + platform_llm: bool | None = None, + platform_storage: bool | None = None, + caller_type: str | None = None, + duration_ms: int | None = None, + error_kind: str | None = None, + metadata: Mapping[str, Any] | None = None, + created_at: float | None = None, +) -> UsageEventDeliveryStatus: + """Deliver one event and fail unless the recorder accepted it durably. + + Only an explicit append, duplicate, or deployment exemption proves the + receipt's billing obligation was handled. Missing and legacy recorders are + unknown delivery, which strict receipt-backed callers must retry. """ recorder = _recorder if recorder is None: - return - try: - recorder( - UsageEvent( - org_id=str(org_id), - event_name=event_name, - event_category=event_category, - user_id=user_id, - request_id=request_id, - session_id=session_id, - pipeline=pipeline, - entity_type=entity_type, - entity_id=entity_id, - event_key=event_key, - extractor_name=extractor_name, - playbook_name=playbook_name, - source=source, - agent_version=agent_version, - backend=backend, - outcome=outcome, - count_value=count_value, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - billing_input_tokens=billing_input_tokens, - platform_llm=platform_llm, - platform_storage=platform_storage, - caller_type=caller_type, - duration_ms=duration_ms, - error_kind=error_kind, - metadata=metadata or {}, - ) + raise UsageEventDeliveryError(UsageEventDeliveryStatus.UNKNOWN) + delivery_outcome = recorder( + UsageEvent( + org_id=str(org_id), + event_name=event_name, + event_category=event_category, + user_id=user_id, + request_id=request_id, + session_id=session_id, + pipeline=pipeline, + entity_type=entity_type, + entity_id=entity_id, + event_key=event_key, + extractor_name=extractor_name, + playbook_name=playbook_name, + source=source, + agent_version=agent_version, + backend=backend, + outcome=outcome, + count_value=count_value, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + billing_input_tokens=billing_input_tokens, + platform_llm=platform_llm, + platform_storage=platform_storage, + caller_type=caller_type, + duration_ms=duration_ms, + error_kind=error_kind, + metadata=metadata or {}, + created_at=time.time() if created_at is None else created_at, + ) + ) + status = ( + UsageEventDeliveryStatus.UNKNOWN + if delivery_outcome is None + else delivery_outcome + ) + if not isinstance(status, UsageEventDeliveryStatus): + raise TypeError( + f"usage event recorder returned an invalid delivery status: {status!r}" ) - except Exception as exc: # noqa: BLE001 - logger.warning("Usage metrics recorder failed: %s", exc) + if status not in { + UsageEventDeliveryStatus.APPENDED, + UsageEventDeliveryStatus.DUPLICATE, + UsageEventDeliveryStatus.EXEMPT, + }: + raise UsageEventDeliveryError(status) + return status diff --git a/tests/client/test_search.py b/tests/client/test_search.py index d0d55b354..adf31957d 100644 --- a/tests/client/test_search.py +++ b/tests/client/test_search.py @@ -1,6 +1,23 @@ +import inspect +from pathlib import Path from typing import Any from reflexio import ReflexioClient +from reflexio.models.api_schema.retriever_schema import ( + SearchUserPlaybookRequest, + UnifiedSearchRequest, +) + + +def _non_null_schema( + field_name: str, *, request_model: type[Any] = UnifiedSearchRequest +) -> dict[str, Any]: + field_schema = request_model.model_json_schema()["properties"][field_name] + return next( + option + for option in field_schema.get("anyOf", [field_schema]) + if option.get("type") != "null" + ) def test_unified_search_serializes_tag_filter(monkeypatch) -> None: @@ -38,3 +55,37 @@ def fake_make_request(method: str, path: str, **kwargs: Any) -> dict[str, Any]: client.search_agent_playbooks(query="billing", user_id="user-1") assert captured["json"]["user_id"] == "user-1" + + +def test_unified_search_docs_track_schema_contract() -> None: + top_k_schema = _non_null_schema("top_k") + identifier_limit = _non_null_schema("request_id")["maxLength"] + interaction_minimum = _non_null_schema("interaction_id")["exclusiveMinimum"] + 1 + client_docs = inspect.getdoc(ReflexioClient.search) or "" + registry_docs = ( + Path(__file__).parents[2] / "docs/lib/methods/unified-search.ts" + ).read_text(encoding="utf-8") + + assert f"1 to {top_k_schema['maximum']}" in client_docs + assert client_docs.count(f"at most {identifier_limit} characters") >= 3 + assert f"positive integer (minimum {interaction_minimum})" in client_docs + + assert f"1 to {top_k_schema['maximum']}" in registry_docs + for field_name in ("user_id", "request_id", "session_id"): + assert f'name: "{field_name}"' in registry_docs + assert registry_docs.count(f"at most {identifier_limit} characters") >= 3 + assert 'name: "interaction_id"' in registry_docs + assert f"positive integer (minimum {interaction_minimum})" in registry_docs + + +def test_user_playbook_search_docs_track_top_k_schema_contract() -> None: + top_k_limit = _non_null_schema("top_k", request_model=SearchUserPlaybookRequest)[ + "maximum" + ] + client_docs = inspect.getdoc(ReflexioClient.search_user_playbooks) or "" + registry_docs = ( + Path(__file__).parents[2] / "docs/lib/methods/user-playbooks.ts" + ).read_text(encoding="utf-8") + + assert f"1 to {top_k_limit}" in client_docs + assert f"1 to {top_k_limit}" in registry_docs diff --git a/tests/client/test_session_outcomes_client.py b/tests/client/test_session_outcomes_client.py index 662a7436b..9b72405e1 100644 --- a/tests/client/test_session_outcomes_client.py +++ b/tests/client/test_session_outcomes_client.py @@ -3,6 +3,9 @@ import inspect from typing import Any +import pytest +from pydantic import ValidationError + from reflexio import ReflexioClient @@ -53,3 +56,23 @@ def test_get_session_outcomes_has_no_untyped_filter_kwargs() -> None: parameter.kind is not inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() ) + + +def test_mark_session_outcome_rejects_partial_server_identity(monkeypatch) -> None: + client = ReflexioClient(api_key="test-key", url_endpoint="http://localhost:8000") + monkeypatch.setattr( + client, + "_make_request", + lambda *_args, **_kwargs: { + "success": True, + "recorded": False, + "outcome_id": "partial-outcome", + }, + ) + + with pytest.raises(ValidationError, match="all populated or all null"): + client.mark_session_outcome( + session_id="session-1", + outcome="success", + occurred_at=1, + ) diff --git a/tests/e2e_tests/conftest.py b/tests/e2e_tests/conftest.py index 4085fbc39..af76368ae 100644 --- a/tests/e2e_tests/conftest.py +++ b/tests/e2e_tests/conftest.py @@ -24,6 +24,7 @@ StorageConfigSQLite, ToolUseConfig, ) +from reflexio.server import usage_metrics from reflexio.server.services.configurator.configurator import DefaultConfigurator from reflexio.server.services.tagging.tagging_scheduler import drain_tagging from reflexio.test_support.llm_mock import ( @@ -58,6 +59,18 @@ """ +@pytest.fixture(autouse=True) +def _install_oss_usage_event_exemption() -> Iterator[None]: + """Match the explicit billing exemption installed by the OSS app lifespan.""" + usage_metrics.configure_usage_event_recorder( + usage_metrics.exempt_usage_event_recorder + ) + try: + yield + finally: + usage_metrics.configure_usage_event_recorder(None) + + @pytest.fixture(autouse=True) def mock_llm(request: pytest.FixtureRequest) -> Iterator[None]: """Keep the standard E2E tier deterministic and credential-free. diff --git a/tests/models/api_schema/test_input_size_bounds.py b/tests/models/api_schema/test_input_size_bounds.py index 911294744..ee35fe46c 100644 --- a/tests/models/api_schema/test_input_size_bounds.py +++ b/tests/models/api_schema/test_input_size_bounds.py @@ -138,10 +138,11 @@ def test_interaction_data_nested_list_bounds() -> None: def test_publish_request_string_bounds() -> None: - cap = 1_000 - _publish(source="s" * cap) - _publish(agent_version="v" * cap) + source_cap = 128 + agent_version_cap = 1_000 + _publish(source="s" * source_cap) + _publish(agent_version="v" * agent_version_cap) with pytest.raises(ValidationError): - _publish(source="s" * (cap + 1)) + _publish(source="s" * (source_cap + 1)) with pytest.raises(ValidationError): - _publish(agent_version="v" * (cap + 1)) + _publish(agent_version="v" * (agent_version_cap + 1)) diff --git a/tests/models/test_open_world_optimization_identity.py b/tests/models/test_open_world_optimization_identity.py new file mode 100644 index 000000000..09c58d9d1 --- /dev/null +++ b/tests/models/test_open_world_optimization_identity.py @@ -0,0 +1,24 @@ +from hashlib import sha256 + +from reflexio.models.api_schema.domain.entities import ( + PlaybookOptimizationArtifact, + PlaybookOptimizationJob, +) + + +def test_open_world_optimization_identity_is_accepted() -> None: + content_json = '{"schema_version":"offline-tuner-open-world-evidence-v1"}' + job = PlaybookOptimizationJob( + optimizer_kind="offline_tuner_open_world", + target_kind="user_playbook", + target_id=7, + ) + artifact = PlaybookOptimizationArtifact( + job_id=1, + artifact_kind="open_world_evidence_bundle", + content_json=content_json, + content_digest=sha256(content_json.encode()).hexdigest(), + ) + + assert job.optimizer_kind == "offline_tuner_open_world" + assert artifact.artifact_kind == "open_world_evidence_bundle" diff --git a/tests/models/test_session_outcome_identity.py b/tests/models/test_session_outcome_identity.py new file mode 100644 index 000000000..cb77b8d35 --- /dev/null +++ b/tests/models/test_session_outcome_identity.py @@ -0,0 +1,620 @@ +from collections.abc import Callable +from datetime import UTC, datetime +from hashlib import sha256 +from itertools import combinations +from typing import Any, Protocol + +import pytest +from pydantic import ValidationError + +from reflexio.models.api_schema.domain.entities import ( + GetSessionOutcomesRequest, + GetSessionOutcomesResponse, + InteractionData, + ManualPlaybookGenerationRequest, + ManualProfileGenerationRequest, + PublishUserInteractionRequest, + Request, + RerunPlaybookGenerationRequest, + RerunProfileGenerationRequest, + SessionOutcomeRecord, + SetSessionOutcomeResponse, +) +from reflexio.models.api_schema.domain.enums import SessionOutcomeKind +from reflexio.models.api_schema.retriever_schema import ( + GetRequestsRequest, + GetUserProfilesRequest, + SearchUserProfileRequest, +) +from reflexio.server.services.storage import session_outcome_identity +from reflexio.server.services.storage.session_outcome_identity import ( + canonical_json_bytes, + canonical_session_trajectory, + canonical_trajectory_bytes, + 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") + + +class _HasSource(Protocol): + @property + def source(self) -> str | None: ... + + +_OUTCOME_SOURCE_MODEL_FACTORIES: tuple[Callable[[str], _HasSource], ...] = ( + lambda source: Request( + request_id="request-1", + user_id="user-1", + session_id="session-1", + source=source, + ), + lambda source: PublishUserInteractionRequest( + user_id="user-1", + session_id="session-1", + interaction_data_list=[InteractionData(content="hello")], + source=source, + ), + lambda source: SessionOutcomeRecord( + user_id="user-1", + session_id="session-1", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source=source, + created_at=2, + ), + lambda source: SetSessionOutcomeResponse(success=True, source=source), + lambda source: GetSessionOutcomesRequest(source=source), + lambda source: GetRequestsRequest(source=source), + lambda source: SearchUserProfileRequest(user_id="user-1", source=source), + lambda source: GetUserProfilesRequest(user_id="user-1", source=source), + lambda source: RerunProfileGenerationRequest(source=source), + lambda source: ManualProfileGenerationRequest(source=source), + lambda source: ManualPlaybookGenerationRequest(source=source), + lambda source: RerunPlaybookGenerationRequest(source=source), +) +_OUTCOME_SOURCE_MODEL_IDS = ( + "request", + "publish", + "record", + "set-response", + "get-outcomes", + "get-requests", + "search-profiles", + "get-profiles", + "rerun-profiles", + "manual-profiles", + "manual-playbooks", + "rerun-playbooks", +) + + +@pytest.mark.parametrize( + "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS +) +def test_outcome_source_models_accept_machine_label( + factory: Callable[[str], _HasSource], +) -> None: + assert factory("support-agent:v2").source == "support-agent:v2" + assert factory("a" * 128).source == "a" * 128 + + +@pytest.mark.parametrize( + "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS +) +@pytest.mark.parametrize( + "source", + [ + "alice@example.com", + "https://example.com/hook", + "support agent", + " support-agent", + "support/agent", + "Support-Agent", + "support-agént", + "a" * 129, + ], +) +def test_outcome_source_models_reject_sensitive_or_free_form_values( + factory: Callable[[str], _HasSource], + source: str, +) -> None: + with pytest.raises(ValidationError): + factory(source) + + +@pytest.mark.parametrize( + "factory", _OUTCOME_SOURCE_MODEL_FACTORIES, ids=_OUTCOME_SOURCE_MODEL_IDS +) +def test_outcome_source_models_preserve_empty_source( + factory: Callable[[str], _HasSource], +) -> None: + assert factory("").source == "" + + +def test_optional_outcome_source_models_preserve_absence() -> None: + assert SetSessionOutcomeResponse(success=False).source is None + assert GetSessionOutcomesRequest().source is None + assert GetRequestsRequest().source is None + assert SearchUserProfileRequest(user_id="user-1").source is None + assert GetUserProfilesRequest(user_id="user-1").source is None + assert RerunProfileGenerationRequest().source is None + assert ManualProfileGenerationRequest().source is None + assert ManualPlaybookGenerationRequest().source is None + assert RerunPlaybookGenerationRequest().source is None + + +def test_outcome_source_json_schema_exposes_exact_constraints() -> None: + source_schema = PublishUserInteractionRequest.model_json_schema()["properties"][ + "source" + ] + + assert { + "maxLength": 128, + "pattern": "^[a-z0-9][a-z0-9._:-]{0,127}$", + "type": "string", + } in source_schema["anyOf"] + + +def test_outcome_contract_digest_is_stable_for_valid_machine_label() -> None: + payload = { + "allowed_values": ["failure", "success", "unknown"], + "finalization_rule": "first_write", + "schema_version": 1, + "source": "support-agent:v2", + } + + assert ( + _outcome_contract_digest(source="support-agent:v2") + == sha256(canonical_json_bytes(payload)).hexdigest() + ) + + +def test_outcome_contract_digest_changes_when_schema_version_changes() -> None: + assert _outcome_contract_digest(schema_version=1) != _outcome_contract_digest( + schema_version=2 + ) + + +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_trajectory_bytes_returns_exact_utf8_representation() -> None: + assert canonical_trajectory_bytes( + {"value": 1.5, "message": "caf\N{LATIN SMALL LETTER E WITH ACUTE}"} + ) == (b'{"message":"caf\xc3\xa9","value":1.5}') + + +def test_canonical_trajectory_bytes_ignores_mapping_insertion_order() -> None: + first = {"messages": [{"content": "hello", "role": "user"}], "session": "s1"} + second = {"session": "s1", "messages": [{"role": "user", "content": "hello"}]} + + assert canonical_trajectory_bytes(first) == canonical_trajectory_bytes(second) + + +@pytest.mark.parametrize( + "value", + [ + {"value": float("nan")}, + {"value": 2**53}, + {"value": object()}, + {1: "non-string-key"}, + {"value": "\ud800"}, + ], + ids=[ + "non-finite-float", + "non-ijson-integer", + "unsupported-object", + "key", + "surrogate", + ], +) +def test_canonical_trajectory_bytes_rejects_the_same_values_as_digest( + value: object, +) -> None: + with pytest.raises((TypeError, ValueError)) as digest_error: + trajectory_digest(value) + + with pytest.raises(type(digest_error.value)) as bytes_error: + canonical_trajectory_bytes(value) + + assert str(bytes_error.value) == str(digest_error.value) + + +def test_trajectory_digest_hashes_canonical_trajectory_bytes() -> None: + trajectory = {"messages": [{"role": "user", "content": "hello"}]} + + assert ( + trajectory_digest(trajectory) + == sha256(canonical_trajectory_bytes(trajectory)).hexdigest() + ) + + +def test_canonical_session_trajectory_normalizes_sqlite_and_postgres_rows() -> None: + sqlite_request = { + "request_id": "parity-request", + "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}]}) + + +@pytest.mark.parametrize( + "container_factory", + [ + lambda value: {"nested": value}, + lambda value: [value], + lambda value: (value,), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_accepts_maximum_canonical_json_depth( + container_factory, +) -> None: + value: object = "leaf" + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH): + value = container_factory(value) + + assert trajectory_digest(value) + + +@pytest.mark.parametrize( + "container_factory", + [ + lambda value: {"nested": value}, + lambda value: [value], + lambda value: (value,), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_rejects_over_maximum_canonical_json_depth( + container_factory, +) -> None: + value: object = "leaf" + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH + 1): + value = container_factory(value) + + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + trajectory_digest(value) + + +@pytest.mark.parametrize( + ("empty_container", "container_factory"), + [ + ({}, lambda value: {"nested": value}), + ([], lambda value: [value]), + ((), lambda value: (value,)), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_accepts_maximum_empty_container_depth( + empty_container, + container_factory, +) -> None: + value: object = empty_container + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH - 1): + value = container_factory(value) + + assert trajectory_digest(value) + + +@pytest.mark.parametrize( + ("empty_container", "container_factory"), + [ + ({}, lambda value: {"nested": value}), + ([], lambda value: [value]), + ((), lambda value: (value,)), + ], + ids=["mapping", "list", "tuple"], +) +def test_trajectory_digest_rejects_over_maximum_empty_container_depth( + empty_container, + container_factory, +) -> None: + value: object = empty_container + for _ in range(session_outcome_identity.MAX_CANONICAL_TRAJECTORY_JSON_DEPTH): + value = container_factory(value) + + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + trajectory_digest(value) + + +def test_session_outcome_record_accepts_unknown_and_serializes_identities() -> None: + record = SessionOutcomeRecord( + outcome_id="outcome-1", + outcome_revision=1, + user_id="user-1", + session_id="session-1", + outcome=SessionOutcomeKind.UNKNOWN, + occurred_at=1, + source="customer_webhook", + outcome_contract_digest="a" * 64, + finalized_trajectory_digest="b" * 64, + created_at=2, + ) + + response = GetSessionOutcomesResponse(success=True, session_outcomes=[record]) + + assert response.model_dump(mode="json")["session_outcomes"] == [ + { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "user_id": "user-1", + "session_id": "session-1", + "outcome": "unknown", + "occurred_at": 1, + "source": "customer_webhook", + "label": None, + "value": None, + "metadata": None, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + "created_at": 2, + } + ] + + +def test_session_outcome_record_accepts_legacy_all_null_identity() -> None: + record = SessionOutcomeRecord( + outcome_id=None, + outcome_revision=None, + user_id="legacy-user", + session_id="legacy-session", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source="customer_webhook", + outcome_contract_digest=None, + finalized_trajectory_digest=None, + created_at=2, + ) + + assert record.outcome_id is None + assert record.outcome_revision is None + assert record.outcome_contract_digest is None + assert record.finalized_trajectory_digest is None + + +def test_session_outcome_record_rejects_partial_legacy_identity() -> None: + with pytest.raises(ValidationError, match="all populated or all null"): + SessionOutcomeRecord( + outcome_id="outcome-1", + outcome_revision=None, + user_id="legacy-user", + session_id="legacy-session", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source="customer_webhook", + outcome_contract_digest=None, + finalized_trajectory_digest=None, + created_at=2, + ) + + +@pytest.mark.parametrize( + "populated_fields", + [ + fields + for populated_count in range(1, 4) + for fields in combinations( + ( + "outcome_id", + "outcome_revision", + "outcome_contract_digest", + "finalized_trajectory_digest", + ), + populated_count, + ) + ], +) +def test_session_outcome_write_result_rejects_partial_identity( + populated_fields: tuple[str, ...], +) -> None: + identity = { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + } + + with pytest.raises(ValueError, match="all populated or all null"): + SessionOutcomeWriteResult( + recorded=False, + **{ + field_name: value + for field_name, value in identity.items() + if field_name in populated_fields + }, + ) + + +@pytest.mark.parametrize( + "populated_fields", + [ + fields + for populated_count in range(1, 4) + for fields in combinations( + ( + "outcome_id", + "outcome_revision", + "outcome_contract_digest", + "finalized_trajectory_digest", + ), + populated_count, + ) + ], +) +def test_set_session_outcome_response_rejects_partial_identity( + populated_fields: tuple[str, ...], +) -> None: + identity = { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + } + + with pytest.raises(ValidationError, match="all populated or all null"): + SetSessionOutcomeResponse( + success=True, + **{ + field_name: value + for field_name, value in identity.items() + if field_name in populated_fields + }, + ) + + +@pytest.mark.parametrize( + "identity", + [ + {}, + { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + }, + ], +) +def test_write_result_and_response_accept_complete_identity_shapes( + identity: dict[str, Any], +) -> None: + write_result = SessionOutcomeWriteResult(recorded=False, **identity) + response = SetSessionOutcomeResponse(success=True, **identity) + + assert write_result.outcome_id == response.outcome_id + + +@pytest.mark.parametrize( + "field_name", ["outcome_contract_digest", "finalized_trajectory_digest"] +) +@pytest.mark.parametrize("digest", ["A" * 64, "a" * 63, "not-a-digest"]) +def test_session_outcome_record_rejects_non_sha256_identity_digests( + field_name: str, digest: str +) -> None: + payload: dict[str, object] = { + "outcome_id": "outcome-1", + "outcome_revision": 1, + "user_id": "user-1", + "session_id": "session-1", + "outcome": "success", + "occurred_at": 1, + "source": "customer_webhook", + "outcome_contract_digest": "a" * 64, + "finalized_trajectory_digest": "b" * 64, + "created_at": 2, + } + payload[field_name] = digest + + with pytest.raises(ValidationError, match="lowercase SHA-256 hex"): + SessionOutcomeRecord(**payload) # type: ignore[arg-type] + + +def test_trajectory_digest_matches_sha256_of_canonical_json() -> None: + assert trajectory_digest({"b": 2, "a": 1}) == sha256(b'{"a":1,"b":2}').hexdigest() diff --git a/tests/server/api_endpoints/test_session_outcomes_integration.py b/tests/server/api_endpoints/test_session_outcomes_integration.py index 5471dd743..645c0ba26 100644 --- a/tests/server/api_endpoints/test_session_outcomes_integration.py +++ b/tests/server/api_endpoints/test_session_outcomes_integration.py @@ -54,13 +54,18 @@ def test_source_is_derived_from_tiebroken_first_request( ) assert response.status_code == 200 - assert response.json() == { + body = response.json() + assert { "success": True, "recorded": True, "message": "Outcome recorded", "user_id": "u1", "source": "canonical-source", - } + }.items() <= body.items() + assert body["outcome_id"] + assert body["outcome_revision"] == 1 + assert len(body["outcome_contract_digest"]) == 64 + assert len(body["finalized_trajectory_digest"]) == 64 assert "stripped unknown fields: source" in caplog.text assert "multiple sources for session source-session" in caplog.text @@ -105,6 +110,169 @@ def test_retry_survives_ordinary_session_deletion( assert retry.json()["source"] == "published" +def test_conflicting_retry_is_not_accepted_only_by_session_identity( + client_with_org: tuple[TestClient, str], +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + storage.add_request( + Request( + request_id="conflict-r1", + user_id="u1", + session_id="conflict-session", + source="published", + created_at=100, + ) + ) + first = client.post( + "/api/session_outcome", + json={ + "session_id": "conflict-session", + "outcome": "success", + "occurred_at": 101, + }, + ) + conflict = client.post( + "/api/session_outcome", + json={ + "session_id": "conflict-session", + "outcome": "failure", + "occurred_at": 101, + }, + ) + + assert first.status_code == 200 + assert first.json()["recorded"] is True + assert conflict.status_code == 200 + assert conflict.json()["success"] is False + assert conflict.json()["recorded"] is False + assert conflict.json()["reason"] == "conflicting_finalization" + + +def test_retry_compares_metadata_as_json_values( + client_with_org: tuple[TestClient, str], +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + storage.add_request( + Request( + request_id="semantic-metadata-r1", + user_id="u1", + session_id="semantic-metadata-session", + source="published", + created_at=100, + ) + ) + payload = { + "session_id": "semantic-metadata-session", + "outcome": "success", + "occurred_at": 101, + "metadata": {"label": "same", "nested": {"one": 1, "two": 2}}, + } + first = client.post("/api/session_outcome", json=payload) + assert first.json()["recorded"] is True + storage.conn.execute( # type: ignore[attr-defined] + "UPDATE session_outcomes SET metadata = ? WHERE session_id = ?", + ( + '{ "nested": { "two": 2, "one": 1 }, "label": "same" }', + "semantic-metadata-session", + ), + ) + storage.conn.commit() # type: ignore[attr-defined] + + retry = client.post("/api/session_outcome", json=payload) + + assert retry.status_code == 200 + assert retry.json()["success"] is True + assert retry.json()["recorded"] is False + assert retry.json().get("reason") is None + + +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: @@ -241,6 +409,41 @@ def test_outcome_warning_values_are_sanitized_and_bounded( assert source_warning.endswith("unsafe?session") +def test_unregistered_acceptance_provider_accepts_long_lived_session_marker( + client_with_org: tuple[TestClient, str], +) -> None: + client, org_id = client_with_org + storage = get_reflexio(org_id=org_id).get_storage() + session_started_at = int(time.time()) - 90 * 86400 + storage.add_request( + Request( + request_id="long-lived-r1", + user_id="u1", + session_id="long-lived-session", + source="published", + created_at=session_started_at, + ) + ) + + response = client.post( + "/api/session_outcome", + json={ + "session_id": "long-lived-session", + "outcome": "success", + "occurred_at": session_started_at + 1, + }, + ) + + assert response.status_code == 200 + assert response.json()["success"] is True + assert response.json()["recorded"] is True + outcomes = storage.get_session_outcomes( + GetSessionOutcomesRequest(session_ids=["long-lived-session"]) + ) + assert len(outcomes) == 1 + assert outcomes[0].occurred_at == session_started_at + 1 + + def test_acceptance_hook_runs_before_persistence_at_exact_deadline( client_with_org: tuple[TestClient, str], monkeypatch ) -> None: diff --git a/tests/server/routes/test_search_exposure_boundary.py b/tests/server/routes/test_search_exposure_boundary.py new file mode 100644 index 000000000..1530aef7f --- /dev/null +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -0,0 +1,363 @@ +"""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 + +import pytest +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[MagicMock]: + 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 + reflexio.search_user_playbooks.return_value = result + with patch( + "reflexio.server.routes.search.reflexio_cache.get_reflexio", + return_value=reflexio, + ): + yield reflexio + + +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 + order: list[str] | None = None + + def record(self, batch: Any) -> None: + if self.order is not None: + self.order.append("record") + self.batches.append(batch) + self.completed = True + + +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_direct_user_playbook_search_records_final_results_before_metering() -> None: + playbooks = [_playbook(21, "Direct first"), _playbook(22, "Direct second")] + order: list[str] = [] + recorder = _Recorder(order=order) + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with ( + _search_results(playbooks), + patch( + "reflexio.server.routes.search.enqueue_search_metering", + side_effect=lambda **_kwargs: order.append("meter_enqueue"), + ), + ): + response = _client().post( + "/api/search_user_playbooks", + json={ + "query": "direct answer", + "user_id": "user-1", + "request_id": "request-direct-1", + "session_id": "session-direct-1", + }, + ) + + assert response.status_code == 200, response.text + assert order == ["record", "meter_enqueue"] + assert len(recorder.batches) == 1 + batch = recorder.batches[0] + assert batch.org_id == "org-1" + assert batch.request_id == "request-direct-1" + assert batch.session_id == "session-direct-1" + assert batch.interaction_id is None + assert batch.user_id == "user-1" + assert batch.user_playbooks == tuple(playbooks) + + +def test_direct_user_playbook_recorder_failure_prevents_metering_and_success() -> None: + order: list[str] = [] + + class _FailingRecorder: + def record(self, _batch: Any) -> None: + order.append("record") + raise RuntimeError("ledger unavailable") + + register_service(SEARCH_EXPOSURE_RECORDER, _FailingRecorder()) + + with ( + _search_results([_playbook(21, "Direct first")]), + patch( + "reflexio.server.routes.search.enqueue_search_metering", + side_effect=lambda **_kwargs: order.append("meter_enqueue"), + ), + ): + response = _client().post( + "/api/search_user_playbooks", + json={ + "query": "direct answer", + "user_id": "user-1", + "request_id": "request-direct-1", + }, + ) + + assert response.status_code == 500 + assert order == ["record"] + + +def test_direct_user_playbook_search_does_not_record_empty_results() -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([]): + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", "user_id": "user-1"}, + ) + + assert response.status_code == 200, response.text + assert recorder.batches == [] + + +def test_direct_user_playbook_search_rejects_101_before_search_or_recording() -> None: + playbooks = [ + _playbook(playbook_id, f"Playbook {playbook_id}") + for playbook_id in range(1, 102) + ] + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results(playbooks) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", "user_id": "user-1", "top_k": 101}, + ) + + assert response.status_code == 422 + reflexio.search_user_playbooks.assert_not_called() + assert recorder.batches == [] + + +def test_direct_user_playbook_search_returns_and_records_exactly_100() -> None: + playbooks = [ + _playbook(playbook_id, f"Playbook {playbook_id}") + for playbook_id in range(1, 101) + ] + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results(playbooks) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", "user_id": "user-1", "top_k": 100}, + ) + + assert response.status_code == 200, response.text + reflexio.search_user_playbooks.assert_called_once() + assert len(response.json()["user_playbooks"]) == 100 + assert len(recorder.batches) == 1 + assert recorder.batches[0].user_playbooks == tuple(playbooks) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("user_id", "u" * 256), + ("request_id", "r" * 256), + ("session_id", "s" * 256), + ], +) +def test_direct_user_playbook_search_rejects_oversized_identifiers_before_search_or_recording( + field: str, + value: str, +) -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([]) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={"query": "direct answer", field: value}, + ) + + assert response.status_code == 422 + reflexio.search_user_playbooks.assert_not_called() + assert recorder.batches == [] + + +def test_direct_user_playbook_search_accepts_255_character_identifiers() -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([]) as reflexio: + response = _client().post( + "/api/search_user_playbooks", + json={ + "query": "direct answer", + "user_id": "u" * 255, + "request_id": "r" * 255, + "session_id": "s" * 255, + }, + ) + + assert response.status_code == 200, response.text + reflexio.search_user_playbooks.assert_called_once() + assert len(recorder.batches) == 0 + + +def test_oss_search_succeeds_when_no_recorder_is_registered() -> None: + with _search_results([_playbook(11, "First")]): + response = _client().post( + "/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 + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("top_k", 101), + ("request_id", "r" * 256), + ("session_id", "s" * 256), + ("user_id", "u" * 256), + ], +) +def test_unified_search_rejects_oversized_work_before_search_execution( + field: str, + value: object, +) -> None: + with _search_results([]) as reflexio: + response = _client().post( + "/api/search", + json={"query": "answer", field: value}, + ) + + assert response.status_code == 422 + reflexio.unified_search.assert_not_called() + + +def test_unified_search_accepts_exact_workload_and_identifier_limits() -> None: + with _search_results([]) as reflexio: + response = _client().post( + "/api/search", + json={ + "query": "answer", + "top_k": 100, + "request_id": "r" * 255, + "session_id": "s" * 255, + "user_id": "u" * 255, + }, + ) + + assert response.status_code == 200, response.text + reflexio.unified_search.assert_called_once() diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index f26a88d00..40cead5fa 100644 --- a/tests/server/services/extraction/test_resume_worker.py +++ b/tests/server/services/extraction/test_resume_worker.py @@ -2,24 +2,49 @@ import json import tempfile +import threading +from collections import Counter +from collections.abc import Callable +from contextlib import ExitStack from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock, patch import pytest -from reflexio.models.api_schema.service_schemas import Interaction, Request +from reflexio.models.api_schema.service_schemas import ( + Interaction, + Request, + UserPlaybook, + UserProfile, +) from reflexio.models.config_schema import ( Config, PendingToolCallConfig, + PlaybookConfig, ProfileExtractorConfig, StorageConfigSQLite, ) from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.billing_meter import ReceiptBillingDeliveryError +from reflexio.server.services.deferred_learning_plan import FinalizationResult from reflexio.server.services.extraction.resume_worker import ( ExtractionResumeWorker, + _finalization_failure_status, _run_playbook_contract_selection, _run_uses_strict_playbook_evidence, ) +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, @@ -32,6 +57,12 @@ build_scope_hash, human_feedback_scope, ) +from reflexio.server.usage_metrics import ( + UsageEvent, + UsageEventDeliveryStatus, + configure_usage_event_recorder, + exempt_usage_event_recorder, +) @pytest.fixture @@ -43,6 +74,15 @@ def storage(): yield SQLiteStorage(org_id="org_1", db_path=f"{temp_dir}/reflexio.db") +@pytest.fixture(autouse=True) +def explicit_test_billing_exemption(): + configure_usage_event_recorder(exempt_usage_event_recorder) + try: + yield + finally: + configure_usage_event_recorder(None) + + @pytest.fixture def request_context(storage): ctx = RequestContext.__new__(RequestContext) @@ -55,7 +95,10 @@ def request_context(storage): profile_extractor_config=ProfileExtractorConfig( extraction_definition_prompt="Extract durable user deployment facts.", ), - pending_tool_call_config=PendingToolCallConfig(enabled=True), + pending_tool_call_config=PendingToolCallConfig( + enabled=True, + max_finalization_attempts=3, + ), ) ctx.configurator.get_agent_context.return_value = "Test agent context" ctx.prompt_manager = MagicMock() @@ -65,6 +108,159 @@ 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 _persist_and_load_run( + storage: SQLiteStorage, run: AgentRunRecord +) -> AgentRunRecord: + storage.create_agent_run(run) + persisted = storage.get_agent_run(run.id) + assert persisted is not None + assert persisted.created_at is not None + return persisted + + +def _gate_initial_receipt_reads( + storages: tuple[SQLiteStorage, SQLiteStorage], +) -> None: + 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) + ] + + +class _DeduplicatingUsageRecorder: + def __init__(self) -> None: + self.attempts: list[UsageEvent] = [] + self.events: list[UsageEvent] = [] + self._accepted_keys: set[str] = set() + self._lock = threading.Lock() + + def __call__(self, event: UsageEvent) -> UsageEventDeliveryStatus: + with self._lock: + self.attempts.append(event) + assert event.event_key is not None + if event.event_key in self._accepted_keys: + return UsageEventDeliveryStatus.DUPLICATE + self._accepted_keys.add(event.event_key) + self.events.append(event) + return UsageEventDeliveryStatus.APPENDED + + @pytest.mark.parametrize( ("schema_name", "expected"), [ @@ -366,7 +562,7 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( patch("litellm.completion", side_effect=[response]), patch( "reflexio.server.services.profile.service." - "ProfileGenerationService._finalize_extracted_items", + "ProfileGenerationService._finalize_extracted_items_with_outcome", side_effect=RuntimeError("storage write failed"), ), ): @@ -393,8 +589,8 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( ), patch( "reflexio.server.services.profile.service." - "ProfileGenerationService._finalize_extracted_items", - return_value=None, + "ProfileGenerationService._finalize_extracted_items_with_outcome", + return_value=FinalizationResult([], won_receipt=False), ) as finalize, ): resumed = worker.drain(max_runs=1) @@ -413,13 +609,13 @@ def test_resume_worker_retries_finalization_without_rerunning_agent( ( "profile", "reflexio.server.services.profile.service." - "ProfileGenerationService._finalize_extracted_items", + "ProfileGenerationService._finalize_extracted_items_with_outcome", "profile", ), ( "playbook", "reflexio.server.services.playbook.service." - "PlaybookGenerationService._finalize_extracted_items", + "PlaybookGenerationService._finalize_extracted_items_with_outcome", "user_playbook", ), ], @@ -442,15 +638,19 @@ def test_resume_bills_only_items_that_survive_finalization( ) dropped = object() survivor = object() + survivor_id = "durable-survivor-id" worker = ExtractionResumeWorker(request_context=request_context) with ( - patch(finalize_path, return_value=[survivor]) as finalize, + patch( + finalize_path, + return_value=FinalizationResult([survivor_id], won_receipt=True), + ) as finalize, patch.object(worker, "_record_finalized_learnings") as record, ): worker._finalize_items(run, [dropped, survivor]) - record.assert_called_once_with(run, [survivor], entity_type=entity_type) + record.assert_called_once_with(run, [survivor_id], entity_type=entity_type) if extractor_kind == "playbook": assert finalize.call_args.kwargs["extraction_run"] is run else: @@ -482,6 +682,1147 @@ def test_resume_worker_tagging_schedule_failure_is_best_effort( worker._schedule_finalized_tagging(run) +def test_resumable_finalization_bills_only_durable_ids_idempotently_on_retry( + request_context, +): + """A mixed batch charges its persisted profile once across finalization retries.""" + run_created_at = datetime(2026, 8, 1, 12, 0, tzinfo=UTC) + agent_completed_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) + run = AgentRunRecord( + id="run_mixed_billing", + binding=AgentBinding( + org_id="org_1", + extractor_kind="profile", + user_id="user_1", + request_id="request_1", + agent_version="v1", + source="api", + ), + status=AgentRunStatus.FINALIZATION_FAILED, + generation_request_snapshot={"request_id": "request_1"}, + agent_completed_at=agent_completed_at, + created_at=run_created_at, + ) + learning_ids = ["profile_1"] + worker = ExtractionResumeWorker(request_context=request_context) + + with ( + patch.object( + worker.storage, + "get_agent_run", + side_effect=AssertionError("billing must not re-read the run"), + ), + patch( + "reflexio.server.billing_meter.record_usage_event_strict", + return_value=UsageEventDeliveryStatus.APPENDED, + ) as record_event, + ): + worker._record_finalized_learnings(run, learning_ids, entity_type="profile") + worker._record_finalized_learnings(run, learning_ids, entity_type="profile") + + 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", + ] + assert [call.kwargs["created_at"] for call in record_event.call_args_list] == [ + run_created_at.timestamp(), + run_created_at.timestamp(), + ] + + +@pytest.mark.parametrize( + ("extractor_kind", "entity_type"), + [("profile", "profile"), ("playbook", "user_playbook")], +) +def test_delivery_failure_after_receipt_commit_retries_billing_without_recompute( + request_context, + storage, + extractor_kind, + entity_type, +): + _seed_interactions(storage) + request_context.configurator.get_config.return_value = Config( + storage_config=StorageConfigSQLite(), + profile_extractor_config=ProfileExtractorConfig( + extraction_definition_prompt="Extract durable user facts.", + ), + user_playbook_extractor_config=PlaybookConfig( + extraction_definition_prompt="Extract durable operating rules.", + ), + pending_tool_call_config=PendingToolCallConfig(enabled=True), + ) + request_id = f"request_{extractor_kind}_delivery_retry" + run_id = f"run_{extractor_kind}_delivery_retry" + committed_output = ( + { + "profiles": [ + { + "content": "User deployment target is AWS ECS.", + "time_to_live": "infinity", + } + ] + } + if extractor_kind == "profile" + else { + "playbooks": [ + { + "content": "Prefer AWS ECS as the deployment target.", + "trigger": "when selecting a deployment target", + "rationale": "The team standardizes on AWS.", + } + ] + } + ) + run_created_at = datetime(2026, 8, 1, 12, 0, tzinfo=UTC) + agent_completed_at = datetime(2026, 8, 31, 23, 59, tzinfo=UTC) + storage.create_agent_run( + AgentRunRecord( + id=run_id, + binding=AgentBinding( + org_id="org_1", + extractor_kind=extractor_kind, + user_id="user_1", + request_id=request_id, + agent_version="v1", + source="api", + source_interaction_ids=[1, 2], + ), + status=AgentRunStatus.FINALIZATION_FAILED, + generation_request_snapshot={ + "request_id": request_id, + "output_schema_name": "StructuredPlaybookList", + }, + committed_output=committed_output, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + finalization_attempts=2, + agent_completed_at=agent_completed_at, + created_at=run_created_at, + ) + ) + stored_run = storage.get_agent_run(run_id) + assert stored_run is not None + assert stored_run.created_at is not None + durable_run_created_at = stored_run.created_at + worker = ExtractionResumeWorker( + request_context=request_context, + llm_client=MagicMock(), + ) + attempts: list[UsageEvent] = [] + accepted: dict[str, UsageEvent] = {} + fail_next = True + + def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: + nonlocal fail_next + attempts.append(event) + if fail_next: + fail_next = False + raise RuntimeError("billing sink unavailable") + assert event.event_key is not None + if event.event_key in accepted: + return UsageEventDeliveryStatus.DUPLICATE + accepted[event.event_key] = event + return UsageEventDeliveryStatus.APPENDED + + configure_usage_event_recorder(recorder) + try: + with ExitStack() as stack: + schedule_tagging = stack.enter_context( + patch.object(worker, "_schedule_finalized_tagging") + ) + if extractor_kind == "profile": + stack.enter_context( + patch( + "reflexio.server.services.profile.components.consolidator." + "ProfileConsolidator.deduplicate", + side_effect=lambda profiles, _user_id, _request_id: ( + profiles, + [], + [], + ), + ) + ) + optimize = aggregate = None + else: + stack.enter_context( + patch.object( + PlaybookGenerationService, + "_configured_playbook_config", + return_value=None, + ) + ) + stack.enter_context( + patch( + "reflexio.server.services.playbook.components.consolidator." + "PlaybookConsolidator.deduplicate", + side_effect=lambda results, *_args, **_kwargs: ( + [playbook for result in results for playbook in result], + [], + [], + ), + ) + ) + optimize = stack.enter_context( + patch.object( + PlaybookGenerationService, + "_enqueue_user_playbook_optimization", + ) + ) + aggregate = stack.enter_context( + patch.object( + PlaybookGenerationService, + "_trigger_playbook_aggregation", + ) + ) + + failed = worker.run_once() + assert failed is not None + assert failed.status == AgentRunStatus.FINALIZATION_FAILED + assert failed.finalization_attempts == 3 + receipt_after_failure = storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type=entity_type, + ) + assert receipt_after_failure + rows_after_failure = ( + [profile.profile_id for profile in storage.get_user_profile("user_1")] + if extractor_kind == "profile" + else [ + str(row[0]) + for row in storage.conn.execute( + "SELECT user_playbook_id FROM user_playbooks " + "WHERE request_id = ? ORDER BY user_playbook_id", + (request_id,), + ).fetchall() + ] + ) + + storage.update_agent_run_status( + run_id, + AgentRunStatus.FINALIZATION_FAILED, + next_resume_at=datetime(2000, 1, 1, tzinfo=UTC), + ) + retry = worker.run_once() + + assert retry is not None + assert retry.status == AgentRunStatus.FINALIZED + assert ( + storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type=entity_type, + ) + == receipt_after_failure + ) + assert rows_after_failure == receipt_after_failure + assert [event.event_key for event in attempts] == [ + f"learn:{entity_type}:{receipt_after_failure[0]}", + f"learn:{entity_type}:{receipt_after_failure[0]}", + ] + assert list(accepted) == [f"learn:{entity_type}:{receipt_after_failure[0]}"] + assert [event.created_at for event in attempts] == [ + durable_run_created_at.timestamp(), + durable_run_created_at.timestamp(), + ] + schedule_tagging.assert_not_called() + if optimize is not None and aggregate is not None: + optimize.assert_called_once() + aggregate.assert_called_once() + finally: + configure_usage_event_recorder(None) + + +@pytest.mark.parametrize("next_attempt_count", [1, 3, 4]) +@pytest.mark.parametrize( + ("delivery_status", "expected_status"), + [ + (UsageEventDeliveryStatus.FAILED, AgentRunStatus.FINALIZATION_FAILED), + (UsageEventDeliveryStatus.UNKNOWN, AgentRunStatus.FINALIZATION_FAILED), + (UsageEventDeliveryStatus.REJECTED, AgentRunStatus.FAILED), + ], +) +def test_receipt_delivery_failure_status_distinguishes_transient_and_permanent( + delivery_status, + next_attempt_count, + expected_status, +): + error = ReceiptBillingDeliveryError(delivery_status) + + assert ( + _finalization_failure_status( + error, + next_attempt_count=next_attempt_count, + max_finalization_attempts=3, + ) + is expected_status + ) + + +@pytest.mark.parametrize("next_attempt_count", [3, 4]) +def test_ordinary_finalization_failure_stops_at_attempt_ceiling(next_attempt_count): + assert ( + _finalization_failure_status( + RuntimeError("ordinary finalization failed"), + next_attempt_count=next_attempt_count, + max_finalization_attempts=3, + ) + is AgentRunStatus.FAILED + ) + + +def test_retry_after_billing_reuses_ids_without_replaying_playbook_schedulers( + request_context, + storage, + monkeypatch, +): + """A post-billing retry reuses durable IDs and skips derived schedulers.""" + monkeypatch.setenv("MOCK_LLM_RESPONSE", "true") + _seed_interactions(storage) + request_context.configurator.get_config.return_value = Config( + storage_config=StorageConfigSQLite(), + 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] = [] + delivery_attempts: list[UsageEvent] = [] + accepted_keys: set[str] = set() + + def deduplicating_recorder(event: UsageEvent) -> UsageEventDeliveryStatus: + delivery_attempts.append(event) + assert event.event_key is not None + if event.event_key in accepted_keys: + return UsageEventDeliveryStatus.DUPLICATE + accepted_keys.add(event.event_key) + events.append(event) + return UsageEventDeliveryStatus.APPENDED + + configure_usage_event_recorder(deduplicating_recorder) + + def _run_failed_then_retried( + run_id: str, + *, + 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]}"] + assert [event.event_key for event in delivery_attempts] == [ + profile_keys[0], + profile_keys[0], + playbook_keys[0], + playbook_keys[0], + ] + + +@pytest.mark.parametrize("failing_scheduler", ["optimizer", "aggregation"]) +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}", + ) + run = _persist_and_load_run(storage, run) + worker = ExtractionResumeWorker( + request_context=_finalization_context(storage), + llm_client=MagicMock(), + ) + candidates = _playbook_candidates( + request_id=run.binding.request_id, + prefix=failing_scheduler, + ) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) + 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 recorder.events + if event.event_name == "learnings_generated" + and event.entity_type == "user_playbook" + ] + assert billing_ids == winner.learning_ids + assert [event.entity_id for event in recorder.attempts] == [ + *winner.learning_ids, + *winner.learning_ids, + ] + + +def test_empty_profile_receipt_wins_once_without_billing_or_recompute(storage): + run = _finalizing_run( + run_id="run_empty_profile", + extractor_kind="profile", + request_id="request_empty_profile", + ) + run = _persist_and_load_run(storage, run) + context = _finalization_context(storage) + worker = ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) + 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 recorder.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), + ) + run = _persist_and_load_run(storage, 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 + + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) + 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 recorder.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", + ) + run = _persist_and_load_run(storage_a, run) + contexts = [_finalization_context(item) for item in (storage_a, storage_b)] + workers = [ + ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + 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) + ] + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) + 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 recorder.events + if event.event_name == "learnings_generated" and event.entity_type == "profile" + ] == ["profile-shared"] + assert [event.entity_id for event in recorder.attempts] == [ + "profile-shared", + "profile-shared", + ] + + +def test_identical_playbook_ids_use_atomic_receipt_owner(tmp_path): + 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", + ) + run = _persist_and_load_run(storage_a, run) + contexts = [_finalization_context(item) for item in (storage_a, storage_b)] + workers = [ + ExtractionResumeWorker(request_context=context, llm_client=MagicMock()) + 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), + ) + + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) + 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 recorder.events + if event.event_name == "learnings_generated" + and event.entity_type == "user_playbook" + ] == ["88", "89"] + assert [event.entity_id for event in recorder.attempts] == ["88", "89", "88", "89"] + + +def test_two_stale_profile_workers_preserve_order_and_bill_only_winner(tmp_path): + 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", + ) + run = _persist_and_load_run(storage_a, run) + _gate_initial_receipt_reads((storage_a, storage_b)) + contexts = [_finalization_context(storage) for storage in (storage_a, storage_b)] + resume_workers = [ + 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) + + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) + 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 recorder.events + if event.event_name == "learnings_generated" and event.entity_type == "profile" + ] + assert [event.entity_id for event in billing_events] == persisted_ids + assert Counter(event.entity_id for event in recorder.attempts) == Counter( + persisted_ids * 2 + ) + + +def test_two_stale_playbook_workers_preserve_order_and_dispatch_once(tmp_path): + 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", + ) + run = _persist_and_load_run(storage_a, run) + _gate_initial_receipt_reads((storage_a, storage_b)) + contexts = [_finalization_context(storage) for storage in (storage_a, storage_b)] + resume_workers = [ + 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) + + recorder = _DeduplicatingUsageRecorder() + configure_usage_event_recorder(recorder) + 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 recorder.events + if event.event_name == "learnings_generated" + and event.entity_type == "user_playbook" + ] + assert billing_ids == receipt_ids + assert Counter(event.entity_id for event in recorder.attempts) == Counter( + receipt_ids * 2 + ) + finally: + configure_usage_event_recorder(None) + + def test_resume_worker_fails_run_when_step_budget_exhausted( monkeypatch, request_context, diff --git a/tests/server/services/governance/test_governance_local_e2e.py b/tests/server/services/governance/test_governance_local_e2e.py index d8ee06b06..e717a807d 100644 --- a/tests/server/services/governance/test_governance_local_e2e.py +++ b/tests/server/services/governance/test_governance_local_e2e.py @@ -1,9 +1,12 @@ from __future__ import annotations +import json +import threading from collections.abc import Generator from datetime import UTC, datetime from pathlib import Path from types import SimpleNamespace +from typing import Any from unittest.mock import patch import pytest @@ -18,12 +21,14 @@ UserProfile, ) from reflexio.models.api_schema.domain.enums import PlaybookStatus +from reflexio.models.api_schema.domain.governance import UserEraseResult from reflexio.models.api_schema.retriever_schema import SearchAgentPlaybookRequest from reflexio.models.config_schema import SearchMode from reflexio.server.services.governance import service as governance_service_module from reflexio.server.services.governance.config import governance_subject_ref from reflexio.server.services.governance.service import GovernanceService from reflexio.server.services.storage.error import SubjectWriteBarrierError +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.sqlite_storage import SQLiteStorage pytestmark = pytest.mark.integration @@ -117,6 +122,154 @@ 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() + + +def test_purge_execution_heartbeat_serializes_concurrent_renewals() -> None: + initial_claim = PurgeExecutionClaim( + purge_id="purge-1", + owner="owner-1", + fence=1, + expires_at=300, + ) + first_renewed_claim = PurgeExecutionClaim( + purge_id="purge-1", + owner="owner-1", + fence=1, + expires_at=400, + ) + second_renewed_claim = PurgeExecutionClaim( + purge_id="purge-1", + owner="owner-1", + fence=1, + expires_at=401, + ) + + class RenewalProgress: + def __init__(self) -> None: + self.event = threading.Event() + self._lock = threading.Lock() + self.source: str | None = None + + def record(self, source: str) -> None: + with self._lock: + if self.event.is_set(): + return + self.source = source + self.event.set() + + class BlockingRenewalStorage: + def __init__(self, progress: RenewalProgress) -> None: + self._progress = progress + self.claims: list[PurgeExecutionClaim] = [] + self._calls_lock = threading.Lock() + self.first_renewal_started = threading.Event() + self.release_first_renewal = threading.Event() + self.second_renewal_started = threading.Event() + + def renew_purge_operation_execution_claim( + self, + _purge_id: str, + claim: PurgeExecutionClaim, + *, + lease_ttl_seconds: int, + ) -> PurgeExecutionClaim: + assert lease_ttl_seconds == 300 + with self._calls_lock: + self.claims.append(claim) + call_number = len(self.claims) + if call_number == 1: + self.first_renewal_started.set() + assert self.release_first_renewal.wait(timeout=5) + return first_renewed_claim + self._progress.record("storage") + self.second_renewal_started.set() + return second_renewed_claim + + class TrackingRenewalLock: + def __init__(self, progress: RenewalProgress) -> None: + self._progress = progress + self._lock = threading.Lock() + self._attempts_lock = threading.Lock() + self._attempts = 0 + self.second_renewal_attempted = threading.Event() + + def __enter__(self) -> TrackingRenewalLock: + with self._attempts_lock: + self._attempts += 1 + if self._attempts == 2: + self._progress.record("renewal lock") + self._lock.acquire() + return self + + def __exit__(self, *_exc: object) -> None: + self._lock.release() + + progress = RenewalProgress() + storage = BlockingRenewalStorage(progress) + heartbeat = governance_service_module._PurgeExecutionHeartbeat( + storage=storage, + purge_id="purge-1", + execution_claim=initial_claim, + ) + renewal_lock = TrackingRenewalLock(progress) + cast_heartbeat: Any = heartbeat + cast_heartbeat._renewal_lock = renewal_lock + errors: list[Exception] = [] + + def renew() -> None: + try: + heartbeat.renew_now() + except Exception as exc: + errors.append(exc) + + allow_second_renewal = threading.Event() + + def renew_second() -> None: + assert allow_second_renewal.wait(timeout=5) + renew() + + first = threading.Thread(target=renew) + second = threading.Thread(target=renew_second) + first.start() + second.start() + try: + assert storage.first_renewal_started.wait(timeout=5) + allow_second_renewal.set() + assert progress.event.wait(timeout=5) + assert progress.source == "renewal lock" + assert not storage.second_renewal_started.is_set() + finally: + allow_second_renewal.set() + storage.release_first_renewal.set() + first.join(timeout=5) + second.join(timeout=5) + + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert storage.claims == [initial_claim, first_renewed_claim] + assert heartbeat.claim() == second_renewed_claim + + @pytest.fixture def storage( tmp_path: Path, @@ -306,6 +459,7 @@ def test_local_governance_e2e_erases_exports_audits_and_preserves_org_agent_play assert erased.deleted_counts["requests"] == 1 assert erased.deleted_counts["user_playbooks"] == 2 assert erased.deleted_counts["agent_success_evaluation_results"] == 1 + assert erased.deleted_counts["session_outcomes"] == 0 assert erased.rebuilt_agent_playbook_ids == [] assert storage.get_user_interaction("alice") == [] @@ -437,6 +591,55 @@ def test_local_governance_e2e_erases_exports_audits_and_preserves_org_agent_play assert len(erase_events_after_retry) == 1 +def test_governance_erasure_uses_authoritative_user_for_session_outcomes_and_receipts( + storage: SQLiteStorage, +) -> None: + alice_ref = governance_subject_ref( + storage.org_id, "alice", "test-governance-secret" + ) + bob_ref = governance_subject_ref(storage.org_id, "bob", "test-governance-secret") + _insert_session_outcome( + storage, + outcome_id="alice-stale-ref", + user_id="alice", + session_id="alice-session", + subject_ref=bob_ref, + ) + _insert_session_outcome( + storage, + outcome_id="bob-conflicting-ref", + user_id="bob", + session_id="bob-session", + subject_ref=alice_ref, + ) + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + + erased = service.erase_user(user_id="alice", request_id="erase-outcomes") + retried = service.erase_user(user_id="alice", request_id="erase-outcomes") + + remaining = storage.conn.execute( + "SELECT outcome_id, user_id FROM session_outcomes ORDER BY outcome_id" + ).fetchall() + assert [(row["outcome_id"], row["user_id"]) for row in remaining] == [ + ("bob-conflicting-ref", "bob") + ] + assert erased.deleted_counts["session_outcomes"] == 1 + assert retried.deleted_counts == erased.deleted_counts + audit = next( + event + for event in storage.list_audit_events(subject_ref=alice_ref) + if event.operation == "ERASE" + ) + assert audit.detail is not None + deleted_counts = audit.detail["deleted_counts"] + assert isinstance(deleted_counts, dict) + assert deleted_counts["session_outcomes"] == 1 + + def test_governance_service_persists_actor_context_in_audit( storage: SQLiteStorage, ) -> None: @@ -517,6 +720,54 @@ def test_completed_erase_retry_reconstructs_response( assert second.rebuilt_agent_playbook_ids == first.rebuilt_agent_playbook_ids +@pytest.mark.parametrize("corrupt_binding", ["purge", "snapshot"]) +def test_completed_erase_retry_rejects_corrupt_authoritative_binding( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, + corrupt_binding: str, +) -> None: + monkeypatch.setenv("REFLEXIO_GOVERNANCE_REF_SECRET", "test-governance-secret") + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + completed = service.erase_user( + user_id="alice", + request_id=f"erase-retry-corrupt-{corrupt_binding}", + ) + + if corrupt_binding == "purge": + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + ("a" * 64, storage.org_id, completed.purge_id), + ) + else: + snapshot = next( + target + for target in storage.list_purge_targets( + completed.purge_id, phase="prepare_targets" + ) + if target.target_name == "target_snapshot" + ) + detail = dict(snapshot.detail or {}) + detail["authoritative_user_digest"] = "a" * 64 + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(detail), storage.org_id, completed.purge_id), + ) + storage.conn.commit() + + with pytest.raises(ValueError, match="authoritative user identity"): + service.erase_user( + user_id="alice", + request_id=f"erase-retry-corrupt-{corrupt_binding}", + ) + + def test_erase_fails_fast_when_service_and_storage_ref_secrets_differ( storage: SQLiteStorage, ) -> None: @@ -629,10 +880,21 @@ def test_second_erase_conflict_preserves_original_barrier_and_write_block( idempotency_key="idem_conflict_first", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000061", ) - storage.begin_subject_erasure_barrier(subject_ref, first_purge.purge_id) + first_claim = storage.claim_purge_operation_execution( + first_purge.purge_id, + lease_owner="test-conflict-first", + lease_ttl_seconds=30, + ) + assert first_claim is not None + storage.begin_subject_erasure_barrier( + subject_ref, + first_purge.purge_id, + execution_claim=first_claim, + ) service = GovernanceService( storage=storage, org_id=storage.org_id, @@ -721,6 +983,7 @@ def erase_subject( storage: SQLiteStorage, subject_ref: str, purge_id: str, + execution_claim: PurgeExecutionClaim, ) -> None: del subject_ref self.calls += 1 @@ -740,6 +1003,7 @@ def erase_subject( status="complete", detail={"count": 2}, deleted_count=2, + execution_claim=execution_claim, ) lifecycle = RetrySafeLifecycle() @@ -785,6 +1049,571 @@ def fail_after_first_lifecycle(*args, **kwargs): assert snapshot.detail["status"] == "complete" +def test_duplicate_erase_waits_for_lifecycle_winner_beyond_old_deadline( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + class SlowSingleUseLifecycle: + def __init__(self) -> None: + self.first_call_started = threading.Event() + self.release_first_call = threading.Event() + self._lock = threading.Lock() + self.calls = 0 + self.duplicate_rejections = 0 + + def erase_subject( + self, + *, + storage: SQLiteStorage, + subject_ref: str, + purge_id: str, + execution_claim: object, + ) -> None: + del storage, subject_ref, purge_id, execution_claim + with self._lock: + self.calls += 1 + call_number = self.calls + if call_number == 1: + self.first_call_started.set() + assert self.release_first_call.wait(timeout=5) + return + self.duplicate_rejections += 1 + raise RuntimeError("provider lifecycle already running") + + lifecycle = SlowSingleUseLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + + def release_winner_on_duplicate_poll(_seconds: float) -> None: + lifecycle.release_first_call.set() + threading.Event().wait(0.001) + + monkeypatch.setattr(service, "_sleep", release_winner_on_duplicate_poll) + winner_results: list[UserEraseResult] = [] + winner_errors: list[BaseException] = [] + + def run_winner() -> None: + try: + winner_results.append( + service.erase_user(user_id="alice", request_id="erase-slow-duplicate") + ) + except BaseException as exc: + winner_errors.append(exc) + + winner = threading.Thread(target=run_winner) + winner.start() + assert lifecycle.first_call_started.wait(timeout=5) + + try: + duplicate = service.erase_user( + user_id="alice", request_id="erase-slow-duplicate" + ) + finally: + lifecycle.release_first_call.set() + winner.join(timeout=5) + + assert not winner.is_alive() + assert winner_errors == [] + assert len(winner_results) == 1 + winner_result = winner_results[0] + assert duplicate.status == "complete" + assert duplicate.purge_id == winner_result.purge_id + assert storage.get_purge_operation(duplicate.purge_id).status == "complete" + barrier = storage.get_subject_write_barrier(duplicate.subject_ref) + assert barrier is not None + assert barrier.status == "erased" + assert lifecycle.calls == 1 + assert lifecycle.duplicate_rejections == 0 + + +def test_duplicate_erase_wait_is_bounded_with_exponential_backoff( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + class ControlledWaitGovernanceService(GovernanceService): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.now = 0.0 + self.sleep_delays: list[float] = [] + + def _monotonic(self) -> float: + return self.now + + def _sleep(self, seconds: float) -> None: + self.sleep_delays.append(seconds) + self.now += seconds + + service = ControlledWaitGovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + claim_attempts = 0 + + def reject_duplicate_claim(*args, **kwargs): + nonlocal claim_attempts + del args, kwargs + claim_attempts += 1 + if claim_attempts > 100: + raise AssertionError("duplicate claim wait exceeded its attempt bound") + + monkeypatch.setattr( + storage, + "claim_purge_operation_execution", + reject_duplicate_claim, + ) + with pytest.raises(RuntimeError, match="retry later") as exc_info: + service.erase_user(user_id="alice", request_id="erase-bounded-duplicate") + + assert ( + type(exc_info.value) is governance_service_module.GovernanceEraseRetryLaterError + ) + assert service.sleep_delays[:5] == pytest.approx([0.05, 0.1, 0.2, 0.4, 0.8]) + assert max(service.sleep_delays) == 1.0 + assert sum(service.sleep_delays) == pytest.approx(5.0) + + +def test_healthy_slow_lifecycle_renews_lease_and_duplicate_converges( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_now = {"value": 100} + monkeypatch.setattr( + governance_service_module, + "_PURGE_EXECUTION_HEARTBEAT_SECONDS", + 0.01, + ) + + from reflexio.server.services.storage.sqlite_storage.governance import ( + _purge as sqlite_purge_module, + ) + + monkeypatch.setattr(sqlite_purge_module, "_epoch_now", lambda: fake_now["value"]) + original_renew = storage.renew_purge_operation_execution_claim + renewal_confirmed = threading.Event() + renewals_by_owner: dict[str, int] = {} + + def observed_renew(*args, **kwargs): + claim = args[1] + renewals_by_owner[claim.owner] = renewals_by_owner.get(claim.owner, 0) + 1 + renewed = original_renew(*args, **kwargs) + if renewals_by_owner[claim.owner] >= 3: + renewal_confirmed.set() + return renewed + + monkeypatch.setattr( + storage, + "renew_purge_operation_execution_claim", + observed_renew, + ) + + class SlowSingleUseLifecycle: + def __init__(self) -> None: + self.first_call_started = threading.Event() + self.release_first_call = threading.Event() + self._lock = threading.Lock() + self.calls = 0 + + def erase_subject( + self, + *, + storage: SQLiteStorage, + subject_ref: str, + purge_id: str, + execution_claim: object, + ) -> None: + del storage, subject_ref, purge_id, execution_claim + with self._lock: + self.calls += 1 + call_number = self.calls + if call_number != 1: + raise RuntimeError("provider lifecycle already running") + self.first_call_started.set() + fake_now["value"] = 350 + assert renewal_confirmed.wait(timeout=5) + fake_now["value"] = 401 + assert self.release_first_call.wait(timeout=5) + + lifecycle = SlowSingleUseLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + winner_results: list[UserEraseResult] = [] + winner_errors: list[BaseException] = [] + duplicate_results: list[UserEraseResult] = [] + duplicate_errors: list[BaseException] = [] + + def run_winner() -> None: + try: + winner_results.append( + service.erase_user(user_id="alice", request_id="erase-renewed-slow") + ) + except BaseException as exc: + winner_errors.append(exc) + + def run_duplicate() -> None: + try: + duplicate_results.append( + service.erase_user(user_id="alice", request_id="erase-renewed-slow") + ) + except BaseException as exc: + duplicate_errors.append(exc) + + winner = threading.Thread(target=run_winner) + winner.start() + assert lifecycle.first_call_started.wait(timeout=5) + assert renewal_confirmed.wait(timeout=5) + + duplicate = threading.Thread(target=run_duplicate) + duplicate.start() + lifecycle.release_first_call.set() + winner.join(timeout=5) + duplicate.join(timeout=5) + + assert not winner.is_alive() + assert not duplicate.is_alive() + assert winner_errors == [] + assert duplicate_errors == [] + assert len(winner_results) == 1 + assert len(duplicate_results) == 1 + assert duplicate_results[0].status == "complete" + assert duplicate_results[0].purge_id == winner_results[0].purge_id + assert lifecycle.calls == 1 + + +def test_heartbeat_renewal_loss_fences_external_lifecycle_and_retry_converges( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_now = {"value": 100} + monkeypatch.setattr( + governance_service_module, + "_PURGE_EXECUTION_HEARTBEAT_SECONDS", + 0.01, + ) + from reflexio.server.services.storage.sqlite_storage.governance import ( + _purge as sqlite_purge_module, + ) + + monkeypatch.setattr(sqlite_purge_module, "_epoch_now", lambda: fake_now["value"]) + original_renew = storage.renew_purge_operation_execution_claim + first_owner: list[str] = [] + renewals_by_owner: dict[str, int] = {} + renewal_lost = threading.Event() + + def fail_first_owner_heartbeat(*args, **kwargs): + claim = args[1] + if not first_owner: + first_owner.append(claim.owner) + renewals_by_owner[claim.owner] = renewals_by_owner.get(claim.owner, 0) + 1 + if claim.owner == first_owner[0] and renewals_by_owner[claim.owner] == 2: + renewal_lost.set() + raise RuntimeError("simulated heartbeat renewal loss") + return original_renew(*args, **kwargs) + + monkeypatch.setattr( + storage, + "renew_purge_operation_execution_claim", + fail_first_owner_heartbeat, + ) + original_apply = storage.apply_governance_user_data_delete + + def apply_then_wait_for_renewal_loss(*args, **kwargs): + result = original_apply(*args, **kwargs) + assert renewal_lost.wait(timeout=5) + return result + + monkeypatch.setattr( + storage, + "apply_governance_user_data_delete", + apply_then_wait_for_renewal_loss, + ) + + class CountingLifecycle: + def __init__(self) -> None: + self.calls = 0 + + def erase_subject(self, **_kwargs) -> None: + self.calls += 1 + + lifecycle = CountingLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + + with pytest.raises(ValueError, match="heartbeat renewal was lost"): + service.erase_user(user_id="alice", request_id="erase-renewal-loss") + + assert lifecycle.calls == 0 + purge_id = str( + storage.conn.execute("SELECT purge_id FROM purge_operations").fetchone()[ + "purge_id" + ] + ) + assert storage.get_purge_operation(purge_id).status == "running" + assert storage.list_audit_events() == [] + + fake_now["value"] = 401 + recovered = service.erase_user( + user_id="alice", + request_id="erase-renewal-loss", + ) + + assert recovered.status == "complete" + assert lifecycle.calls == 1 + assert storage.get_purge_operation(purge_id).status == "complete" + + +def test_synchronous_renewal_loss_skips_lifecycle_and_retry_converges( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_renew = storage.renew_purge_operation_execution_claim + renewal_attempts = 0 + renewal_failed = threading.Event() + lifecycle_called = threading.Event() + + def fail_mandatory_lifecycle_renewal(*args, **kwargs): + nonlocal renewal_attempts + renewal_attempts += 1 + if renewal_attempts == 2: + renewal_failed.set() + raise RuntimeError("simulated synchronous renewal loss") + return original_renew(*args, **kwargs) + + monkeypatch.setattr( + storage, + "renew_purge_operation_execution_claim", + fail_mandatory_lifecycle_renewal, + ) + + class CountingLifecycle: + def __init__(self) -> None: + self.calls = 0 + + def erase_subject(self, **_kwargs) -> None: + self.calls += 1 + lifecycle_called.set() + + lifecycle = CountingLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + + with pytest.raises(ValueError, match="heartbeat renewal was lost"): + service.erase_user(user_id="alice", request_id="erase-sync-renewal-loss") + + purge_id = str( + storage.conn.execute("SELECT purge_id FROM purge_operations").fetchone()[ + "purge_id" + ] + ) + purge = storage.get_purge_operation(purge_id) + assert purge.subject_ref is not None + barrier = storage.get_subject_write_barrier(purge.subject_ref) + assert renewal_attempts == 2 + assert renewal_failed.is_set() + assert not lifecycle_called.is_set() + assert lifecycle.calls == 0 + assert purge.status == "running" + assert barrier is not None + assert barrier.status == "erasing" + assert storage.list_audit_events() == [] + + storage.conn.execute( + "UPDATE purge_operations SET execution_claim_expires_at = 0 WHERE purge_id = ?", + (purge_id,), + ) + storage.conn.commit() + + def fail_if_recovery_polls(_seconds: float) -> None: + raise AssertionError("expired synchronous-renewal claim did not recover") + + monkeypatch.setattr(service, "_sleep", fail_if_recovery_polls) + + recovered = service.erase_user( + user_id="alice", + request_id="erase-sync-renewal-loss", + ) + + assert recovered.status == "complete" + assert lifecycle_called.is_set() + assert lifecycle.calls == 1 + assert storage.get_purge_operation(purge_id).status == "complete" + + +def test_stale_running_erase_claim_recovers_after_crash( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + original_begin_barrier = storage.begin_subject_erasure_barrier + crash_once = True + + def crash_after_claim(subject_ref: str, purge_id: str, **kwargs): + nonlocal crash_once + if crash_once: + crash_once = False + raise SystemExit("simulated crash after claim") + return original_begin_barrier(subject_ref, purge_id, **kwargs) + + monkeypatch.setattr( + storage, + "begin_subject_erasure_barrier", + crash_after_claim, + ) + + with pytest.raises(SystemExit, match="simulated crash after claim"): + service.erase_user(user_id="alice", request_id="erase-crash-after-claim") + + storage.conn.execute("UPDATE purge_operations SET execution_claim_expires_at = 0") + storage.conn.commit() + + def fail_if_duplicate_polls(_seconds: float) -> None: + raise AssertionError("stale running claim did not recover") + + monkeypatch.setattr(service, "_sleep", fail_if_duplicate_polls) + + recovered = service.erase_user( + user_id="alice", + request_id="erase-crash-after-claim", + ) + + assert recovered.status == "complete" + assert storage.get_purge_operation(recovered.purge_id).status == "complete" + barrier = storage.get_subject_write_barrier(recovered.subject_ref) + assert barrier is not None + assert barrier.status == "erased" + + +def test_delete_committed_before_completion_retry_converges_idempotently( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + original_complete = storage.complete_subject_erasure_barrier_after_empty_check + completion_attempts = 0 + + def crash_after_delete_targets(*args, **kwargs): + nonlocal completion_attempts + completion_attempts += 1 + if completion_attempts == 1: + raise SystemExit("simulated crash after delete") + return original_complete(*args, **kwargs) + + monkeypatch.setattr( + storage, + "complete_subject_erasure_barrier_after_empty_check", + crash_after_delete_targets, + ) + + with pytest.raises(SystemExit, match="simulated crash after delete"): + service.erase_user(user_id="alice", request_id="erase-crash-after-delete") + + storage.conn.execute("UPDATE purge_operations SET execution_claim_expires_at = 0") + storage.conn.commit() + + delete_targets_after_crash = storage.list_purge_targets( + storage.conn.execute("SELECT purge_id FROM purge_operations").fetchone()[ + "purge_id" + ], + phase="delete", + ) + assert delete_targets_after_crash + assert all(target.status == "complete" for target in delete_targets_after_crash) + + def fail_if_duplicate_polls(_seconds: float) -> None: + raise AssertionError("stale post-delete claim did not recover") + + monkeypatch.setattr(service, "_sleep", fail_if_duplicate_polls) + + recovered = service.erase_user( + user_id="alice", + request_id="erase-crash-after-delete", + ) + + assert recovered.status == "complete" + assert storage.get_purge_operation(recovered.purge_id).status == "complete" + assert completion_attempts == 2 + + +def test_pending_duplicate_erase_has_one_durable_execution_owner( + storage: SQLiteStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_begin = storage.begin_purge_operation + both_pending = threading.Barrier(2) + + def synchronized_begin(*args, **kwargs): + purge = original_begin(*args, **kwargs) + both_pending.wait(timeout=5) + return purge + + monkeypatch.setattr(storage, "begin_purge_operation", synchronized_begin) + + class CountingLifecycle: + def __init__(self) -> None: + self.calls = 0 + self._lock = threading.Lock() + + def erase_subject(self, **_kwargs) -> None: + with self._lock: + self.calls += 1 + + lifecycle = CountingLifecycle() + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + subject_erasure_lifecycle=lifecycle, + ) + results: list[UserEraseResult] = [] + errors: list[BaseException] = [] + + def erase() -> None: + try: + results.append( + service.erase_user(user_id="alice", request_id="erase-pending-race") + ) + except BaseException as exc: + errors.append(exc) + + callers = [threading.Thread(target=erase) for _ in range(2)] + for caller in callers: + caller.start() + for caller in callers: + caller.join(timeout=5) + + assert all(not caller.is_alive() for caller in callers) + assert errors == [] + assert len(results) == 2 + assert results[0].purge_id == results[1].purge_id + assert all(result.status == "complete" for result in results) + assert lifecycle.calls == 1 + assert storage.get_purge_operation(results[0].purge_id).status == "complete" + + def test_session_export_paginates_by_returned_rows_when_requests_are_missing() -> None: class _Storage: def __init__(self) -> None: @@ -809,3 +1638,58 @@ def get_sessions(self, *, user_id: str, top_k: int, offset: int): assert storage.calls == [0, 1000] assert [request.request_id for request in requests] == ["req-1"] assert sessions == [{"session_id": "session-a", "request_ids": ["req-1"]}] + + +def test_rebuild_agent_playbooks_forwards_the_active_execution_claim() -> None: + target = SimpleNamespace( + target_name="agent_playbook", + target_ref="17", + status="running", + detail={"remaining_source_windows": []}, + ) + + class _Storage: + def __init__(self) -> None: + self.applied: list[dict[str, object]] = [] + + def list_purge_targets(self, purge_id: str, *, phase: str): + assert purge_id == "purge_claimed_rebuild" + assert phase == "rebuild_without_erased_sources" + return [target] + + def get_user_playbooks_by_ids_any_user(self, ids: list[int]): + assert ids == [] + return [] + + def apply_governance_agent_playbook_rebuild(self, **kwargs: object) -> None: + self.applied.append(kwargs) + + storage = _Storage() + service = GovernanceService(storage=storage, org_id="org", ref_secret="secret") + claim = PurgeExecutionClaim( + purge_id="purge_claimed_rebuild", + owner="worker-a", + fence=1, + expires_at=2_000_000_000, + ) + + rebuilt_ids = service._rebuild_agent_playbooks( + "purge_claimed_rebuild", + execution_claim=claim, + ) + + assert rebuilt_ids == [17] + assert storage.applied == [ + { + "purge_id": "purge_claimed_rebuild", + "agent_playbook_id": 17, + "remaining_source_windows": [], + "content": None, + "trigger": None, + "rationale": None, + "blocking_issue": None, + "expanded_terms": None, + "tags": None, + "execution_claim": claim, + } + ] diff --git a/tests/server/services/governance/test_subject_write_barrier_sqlite.py b/tests/server/services/governance/test_subject_write_barrier_sqlite.py index c649bcaaf..c82547119 100644 --- a/tests/server/services/governance/test_subject_write_barrier_sqlite.py +++ b/tests/server/services/governance/test_subject_write_barrier_sqlite.py @@ -1,5 +1,7 @@ from __future__ import annotations +import hashlib +import hmac from datetime import UTC, datetime from pathlib import Path @@ -14,12 +16,13 @@ UserPlaybook, UserProfile, ) -from reflexio.models.api_schema.domain.governance import AuditEvent +from reflexio.models.api_schema.domain.governance import AuditEvent, SubjectWriteBarrier from reflexio.server.services.governance.config import governance_subject_ref from reflexio.server.services.storage.error import ( StorageError, SubjectWriteBarrierError, ) +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, ) @@ -36,14 +39,78 @@ def _storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> SQLiteStorage: return SQLiteStorage(org_id="org-barrier", db_path=str(tmp_path / "barrier.db")) +def _claim_purge(storage: SQLiteStorage, purge_id: str) -> PurgeExecutionClaim: + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"test-{purge_id}", + lease_ttl_seconds=30, + ) + if claim is None: + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"test-{purge_id}", + lease_ttl_seconds=30, + ) + assert claim is not None + return claim + + +def _typed_test_claim_for_unvalidated_purge_id(purge_id: str) -> PurgeExecutionClaim: + return PurgeExecutionClaim( + purge_id=purge_id, + owner="test-unvalidated", + fence=1, + expires_at=1, + ) + + +def _begin_claimed_subject_erasure_barrier( + storage: SQLiteStorage, + subject_ref: str, + purge_id: str, +) -> SubjectWriteBarrier: + return storage.begin_subject_erasure_barrier( + subject_ref, + purge_id, + execution_claim=_claim_purge(storage, purge_id), + ) + + +def _authoritative_user_digest(storage: SQLiteStorage, purge_id: str) -> str: + return storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + + +def _expected_authoritative_user_digest( + *, secret: str, org_id: str, purge_id: str, user_id: str +) -> str: + material = f"authoritative-user-v1\0{org_id}\0{purge_id}\0{user_id}" + return hmac.new(secret.encode(), material.encode(), hashlib.sha256).hexdigest() + + def _mark_all_completion_targets(storage: SQLiteStorage, purge_id: str) -> None: + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - detail={"prepared": True}, + execution_claim=claim, + detail={ + "prepared": True, + "authoritative_user_digest": _authoritative_user_digest(storage, purge_id), + }, ) # Single source of truth — a stale local copy of the canonical tuple is # exactly how this suite went red when new delete targets landed. @@ -54,6 +121,7 @@ def _mark_all_completion_targets(storage: SQLiteStorage, purge_id: str) -> None: phase="delete", status="complete", target_ref="all", + execution_claim=claim, detail={"count": 0}, ) @@ -77,8 +145,190 @@ def _complete_empty_purge( idempotency_key=purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) + + +def test_begin_purge_operation_keys_authoritative_user_digest( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_keyed_authoritative_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_keyed_authoritative_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000056", + ) + + stored_digest = _authoritative_user_digest(storage, purge_id) + assert stored_digest == _expected_authoritative_user_digest( + secret="barrier-secret", + org_id="org-barrier", + purge_id=purge_id, + user_id="alice", + ) + assert stored_digest != hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + + +@pytest.mark.parametrize("legacy_digest", [None, "unkeyed"]) +def test_begin_purge_operation_upgrades_validated_legacy_authoritative_digest( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + legacy_digest: str | None, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_legacy_authoritative_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_authoritative_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000057", + ) + persisted_digest = ( + hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + if legacy_digest == "unkeyed" + else None + ) + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = NULLIF(?, '') + WHERE org_id = ? AND purge_id = ?""", + (persisted_digest or "", storage.org_id, purge_id), + ) + storage.conn.commit() + + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_authoritative_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000057", + ) + + assert _authoritative_user_digest( + storage, purge_id + ) == _expected_authoritative_user_digest( + secret="barrier-secret", + org_id="org-barrier", + purge_id=purge_id, + user_id="alice", + ) + + +def test_begin_purge_operation_does_not_upgrade_legacy_digest_for_wrong_user( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_legacy_wrong_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_wrong_identity", + operation_type="user_erasure", + scope_type="user", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000058", + authoritative_user_id="alice", + ) + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = NULL + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + + with pytest.raises(ValueError, match="must match subject_ref"): + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_legacy_wrong_identity", + operation_type="user_erasure", + scope_type="user", + subject_ref=subject_ref, + request_ref="reqref_v1_00000000000000000000000000000058", + authoritative_user_id="bob", + ) + + assert _authoritative_user_digest(storage, purge_id) is None + + +def test_completion_checks_session_outcomes_by_authoritative_user_id( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage = _storage(tmp_path, monkeypatch) + purge_id = "purge_exact_outcome_identity" + subject_ref = governance_subject_ref("org-barrier", "alice", "barrier-secret") + request_ref = "reqref_v1_00000000000000000000000000000059" + storage.begin_purge_operation( + purge_id=purge_id, + idempotency_key="idem_exact_outcome_identity", + operation_type="user_erasure", + scope_type="user", + authoritative_user_id="alice", + subject_ref=subject_ref, + request_ref=request_ref, + ) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge_id) + _mark_all_completion_targets(storage, purge_id) + storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES (?, 1, ?, ?, 'success', 1, 'test', ?, ?, ?, 1)""", + ( + "outcome-bob", + "bob", + "session-bob", + "contract-digest", + "trajectory-digest", + governance_subject_ref("org-barrier", "bob", "barrier-secret"), + ), + ) + storage.conn.commit() + original_subject_ref = storage._subject_ref_for_user_id + + def subject_ref_for_authoritative_user_only(user_id: str) -> str: + assert user_id == "alice", "completion enumerated an unrelated outcome user" + return original_subject_ref(user_id) + + monkeypatch.setattr( + storage, + "_subject_ref_for_user_id", + subject_ref_for_authoritative_user_only, + ) + + completed = storage.complete_subject_erasure_barrier_after_empty_check( + purge_id, + AuditEvent( + org_id="org-barrier", + operation="ERASE", + entity_type="request", + subject_ref=subject_ref, + request_ref=request_ref, + idempotency_key=purge_id, + detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, + ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) + assert completed.status == "complete" + def test_barrier_blocks_request_interaction_and_profile_writes( tmp_path: Path, @@ -91,11 +341,14 @@ def test_barrier_blocks_request_interaction_and_profile_writes( idempotency_key="idem_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_11111111111111111111111111111111", ) - barrier = storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + barrier = _begin_claimed_subject_erasure_barrier( + storage, subject_ref, purge.purge_id + ) assert barrier.status == "erasing" with pytest.raises(SubjectWriteBarrierError): @@ -145,6 +398,7 @@ def test_barrier_blocks_playbook_eval_and_source_window_writes( idempotency_key="idem_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_11111111111111111111111111111111", ) @@ -174,7 +428,7 @@ def test_barrier_blocks_playbook_eval_and_source_window_writes( ] )[0] - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.save_user_playbooks( @@ -236,10 +490,11 @@ def test_barrier_blocks_deferred_evaluation_tag_write( idempotency_key="idem_deferred_tag_write", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_11111111111111111111111111111112", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.update_agent_success_evaluation_result_tags( @@ -263,16 +518,21 @@ def test_begin_subject_erasure_barrier_requires_matching_purge( idempotency_key="idem_barrier_match", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=alice_subject_ref, request_ref="reqref_v1_00000000000000000000000000000021", ) with pytest.raises(ValueError, match="subject_ref must match"): - storage.begin_subject_erasure_barrier(bob_subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, bob_subject_ref, purge.purge_id) with pytest.raises(ValueError, match="not found"): storage.begin_subject_erasure_barrier( - alice_subject_ref, "purge_barrier_missing" + alice_subject_ref, + "purge_barrier_missing", + execution_claim=_typed_test_claim_for_unvalidated_purge_id( + "purge_barrier_missing" + ), ) @@ -287,6 +547,7 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( idempotency_key="idem_barrier_first", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000051", ) @@ -295,11 +556,12 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( idempotency_key="idem_barrier_second", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000052", ) - storage.begin_subject_erasure_barrier(subject_ref, first_purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, first_purge.purge_id) with pytest.raises(ValueError, match="matching barrier"): storage.fail_subject_erasure_barrier( @@ -307,6 +569,7 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( second_purge.purge_id, error_code="governance_erase_failed", error_detail="ValueError", + execution_claim=_claim_purge(storage, second_purge.purge_id), ) barrier = storage.get_subject_write_barrier(subject_ref) @@ -326,7 +589,7 @@ def test_fail_subject_erasure_barrier_requires_matching_barrier_row( ) -def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( +def test_begin_subject_erasure_barrier_rejects_inactive_claim_after_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -338,10 +601,11 @@ def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( idempotency_key="idem_barrier_terminal_begin", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref=request_ref, ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -349,7 +613,14 @@ def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( request_ref=request_ref, ) - barrier = storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + with pytest.raises(ValueError, match="purge execution claim"): + storage.begin_subject_erasure_barrier( + subject_ref, + purge.purge_id, + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge.purge_id), + ) + barrier = storage.get_subject_write_barrier(subject_ref) + assert barrier is not None stored_barrier = storage.get_subject_write_barrier(subject_ref) stored_purge = storage.get_purge_operation(purge.purge_id) @@ -360,7 +631,7 @@ def test_begin_subject_erasure_barrier_preserves_terminal_erased_state( assert stored_purge.status == "complete" -def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( +def test_fail_subject_erasure_barrier_rejects_inactive_claim_after_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -372,10 +643,11 @@ def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( idempotency_key="idem_barrier_terminal_fail", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref=request_ref, ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -383,12 +655,13 @@ def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( request_ref=request_ref, ) - with pytest.raises(ValueError, match="matching barrier"): + with pytest.raises(ValueError, match="purge execution claim"): storage.fail_subject_erasure_barrier( subject_ref, purge.purge_id, error_code="governance_erase_failed", error_detail="late_failure", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge.purge_id), ) barrier = storage.get_subject_write_barrier(subject_ref) @@ -401,7 +674,7 @@ def test_fail_subject_erasure_barrier_rejects_terminal_erased_state( assert purge_after_failure.error_code is None -def test_fail_purge_operation_rejects_terminal_complete_state( +def test_fail_purge_operation_rejects_inactive_claim_after_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -413,10 +686,11 @@ def test_fail_purge_operation_rejects_terminal_complete_state( idempotency_key="idem_barrier_terminal_purge_fail", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref=request_ref, ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) _complete_empty_purge( storage, purge_id=purge.purge_id, @@ -424,11 +698,12 @@ def test_fail_purge_operation_rejects_terminal_complete_state( request_ref=request_ref, ) - with pytest.raises(ValueError, match="already complete"): + with pytest.raises(ValueError, match="purge execution claim"): storage.fail_purge_operation( purge.purge_id, error_code="governance_erase_failed", error_detail="late_failure", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge.purge_id), ) barrier = storage.get_subject_write_barrier(subject_ref) @@ -470,10 +745,11 @@ def test_guarded_completion_allows_purged_retained_skeletons( idempotency_key="idem_purged_skeletons", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000061", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) assert storage.purge_content(entity_type="profile", entity_id=profile.profile_id) assert storage.purge_content( @@ -535,6 +811,7 @@ def test_guarded_completion_requires_empty_subject_rows( idempotency_key="idem_guarded_complete", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_0123456789abcdef0123456789abcdef", ) @@ -548,14 +825,20 @@ def test_guarded_completion_requires_empty_subject_rows( created_at=_now(), ) ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) storage.record_purge_target( purge.purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - detail={"prepared": True}, + execution_claim=_claim_purge(storage, purge.purge_id), + detail={ + "prepared": True, + "authoritative_user_digest": _authoritative_user_digest( + storage, purge.purge_id + ), + }, ) with pytest.raises(ValueError, match="same-subject rows remain"): @@ -570,6 +853,8 @@ def test_guarded_completion_requires_empty_subject_rows( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -584,6 +869,7 @@ def test_guarded_completion_requires_empty_legacy_null_subject_rows( idempotency_key="idem_guarded_legacy", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000031", ) @@ -605,14 +891,20 @@ def test_guarded_completion_requires_empty_legacy_null_subject_rows( ) storage.conn.commit() - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) storage.record_purge_target( purge.purge_id, target_name="target_snapshot", phase="prepare_targets", status="complete", target_ref="all", - detail={"prepared": True}, + execution_claim=_claim_purge(storage, purge.purge_id), + detail={ + "prepared": True, + "authoritative_user_digest": _authoritative_user_digest( + storage, purge.purge_id + ), + }, ) with pytest.raises(ValueError, match="same-subject rows remain"): @@ -627,6 +919,8 @@ def test_guarded_completion_requires_empty_legacy_null_subject_rows( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -641,6 +935,7 @@ def test_guarded_completion_requires_existing_erasing_subject_barrier( idempotency_key="idem_missing_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000032", ) @@ -658,6 +953,8 @@ def test_guarded_completion_requires_existing_erasing_subject_barrier( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -672,15 +969,17 @@ def test_guarded_completion_rejects_failed_subject_barrier( idempotency_key="idem_failed_barrier", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000035", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) storage.fail_subject_erasure_barrier( subject_ref, purge.purge_id, error_code="test_failed_barrier", error_detail="RuntimeError", + execution_claim=_claim_purge(storage, purge.purge_id), ) _mark_all_completion_targets(storage, purge.purge_id) @@ -696,6 +995,8 @@ def test_guarded_completion_rejects_failed_subject_barrier( idempotency_key=purge.purge_id, detail={"deleted_counts": {}, "rebuilt_agent_playbook_ids": []}, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) @@ -718,10 +1019,11 @@ def test_barrier_blocks_profile_update_paths( idempotency_key="idem_profile_update", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000033", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.update_user_profile_tags("alice", profile.profile_id, ["blocked"]) @@ -758,10 +1060,11 @@ def test_barrier_blocks_user_playbook_update_paths( idempotency_key="idem_playbook_update", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000034", ) - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.archive_user_playbook_by_id("alice", playbook.user_playbook_id) @@ -789,6 +1092,7 @@ def test_assert_subject_writable_blocks_only_barriered_subject( idempotency_key="idem_assert_writable", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=barriered_subject_ref, request_ref=request_ref, ) @@ -797,7 +1101,9 @@ def test_assert_subject_writable_blocks_only_barriered_subject( storage.assert_subject_writable(barriered_subject_ref) storage.assert_subject_writable(other_subject_ref) - storage.begin_subject_erasure_barrier(barriered_subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier( + storage, barriered_subject_ref, purge.purge_id + ) # The 'erasing' barrier blocks only its own subject. with pytest.raises(SubjectWriteBarrierError, match="blocked by erasure barrier"): @@ -831,6 +1137,7 @@ def test_source_window_write_blocks_legacy_null_subject_ref_user_playbook( idempotency_key="idem_source_window_legacy", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=subject_ref, request_ref="reqref_v1_00000000000000000000000000000041", ) @@ -867,7 +1174,7 @@ def test_source_window_write_blocks_legacy_null_subject_ref_user_playbook( ] )[0] - storage.begin_subject_erasure_barrier(subject_ref, purge.purge_id) + _begin_claimed_subject_erasure_barrier(storage, subject_ref, purge.purge_id) with pytest.raises(SubjectWriteBarrierError): storage.set_source_windows_for_agent_playbook( diff --git a/tests/server/services/lineage/test_gc_scheduler_global_sweep.py b/tests/server/services/lineage/test_gc_scheduler_global_sweep.py index 642d9a125..62ed4cfb6 100644 --- a/tests/server/services/lineage/test_gc_scheduler_global_sweep.py +++ b/tests/server/services/lineage/test_gc_scheduler_global_sweep.py @@ -2,6 +2,7 @@ import pytest +from reflexio.server.auth import DEFAULT_ORG_ID from reflexio.server.services.lineage import gc_scheduler from reflexio.server.services.lineage.gc_scheduler import ( LineageGCScheduler, @@ -26,8 +27,13 @@ def _cfg(enabled: bool): @pytest.fixture(autouse=True) def _isolate_hooks(): clear_global_sweeps() + clear_always = getattr(gc_scheduler, "clear_always_global_sweeps", None) + if clear_always is not None: + clear_always() yield clear_global_sweeps() + if clear_always is not None: + clear_always() def test_global_sweep_runs_once_when_enabled(): @@ -97,3 +103,69 @@ def test_run_once_invokes_global_sweeps(monkeypatch): scheduler._run_once() assert len(calls) == 1 + + +def test_run_once_invokes_always_global_sweep_when_expiry_disabled(monkeypatch): + calls: list[int] = [] + register = getattr(gc_scheduler, "register_always_global_sweep", None) + assert callable(register) + register(lambda now: calls.append(now) or 1) + + cfg = types.SimpleNamespace( + lineage_gc=types.SimpleNamespace(poll_interval_seconds=10), + expiry_reclamation=types.SimpleNamespace(enabled=False), + ) + ctx = types.SimpleNamespace( + configurator=types.SimpleNamespace(get_config=lambda: cfg), + ) + scheduler = LineageGCScheduler( + request_context_factory=lambda _org_id: ctx, # type: ignore[arg-type] + bootstrap_org_id="org-boot", + ) + monkeypatch.setattr(scheduler, "_discover_org_ids", lambda _ctx: []) + monkeypatch.setattr(scheduler, "_gc_tick", lambda _org_ids, **_kwargs: None) + + scheduler._run_once() + + assert len(calls) == 1 + + +def test_run_once_invokes_always_global_sweep_during_empty_fleet_retry(): + calls: list[int] = [] + gc_scheduler.register_always_global_sweep(lambda now: calls.append(now) or 1) + factory_calls: list[str] = [] + scheduler = LineageGCScheduler( + request_context_factory=lambda org_id: factory_calls.append(org_id), # type: ignore[arg-type] + bootstrap_org_id=DEFAULT_ORG_ID, + org_id_provider=lambda: [DEFAULT_ORG_ID], + ) + + assert scheduler._run_once() == 5 + assert len(calls) == 1 + assert factory_calls == [] + + +@pytest.mark.parametrize("bootstrap_failure", ["context", "config"]) +def test_run_once_invokes_always_global_sweep_during_bootstrap_failure( + bootstrap_failure: str, +): + calls: list[int] = [] + gc_scheduler.register_always_global_sweep(lambda now: calls.append(now) or 1) + + def get_config(): + raise RuntimeError("config bootstrap failed") + + def request_context_factory(_org_id: str): + if bootstrap_failure == "context": + raise RuntimeError("context bootstrap failed") + return types.SimpleNamespace( + configurator=types.SimpleNamespace(get_config=get_config) + ) + + scheduler = LineageGCScheduler( + request_context_factory=request_context_factory, # type: ignore[arg-type] + bootstrap_org_id="org-boot", + ) + + assert scheduler._run_once() == 86400 + assert len(calls) == 1 diff --git a/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py b/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py index 78300ee39..cb369c7d4 100644 --- a/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_agent_run_storage.py @@ -5,6 +5,8 @@ import pytest +from reflexio.models.api_schema.service_schemas import UserProfile +from reflexio.server.services.storage.error import StorageError from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.sqlite_storage._agent_run import _dt from reflexio.server.services.storage.storage_base import ( @@ -97,6 +99,213 @@ 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 4ff0df379..17d1c385d 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_retrieved_learning.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_retrieved_learning.py @@ -117,11 +117,26 @@ def _erase(storage: SQLiteStorage, user_id: str, purge_id: str) -> dict[str, int idempotency_key=f"idem_{purge_id}", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + authoritative_user_id=user_id, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, ) - storage.prepare_governance_erase_targets(purge_id, user_id) - return storage.apply_governance_user_data_delete(purge_id, user_id) + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"test-{purge_id}", + lease_ttl_seconds=30, + ) + assert claim is not None + storage.prepare_governance_erase_targets( + purge_id, + user_id, + execution_claim=claim, + ) + return storage.apply_governance_user_data_delete( + purge_id, + user_id, + execution_claim=claim, + ) def test_erase_scrubs_rle_rows_and_all_state_namespaces(storage) -> None: diff --git a/tests/server/services/storage/sqlite_storage/test_governance_storage.py b/tests/server/services/storage/sqlite_storage/test_governance_storage.py index 880cbad09..8ed1b9f87 100644 --- a/tests/server/services/storage/sqlite_storage/test_governance_storage.py +++ b/tests/server/services/storage/sqlite_storage/test_governance_storage.py @@ -1,10 +1,15 @@ from __future__ import annotations +import ast +import hashlib +import inspect import json import sqlite3 import threading import time +from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Any, Literal, cast from unittest.mock import patch @@ -24,6 +29,9 @@ ) from reflexio.models.api_schema.retriever_schema import SearchAgentPlaybookRequest from reflexio.models.config_schema import GovernanceRetentionConfig +from reflexio.server.services.governance.config import governance_subject_ref +from reflexio.server.services.governance.service import GovernanceService +from reflexio.server.services.storage.governance_claims import PurgeExecutionClaim from reflexio.server.services.storage.governance_validation import ( _CANONICAL_DELETE_TARGET_NAMES, ) @@ -37,17 +45,55 @@ 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._rebuild_hide import ( + RebuildHideMixin as RebuildHideContractMixin, +) +from reflexio.server.services.storage.storage_base.governance._subject_barrier import ( + SubjectBarrierMixin, +) pytestmark = pytest.mark.integration -SUBJECT_REF = "subref_v1_" + "a" * 32 -OTHER_SUBJECT_REF = "subref_v1_" + "c" * 32 +SUBJECT_REF = governance_subject_ref("org1", "alice", "test-governance-secret") +OTHER_SUBJECT_REF = governance_subject_ref("org1", "bob", "test-governance-secret") REQUEST_REF = "reqref_v1_" + "b" * 32 OTHER_REQUEST_REF = "reqref_v1_" + "d" * 32 ACTOR_REF = "actref_v1_" + "e" * 32 # Single source of truth — a stale local copy of this tuple is exactly how # this suite went red when new canonical targets landed without test updates. CANONICAL_DELETE_TARGET_NAMES = _CANONICAL_DELETE_TARGET_NAMES +CLAIMED_ERASURE_MUTATIONS = { + "record_purge_target", + "prepare_governance_erase_targets", + "fail_purge_operation", + "begin_subject_erasure_barrier", + "complete_subject_erasure_barrier_after_empty_check", + "fail_subject_erasure_barrier", + "apply_governance_user_data_delete", + "hide_governance_agent_playbooks_for_rebuild", + "apply_governance_agent_playbook_rebuild", + "complete_purge_operation_with_audit", +} + + +def _begin_test_purge_operation(storage: SQLiteStorage, **kwargs: Any): + if ( + kwargs.get("operation_type") == "user_erasure" + and kwargs.get("scope_type") == "user" + and "authoritative_user_id" not in kwargs + ): + subject_ref = kwargs.get("subject_ref") + for user_id in ("alice", "bob"): + if storage._subject_ref_for_user_id(user_id) == subject_ref: + kwargs["authoritative_user_id"] = user_id + break + return SQLiteStorage.begin_purge_operation(storage, **kwargs) @pytest.fixture @@ -68,29 +114,155 @@ def _make_storage(org_id: str) -> SQLiteStorage: yield _make_storage -def _begin_purge(storage: SQLiteStorage, purge_id: str) -> str: - purge = storage.begin_purge_operation( +def _begin_purge( + storage: SQLiteStorage, + purge_id: str, + *, + subject_ref: str | None = None, + authoritative_user_id: str = "alice", +) -> str: + subject_ref = subject_ref or storage._subject_ref_for_user_id(authoritative_user_id) + purge = _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key=f"idem_{purge_id}", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=subject_ref, request_ref=REQUEST_REF, + authoritative_user_id=authoritative_user_id, ) + claim = _claim_purge(storage, purge.purge_id) storage.record_purge_target( purge_id=purge.purge_id, target_name="target_snapshot", target_ref="all", phase="prepare_targets", status="complete", + execution_claim=claim, detail={ + "authoritative_user_digest": storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge.purge_id), + ).fetchone()["authoritative_user_digest"], "owned_user_playbook_ids": [11], }, ) return purge.purge_id -def _add_complete_delete_target_matrix(storage: SQLiteStorage, purge_id: str) -> None: +def _begin_raw_user_erasure_purge(storage: SQLiteStorage, purge_id: str) -> str: + return _begin_raw_user_erasure_purge_for_subject( + storage, purge_id, subject_ref=SUBJECT_REF + ) + + +def _begin_raw_user_erasure_purge_for_subject( + storage: SQLiteStorage, + purge_id: str, + *, + subject_ref: str, + authoritative_user_id: str = "alice", +) -> str: + purge = _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=subject_ref, + request_ref=REQUEST_REF, + authoritative_user_id=authoritative_user_id, + ) + return purge.purge_id + + +def _claim_then_take_over(storage: SQLiteStorage, purge_id: str): + first_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-a", + lease_ttl_seconds=30, + ) + if first_claim is None: + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + first_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-a", + lease_ttl_seconds=30, + ) + assert first_claim is not None + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + takeover_claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner="owner-b", + lease_ttl_seconds=30, + ) + assert takeover_claim is not None + assert takeover_claim.fence == first_claim.fence + 1 + return first_claim, takeover_claim + + +def _claim_purge(storage: SQLiteStorage, purge_id: str) -> PurgeExecutionClaim: + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"owner-{purge_id}", + lease_ttl_seconds=30, + ) + if claim is None: + storage.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ) + storage.conn.commit() + claim = storage.claim_purge_operation_execution( + purge_id, + lease_owner=f"owner-{purge_id}", + lease_ttl_seconds=30, + ) + assert claim is not None + return claim + + +def _typed_test_claim_for_unvalidated_purge_id(purge_id: str) -> PurgeExecutionClaim: + return PurgeExecutionClaim( + purge_id=purge_id, + owner="test-invalid-purge-id", + fence=1, + expires_at=1, + ) + + +def _assert_rejects_missing_claim( + omitted_call: Callable[[], object], + none_call: Callable[[], object], +) -> None: + with pytest.raises((TypeError, ValueError)): + omitted_call() + with pytest.raises((TypeError, ValueError)): + none_call() + + +def _add_complete_delete_target_matrix( + storage: SQLiteStorage, + purge_id: str, + *, + execution_claim: PurgeExecutionClaim, +) -> None: for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -98,13 +270,34 @@ def _add_complete_delete_target_matrix(storage: SQLiteStorage, purge_id: str) -> target_ref="all", phase="delete", status="complete", + execution_claim=execution_claim, ) -def _begin_completeable_purge(storage: SQLiteStorage, purge_id: str) -> str: - purge_id = _begin_purge(storage, purge_id) - _add_complete_delete_target_matrix(storage, purge_id) - storage.begin_subject_erasure_barrier(SUBJECT_REF, purge_id) +def _begin_completeable_purge( + storage: SQLiteStorage, + purge_id: str, + *, + subject_ref: str = SUBJECT_REF, + authoritative_user_id: str = "alice", +) -> str: + purge_id = _begin_purge( + storage, + purge_id, + subject_ref=subject_ref, + authoritative_user_id=authoritative_user_id, + ) + claim = _claim_purge(storage, purge_id) + _add_complete_delete_target_matrix( + storage, + purge_id, + execution_claim=claim, + ) + storage.begin_subject_erasure_barrier( + subject_ref, + purge_id, + execution_claim=claim, + ) return purge_id @@ -113,12 +306,13 @@ def _erase_event( purge_id: str, status: AuditStatus = "ok", operation: AuditOperation = "ERASE", + subject_ref: str = SUBJECT_REF, ): return AuditEvent( org_id="org1", operation=operation, entity_type="request", - subject_ref=SUBJECT_REF, + subject_ref=subject_ref, request_ref=REQUEST_REF, idempotency_key=purge_id, status=status, @@ -346,6 +540,7 @@ def _record_agent_playbook_rebuild_target( target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status=status, + execution_claim=_claim_purge(storage, purge_id), detail={ "original_source_windows": original_windows or [ @@ -487,7 +682,8 @@ def test_list_audit_events_rejects_cross_org_override(storage_factory): def test_purge_targets_require_snapshot_marker(storage): - purge = storage.begin_purge_operation( + purge = _begin_test_purge_operation( + storage, purge_id="purge_snapshot_marker", idempotency_key="idem_snapshot_marker", operation_type="user_erasure", @@ -502,10 +698,15 @@ def test_purge_targets_require_snapshot_marker(storage): phase="delete", status="complete", deleted_count=1, + execution_claim=_claim_purge(storage, purge.purge_id), ) assert storage.purge_targets_prepared(purge.purge_id) is False - storage.begin_subject_erasure_barrier(SUBJECT_REF, purge.purge_id) + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge.purge_id, + execution_claim=_claim_purge(storage, purge.purge_id), + ) with pytest.raises(ValueError, match="target snapshot"): storage.complete_purge_operation_with_audit( purge.purge_id, @@ -517,25 +718,293 @@ def test_purge_targets_require_snapshot_marker(storage): request_ref=REQUEST_REF, idempotency_key=purge.purge_id, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge.purge_id), ) def test_complete_purge_operation_with_audit_is_atomic_success_path(storage): purge_id = _begin_completeable_purge(storage, "purge_atomic_success") + complete_claim = _claim_purge(storage, purge_id) complete = storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=complete_claim, ) assert complete.status == "complete" rows = storage.list_audit_events(subject_ref=SUBJECT_REF) assert [row.operation for row in rows] == ["ERASE"] - same = storage.complete_purge_operation_with_audit( + with pytest.raises(ValueError, match="purge execution claim"): + storage.complete_purge_operation_with_audit( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=complete_claim, + ) + assert len(storage.list_audit_events(subject_ref=SUBJECT_REF)) == 1 + + +def test_complete_purge_operation_with_audit_rejects_wrong_authoritative_user(storage): + purge_id = _begin_completeable_purge(storage, "purge_wrong_complete_user") + + with pytest.raises(ValueError, match="authoritative user identity"): + storage.complete_purge_operation_with_audit( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="bob", + execution_claim=_claim_purge(storage, purge_id), + ) + + assert storage.get_purge_operation(purge_id).status == "running" + assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] + + +@pytest.mark.parametrize( + "complete_method_name", + [ + "complete_purge_operation_with_audit", + "complete_subject_erasure_barrier_after_empty_check", + ], +) +def test_user_erasure_completion_contracts_reject_org_purge( + storage, + complete_method_name, +): + purge = _begin_test_purge_operation( + storage, + purge_id=f"purge_org_scope_{complete_method_name}", + idempotency_key=f"idem_org_scope_{complete_method_name}", + operation_type="org_purge", + scope_type="org", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + ) + claim = _claim_purge(storage, purge.purge_id) + storage.record_purge_target( + purge.purge_id, + target_name="target_snapshot", + target_ref="all", + phase="prepare_targets", + status="complete", + detail={"prepared": True}, + execution_claim=claim, + ) + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge.purge_id, + execution_claim=claim, + ) + + complete = getattr(storage, complete_method_name) + with pytest.raises(ValueError, match="user erasure"): + complete( + purge.purge_id, + _erase_event(purge_id=purge.purge_id), + authoritative_user_id="arbitrary-user", + execution_claim=claim, + ) + + assert storage.get_purge_operation(purge.purge_id).status == "running" + assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] + + +@pytest.mark.parametrize( + ("purge_binding", "snapshot_binding"), + [ + pytest.param("legacy", "legacy", id="both-unkeyed"), + pytest.param(None, None, id="both-null"), + pytest.param("current", "legacy", id="interrupted-row-only-upgrade"), + ], +) +def test_idempotent_retry_upgrades_legacy_identity_bindings_and_resumes_completion( + storage, + purge_binding, + snapshot_binding, +): + purge_id = _begin_completeable_purge(storage, "purge_interrupted_legacy_resume") + current_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + legacy_digest = hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + + def resolve_binding(binding): + return ( + legacy_digest + if binding == "legacy" + else current_digest + if binding == "current" + else None + ) + + snapshot_row = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + snapshot_detail = json.loads(snapshot_row["detail"]) + snapshot_detail["authoritative_user_digest"] = resolve_binding(snapshot_binding) + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + (resolve_binding(purge_binding), storage.org_id, purge_id), + ) + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(snapshot_detail), storage.org_id, purge_id), + ) + storage.conn.commit() + + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + + adopted_purge_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + adopted_snapshot = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + assert adopted_purge_digest == current_digest + assert ( + json.loads(adopted_snapshot["detail"])["authoritative_user_digest"] + == current_digest + ) + + completed = storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) + assert completed.status == "complete" + + +def test_idempotent_retry_rolls_back_purge_digest_when_snapshot_binding_mismatches( + storage, +): + purge_id = _begin_completeable_purge(storage, "purge_mismatched_legacy_snapshot") + legacy_digest = hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + snapshot_row = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + snapshot_detail = json.loads(snapshot_row["detail"]) + snapshot_detail["authoritative_user_digest"] = "mismatched-digest" + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + (legacy_digest, storage.org_id, purge_id), + ) + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(snapshot_detail), storage.org_id, purge_id), + ) + storage.conn.commit() + + with pytest.raises(ValueError, match="authoritative user identity"): + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + + persisted_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + assert persisted_digest == legacy_digest + + +def test_idempotent_retry_rolls_back_row_upgrade_when_snapshot_upgrade_fails(storage): + purge_id = _begin_completeable_purge(storage, "purge_snapshot_upgrade_failure") + legacy_digest = hashlib.sha256(f"{purge_id}\0alice".encode()).hexdigest() + snapshot_row = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + snapshot_detail = json.loads(snapshot_row["detail"]) + snapshot_detail["authoritative_user_digest"] = legacy_digest + storage.conn.execute( + """UPDATE purge_operations SET authoritative_user_digest = ? + WHERE org_id = ? AND purge_id = ?""", + (legacy_digest, storage.org_id, purge_id), + ) + storage.conn.execute( + """UPDATE purge_operation_targets SET detail = ? + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (json.dumps(snapshot_detail), storage.org_id, purge_id), + ) + storage.conn.execute( + f"""CREATE TRIGGER fail_snapshot_digest_upgrade + BEFORE UPDATE OF detail ON purge_operation_targets + WHEN OLD.org_id = '{storage.org_id}' AND OLD.purge_id = '{purge_id}' + AND OLD.target_name = 'target_snapshot' + BEGIN + SELECT RAISE(ABORT, 'snapshot upgrade failed'); + END""" + ) + storage.conn.commit() + + with pytest.raises(sqlite3.IntegrityError, match="snapshot upgrade failed"): + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + + persisted_digest = storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (storage.org_id, purge_id), + ).fetchone()["authoritative_user_digest"] + persisted_snapshot = storage.conn.execute( + """SELECT detail FROM purge_operation_targets + WHERE org_id = ? AND purge_id = ? AND target_name = 'target_snapshot' + AND target_ref = 'all' AND phase = 'prepare_targets'""", + (storage.org_id, purge_id), + ).fetchone() + assert persisted_digest == legacy_digest + assert ( + json.loads(persisted_snapshot["detail"])["authoritative_user_digest"] + == legacy_digest ) - assert same.status == "complete" - assert len(storage.list_audit_events(subject_ref=SUBJECT_REF)) == 1 def test_complete_purge_operation_with_audit_begins_immediate_transaction_before_reads( @@ -548,6 +1017,8 @@ def test_complete_purge_operation_with_audit_begins_immediate_transaction_before storage.complete_purge_operation_with_audit( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) finally: storage.conn.set_trace_callback(None) @@ -565,6 +1036,7 @@ def test_complete_purge_operation_with_audit_begins_immediate_transaction_before def test_apply_governance_delete_begins_immediate_transaction_before_reads(storage): purge_id = _begin_purge(storage, "purge_delete_begin_immediate") + claim = _claim_purge(storage, purge_id) for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -572,13 +1044,18 @@ def test_apply_governance_delete_begins_immediate_transaction_before_reads(stora target_ref="all", phase="delete", status="pending", + execution_claim=claim, detail={"count": 0}, ) statements: list[str] = [] storage.conn.set_trace_callback(statements.append) try: with pytest.raises(ValueError, match="prepared purge snapshot"): - storage.apply_governance_user_data_delete(purge_id, "empty-user") + storage.apply_governance_user_data_delete( + purge_id, + "alice", + execution_claim=claim, + ) finally: storage.conn.set_trace_callback(None) @@ -619,6 +1096,8 @@ def test_complete_purge_operation_with_audit_accepts_planned_success_detail(stor "rebuilt_agent_playbook_ids": rebuilt_ids, }, ), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert complete.status == "complete" @@ -648,7 +1127,12 @@ def test_complete_purge_operation_rejects_audit_refs_that_mismatch_persisted_pur event = _erase_event(purge_id=purge_id).model_copy(update=event_kwargs) with pytest.raises(ValueError, match=match): - storage.complete_purge_operation_with_audit(purge_id, event) + storage.complete_purge_operation_with_audit( + purge_id, + event, + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) assert storage.get_purge_operation(purge_id).status == "running" assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] @@ -671,7 +1155,8 @@ def test_complete_purge_operation_rejects_audit_refs_that_mismatch_persisted_pur def test_begin_purge_operation_rejects_mismatched_idempotent_retry( storage, retry_kwargs, match ): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_begin_retry", idempotency_key="idem_begin_retry", operation_type="user_erasure", @@ -681,7 +1166,8 @@ def test_begin_purge_operation_rejects_mismatched_idempotent_retry( ) with pytest.raises(ValueError, match=match): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=retry_kwargs.get("purge_id", "purge_begin_retry"), idempotency_key="idem_begin_retry", operation_type=retry_kwargs.get("operation_type", "user_erasure"), @@ -698,7 +1184,8 @@ def test_begin_purge_operation_rejects_mismatched_idempotent_retry( def test_begin_purge_operation_rejects_numeric_idempotency_key(storage): with pytest.raises(ValueError, match="idempotency_key"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_numeric_idem", idempotency_key="12345", operation_type="user_erasure", @@ -714,7 +1201,8 @@ def test_begin_purge_operation_rejects_numeric_idempotency_key(storage): def test_begin_purge_operation_accepts_code_shaped_idempotency_key_with_content( storage, ): - purge = storage.begin_purge_operation( + purge = _begin_test_purge_operation( + storage, purge_id="purge_content_retry", idempotency_key="content_purge_retry_1", operation_type="user_erasure", @@ -729,7 +1217,8 @@ def test_begin_purge_operation_accepts_code_shaped_idempotency_key_with_content( @pytest.mark.parametrize("purge_id", ["purge_1", "purge_123"]) def test_begin_purge_operation_rejects_raw_numeric_purge_suffix(storage, purge_id): with pytest.raises(ValueError, match="purge_id"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key=f"idem_{purge_id}", operation_type="user_erasure", @@ -787,7 +1276,12 @@ def test_complete_purge_operation_rejects_invalid_audit_event(storage, event, ma purge_id = _begin_completeable_purge(storage, "purge_invalid") with pytest.raises(ValueError, match=match): - storage.complete_purge_operation_with_audit(purge_id, event) + storage.complete_purge_operation_with_audit( + purge_id, + event, + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), + ) assert storage.get_purge_operation(purge_id).status == "running" assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] @@ -816,7 +1310,10 @@ def test_complete_purge_operation_requires_matching_existing_erase_row( with pytest.raises(ValueError, match=match): storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).status == "running" @@ -873,7 +1370,10 @@ def test_complete_purge_operation_rejects_mismatched_existing_erase_row( with pytest.raises(ValueError, match="matching successful ERASE"): storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).status == "running" @@ -902,53 +1402,689 @@ def test_append_audit_event_rejects_successful_erase_without_idempotency_key(sto ) -def test_complete_purge_operation_requires_full_delete_target_matrix(storage): - purge_id = _begin_purge(storage, "purge_snapshot_only") - storage.begin_subject_erasure_barrier(SUBJECT_REF, purge_id) +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] = [] + errors: list[BaseException] = [] + + def claim(storage_instance: SQLiteStorage, owner: str) -> None: + try: + ready.wait(timeout=5) + claims.append( + storage_instance.claim_purge_operation_execution( + purge.purge_id, + lease_owner=owner, + lease_ttl_seconds=30, + ) + ) + except BaseException as exc: # noqa: BLE001 - intentional thread error capture + errors.append(exc) + + callers = [ + threading.Thread(target=claim, args=(storage_a, "owner-a")), + 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) + assert errors == [] + assert len(claims) == len(callers) + live_claims = [claim for claim in claims if claim is not None] + assert len(live_claims) == 1 + first_claim = live_claims[0] + 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", + }, + RebuildHideContractMixin: { + "hide_governance_agent_playbooks_for_rebuild", + "apply_governance_agent_playbook_rebuild", + }, + 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" + assert production_root.exists() + python_modules = list(production_root.rglob("*.py")) + assert python_modules + violations: list[str] = [] + for path in python_modules: + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + function = node.func + if not isinstance(function, ast.Attribute): + continue + if function.attr not in CLAIMED_ERASURE_MUTATIONS: + continue + claim_keywords = [ + keyword for keyword in node.keywords if keyword.arg == "execution_claim" + ] + if not claim_keywords: + violations.append(f"{path.relative_to(production_root)}:{node.lineno}") + continue + if any(keyword.arg is None for keyword in node.keywords): + violations.append( + f"{path.relative_to(production_root)}:{node.lineno}: **kwargs" + ) + claim_value = claim_keywords[0].value + if isinstance(claim_value, ast.Constant) and claim_value.value is None: + violations.append( + f"{path.relative_to(production_root)}:{node.lineno}: None" + ) + assert violations == [] + + +def test_sqlite_claimed_erasure_mutations_reject_omitted_and_none_claim(storage): + target_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_target") + _, target_claim = _claim_then_take_over(storage, target_purge_id) + _assert_rejects_missing_claim( + lambda: storage.record_purge_target( + purge_id=target_purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + ), + lambda: storage.record_purge_target( + purge_id=target_purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.list_purge_targets(target_purge_id, phase="delete") == [] + storage.record_purge_target( + purge_id=target_purge_id, + target_name="interaction", + target_ref="all", + phase="delete", + status="complete", + execution_claim=target_claim, + ) + + barrier_user_id = "no-claim-barrier-user" + barrier_subject_ref = storage._subject_ref_for_user_id(barrier_user_id) + barrier_purge_id = _begin_raw_user_erasure_purge_for_subject( + storage, + "purge_no_claim_barrier", + subject_ref=barrier_subject_ref, + authoritative_user_id=barrier_user_id, + ) + _, barrier_claim = _claim_then_take_over(storage, barrier_purge_id) + _assert_rejects_missing_claim( + lambda: storage.begin_subject_erasure_barrier( + barrier_subject_ref, barrier_purge_id + ), + lambda: storage.begin_subject_erasure_barrier( + barrier_subject_ref, + barrier_purge_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_subject_write_barrier(barrier_subject_ref) is None + storage.begin_subject_erasure_barrier( + barrier_subject_ref, + barrier_purge_id, + execution_claim=barrier_claim, + ) + + prepare_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_prepare") + _, prepare_claim = _claim_then_take_over(storage, prepare_purge_id) + _assert_rejects_missing_claim( + lambda: storage.prepare_governance_erase_targets( + purge_id=prepare_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + ), + lambda: storage.prepare_governance_erase_targets( + purge_id=prepare_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.list_purge_targets(prepare_purge_id) == [] + storage.prepare_governance_erase_targets( + purge_id=prepare_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=prepare_claim, + ) + + delete_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_delete") + delete_claim = _claim_purge(storage, delete_purge_id) + storage.prepare_governance_erase_targets( + purge_id=delete_purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=delete_claim, + ) + _assert_rejects_missing_claim( + lambda: storage.apply_governance_user_data_delete(delete_purge_id, "alice"), + lambda: storage.apply_governance_user_data_delete( + delete_purge_id, + "alice", + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert all( + target.status == "pending" + for target in storage.list_purge_targets(delete_purge_id, phase="delete") + ) + + complete_user_id = "no-claim-complete-user" + complete_subject_ref = storage._subject_ref_for_user_id(complete_user_id) + complete_purge_id = _begin_completeable_purge( + storage, + "purge_no_claim_complete", + subject_ref=complete_subject_ref, + authoritative_user_id=complete_user_id, + ) + _claim_purge(storage, complete_purge_id) + _assert_rejects_missing_claim( + lambda: storage.complete_subject_erasure_barrier_after_empty_check( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + ), + lambda: storage.complete_subject_erasure_barrier_after_empty_check( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_purge_operation(complete_purge_id).status == "running" + _assert_rejects_missing_claim( + lambda: storage.complete_purge_operation_with_audit( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + ), + lambda: storage.complete_purge_operation_with_audit( + complete_purge_id, + _erase_event( + purge_id=complete_purge_id, + subject_ref=complete_subject_ref, + ), + authoritative_user_id=complete_user_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_purge_operation(complete_purge_id).status == "running" + + fail_barrier_user_id = "no-claim-fail-barrier-user" + fail_barrier_subject_ref = storage._subject_ref_for_user_id(fail_barrier_user_id) + fail_barrier_purge_id = _begin_raw_user_erasure_purge_for_subject( + storage, + "purge_no_claim_fail_barrier", + subject_ref=fail_barrier_subject_ref, + authoritative_user_id=fail_barrier_user_id, + ) + fail_barrier_claim = _claim_purge(storage, fail_barrier_purge_id) + storage.begin_subject_erasure_barrier( + fail_barrier_subject_ref, + fail_barrier_purge_id, + execution_claim=fail_barrier_claim, + ) + _assert_rejects_missing_claim( + lambda: storage.fail_subject_erasure_barrier( + fail_barrier_subject_ref, + fail_barrier_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + ), + lambda: storage.fail_subject_erasure_barrier( + fail_barrier_subject_ref, + fail_barrier_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=None, # type: ignore[arg-type] + ), + ) + barrier = storage.get_subject_write_barrier(fail_barrier_subject_ref) + assert barrier is not None + assert barrier.status == "erasing" + + fail_purge_id = _begin_raw_user_erasure_purge(storage, "purge_no_claim_fail") + _claim_purge(storage, fail_purge_id) + _assert_rejects_missing_claim( + lambda: storage.fail_purge_operation( + fail_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + ), + lambda: storage.fail_purge_operation( + fail_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert storage.get_purge_operation(fail_purge_id).status == "running" + + +def test_stale_execution_claim_cannot_start_subject_barrier(storage): + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_barrier_start") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge_id, + execution_claim=first_claim, + ) + assert storage.get_subject_write_barrier(SUBJECT_REF) is None + + barrier = storage.begin_subject_erasure_barrier( + SUBJECT_REF, + purge_id, + execution_claim=takeover_claim, + ) + assert barrier.status == "erasing" + + +def test_stale_execution_claim_cannot_prepare_delete_targets(storage): + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_prepare") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.prepare_governance_erase_targets( + purge_id=purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=first_claim, + ) + assert storage.list_purge_targets(purge_id) == [] + + storage.prepare_governance_erase_targets( + purge_id=purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=takeover_claim, + ) + assert storage.purge_targets_prepared(purge_id) + + +def test_stale_execution_claim_cannot_apply_protected_delete(storage): + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_delete") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + storage.prepare_governance_erase_targets( + purge_id=purge_id, + user_id="alice", + owned_user_playbook_ids=set(), + execution_claim=takeover_claim, + ) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.apply_governance_user_data_delete( + purge_id, + "alice", + execution_claim=first_claim, + ) + assert all( + target.status == "pending" + for target in storage.list_purge_targets(purge_id, phase="delete") + ) + + counts = storage.apply_governance_user_data_delete( + purge_id, + "alice", + execution_claim=takeover_claim, + ) + assert counts["requests"] == 0 + assert all( + target.status == "complete" + for target in storage.list_purge_targets(purge_id, phase="delete") + ) + + +def test_stale_execution_claim_cannot_complete_purge_or_barrier(storage): + purge_id = _begin_completeable_purge(storage, "purge_stale_complete") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) - with pytest.raises(ValueError, match="delete target matrix"): - storage.complete_purge_operation_with_audit( + with pytest.raises(ValueError, match="purge execution claim"): + storage.complete_subject_erasure_barrier_after_empty_check( purge_id, _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=first_claim, ) - assert storage.get_purge_operation(purge_id).status == "running" + barrier = storage.get_subject_write_barrier(SUBJECT_REF) + assert barrier is not None + assert barrier.status == "erasing" assert storage.list_audit_events(subject_ref=SUBJECT_REF) == [] + completed = storage.complete_subject_erasure_barrier_after_empty_check( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=takeover_claim, + ) + assert completed.status == "complete" -def test_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, + +def test_stale_execution_claim_cannot_fail_purge_or_barrier(storage): + barrier_purge_id = _begin_raw_user_erasure_purge( + storage, "purge_stale_barrier_failure" + ) + barrier_claim = _claim_purge(storage, barrier_purge_id) + storage.begin_subject_erasure_barrier( + SUBJECT_REF, + barrier_purge_id, + execution_claim=barrier_claim, + ) + first_claim, takeover_claim = _claim_then_take_over(storage, barrier_purge_id) + + with pytest.raises(ValueError, match="purge execution claim"): + storage.fail_subject_erasure_barrier( + SUBJECT_REF, + barrier_purge_id, error_code="governance_erase_failed", error_detail="RuntimeError", + execution_claim=first_claim, ) - assert failed.completed_at == 111 + barrier = storage.get_subject_write_barrier(SUBJECT_REF) + assert barrier is not None + assert barrier.status == "erasing" - with patch.object(erase_execution_module, "_epoch_now", return_value=222): - completed = storage.complete_purge_operation_with_audit( + failed_barrier = storage.fail_subject_erasure_barrier( + SUBJECT_REF, + barrier_purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=takeover_claim, + ) + assert failed_barrier.status == "failed" + + purge_id = _begin_raw_user_erasure_purge(storage, "purge_stale_purge_failure") + first_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + with pytest.raises(ValueError, match="purge execution claim"): + storage.fail_purge_operation( purge_id, - _erase_event(purge_id=purge_id), + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=first_claim, ) + assert storage.get_purge_operation(purge_id).status == "running" - assert completed.status == "complete" - assert completed.completed_at == 222 + failed_purge = storage.fail_purge_operation( + purge_id, + error_code="governance_erase_failed", + error_detail="RuntimeError", + execution_claim=takeover_claim, + ) + assert failed_purge.status == "failed" def test_prepare_governance_erase_targets_sanitizes_snapshot_detail(storage): - storage.begin_purge_operation( + user_id = "user_123@example.com" + _begin_test_purge_operation( + storage, purge_id="purge_detail", idempotency_key="idem_purge_detail", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) storage.prepare_governance_erase_targets( purge_id="purge_detail", - user_id="user_123@example.com", + user_id=user_id, owned_user_playbook_ids={7}, + execution_claim=_claim_purge(storage, "purge_detail"), ) snapshot = next( @@ -958,23 +2094,33 @@ def test_prepare_governance_erase_targets_sanitizes_snapshot_detail(storage): ) if target.target_name == "target_snapshot" ) - assert snapshot.detail == {"owned_user_playbook_ids": [7]} + assert snapshot.detail == { + "authoritative_user_digest": storage.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = 'purge_detail'""", + (storage.org_id,), + ).fetchone()["authoritative_user_digest"], + "owned_user_playbook_ids": [7], + } def test_apply_governance_user_data_delete_rejects_playbook_snapshot_drift(storage): user_id = "user-snapshot-drift" owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_snapshot_drift", idempotency_key="idem_purge_snapshot_drift", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) storage.prepare_governance_erase_targets( purge_id="purge_snapshot_drift", user_id=user_id, + execution_claim=_claim_purge(storage, "purge_snapshot_drift"), ) storage.conn.execute( """INSERT INTO user_playbooks ( @@ -991,7 +2137,11 @@ def test_apply_governance_user_data_delete_rejects_playbook_snapshot_drift(stora storage.conn.commit() with pytest.raises(ValueError, match="prepared purge snapshot"): - storage.apply_governance_user_data_delete("purge_snapshot_drift", user_id) + storage.apply_governance_user_data_delete( + "purge_snapshot_drift", + user_id, + execution_claim=_claim_purge(storage, "purge_snapshot_drift"), + ) remaining_ids = { int(row["user_playbook_id"]) @@ -1006,13 +2156,16 @@ def test_apply_governance_user_data_delete_rejects_playbook_snapshot_drift(stora def test_prepare_governance_erase_targets_does_not_plan_org_agent_playbook_rebuilds( storage, ): - storage.begin_purge_operation( + user_id = "user-rebuild-windows" + _begin_test_purge_operation( + storage, purge_id="purge_rebuild_windows", idempotency_key="idem_purge_rebuild_windows", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) agent_playbook_id = _seed_agent_playbook( storage, @@ -1026,8 +2179,9 @@ def test_prepare_governance_erase_targets_does_not_plan_org_agent_playbook_rebui storage.prepare_governance_erase_targets( purge_id="purge_rebuild_windows", - user_id="user-rebuild-windows", + user_id=user_id, owned_user_playbook_ids={7}, + execution_claim=_claim_purge(storage, "purge_rebuild_windows"), ) assert ( @@ -1049,19 +2203,22 @@ def test_prepare_governance_erase_targets_records_full_delete_matrix_counts(stor session_id="session_seed", evaluation_name="governance_prepare_counts", ) - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_prepare_counts", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) storage.prepare_governance_erase_targets( purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) delete_targets = { @@ -1095,17 +2252,379 @@ def test_prepare_governance_erase_targets_records_full_delete_matrix_counts(stor } +def test_sqlite_rebuild_mutations_reject_missing_and_stale_claims_without_state_changes( + storage, +): + purge_id = "purge_rebuild_claim_matrix" + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key="idem_purge_rebuild_claim_matrix", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + agent_playbook_id = _seed_agent_playbook( + storage, + status=None, + source_windows=[ + AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), + AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), + ], + ) + _record_agent_playbook_rebuild_target( + storage, + purge_id=purge_id, + agent_playbook_id=agent_playbook_id, + previous_lifecycle_status=None, + ) + + def state_snapshot() -> tuple[object, ...]: + playbook = storage.get_agent_playbook_by_id( + agent_playbook_id, + include_tombstones=True, + ) + assert playbook is not None + return ( + playbook.model_dump( + mode="json", + include={ + "content", + "trigger", + "rationale", + "blocking_issue", + "expanded_terms", + "tags", + "status", + }, + ), + [ + window.model_dump(mode="json") + for window in storage.get_source_windows_for_agent_playbook( + agent_playbook_id + ) + ], + [ + target.model_dump(mode="json") + for target in storage.list_purge_targets(purge_id) + ], + dict( + storage.conn.execute( + "SELECT * FROM purge_operations WHERE org_id = ? AND purge_id = ?", + (storage.org_id, purge_id), + ).fetchone() + ), + ) + + before_missing_hide = state_snapshot() + _assert_rejects_missing_claim( + lambda: storage.hide_governance_agent_playbooks_for_rebuild(purge_id), + lambda: storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert state_snapshot() == before_missing_hide + + stale_claim, takeover_claim = _claim_then_take_over(storage, purge_id) + before_stale_hide = state_snapshot() + with pytest.raises(ValueError, match="purge execution claim"): + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=stale_claim, + ) + assert state_snapshot() == before_stale_hide + + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=takeover_claim, + ) + apply_kwargs = { + "purge_id": purge_id, + "agent_playbook_id": agent_playbook_id, + "remaining_source_windows": [ + {"user_playbook_id": 9, "source_interaction_ids": [201]} + ], + "content": "rebuilt claimed content", + "trigger": "rebuilt claimed trigger", + "rationale": "rebuilt claimed rationale", + "blocking_issue": None, + "expanded_terms": "rebuilt claimed terms", + "tags": ["rebuilt-claimed"], + } + before_missing_apply = state_snapshot() + _assert_rejects_missing_claim( + lambda: storage.apply_governance_agent_playbook_rebuild(**apply_kwargs), + lambda: storage.apply_governance_agent_playbook_rebuild( + **apply_kwargs, + execution_claim=None, # type: ignore[arg-type] + ), + ) + assert state_snapshot() == before_missing_apply + with pytest.raises(ValueError, match="purge execution claim"): + storage.apply_governance_agent_playbook_rebuild( + **apply_kwargs, + execution_claim=stale_claim, + ) + assert state_snapshot() == before_missing_apply + + storage.apply_governance_agent_playbook_rebuild( + **apply_kwargs, + execution_claim=takeover_claim, + ) + rebuilt = storage.get_agent_playbook_by_id(agent_playbook_id) + assert rebuilt is not None + assert rebuilt.content == "rebuilt claimed content" + assert rebuilt.status is None + assert storage.get_source_windows_for_agent_playbook(agent_playbook_id) == [ + AgentPlaybookSourceWindow( + user_playbook_id=9, + source_interaction_ids=[201], + ) + ] + rebuild_target = storage.list_purge_targets( + purge_id, + phase="rebuild_without_erased_sources", + ) + assert len(rebuild_target) == 1 + assert rebuild_target[0].status == "complete" + + +def _prepare_blocking_embedding_rebuild( + storage: SQLiteStorage, + *, + suffix: str, +) -> tuple[int, PurgeExecutionClaim, dict[str, object]]: + purge_id = f"purge_blocking_embedding_{suffix}" + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key=f"idem_{purge_id}", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + agent_playbook_id = _seed_agent_playbook( + storage, + status=None, + source_windows=[ + AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), + AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), + ], + ) + _record_agent_playbook_rebuild_target( + storage, + purge_id=purge_id, + agent_playbook_id=agent_playbook_id, + previous_lifecycle_status=None, + ) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) + return ( + agent_playbook_id, + claim, + { + "purge_id": purge_id, + "agent_playbook_id": agent_playbook_id, + "remaining_source_windows": [ + {"user_playbook_id": 9, "source_interaction_ids": [201]} + ], + "content": "rebuilt after embedding", + "trigger": "embedding trigger", + "rationale": "embedding rationale", + "blocking_issue": None, + "expanded_terms": "embedding terms", + "tags": ["embedding"], + "execution_claim": claim, + }, + ) + + +def test_sqlite_rebuild_embedding_does_not_hold_shared_storage_lock(storage) -> None: + agent_playbook_id, _, apply_kwargs = _prepare_blocking_embedding_rebuild( + storage, + suffix="ordinary_read", + ) + embedding_started = threading.Event() + release_embedding = threading.Event() + rebuild_errors: list[BaseException] = [] + read_finished = threading.Event() + + def blocking_embedding(_text: str, purpose: str = "document") -> list[float]: + assert purpose == "document" + embedding_started.set() + assert release_embedding.wait(timeout=5) + return [0.0] * 512 + + def rebuild() -> None: + try: + storage.apply_governance_agent_playbook_rebuild(**apply_kwargs) + except BaseException as exc: # pragma: no cover - asserted below + rebuild_errors.append(exc) + + def ordinary_read() -> None: + storage.get_agent_playbook_by_id(agent_playbook_id, include_tombstones=True) + read_finished.set() + + with patch.object(storage, "_get_embedding", side_effect=blocking_embedding): + rebuild_thread = threading.Thread(target=rebuild) + rebuild_thread.start() + read_thread = threading.Thread(target=ordinary_read) + try: + assert embedding_started.wait(timeout=5) + read_thread.start() + read_completed_while_embedding_blocked = read_finished.wait(timeout=0.5) + finally: + release_embedding.set() + rebuild_thread.join(timeout=5) + if read_thread.ident is not None: + read_thread.join(timeout=5) + + assert read_completed_while_embedding_blocked is True + assert rebuild_errors == [] + assert not rebuild_thread.is_alive() + assert not read_thread.is_alive() + + +def test_sqlite_rebuild_revalidates_claim_after_embedding_takeover( + storage_factory, +) -> None: + storage = storage_factory("org1") + peer = storage_factory("org1") + agent_playbook_id, stale_claim, apply_kwargs = _prepare_blocking_embedding_rebuild( + storage, + suffix="claim_takeover", + ) + before = storage.get_agent_playbook_by_id( + agent_playbook_id, + include_tombstones=True, + ) + assert before is not None + embedding_started = threading.Event() + release_embedding = threading.Event() + rebuild_errors: list[BaseException] = [] + + def blocking_embedding(_text: str, purpose: str = "document") -> list[float]: + assert purpose == "document" + embedding_started.set() + assert release_embedding.wait(timeout=5) + return [0.0] * 512 + + def rebuild() -> None: + try: + storage.apply_governance_agent_playbook_rebuild(**apply_kwargs) + except BaseException as exc: + rebuild_errors.append(exc) + + takeover_claim: PurgeExecutionClaim | None = None + with patch.object(storage, "_get_embedding", side_effect=blocking_embedding): + rebuild_thread = threading.Thread(target=rebuild) + rebuild_thread.start() + try: + assert embedding_started.wait(timeout=5) + peer.conn.execute("PRAGMA busy_timeout = 500") + peer.conn.execute( + """UPDATE purge_operations + SET execution_claim_expires_at = 0 + WHERE org_id = ? AND purge_id = ?""", + (peer.org_id, apply_kwargs["purge_id"]), + ) + peer.conn.commit() + takeover_claim = peer.claim_purge_operation_execution( + str(apply_kwargs["purge_id"]), + lease_owner="takeover-owner", + lease_ttl_seconds=30, + ) + finally: + release_embedding.set() + rebuild_thread.join(timeout=5) + + assert takeover_claim is not None + assert takeover_claim.fence == stale_claim.fence + 1 + assert len(rebuild_errors) == 1 + assert isinstance(rebuild_errors[0], ValueError) + assert "purge execution claim" in str(rebuild_errors[0]) + after = storage.get_agent_playbook_by_id( + agent_playbook_id, + include_tombstones=True, + ) + assert after == before + [target] = storage.list_purge_targets( + str(apply_kwargs["purge_id"]), + phase="rebuild_without_erased_sources", + ) + assert target.status == "running" + + +def test_governance_service_rebuilds_through_claimed_sqlite_contract(storage): + purge_id = "purge_service_claimed_rebuild" + _begin_test_purge_operation( + storage, + purge_id=purge_id, + idempotency_key="idem_purge_service_claimed_rebuild", + operation_type="user_erasure", + scope_type="user", + subject_ref=SUBJECT_REF, + request_ref=REQUEST_REF, + authoritative_user_id="alice", + ) + agent_playbook_id = _seed_agent_playbook( + storage, + status=None, + source_windows=[ + AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), + AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), + ], + ) + _record_agent_playbook_rebuild_target( + storage, + purge_id=purge_id, + agent_playbook_id=agent_playbook_id, + previous_lifecycle_status=None, + ) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) + service = GovernanceService( + storage=storage, + org_id=storage.org_id, + ref_secret="test-governance-secret", + ) + + rebuilt_ids = service._rebuild_agent_playbooks( + purge_id, + execution_claim=claim, + ) + + assert rebuilt_ids == [agent_playbook_id] + rebuilt = storage.get_agent_playbook_by_id(agent_playbook_id) + assert rebuilt is not None + assert rebuilt.content == "source-playbook-9" + assert rebuilt.status is None + + def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_and_hide_marker( storage, ): purge_id = "purge_hide_rebuild" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_hide_rebuild", operation_type="user_erasure", scope_type="user", subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, + authoritative_user_id="alice", ) agent_playbook_id = _seed_agent_playbook( storage, @@ -1132,7 +2651,10 @@ def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_an ], } - hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=_claim_purge(storage, purge_id), + ) assert hidden_ids == [agent_playbook_id] status = storage.conn.execute( @@ -1161,7 +2683,8 @@ def test_hide_governance_agent_playbooks_for_rebuild_sets_archive_in_progress_an def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage): purge_id = "purge_rebuild_complete" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_complete", operation_type="user_erasure", @@ -1177,12 +2700,14 @@ def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1200,6 +2725,7 @@ def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) expected_detail = { "original_source_windows": [ @@ -1224,6 +2750,7 @@ def test_apply_governance_agent_playbook_rebuild_completes_planned_phase(storage blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) rebuild_target = next( @@ -1261,7 +2788,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_ad_hoc_rebuild_without_ storage, ): purge_id = "purge_rebuild_requires_target" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_requires_target", operation_type="user_erasure", @@ -1299,6 +2827,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_ad_hoc_rebuild_without_ blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=_claim_purge(storage, purge_id), ) assert ( @@ -1324,7 +2853,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_rebuild_before_hide_pha storage, ): purge_id = "purge_rebuild_requires_hide" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_requires_hide", operation_type="user_erasure", @@ -1340,12 +2870,14 @@ def test_apply_governance_agent_playbook_rebuild_rejects_rebuild_before_hide_pha AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1379,6 +2911,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_rebuild_before_hide_pha blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) assert ( @@ -1419,7 +2952,8 @@ def test_apply_governance_agent_playbook_rebuild_succeeds_after_prepare_and_hide storage, ): purge_id = "purge_rebuild_prepare_hide" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_prepare_hide", operation_type="user_erasure", @@ -1441,7 +2975,11 @@ def test_apply_governance_agent_playbook_rebuild_succeeds_after_prepare_and_hide agent_playbook_id=agent_playbook_id, previous_lifecycle_status=None, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, @@ -1455,6 +2993,7 @@ def test_apply_governance_agent_playbook_rebuild_succeeds_after_prepare_and_hide blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) rebuild_target = next( @@ -1485,7 +3024,8 @@ def test_apply_governance_agent_playbook_rebuild_does_not_complete_target_when_s storage, monkeypatch ): purge_id = "purge_rebuild_search_refresh_failure" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_search_refresh_failure", operation_type="user_erasure", @@ -1501,12 +3041,14 @@ def test_apply_governance_agent_playbook_rebuild_does_not_complete_target_when_s AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1524,6 +3066,7 @@ def test_apply_governance_agent_playbook_rebuild_does_not_complete_target_when_s target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) original_row = storage.conn.execute( """SELECT content, trigger, rationale, blocking_issue, expanded_terms, tags, status @@ -1561,6 +3104,7 @@ def fail_search_refresh(*args, **kwargs): blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) assert ( @@ -1598,7 +3142,8 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when storage, ): purge_id = "purge_rebuild_remove_orphan" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_remove_orphan", operation_type="user_erasure", @@ -1613,12 +3158,14 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when AgentPlaybookSourceWindow(user_playbook_id=7, source_interaction_ids=[101]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1633,6 +3180,7 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) storage.apply_governance_agent_playbook_rebuild( @@ -1645,6 +3193,7 @@ def test_apply_governance_agent_playbook_rebuild_removes_orphaned_aggregate_when blocking_issue=None, expanded_terms=None, tags=None, + execution_claim=claim, ) rebuild_target = next( @@ -1699,7 +3248,8 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta storage, ): purge_id = "purge_rebuild_restore_archived" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_restore_archived", operation_type="user_erasure", @@ -1715,12 +3265,14 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta AgentPlaybookSourceWindow(user_playbook_id=9, source_interaction_ids=[201]), ], ) + claim = _claim_purge(storage, purge_id) storage.record_purge_target( purge_id=purge_id, target_name="agent_playbook", target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=claim, detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -1738,6 +3290,7 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=claim, ) storage.apply_governance_agent_playbook_rebuild( @@ -1752,6 +3305,7 @@ def test_apply_governance_agent_playbook_rebuild_restores_previous_lifecycle_sta blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) rebuilt_status = storage.conn.execute( @@ -1765,7 +3319,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl storage, ): purge_id = "purge_rebuild_second_call" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_second_call", operation_type="user_erasure", @@ -1787,7 +3342,11 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl agent_playbook_id=agent_playbook_id, previous_lifecycle_status=Status.ARCHIVED.value, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, agent_playbook_id=agent_playbook_id, @@ -1800,6 +3359,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) before_row = storage.conn.execute( @@ -1838,6 +3398,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_second_call_after_compl blocking_issue={"issue": "should not persist"}, expanded_terms="mutated terms", tags=["mutated"], + execution_claim=claim, ) after_row = storage.conn.execute( @@ -1877,7 +3438,8 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet storage, ): purge_id = "purge_hide_after_complete" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_hide_after_complete", operation_type="user_erasure", @@ -1899,7 +3461,11 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet agent_playbook_id=agent_playbook_id, previous_lifecycle_status=Status.ARCHIVED.value, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, agent_playbook_id=agent_playbook_id, @@ -1912,6 +3478,7 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) before_status = storage.conn.execute( @@ -1934,7 +3501,10 @@ def test_hide_governance_agent_playbooks_for_rebuild_is_idempotent_after_complet and target.target_ref == str(agent_playbook_id) ) - hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) after_status = storage.conn.execute( "SELECT status FROM agent_playbooks WHERE agent_playbook_id = ?", @@ -1969,7 +3539,8 @@ def test_hide_governance_agent_playbooks_for_rebuild_does_not_reopen_complete_ta storage, monkeypatch ): purge_id = "purge_hide_stale_prelock" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_hide_stale_prelock", operation_type="user_erasure", @@ -1991,7 +3562,11 @@ def test_hide_governance_agent_playbooks_for_rebuild_does_not_reopen_complete_ta agent_playbook_id=agent_playbook_id, previous_lifecycle_status=Status.ARCHIVED.value, ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + claim = _claim_purge(storage, purge_id) + storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) storage.apply_governance_agent_playbook_rebuild( purge_id=purge_id, agent_playbook_id=agent_playbook_id, @@ -2004,6 +3579,7 @@ def test_hide_governance_agent_playbooks_for_rebuild_does_not_reopen_complete_ta blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=claim, ) before_status = storage.conn.execute( @@ -2036,7 +3612,10 @@ def stale_list_purge_targets(*_args, **_kwargs): monkeypatch.setattr(storage, "list_purge_targets", stale_list_purge_targets) - hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild(purge_id) + hidden_ids = storage.hide_governance_agent_playbooks_for_rebuild( + purge_id, + execution_claim=claim, + ) assert hidden_ids == [] assert ( @@ -2080,19 +3659,22 @@ def test_prepare_governance_erase_targets_is_idempotent_after_completed_snapshot ): purge_id = "purge_prepare_idempotent_after_snapshot" user_id = "user-prepare-idempotent-after-snapshot" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_prepare_idempotent_after_snapshot", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) storage.prepare_governance_erase_targets( purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) before_targets = [ @@ -2112,6 +3694,7 @@ def test_prepare_governance_erase_targets_is_idempotent_after_completed_snapshot purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) after_targets = [ @@ -2144,8 +3727,9 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): idempotency_key=f"idem_{storage_instance.org_id}_{purge_id}", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage_instance._subject_ref_for_user_id("alice"), request_ref=request_ref, + authoritative_user_id="alice", ) storage_org1.record_purge_target( @@ -2154,6 +3738,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="prepare_targets", status="complete", + execution_claim=_claim_purge(storage_org1, purge_id), detail={"prepared": True}, ) storage_org1.record_purge_target( @@ -2162,6 +3747,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="delete", status="pending", + execution_claim=_claim_purge(storage_org1, purge_id), detail={"count": 1}, ) storage_org2.record_purge_target( @@ -2170,6 +3756,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage_org2, purge_id), detail={"count": 2}, deleted_count=2, ) @@ -2197,6 +3784,7 @@ def test_purge_targets_are_scoped_by_org_for_same_purge_id(storage_factory): target_ref="all", phase="delete", status="running", + execution_claim=_claim_purge(storage_org2, purge_id), detail={"count": 3}, ) @@ -2293,6 +3881,7 @@ def test_record_purge_target_validates_governance_fields(storage, kwargs, match) "phase": "delete", "status": "running", "target_ref": "all", + "execution_claim": _claim_purge(storage, purge_id), } params.update(kwargs) @@ -2331,6 +3920,7 @@ def test_record_purge_target_rejects_invalid_deleted_count( phase="delete", status="complete", deleted_count=deleted_count, + execution_claim=_claim_purge(storage, purge_id), ) @@ -2348,6 +3938,7 @@ def test_record_purge_target_accepts_nonnegative_detail_deleted_count( target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail={"deleted_count": detail_deleted_count}, ) @@ -2370,6 +3961,7 @@ def test_record_purge_target_rejects_negative_detail_deleted_count(storage): phase="delete", status="complete", detail={"deleted_count": -1}, + execution_claim=_claim_purge(storage, purge_id), ) @@ -2428,6 +4020,16 @@ def test_record_purge_target_rejects_negative_detail_deleted_count(storage): None, id="audit-detail-allowed-deleted-counts", ), + pytest.param( + {"deleted_counts": {"session_outcomes": 1}}, + None, + id="audit-detail-allowed-session-outcome-counts", + ), + pytest.param( + {"deleted_counts": {"session_outcome": 1}}, + "session_outcome", + id="audit-detail-rejects-unknown-deleted-count-key", + ), pytest.param( {"agent_playbook_id": 7}, None, id="audit-detail-allowed-agent-playbook-id" ), @@ -2464,6 +4066,7 @@ def test_record_purge_target_accepts_target_detail_shapes(storage): purge_id = _begin_purge(storage, "purge_target_detail_shapes") detail = { + "authoritative_user_digest": "a" * 64, "owned_user_playbook_ids": [7], "source_interaction_ids": [11, 12], "original_source_windows": [ @@ -2481,6 +4084,7 @@ def test_record_purge_target_accepts_target_detail_shapes(storage): target_ref="7", phase="rebuild_without_erased_sources", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -2532,6 +4136,7 @@ def test_fail_purge_operation_rejects_raw_error_detail(storage): purge_id, error_code="boom", error_detail="RuntimeError: request reqref_123 for alice@example.com", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_detail is None @@ -2545,6 +4150,7 @@ def test_fail_purge_operation_rejects_freeform_error_detail(storage): purge_id, error_code="PURGE_TARGET_FAILED", error_detail="stable failure detail", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_detail is None @@ -2557,6 +4163,7 @@ def test_fail_purge_operation_persists_code_shaped_error_detail(storage): purge_id, error_code="PURGE_TARGET_FAILED", error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert failed.status == "failed" @@ -2569,6 +4176,7 @@ def test_fail_missing_purge_rolls_back_implicit_transaction(storage): "purge_missing", error_code="PURGE_TARGET_FAILED", error_detail="target_delete_failed", + execution_claim=_typed_test_claim_for_unvalidated_purge_id("purge_missing"), ) assert storage.conn.in_transaction is False @@ -2577,6 +4185,7 @@ def test_fail_missing_purge_rolls_back_implicit_transaction(storage): def test_record_purge_target_rolls_back_after_write_failure(storage, monkeypatch): purge_id = _begin_purge(storage, "purge_target_write_failure") + claim = _claim_purge(storage, purge_id) def _write_then_raise(**_kwargs: object) -> None: storage.conn.execute( @@ -2594,6 +4203,7 @@ def _write_then_raise(**_kwargs: object) -> None: target_ref="all", phase="delete", status="running", + execution_claim=claim, ) assert storage.conn.in_transaction is False @@ -2633,6 +4243,7 @@ def _trace(statement: str) -> None: "user", SUBJECT_REF, REQUEST_REF, + authoritative_user_id="alice", ) assert entered.wait(timeout=1) time.sleep(0.05) @@ -2653,9 +4264,11 @@ def test_prepare_targets_rechecks_snapshot_after_two_connection_write_lock( idempotency_key="idem_two_connection_prepare", operation_type="user_erasure", scope_type="user", + authoritative_user_id="alice", subject_ref=SUBJECT_REF, request_ref=REQUEST_REF, ) + claim = _claim_purge(first, purge_id) first.conn.execute("BEGIN IMMEDIATE") first._record_purge_target_locked( purge_id=purge_id, @@ -2673,7 +4286,14 @@ def test_prepare_targets_rechecks_snapshot_after_two_connection_write_lock( target_ref="all", phase="prepare_targets", status="complete", - detail={"owned_user_playbook_ids": []}, + detail={ + "authoritative_user_digest": first.conn.execute( + """SELECT authoritative_user_digest FROM purge_operations + WHERE org_id = ? AND purge_id = ?""", + (first.org_id, purge_id), + ).fetchone()["authoritative_user_digest"], + "owned_user_playbook_ids": [], + }, deleted_count=0, error_detail=None, ) @@ -2691,8 +4311,9 @@ def _trace(statement: str) -> None: future = executor.submit( second.prepare_governance_erase_targets, purge_id, - "two-connection-user", - set(), + "alice", + execution_claim=claim, + owned_user_playbook_ids=set(), ) assert entered.wait(timeout=1) time.sleep(0.05) @@ -2719,6 +4340,7 @@ def test_fail_purge_operation_accepts_code_shaped_error_code_with_prompt_or_cont purge_id, error_code=error_code, error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert failed.status == "failed" @@ -2733,6 +4355,7 @@ def test_fail_purge_operation_rejects_prompt_content_prose_error_detail(storage) purge_id, error_code="PURGE_TARGET_FAILED", error_detail="prompt content leaked from upstream", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_detail is None @@ -2755,6 +4378,7 @@ def test_fail_purge_operation_validates_error_code(storage, error_code, match): purge_id, error_code=error_code, error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert failed.status == "failed" assert failed.error_code == error_code @@ -2765,6 +4389,7 @@ def test_fail_purge_operation_validates_error_code(storage, error_code, match): purge_id, error_code=error_code, error_detail="target_delete_failed", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.get_purge_operation(purge_id).error_code is None @@ -2903,7 +4528,8 @@ def test_begin_purge_operation_validates_top_level_refs( storage, subject_ref, request_ref, match ): with pytest.raises(ValueError, match=match): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_top_level_refs", idempotency_key="idem_purge_top_level_refs", operation_type="user_erasure", @@ -2928,7 +4554,8 @@ def test_begin_purge_operation_rejects_invalid_enum_values( storage, operation_type, scope_type, match ): with pytest.raises(ValueError, match=match): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_invalid_enum", idempotency_key="idem_purge_invalid_enum", operation_type=operation_type, @@ -2949,7 +4576,8 @@ def test_begin_purge_operation_rejects_invalid_enum_values( ) def test_begin_purge_operation_rejects_unsafe_purge_id(storage, purge_id): with pytest.raises(ValueError, match="purge_id"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_invalid_id", operation_type="user_erasure", @@ -2990,6 +4618,7 @@ def test_record_purge_target_rejects_mixed_case_window_keys(storage, detail_key) target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={detail_key: [{"User_Playbook_Id": "alice@example.com"}]}, ) @@ -3025,6 +4654,7 @@ def test_record_purge_target_requires_window_user_playbook_id(storage, detail_ke target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={detail_key: [{"source_interaction_ids": [1, 2]}]}, ) @@ -3046,6 +4676,7 @@ def test_record_purge_target_accepts_previous_lifecycle_status_for_rebuild_targe target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [11, 12]} @@ -3124,6 +4755,7 @@ def test_record_purge_target_rejects_invalid_previous_lifecycle_status_detail( target_ref="7", phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -3151,6 +4783,7 @@ def test_record_purge_target_validates_target_ref_contract(storage, target_ref, target_ref=target_ref, phase="delete", status="running", + execution_claim=_claim_purge(storage, purge_id), ) return @@ -3161,6 +4794,7 @@ def test_record_purge_target_validates_target_ref_contract(storage, target_ref, target_ref=target_ref, phase="delete", status="running", + execution_claim=_claim_purge(storage, purge_id), ) @@ -3195,6 +4829,7 @@ def test_persistence_paths_reject_unsafe_purge_id(storage, purge_id): target_ref="all", phase="delete", status="running", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge_id), ) with pytest.raises(ValueError, match="purge_id"): @@ -3208,6 +4843,8 @@ def test_persistence_paths_reject_unsafe_purge_id(storage, purge_id): request_ref=REQUEST_REF, idempotency_key=purge_id, ), + authoritative_user_id="alice", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge_id), ) with pytest.raises(ValueError, match="purge_id"): @@ -3225,6 +4862,9 @@ def test_apply_governance_user_data_delete_rejects_unsafe_purge_id_before_side_e storage.apply_governance_user_data_delete( purge_id="alice@example.com", user_id=user_id, + execution_claim=_typed_test_claim_for_unvalidated_purge_id( + "alice@example.com" + ), ) remaining = _user_scoped_row_counts(storage, user_id=user_id) @@ -3239,7 +4879,11 @@ def test_apply_governance_user_data_delete_rejects_unsafe_purge_id_before_side_e def test_apply_governance_user_data_delete_rejects_unexpected_target_name_from_internal_counts( storage, monkeypatch ): - purge_id = _begin_purge(storage, "purge_internal_target_name") + purge_id = _begin_purge( + storage, + "purge_internal_target_name", + authoritative_user_id="user-delete-seed", + ) for target_name in CANONICAL_DELETE_TARGET_NAMES: storage.record_purge_target( purge_id=purge_id, @@ -3247,6 +4891,7 @@ def test_apply_governance_user_data_delete_rejects_unexpected_target_name_from_i target_ref="all", phase="delete", status="pending", + execution_claim=_claim_purge(storage, purge_id), detail={"count": 0}, ) @@ -3269,6 +4914,7 @@ def _stub_clear_user_data_for_governance_locked( storage.apply_governance_user_data_delete( purge_id=purge_id, user_id="user-delete-seed", + execution_claim=_claim_purge(storage, purge_id), ) delete_targets = storage.list_purge_targets(purge_id, phase="delete") @@ -3278,8 +4924,12 @@ def _stub_clear_user_data_for_governance_locked( def test_apply_governance_user_data_delete_requires_complete_prepared_delete_matrix( storage, monkeypatch ): - purge_id = _begin_purge(storage, "purge_delete_requires_prepared_matrix") user_id = "user-delete-seed" + purge_id = _begin_purge( + storage, + "purge_delete_requires_prepared_matrix", + authoritative_user_id=user_id, + ) expected_user_id = user_id _seed_user_scoped_rows(storage, user_id=user_id) baseline_counts = _user_scoped_row_counts(storage, user_id=user_id) @@ -3289,6 +4939,7 @@ def test_apply_governance_user_data_delete_requires_complete_prepared_delete_mat target_ref="all", phase="delete", status="pending", + execution_claim=_claim_purge(storage, purge_id), detail={"count": 1}, ) storage.record_purge_target( @@ -3297,6 +4948,7 @@ def test_apply_governance_user_data_delete_requires_complete_prepared_delete_mat target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail={"count": 0}, deleted_count=0, ) @@ -3322,6 +4974,7 @@ def _stub_clear_user_data_for_governance_locked( storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert clear_locked_called is False @@ -3338,13 +4991,15 @@ def test_apply_governance_user_data_delete_preserves_org_agent_playbooks_without ): purge_id = "purge_delete_requires_hide" user_id = "user-delete-hide-required" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_delete_requires_hide", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) _seed_eval_result( @@ -3368,11 +5023,13 @@ def test_apply_governance_user_data_delete_preserves_org_agent_playbooks_without purge_id=purge_id, user_id=user_id, owned_user_playbook_ids=owned_user_playbook_ids, + execution_claim=_claim_purge(storage, purge_id), ) counts = storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert counts["user_playbooks"] == 1 @@ -3400,13 +5057,15 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( """ purge_id = "purge_delete_after_hide" user_id = "user-delete-hide-complete" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_delete_after_hide", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) _seed_eval_result( @@ -3440,12 +5099,17 @@ 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, + execution_claim=_claim_purge(storage, purge_id), ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) counts = storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert counts == { @@ -3526,13 +5190,15 @@ def test_apply_governance_user_data_delete_retains_lineage_skeleton( def test_apply_governance_user_data_delete_is_failure_atomic(storage, monkeypatch): purge_id = "purge_delete_atomic" user_id = "user-delete-atomic" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_delete_atomic", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id(user_id), request_ref=REQUEST_REF, + authoritative_user_id=user_id, ) owned_user_playbook_ids = _seed_prepare_counts_user_data(storage, user_id=user_id) affected_user_playbook_id = min(owned_user_playbook_ids) @@ -3550,8 +5216,12 @@ 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, + execution_claim=_claim_purge(storage, purge_id), ) - storage.hide_governance_agent_playbooks_for_rebuild(purge_id) before_counts = _user_scoped_row_counts(storage, user_id=user_id) before_profile_rows = storage.conn.execute( @@ -3605,6 +5275,7 @@ def _raising_record_purge_target_locked( storage.apply_governance_user_data_delete( purge_id=purge_id, user_id=user_id, + execution_claim=_claim_purge(storage, purge_id), ) assert _user_scoped_row_counts(storage, user_id=user_id) == before_counts @@ -3655,6 +5326,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_unsafe_purge_id_before_ blocking_issue=None, expanded_terms="updated terms", tags=["updated"], + execution_claim=_typed_test_claim_for_unvalidated_purge_id("request_12345"), ) after_row = storage.conn.execute( @@ -3676,6 +5348,7 @@ def test_fail_purge_operation_rejects_unsafe_purge_id_before_side_effects(storag SUBJECT_REF, "governance.error", "detail.code", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(SUBJECT_REF), ) failed = storage.get_purge_operation(purge_id) @@ -3688,7 +5361,8 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so storage, ): purge_id = "purge_rebuild_windows_mismatch" - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_rebuild_windows_mismatch", operation_type="user_erasure", @@ -3710,6 +5384,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so target_ref=str(agent_playbook_id), phase="rebuild_without_erased_sources", status="running", + execution_claim=_claim_purge(storage, purge_id), detail={ "original_source_windows": [ {"user_playbook_id": 7, "source_interaction_ids": [101]}, @@ -3727,6 +5402,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so target_ref=str(agent_playbook_id), phase="hide_for_rebuild", status="complete", + execution_claim=_claim_purge(storage, purge_id), ) original_row = storage.conn.execute( """SELECT content, trigger, rationale, blocking_issue, expanded_terms, tags, status @@ -3749,6 +5425,7 @@ def test_apply_governance_agent_playbook_rebuild_rejects_mismatched_remaining_so blocking_issue=None, expanded_terms="rebuilt terms", tags=["rebuilt"], + execution_claim=_claim_purge(storage, purge_id), ) rebuilt_row = storage.conn.execute( @@ -3867,6 +5544,7 @@ def test_record_purge_target_rejects_invalid_enum_values( target_ref="all", phase=phase, status=status, + execution_claim=_claim_purge(storage, purge_id), ) @@ -3919,7 +5597,8 @@ def test_governance_persistence_rejects_unsafe_idempotency_keys( storage.append_audit_event(event) with pytest.raises(ValueError, match="idempotency_key"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id="purge_unsafe_idem", idempotency_key=idempotency_key, operation_type="user_erasure", @@ -4024,6 +5703,7 @@ def test_governance_detail_rejects_noncanonical_status_and_route( target_ref="all", phase="delete", status="running", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -4034,7 +5714,8 @@ def test_governance_detail_rejects_noncanonical_status_and_route( ) def test_begin_purge_operation_rejects_identifier_like_purge_suffix(storage, purge_id): with pytest.raises(ValueError, match="purge_id"): - storage.begin_purge_operation( + _begin_test_purge_operation( + storage, purge_id=purge_id, idempotency_key="idem_purge_identifier_suffix", operation_type="user_erasure", @@ -4070,6 +5751,7 @@ def test_record_purge_target_canonicalizes_detail_keys_before_persistence(storag target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail={" Deleted_Counts ": {"requests": 2}}, deleted_count=2, ) @@ -4105,6 +5787,7 @@ def test_governance_detail_rejects_duplicate_normalized_keys(storage, persistenc target_ref="all", phase="delete", status="complete", + execution_claim=_claim_purge(storage, purge_id), detail=detail, ) @@ -4232,6 +5915,7 @@ def test_record_purge_target_validates_target_ref_by_phase_and_name( target_ref=target_ref, phase=phase, status="running", + execution_claim=_claim_purge(storage, purge_id), ) return @@ -4242,6 +5926,7 @@ def test_record_purge_target_validates_target_ref_by_phase_and_name( target_ref=target_ref, phase=phase, status="running", + execution_claim=_claim_purge(storage, purge_id), ) @@ -4302,13 +5987,15 @@ def test_init_governance_tables_upgrades_legacy_purge_target_table(tmp_path): with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): storage = SQLiteStorage(org_id="org1", db_path=str(db_path)) - purge_id = storage.begin_purge_operation( + purge_id = _begin_test_purge_operation( + storage, purge_id="purge_legacy_upgrade", idempotency_key="idem_legacy_upgrade", operation_type="user_erasure", scope_type="user", - subject_ref=SUBJECT_REF, + subject_ref=storage._subject_ref_for_user_id("alice"), request_ref=REQUEST_REF, + authoritative_user_id="alice", ).purge_id storage.record_purge_target( purge_id=purge_id, @@ -4316,6 +6003,7 @@ def test_init_governance_tables_upgrades_legacy_purge_target_table(tmp_path): target_ref="all", phase="prepare_targets", status="complete", + execution_claim=_claim_purge(storage, purge_id), ) assert storage.purge_targets_prepared(purge_id) is True @@ -4522,18 +6210,23 @@ def test_successful_erase_audit_row_exists_only_after_complete_purge(storage): # (2) The one legitimate writer produces exactly one successful-ERASE row, # and only after the purge_operation transitions to 'complete'. completed = storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_claim_purge(storage, purge_id), ) assert completed.status == "complete" erase_rows = _successful_erase_audit_rows(storage) assert [event.idempotency_key for event in erase_rows] == [purge_id] _assert_successful_erase_rows_only_for_complete_purges(storage) - # Idempotent re-completion neither duplicates the row nor breaks the - # invariant. - storage.complete_purge_operation_with_audit( - purge_id, _erase_event(purge_id=purge_id) - ) + with pytest.raises(ValueError, match="purge execution claim"): + storage.complete_purge_operation_with_audit( + purge_id, + _erase_event(purge_id=purge_id), + authoritative_user_id="alice", + execution_claim=_typed_test_claim_for_unvalidated_purge_id(purge_id), + ) assert [ event.idempotency_key for event in _successful_erase_audit_rows(storage) ] == [purge_id] diff --git a/tests/server/services/storage/sqlite_storage/test_session_id_migration.py b/tests/server/services/storage/sqlite_storage/test_session_id_migration.py index fb40107bd..2a6512833 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_id_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_id_migration.py @@ -20,6 +20,7 @@ source TEXT NOT NULL DEFAULT '', agent_version TEXT NOT NULL DEFAULT '', session_id TEXT, + governance_subject_ref TEXT, metadata TEXT NOT NULL DEFAULT '{}' ); """ @@ -29,12 +30,35 @@ def _seed_legacy_db(db_path: str) -> None: conn = sqlite3.connect(db_path) conn.executescript(_LEGACY_REQUESTS_DDL) conn.executemany( - "INSERT INTO requests (request_id, user_id, created_at, source, session_id) " - "VALUES (?, ?, ?, ?, ?)", + """INSERT INTO requests ( + request_id, user_id, created_at, source, session_id, + governance_subject_ref + ) VALUES (?, ?, ?, ?, ?, ?)""", [ - ("r-null", "u1", "2026-01-01T00:00:00+00:00", "web", None), - ("r-blank", "u1", "2026-01-01T00:00:01+00:00", "web", " "), - ("r-valid", "u1", "2026-01-01T00:00:02+00:00", "web", "s-valid"), + ( + "r-null", + "u1", + "2026-01-01T00:00:00+00:00", + "web", + None, + "subject-null", + ), + ( + "r-blank", + "u1", + "2026-01-01T00:00:01+00:00", + "web", + " ", + "subject-blank", + ), + ( + "r-valid", + "u1", + "2026-01-01T00:00:02+00:00", + "web", + "s-valid", + "subject-valid", + ), ], ) conn.commit() @@ -89,6 +113,28 @@ def test_migration_enforces_not_null_and_non_empty(tmp_path): conn.close() +def test_migration_preserves_governance_subject_ref(tmp_path): + db_path = str(tmp_path / "legacy.db") + _seed_legacy_db(db_path) + + SQLiteStorage(org_id="0", db_path=db_path) + + conn = sqlite3.connect(db_path) + try: + subject_refs = dict( + conn.execute( + "SELECT request_id, governance_subject_ref FROM requests" + ).fetchall() + ) + finally: + conn.close() + assert subject_refs == { + "r-null": "subject-null", + "r-blank": "subject-blank", + "r-valid": "subject-valid", + } + + def test_migration_is_idempotent(tmp_path): db_path = str(tmp_path / "legacy.db") _seed_legacy_db(db_path) diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py index 54a8d5d13..4e484c818 100644 --- a/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_downgrade_migration.py @@ -1,4 +1,4 @@ -"""Regression coverage for reverting the session-outcome identity schema.""" +"""Regression coverage for rebuilding session-outcome identity after rollback.""" from unittest.mock import patch @@ -42,7 +42,7 @@ def _storage(db_path: str) -> SQLiteStorage: return SQLiteStorage(org_id="downgrade-org", db_path=db_path) -def test_identity_schema_is_downgraded_and_current_writes_resume(tmp_path) -> None: +def test_identity_schema_is_preserved_and_current_writes_resume(tmp_path) -> None: db_path = str(tmp_path / "identity-schema.db") storage = _storage(db_path) for session_id in ("kept", "unknown", "new"): @@ -84,6 +84,8 @@ def test_identity_schema_is_downgraded_and_current_writes_resume(tmp_path) -> No row["name"] for row in migrated.conn.execute("PRAGMA table_info(session_outcomes)") } == { + "outcome_id", + "outcome_revision", "user_id", "session_id", "outcome", @@ -92,13 +94,17 @@ def test_identity_schema_is_downgraded_and_current_writes_resume(tmp_path) -> No "label", "value", "metadata", + "outcome_contract_digest", + "finalized_trajectory_digest", "governance_subject_ref", "created_at", } records = migrated.get_session_outcomes(GetSessionOutcomesRequest()) - assert [(record.session_id, record.outcome) for record in records] == [ - ("kept", SessionOutcomeKind.SUCCESS) - ] + records_by_session = {record.session_id: record for record in records} + assert records_by_session["kept"].outcome is SessionOutcomeKind.SUCCESS + assert records_by_session["kept"].outcome_id == "outcome-kept" + assert records_by_session["unknown"].outcome is SessionOutcomeKind.UNKNOWN + assert records_by_session["unknown"].outcome_id == "outcome-unknown" request = SetSessionOutcomeRequest( session_id="new", outcome=SessionOutcomeKind.FAILURE, occurred_at=101 @@ -165,29 +171,28 @@ def test_legacy_schema_backfills_required_subject_ref( ) -def test_legacy_empty_subject_default_is_rebuilt_and_backfilled(tmp_path) -> None: +def test_identity_empty_subject_default_is_rebuilt_without_changing_identity( + tmp_path, +) -> None: db_path = str(tmp_path / "legacy-empty-subject-default.db") storage = _storage(db_path) storage.conn.execute("DROP TABLE session_outcomes") storage.conn.executescript( - """CREATE TABLE session_outcomes ( - user_id TEXT NOT NULL, - session_id TEXT NOT NULL, - outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), - occurred_at INTEGER NOT NULL, - source TEXT NOT NULL, - label TEXT, - value REAL, - metadata TEXT, - governance_subject_ref TEXT NOT NULL DEFAULT '', - created_at INTEGER NOT NULL, - PRIMARY KEY (user_id, session_id) - );""" + _IDENTITY_SCHEMA.replace( + "governance_subject_ref TEXT NOT NULL,", + "governance_subject_ref TEXT NOT NULL DEFAULT '',", + ) ) storage.conn.execute( """INSERT INTO session_outcomes ( - user_id, session_id, outcome, occurred_at, source, created_at - ) VALUES ('legacy-user', 'legacy-session', 'success', 101, 'legacy', 102)""" + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, created_at + ) VALUES ( + 'stable-outcome-id', 3, 'legacy-user', 'legacy-session', 'unknown', + 101, 'legacy', ?, ?, 102 + )""", + ("a" * 64, "b" * 64), ) storage.conn.commit() storage.conn.close() @@ -195,12 +200,25 @@ def test_legacy_empty_subject_default_is_rebuilt_and_backfilled(tmp_path) -> Non migrated = _storage(db_path) row = migrated.conn.execute( - "SELECT governance_subject_ref FROM session_outcomes" + """SELECT outcome_id, outcome_revision, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref + FROM session_outcomes""" ).fetchone() assert row is not None + assert row["outcome_id"] == "stable-outcome-id" + assert row["outcome_revision"] == 3 + assert row["outcome_contract_digest"] == "a" * 64 + assert row["finalized_trajectory_digest"] == "b" * 64 assert row["governance_subject_ref"] == migrated._subject_ref_for_user_id( "legacy-user" ) + subject_column = next( + item + for item in migrated.conn.execute("PRAGMA table_info(session_outcomes)") + if item["name"] == "governance_subject_ref" + ) + assert subject_column["notnull"] == 1 + assert subject_column["dflt_value"] is None def test_sqlite_versions_without_returning_and_drop_column_are_rejected( diff --git a/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py new file mode 100644 index 000000000..c2602c6d0 --- /dev/null +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -0,0 +1,608 @@ +"""Regression coverage for the SQLite session-outcome identity migration.""" + +import json +import sqlite3 +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, + _prefetch_canonical_session_trajectory_digests, +) + +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) +); +""" + + +_IDENTITY_COMPLETE_UNCONSTRAINED_REVISION_DDL = """ +CREATE TABLE session_outcomes ( + outcome_id TEXT NOT NULL UNIQUE, + outcome_revision INTEGER, + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + outcome TEXT NOT NULL CHECK (outcome IN ('success', 'failure')), + occurred_at INTEGER NOT NULL, + source TEXT NOT NULL, + label TEXT, + value REAL, + metadata TEXT, + outcome_contract_digest TEXT NOT NULL, + finalized_trajectory_digest TEXT NOT NULL, + governance_subject_ref TEXT NOT NULL, + created_at INTEGER NOT NULL, + PRIMARY KEY (user_id, session_id) +); +""" + + +def _replace_session_outcomes_with_unconstrained_revision( + storage: SQLiteStorage, revision: int | None +) -> None: + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_IDENTITY_COMPLETE_UNCONSTRAINED_REVISION_DDL) + storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + "legacy-outcome-id", + revision, + "legacy-user", + "legacy-session", + "success", + 101, + "legacy-source", + "a" * 64, + "b" * 64, + "subject:legacy-user", + 102, + ), + ) + storage.conn.commit() + + +def test_migration_preserves_populated_legacy_outcomes_with_unambiguous_ids( + tmp_path, +) -> None: + 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"]) + ) + + +def test_migration_defaults_null_legacy_outcome_revision_to_one(tmp_path) -> None: + db_path = str(tmp_path / "legacy-null-revision.db") + storage = SQLiteStorage(org_id="legacy-null-revision", db_path=db_path) + _replace_session_outcomes_with_unconstrained_revision(storage, None) + storage.conn.close() + + migrated = SQLiteStorage(org_id="legacy-null-revision", db_path=db_path) + + row = migrated.conn.execute( + "SELECT outcome_revision FROM session_outcomes WHERE session_id = ?", + ("legacy-session",), + ).fetchone() + assert row is not None + assert row["outcome_revision"] == 1 + + +def test_migration_rejects_zero_legacy_outcome_revision_atomically(tmp_path) -> None: + db_path = str(tmp_path / "legacy-zero-revision.db") + storage = SQLiteStorage(org_id="legacy-zero-revision", db_path=db_path) + _replace_session_outcomes_with_unconstrained_revision(storage, 0) + storage.conn.close() + + with pytest.raises(sqlite3.IntegrityError, match="outcome_revision"): + SQLiteStorage(org_id="legacy-zero-revision", db_path=db_path) + + with sqlite3.connect(db_path) as probe: + stored_revision = probe.execute( + "SELECT outcome_revision FROM session_outcomes WHERE session_id = ?", + ("legacy-session",), + ).fetchone()[0] + stranded_legacy_table = probe.execute( + """SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes_legacy'""" + ).fetchone() + + assert stored_revision == 0 + assert stranded_legacy_table is None + + +def test_migration_prefetches_trajectory_inputs_in_bounded_chunks(tmp_path) -> None: + db_path = str(tmp_path / "legacy-session-outcome-query-scaling.db") + storage = SQLiteStorage(org_id="legacy-query-scaling", db_path=db_path) + row_count = 501 + storage.conn.executemany( + """INSERT INTO requests ( + request_id, user_id, created_at, source, session_id + ) VALUES (?, ?, ?, ?, ?)""", + [ + ( + f"request-{index}", + f"user-{index}", + str(index), + "legacy-source", + f"session-{index}", + ) + for index in range(row_count) + ], + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + storage.conn.executemany( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?)""", + [ + ( + f"user-{index}", + f"session-{index}", + "success", + index, + "legacy-source", + f"subject-{index}", + index, + ) + for index in range(row_count) + ], + ) + storage.conn.commit() + + statements: list[str] = [] + storage.conn.set_trace_callback(statements.append) + try: + storage._migrate_session_outcomes_schema() + finally: + storage.conn.set_trace_callback(None) + + trajectory_input_queries = [ + statement + for statement in statements + if statement.lstrip().upper().startswith("SELECT") + and ("FROM requests" in statement or "FROM interactions" in statement) + ] + assert len(trajectory_input_queries) == 4 + assert ( + storage.conn.execute("SELECT COUNT(*) FROM session_outcomes").fetchone()[0] + == row_count + ) + + +def test_migration_prefetch_retains_only_trajectory_digests(tmp_path) -> None: + storage = SQLiteStorage( + org_id="legacy-digest-retention", + db_path=str(tmp_path / "legacy-digest-retention.db"), + ) + storage.add_request( + Request( + request_id="request-digest-retention", + user_id="digest-user", + session_id="digest-session", + source="legacy-source", + created_at=100, + ) + ) + + digests = _prefetch_canonical_session_trajectory_digests( + storage.conn, ["digest-session"] + ) + + assert digests == { + "digest-session": trajectory_digest( + _canonical_session_snapshot(storage.conn, "digest-session") + ) + } + assert all(isinstance(digest, str) for digest in digests.values()) + + +def test_session_outcome_rebuild_failure_rolls_back_renamed_legacy_table( + tmp_path, monkeypatch +) -> None: + db_path = str(tmp_path / "legacy-session-outcome-atomicity.db") + storage = SQLiteStorage(org_id="legacy-atomicity", db_path=db_path) + storage.add_request( + Request( + request_id="request-atomicity-session", + user_id="atomicity-user", + session_id="atomicity-session", + source="atomicity-source", + created_at=100, + ) + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + legacy_rows = [ + ( + "atomicity-user", + "atomicity-session", + "success", + 101, + "atomicity-source", + "atomicity-label", + 4.5, + json.dumps({"kept": True}, sort_keys=True), + "subject:atomicity-user", + 102, + ), + ( + "atomicity-user-2", + "atomicity-session-2", + "failure", + 201, + "atomicity-source-2", + None, + None, + None, + "subject:atomicity-user-2", + 202, + ), + ] + storage.conn.executemany( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, label, value, + metadata, governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + legacy_rows, + ) + storage.conn.commit() + storage.conn.close() + + with sqlite3.connect(db_path) as probe: + original_schema = probe.execute( + """SELECT sql FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes'""" + ).fetchone()[0] + original_rows = probe.execute( + "SELECT * FROM session_outcomes ORDER BY user_id, session_id" + ).fetchall() + + real_connect = sqlite3.connect + failing_connections = [] + + class _FailAfterSessionOutcomeRename(sqlite3.Connection): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.saw_session_outcome_rename = False + self.injected_failure = False + + def execute(self, sql, parameters=(), /): + normalized = " ".join(str(sql).split()) + if ( + normalized + == "ALTER TABLE session_outcomes RENAME TO session_outcomes_legacy" + ): + self.saw_session_outcome_rename = True + if self.saw_session_outcome_rename and normalized.startswith( + "CREATE TABLE session_outcomes (" + ): + self.injected_failure = True + raise RuntimeError("injected session_outcomes rebuild failure") + return super().execute(sql, parameters) + + def _connect_with_failure(*args, **kwargs): + kwargs["factory"] = _FailAfterSessionOutcomeRename + conn = real_connect(*args, **kwargs) + failing_connections.append(conn) + return conn + + monkeypatch.setattr(sqlite3, "connect", _connect_with_failure) + + with pytest.raises(RuntimeError, match="injected session_outcomes rebuild failure"): + SQLiteStorage(org_id="legacy-atomicity", db_path=db_path) + + assert any( + conn.saw_session_outcome_rename and conn.injected_failure + for conn in failing_connections + ) + for conn in failing_connections: + conn.close() + + with sqlite3.connect(db_path) as probe: + restored_schema = probe.execute( + """SELECT sql FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes'""" + ).fetchone()[0] + restored_rows = probe.execute( + "SELECT * FROM session_outcomes ORDER BY user_id, session_id" + ).fetchall() + stranded_legacy_table = probe.execute( + """SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'session_outcomes_legacy'""" + ).fetchone() + + assert restored_schema == original_schema + assert restored_rows == original_rows + assert stranded_legacy_table is None + + +@pytest.mark.parametrize( + ("with_subject_column", "stored_subject_ref"), + [(False, None), (True, None), (True, " ")], + 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 557326342..0dd748a1f 100644 --- a/tests/server/services/storage/test_playbook_aggregation_state_integration.py +++ b/tests/server/services/storage/test_playbook_aggregation_state_integration.py @@ -1102,23 +1102,9 @@ def test_incremental_run_refreshes_agent_and_centroid_after_match( request_context=context, agent_version="v1", ) - learning_meter = MagicMock() - monkeypatch.setattr(aggregator, "_record_learnings_generated", learning_meter) - first = aggregator.run( - PlaybookAggregatorRequest(agent_version="v1", operation_key="test-run-1") - ) + first = aggregator.run(PlaybookAggregatorRequest(agent_version="v1")) assert first["playbooks_generated"] == 1 - first_saved_id = store.conn.execute( - "SELECT agent_playbook_id FROM agent_playbooks" - ).fetchone()[0] - learning_meter.assert_called_once_with( - learning_ids=[str(first_saved_id)], - playbook_name="playbook", - request_id="test-run-1", - metadata=first, - total_count=1, - ) assert store.conn.execute("SELECT count(*) FROM agent_playbooks").fetchone()[0] == 1 _insert_current(store, 3, embedding=encoded, trigger="same trigger") diff --git a/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py b/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py index d6439ec30..0a6c64f16 100644 --- a/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py +++ b/tests/server/services/storage/test_playbook_optimization_replay_contract_integration.py @@ -14,6 +14,8 @@ from reflexio.models.api_schema import service_schemas as schemas from reflexio.server.services.playbook.publication import canonical_json_bytes from reflexio.server.services.storage.error import ( + OptimizationArtifactIntegrityError, + OptimizationJobIdentityConflictError, OptimizationJobLeaseLiveError, StorageError, ) @@ -113,10 +115,21 @@ def test_same_attempt_key_returns_one_active_job(storage: BaseStorage) -> None: def test_conflicting_active_identity_is_rejected(storage: BaseStorage) -> None: storage.create_or_get_playbook_optimization_job(_replay_job("d1", "a1")) - with pytest.raises(StorageError, match="immutable optimizer job identity"): + with pytest.raises( + OptimizationJobIdentityConflictError, + match="immutable optimizer job identity", + ): storage.create_or_get_playbook_optimization_job(_replay_job("d1", "a2")) +def test_sqlite_rejects_open_world_optimizer_jobs(storage: BaseStorage) -> None: + open_world_job = _replay_job("open-world-discovery", "open-world-attempt") + open_world_job.optimizer_kind = "offline_tuner_open_world" + + with pytest.raises(StorageError, match="CHECK constraint failed"): + storage.create_or_get_playbook_optimization_job(open_world_job) + + def test_gepa_publication_reclaim_contract_has_none_live_and_reclaimed_outcomes( storage: BaseStorage, ) -> None: @@ -348,7 +361,10 @@ def test_artifact_upsert_canonicalizes_equivalent_json_and_requires_digest_and_c fence=claim.fence, now=5_001, ) - with pytest.raises(StorageError, match="artifact digest conflict"): + with pytest.raises( + OptimizationArtifactIntegrityError, + match="artifact digest conflict", + ): storage.upsert_playbook_optimization_artifact( _artifact( job_id=job.job_id, @@ -359,6 +375,52 @@ def test_artifact_upsert_canonicalizes_equivalent_json_and_requires_digest_and_c ) +def test_malformed_persisted_artifact_raises_typed_integrity_error( + storage: BaseStorage, +) -> None: + job = storage.create_or_get_playbook_optimization_job(_replay_job("d1", "a1")) + claim = storage.claim_playbook_optimization_job( + job_id=job.job_id, + owner="worker-a", + lease_seconds=60, + now=5_000, + ) + saved = storage.upsert_playbook_optimization_artifact( + _artifact(job_id=job.job_id), + fence=claim.fence, + now=5_001, + ) + assert isinstance(storage, SQLiteStorage) + storage.conn.execute( + "UPDATE playbook_optimization_artifacts SET content_digest = ? " + "WHERE artifact_id = ?", + ("f" * 64, saved.artifact_id), + ) + storage.conn.commit() + + with pytest.raises( + OptimizationArtifactIntegrityError, + match="optimizer artifact row is malformed", + ): + storage.get_playbook_optimization_artifact( + job.job_id, + "expected_population_manifest", + ) + + +def test_artifact_storage_failures_remain_generic(storage: BaseStorage) -> None: + assert isinstance(storage, SQLiteStorage) + storage.conn.execute("DROP TABLE playbook_optimization_artifacts") + + with pytest.raises(StorageError) as raised: + storage.get_playbook_optimization_artifact( + 1, + "expected_population_manifest", + ) + + assert type(raised.value) is StorageError + + def test_artifact_model_rejects_malformed_json() -> None: with pytest.raises( ValidationError, match="artifact content_json must be valid JSON" @@ -429,6 +491,240 @@ def test_artifact_model_uses_publication_numeric_contract() -> None: ) +def test_previous_artifact_schema_is_upgraded_without_losing_constraints( + tmp_path: Path, +) -> None: + db_path = tmp_path / "previous-artifact-schema.db" + initial_store = SQLiteStorage( + org_id="previous-artifact-schema", db_path=str(db_path) + ) + initial_store.conn.close() + + legacy_content = '{"source":"legacy"}' + legacy_digest = sha256(legacy_content.encode()).hexdigest() + conn = sqlite3.connect(db_path) + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("DROP INDEX idx_poa_job") + conn.execute("DROP TABLE playbook_optimization_artifacts") + conn.executescript( + """ + CREATE TABLE playbook_optimization_artifacts ( + artifact_id INTEGER PRIMARY KEY AUTOINCREMENT, + job_id INTEGER NOT NULL, + artifact_kind TEXT NOT NULL CHECK (artifact_kind IN ( + 'expected_population_manifest', + 'generation_selection', + 'replay_manifest', + 'candidate', + 'candidate_search_projection' + )), + content_json TEXT NOT NULL, + content_digest TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (job_id, artifact_kind), + FOREIGN KEY (job_id) REFERENCES playbook_optimization_jobs(job_id) + ON DELETE CASCADE + ); + CREATE INDEX idx_poa_job ON playbook_optimization_artifacts(job_id); + """ + ) + conn.execute( + """INSERT INTO playbook_optimization_jobs ( + job_id, optimizer_kind, target_kind, target_id, status, + lease_owner, lease_fence, lease_expires_at, created_at, updated_at + ) VALUES ( + 41, 'offline_tuner_replay', 'user_playbook', 9, 'running', + 'worker-a', 3, 1000, 101, 102 + )""" + ) + conn.execute( + """INSERT INTO playbook_optimization_artifacts ( + artifact_id, job_id, artifact_kind, content_json, content_digest, + created_at, updated_at + ) VALUES (77, 41, 'candidate', ?, ?, 103, 104)""", + (legacy_content, legacy_digest), + ) + conn.commit() + conn.close() + + store = SQLiteStorage(org_id="previous-artifact-schema", db_path=str(db_path)) + try: + legacy_row = store.conn.execute( + "SELECT * FROM playbook_optimization_artifacts WHERE artifact_id = 77" + ).fetchone() + assert dict(legacy_row) == { + "artifact_id": 77, + "job_id": 41, + "artifact_kind": "candidate", + "content_json": legacy_content, + "content_digest": legacy_digest, + "created_at": 103, + "updated_at": 104, + } + assert store.conn.execute("PRAGMA foreign_key_check").fetchall() == [] + assert store.conn.execute("PRAGMA foreign_keys").fetchone()[0] == 1 + + with pytest.raises(sqlite3.IntegrityError): + store.conn.execute( + """INSERT INTO playbook_optimization_artifacts VALUES + (78, 41, 'candidate', '{}', ?, 105, 106)""", + (sha256(b"{}").hexdigest(),), + ) + store.conn.rollback() + assert ( + store.conn.execute( + """SELECT 1 FROM sqlite_master + WHERE type = 'index' AND name = 'idx_poa_job'""" + ).fetchone() + is not None + ) + with pytest.raises(sqlite3.IntegrityError): + store.conn.execute( + """INSERT INTO playbook_optimization_artifacts VALUES + (79, 41, 'unknown_kind', '{}', ?, 105, 106)""", + (sha256(b"{}").hexdigest(),), + ) + store.conn.rollback() + + evidence_json = '{"cases":[1]}' + evidence = schemas.PlaybookOptimizationArtifact( + job_id=41, + artifact_kind="open_world_evidence_bundle", + content_json=evidence_json, + content_digest=sha256(evidence_json.encode()).hexdigest(), + created_at=107, + updated_at=108, + ) + saved = store.upsert_playbook_optimization_artifact(evidence, fence=3, now=500) + assert ( + store.get_playbook_optimization_artifact(41, "open_world_evidence_bundle") + == saved + ) + + assert store.migrate() is True + store.conn.execute("DELETE FROM playbook_optimization_jobs WHERE job_id = 41") + store.conn.commit() + assert ( + store.conn.execute( + "SELECT COUNT(*) FROM playbook_optimization_artifacts WHERE job_id = 41" + ).fetchone()[0] + == 0 + ) + finally: + store.conn.close() + + assert not hasattr(BaseStorage, "load_open_world_evidence_snapshot") + assert not hasattr(SQLiteStorage, "load_open_world_evidence_snapshot") + + +def test_optimizer_job_rebuild_preserves_deleted_id_high_water_and_repeats( + tmp_path: Path, +) -> None: + db_path = tmp_path / "legacy-job-sequence.db" + _create_legacy_optimizer_schema(db_path) + conn = sqlite3.connect(db_path) + conn.execute("DELETE FROM playbook_optimization_jobs WHERE job_id IN (5, 6)") + conn.commit() + conn.close() + + first_store = SQLiteStorage(org_id="job-sequence-first", db_path=str(db_path)) + first = first_store.create_playbook_optimization_job(_replay_job("d1", "a1")) + first_store.conn.close() + second_store = SQLiteStorage(org_id="job-sequence-second", db_path=str(db_path)) + second = second_store.create_playbook_optimization_job( + _replay_job("d2", "a2").model_copy(update={"target_id": 42}) + ) + second_store.conn.close() + + assert first.job_id == 7 + assert second.job_id == 8 + + +def test_empty_optimizer_job_rebuild_preserves_deleted_id_high_water( + tmp_path: Path, +) -> None: + db_path = tmp_path / "empty-legacy-job-sequence.db" + _create_legacy_optimizer_schema(db_path) + conn = sqlite3.connect(db_path) + conn.execute("DELETE FROM playbook_optimization_jobs") + conn.commit() + conn.close() + + store = SQLiteStorage(org_id="empty-job-sequence", db_path=str(db_path)) + job = store.create_playbook_optimization_job(_replay_job("d1", "a1")) + store.conn.close() + + assert job.job_id == 7 + + +def test_artifact_rebuild_preserves_deleted_id_high_water_and_repeats( + tmp_path: Path, +) -> None: + db_path = tmp_path / "legacy-artifact-sequence.db" + store = SQLiteStorage(org_id="artifact-sequence-setup", db_path=str(db_path)) + parent = store.create_playbook_optimization_job(_replay_job("d1", "a1")) + store.conn.execute("DROP INDEX idx_poa_job") + store.conn.execute("DROP TABLE playbook_optimization_artifacts") + store.conn.executescript( + """ + CREATE TABLE playbook_optimization_artifacts ( + artifact_id INTEGER PRIMARY KEY AUTOINCREMENT, + job_id INTEGER NOT NULL, + artifact_kind TEXT NOT NULL CHECK (artifact_kind IN ('candidate')), + content_json TEXT NOT NULL, + content_digest TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (job_id, artifact_kind), + FOREIGN KEY (job_id) REFERENCES playbook_optimization_jobs(job_id) + ON DELETE CASCADE + ); + CREATE INDEX idx_poa_job ON playbook_optimization_artifacts(job_id); + """ + ) + digest = sha256(b"{}").hexdigest() + store.conn.execute( + "INSERT INTO playbook_optimization_artifacts VALUES " + "(9, ?, 'candidate', '{}', ?, 1, 1)", + (parent.job_id, digest), + ) + store.conn.execute("DELETE FROM playbook_optimization_artifacts") + store.conn.commit() + store.conn.close() + + first_store = SQLiteStorage(org_id="artifact-sequence-first", db_path=str(db_path)) + first_store.conn.execute( + "INSERT INTO playbook_optimization_artifacts " + "(job_id, artifact_kind, content_json, content_digest, created_at, updated_at) " + "VALUES (?, 'candidate', '{}', ?, 1, 1)", + (parent.job_id, digest), + ) + first_id = first_store.conn.execute( + "SELECT artifact_id FROM playbook_optimization_artifacts" + ).fetchone()[0] + first_store.conn.execute("DELETE FROM playbook_optimization_artifacts") + first_store.conn.commit() + first_store.conn.close() + + second_store = SQLiteStorage( + org_id="artifact-sequence-second", db_path=str(db_path) + ) + second_store.conn.execute( + "INSERT INTO playbook_optimization_artifacts " + "(job_id, artifact_kind, content_json, content_digest, created_at, updated_at) " + "VALUES (?, 'candidate', '{}', ?, 1, 1)", + (parent.job_id, digest), + ) + second_id = second_store.conn.execute( + "SELECT artifact_id FROM playbook_optimization_artifacts" + ).fetchone()[0] + second_store.conn.close() + + assert first_id == 10 + assert second_id == 11 + + @pytest.mark.parametrize( "current_stage", [ diff --git a/tests/server/services/storage/test_sqlite_storage.py b/tests/server/services/storage/test_sqlite_storage.py index 81584e44a..e187523cc 100644 --- a/tests/server/services/storage/test_sqlite_storage.py +++ b/tests/server/services/storage/test_sqlite_storage.py @@ -30,6 +30,7 @@ _true_rrf_merge, _vector_rank_rows, ) +from reflexio.server.services.storage.sqlite_storage import _base as sqlite_storage_base from reflexio.server.services.storage.sqlite_storage._base import ( _epoch_to_iso, _iso_to_epoch, @@ -57,6 +58,42 @@ def storage(): yield SQLiteStorage(org_id="0", db_path=f"{temp_dir}/reflexio.db") +def test_sqlite_storage_rejects_sqlite_before_returning_support( + tmp_path, monkeypatch +) -> None: + monkeypatch.setattr(sqlite_storage_base.sqlite3, "sqlite_version_info", (3, 34, 99)) + db_path = tmp_path / "missing" / "nested" / "old.db" + + with ( + patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512), + pytest.raises( + RuntimeError, + match=r"SQLite 3\.35\.0 or newer is required; detected 3\.34\.99", + ), + ): + SQLiteStorage(org_id="version-check", db_path=str(db_path)) + + assert not db_path.parent.exists() + assert not db_path.exists() + + +def test_sqlite_storage_accepts_sqlite_with_returning_support( + tmp_path, monkeypatch +) -> None: + monkeypatch.setattr(sqlite_storage_base.sqlite3, "sqlite_version_info", (3, 35, 0)) + + with patch.object(SQLiteStorage, "_get_embedding", return_value=[0.0] * 512): + storage = SQLiteStorage( + org_id="version-check", db_path=str(tmp_path / "supported.db") + ) + try: + assert storage.conn.execute( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'purge_operations'" + ).fetchone() + finally: + storage.conn.close() + + # --------------------------------------------------------------------------- # _sanitize_fts_query tests # --------------------------------------------------------------------------- diff --git a/tests/server/services/storage/test_storage_contract_clear_user_data.py b/tests/server/services/storage/test_storage_contract_clear_user_data.py index cf550ae22..1ef977535 100644 --- a/tests/server/services/storage/test_storage_contract_clear_user_data.py +++ b/tests/server/services/storage/test_storage_contract_clear_user_data.py @@ -24,6 +24,8 @@ UserPlaybook, UserProfile, ) +from reflexio.server.services.governance.config import governance_subject_ref +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.services.storage.storage_base import BaseStorage pytestmark = pytest.mark.integration @@ -170,6 +172,69 @@ def test_clear_unknown_user_is_noop(self, storage: BaseStorage) -> None: assert len(storage.get_user_playbooks(user_id="userA")) == 1 assert storage.get_request("req_a") is not None + def test_session_outcomes_use_authoritative_user_and_report_stable_zero( + self, storage: BaseStorage, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("REFLEXIO_GOVERNANCE_REF_SECRET", "test-governance-secret") + sqlite_storage = storage + assert isinstance(sqlite_storage, SQLiteStorage) + alice_ref = governance_subject_ref( + sqlite_storage.org_id, "alice", "test-governance-secret" + ) + bob_ref = governance_subject_ref( + sqlite_storage.org_id, "bob", "test-governance-secret" + ) + for outcome_id, user_id, subject_ref in ( + ("alice-stale", "alice", bob_ref), + ("bob-conflict", "bob", alice_ref), + ): + sqlite_storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES (?, 1, ?, ?, 'success', 100, 'test', ?, ?, ?, 101)""", + ( + outcome_id, + user_id, + f"session-{outcome_id}", + "a" * 64, + "b" * 64, + subject_ref, + ), + ) + sqlite_storage.conn.commit() + + counts = storage.clear_user_data("alice") + zero_counts = storage.clear_user_data("missing-user") + + remaining = sqlite_storage.conn.execute( + "SELECT outcome_id FROM session_outcomes ORDER BY outcome_id" + ).fetchall() + assert [row["outcome_id"] for row in remaining] == ["bob-conflict"] + assert counts["session_outcomes"] == 1 + assert zero_counts["session_outcomes"] == 0 + + def test_default_clear_user_data_preserves_session_outcome_count( + self, storage: BaseStorage + ) -> None: + sqlite_storage = storage + assert isinstance(sqlite_storage, SQLiteStorage) + sqlite_storage.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, outcome, + occurred_at, source, outcome_contract_digest, + finalized_trajectory_digest, governance_subject_ref, created_at + ) VALUES ('default-clear', 1, 'alice', 'default-clear-session', + 'success', 100, 'test', ?, ?, 'subref_v1_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', 101)""", + ("a" * 64, "b" * 64), + ) + sqlite_storage.conn.commit() + + counts = BaseStorage.clear_user_data(sqlite_storage, "alice") + + assert counts["session_outcomes"] == 1 + def test_returned_counts_match_seeded_rows(self, storage: BaseStorage) -> None: """Per-entity counts must reflect actual seeded row counts for the user.""" # Seed userA with two of each. diff --git a/tests/server/services/storage/test_storage_contract_retention.py b/tests/server/services/storage/test_storage_contract_retention.py index 6cdd17083..2612dc1c4 100644 --- a/tests/server/services/storage/test_storage_contract_retention.py +++ b/tests/server/services/storage/test_storage_contract_retention.py @@ -1,6 +1,8 @@ """Contract tests for generic row-retention storage methods.""" from datetime import UTC, datetime +from typing import Any, cast +from unittest.mock import patch import pytest @@ -188,6 +190,50 @@ def test_retention_request_cascade_cleans_interaction_fts( assert len(fts_kept) == 1, "fts row for surviving interaction must remain" +def test_retention_exposure_age_boundary_is_strict(storage: BaseStorage) -> None: + """Only exposure evidence strictly older than 14 days is row-cap eligible.""" + from reflexio.server.services.storage.retention import ( + OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS, + ) + + now = 2_000_000_000 + cutoff = now - OPEN_WORLD_EVIDENCE_RETENTION_WINDOW_SECONDS + conn = storage.conn # type: ignore[attr-defined] + conn.execute( + """CREATE TABLE user_playbook_exposure_events ( + exposure_event_id TEXT PRIMARY KEY, + ingested_at INTEGER NOT NULL + )""" + ) + conn.executemany( + """INSERT INTO user_playbook_exposure_events + (exposure_event_id, ingested_at) + VALUES (?, ?)""", + [ + ("older", cutoff - 1), + ("exact", cutoff), + ("newer", cutoff + 1), + ], + ) + conn.commit() + + retention_storage = cast(Any, storage) + with patch( + "reflexio.server.services.storage.retention_mixin.time.time", + return_value=now, + ): + deleted = retention_storage.delete_oldest_retention_target_rows( + "user_playbook_exposure_events", 3 + ) + + assert deleted == 1 + remaining = conn.execute( + "SELECT exposure_event_id FROM user_playbook_exposure_events " + "ORDER BY exposure_event_id" + ).fetchall() + assert [row["exposure_event_id"] for row in remaining] == ["exact", "newer"] + + # --------------------------------------------------------------------------- # Playbook retention FTS + vec cleanup (B3h) # --------------------------------------------------------------------------- diff --git a/tests/server/services/storage/test_storage_contract_session_outcomes.py b/tests/server/services/storage/test_storage_contract_session_outcomes.py index b07a2223b..6716214b4 100644 --- a/tests/server/services/storage/test_storage_contract_session_outcomes.py +++ b/tests/server/services/storage/test_storage_contract_session_outcomes.py @@ -1,11 +1,18 @@ """Session outcome storage contract.""" +from typing import cast + from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, Request, + SessionOutcomeFailureReason, SessionOutcomeKind, SetSessionOutcomeRequest, ) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._base import ( + _canonical_session_snapshot, +) from reflexio.server.services.storage.storage_base import BaseStorage @@ -39,14 +46,345 @@ def test_first_write_preserves_outcome_fields(storage: BaseStorage) -> None: assert first.recorded is True assert duplicate.recorded is False + assert duplicate.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION assert duplicate.source == "published" records = storage.get_session_outcomes(GetSessionOutcomesRequest(label="booked")) assert len(records) == 1 + assert records[0].outcome_id + assert records[0].outcome_revision == 1 + outcome_contract_digest = records[0].outcome_contract_digest + finalized_trajectory_digest = records[0].finalized_trajectory_digest + assert outcome_contract_digest is not None + assert finalized_trajectory_digest is not None + assert len(outcome_contract_digest) == 64 + assert len(finalized_trajectory_digest) == 64 assert records[0].outcome == SessionOutcomeKind.SUCCESS assert records[0].value == 12.0 assert records[0].metadata == {"crm": "test"} +def test_exact_finalization_retry_is_idempotent(storage: BaseStorage) -> None: + storage.add_request( + Request( + request_id="retry-r1", + user_id="u1", + session_id="exact-retry", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="exact-retry", + outcome=SessionOutcomeKind.UNKNOWN, + occurred_at=101, + metadata={"reason": "not enough information"}, + ) + context = storage.get_session_outcome_context("exact-retry") + + first = storage.record_session_outcome( + request, created_at=102, expected_context=context + ) + retry = storage.record_session_outcome( + request, created_at=103, expected_context=context + ) + + assert first.recorded is True + assert retry.recorded is False + assert retry.reason is None + assert retry.outcome_id == first.outcome_id + assert retry.outcome_revision == first.outcome_revision == 1 + assert retry.outcome_contract_digest == first.outcome_contract_digest + assert retry.finalized_trajectory_digest == first.finalized_trajectory_digest + + +def test_legacy_all_null_identity_exact_retry_uses_available_context( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="legacy-r1", + user_id="legacy-user", + session_id="legacy-retry", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="legacy-retry", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + metadata={"legacy": True}, + ) + first = storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("legacy-retry"), + ) + assert first.recorded is True + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + "CREATE TABLE legacy_session_outcomes AS SELECT * FROM session_outcomes" + ) + sqlite_storage.conn.execute("DROP TABLE session_outcomes") + sqlite_storage.conn.execute( + "ALTER TABLE legacy_session_outcomes RENAME TO session_outcomes" + ) + sqlite_storage.conn.execute( + """UPDATE session_outcomes + SET outcome_id = NULL, outcome_revision = NULL, + outcome_contract_digest = NULL, + finalized_trajectory_digest = NULL + WHERE session_id = ?""", + ("legacy-retry",), + ) + sqlite_storage.conn.commit() + + retry = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context("legacy-retry"), + ) + + assert retry.recorded is False + assert retry.reason is None + assert retry.outcome_id is None + assert retry.outcome_revision is None + assert retry.outcome_contract_digest is None + assert retry.finalized_trajectory_digest is None + + +def test_legacy_all_null_identity_changed_payload_still_conflicts( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="legacy-conflict-r1", + user_id="legacy-user", + session_id="legacy-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="legacy-conflict", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("legacy-conflict"), + ) + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + "CREATE TABLE legacy_session_outcomes AS SELECT * FROM session_outcomes" + ) + sqlite_storage.conn.execute("DROP TABLE session_outcomes") + sqlite_storage.conn.execute( + "ALTER TABLE legacy_session_outcomes RENAME TO session_outcomes" + ) + sqlite_storage.conn.execute( + """UPDATE session_outcomes + SET outcome_id = NULL, outcome_revision = NULL, + outcome_contract_digest = NULL, + finalized_trajectory_digest = NULL + WHERE session_id = ?""", + ("legacy-conflict",), + ) + sqlite_storage.conn.commit() + + retry = storage.record_session_outcome( + request.model_copy(update={"outcome": SessionOutcomeKind.FAILURE}), + created_at=103, + expected_context=storage.get_session_outcome_context("legacy-conflict"), + ) + + assert retry.recorded is False + assert retry.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + + +def test_legacy_all_null_identity_changed_governance_context_conflicts( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="legacy-governance-r1", + user_id="legacy-user", + session_id="legacy-governance-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="legacy-governance-conflict", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context( + "legacy-governance-conflict" + ), + ) + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + "CREATE TABLE legacy_session_outcomes AS SELECT * FROM session_outcomes" + ) + sqlite_storage.conn.execute("DROP TABLE session_outcomes") + sqlite_storage.conn.execute( + "ALTER TABLE legacy_session_outcomes RENAME TO session_outcomes" + ) + sqlite_storage.conn.execute( + """UPDATE session_outcomes + SET outcome_id = NULL, outcome_revision = NULL, + outcome_contract_digest = NULL, + finalized_trajectory_digest = NULL + WHERE session_id = ?""", + ("legacy-governance-conflict",), + ) + sqlite_storage.conn.execute( + "UPDATE requests SET governance_subject_ref = ? WHERE session_id = ?", + ( + sqlite_storage._subject_ref_for_user_id("different-user"), + "legacy-governance-conflict", + ), + ) + sqlite_storage.conn.commit() + + retry = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context( + "legacy-governance-conflict" + ), + ) + + assert retry.recorded is False + assert retry.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + + +def test_sqlite_canonical_snapshot_loads_interactions_in_one_query( + storage: BaseStorage, +) -> None: + sqlite_storage = cast(SQLiteStorage, storage) + for index in range(3): + storage.add_request( + Request( + request_id=f"snapshot-r{index}", + user_id="snapshot-user", + session_id="snapshot-session", + source="published", + created_at=100 + index, + ) + ) + statements: list[str] = [] + sqlite_storage.conn.set_trace_callback(statements.append) + try: + snapshot = _canonical_session_snapshot(sqlite_storage.conn, "snapshot-session") + finally: + sqlite_storage.conn.set_trace_callback(None) + + interaction_queries = [ + statement for statement in statements if "FROM interactions" in statement + ] + assert [item["request"]["request_id"] for item in snapshot["requests"]] == [ + "snapshot-r0", + "snapshot-r1", + "snapshot-r2", + ] + assert len(interaction_queries) == 1 + + +def test_changed_contract_identity_is_conflicting_finalization( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="contract-r1", + user_id="u1", + session_id="contract-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="contract-conflict", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + first = storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("contract-conflict"), + ) + sqlite_storage = cast(SQLiteStorage, storage) + sqlite_storage.conn.execute( + """UPDATE session_outcomes SET outcome_contract_digest = ? + WHERE session_id = ?""", + ("0" * 64, "contract-conflict"), + ) + sqlite_storage.conn.commit() + + changed_contract = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context("contract-conflict"), + ) + + assert first.recorded is True + assert changed_contract.recorded is False + assert ( + changed_contract.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + ) + + +def test_changed_session_trajectory_is_conflicting_finalization( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="trajectory-r1", + user_id="u1", + session_id="trajectory-conflict", + source="published", + created_at=100, + ) + ) + request = SetSessionOutcomeRequest( + session_id="trajectory-conflict", + outcome=SessionOutcomeKind.FAILURE, + occurred_at=101, + ) + first = storage.record_session_outcome( + request, + created_at=102, + expected_context=storage.get_session_outcome_context("trajectory-conflict"), + ) + storage.add_request( + Request( + request_id="trajectory-r2", + user_id="u1", + session_id="trajectory-conflict", + source="published", + created_at=103, + ) + ) + + changed_trajectory = storage.record_session_outcome( + request, + created_at=104, + expected_context=storage.get_session_outcome_context("trajectory-conflict"), + ) + + assert first.recorded is True + assert changed_trajectory.recorded is False + assert ( + changed_trajectory.reason + == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + ) + + def test_unknown_session_is_rejected(storage: BaseStorage) -> None: context = storage.get_session_outcome_context("missing") result = storage.record_session_outcome( diff --git a/tests/server/services/test_generation_billing_emission.py b/tests/server/services/test_generation_billing_emission.py index 6138f19f4..ec622bc1c 100644 --- a/tests/server/services/test_generation_billing_emission.py +++ b/tests/server/services/test_generation_billing_emission.py @@ -15,7 +15,14 @@ from typing import Any from unittest.mock import MagicMock, patch -from reflexio.models.api_schema.service_schemas import Interaction, Request +import pytest + +from reflexio.models.api_schema.domain.entities import UserProfile +from reflexio.models.api_schema.service_schemas import ( + Interaction, + Request, + UserPlaybook, +) from reflexio.models.config_schema import ( Config, ProfileExtractorConfig, @@ -28,6 +35,14 @@ BaseGenerationService, PreparedGenerationRun, ) +from reflexio.server.services.deferred_learning_plan import ( + PlaybookWritePlan, + ProfileWritePlan, +) +from reflexio.server.services.playbook.service import ( + PlaybookGenerationService, + PlaybookGenerationServiceConfig, +) from reflexio.server.services.profile.profile_generation_service_utils import ( ProfileGenerationRequest, ) @@ -38,6 +53,7 @@ from reflexio.server.services.service_utils import format_sessions_to_history_string from reflexio.server.services.storage.sqlite_storage import SQLiteStorage from reflexio.server.usage_metrics import UsageEvent, configure_usage_event_recorder +from reflexio.test_support.llm_mock import patched_litellm # --------------------------------------------------------------------------- # Helpers: build a RequestContext backed by a real SQLiteStorage, mirroring @@ -137,9 +153,9 @@ def _run_profile_generation(storage: SQLiteStorage, *, auto_run: bool = True) -> def test_real_extraction_emits_tokens_and_learnings(tmp_path): """A successful extraction emits extraction_tokens + learnings_generated. - MOCK_LLM_RESPONSE=true is active (autouse conftest), so the profile extractor - takes the deterministic mock path and we verify the billing events it triggers. - auto_run=False bypasses the should_run gate so extraction always fires. + The local LiteLLM patch keeps this deterministic even when the test is selected + alongside E2E paths. auto_run=False bypasses the should_run gate so extraction + always fires. """ events: list[UsageEvent] = [] configure_usage_event_recorder(events.append) @@ -149,7 +165,8 @@ def test_real_extraction_emits_tokens_and_learnings(tmp_path): org_id=_ORG_ID, db_path=str(tmp_path / "reflexio.db") ) _seed_interactions(storage) - _run_profile_generation(storage, auto_run=False) + with patched_litellm(): + _run_profile_generation(storage, auto_run=False) finally: configure_usage_event_recorder(None) @@ -165,12 +182,123 @@ def test_real_extraction_emits_tokens_and_learnings(tmp_path): assert tok.billing_input_tokens == tok.count_value assert tok.platform_llm is True # no api_key_config in the seeded Config - # learnings_generated.count_value must equal the existing generation_succeeded count. + # This fixture retains every generated output, so billing and success telemetry + # have equal counts here. gen = next(e for e in learning if e.event_name == "learnings_generated") succ = next(e for e in events if e.event_name == "generation_succeeded") assert gen.count_value == succ.count_value +def test_online_learning_bills_survivors_but_telemetry_counts_raw_results(tmp_path): + """Telemetry counts raw output while billing counts retained write-plan items.""" + storage = _build_sqlite_storage(tmp_path) + service = _build_profile_service(storage) + extracted_profiles = [ + UserProfile( + profile_id=profile_id, + user_id=_USER_ID, + content=profile_id, + last_modified_timestamp=1_000, + generated_from_request_id=_REQUEST_ID, + ) + for profile_id in ("retained", "dropped") + ] + write_plan = ProfileWritePlan( + user_id=_USER_ID, + request_id=_REQUEST_ID, + new_profiles=extracted_profiles[:1], + superseded_ids=[], + ) + + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object(service, "_prepare_generation_run", return_value=_prepared()), + patch.object( + service, "_execute_extractor", return_value=extracted_profiles + ), + patch.object(service, "_resolve_write_plan", return_value=write_plan), + patch.object(service, "_finalize_extraction_runs"), + patch.object(service, "_persist_write_plan"), + patch.object(service, "_extraction_input_text", return_value=""), + ): + plan = service.compute_generation(MagicMock()) + assert plan is not None + assert plan.generated_count == 2 + assert plan.billable_count == 1 + service.persist_generation(plan) + service.emit_generation_side_effects(plan) + finally: + configure_usage_event_recorder(None) + + billed = [event for event in events if event.event_name == "learnings_generated"] + assert [event.count_value for event in billed] == [1] + succeeded = [ + event for event in events if event.event_name == "generation_succeeded" + ] + assert [event.count_value for event in succeeded] == [2] + + +def test_online_playbook_bills_only_write_plan_survivors(tmp_path): + """Playbook candidates removed during write-plan resolution are not billable.""" + storage = _build_sqlite_storage(tmp_path) + context = _request_context(storage) + service = PlaybookGenerationService( + llm_client=LiteLLMClient(LiteLLMConfig(model="gpt-4o-mini")), + request_context=context, + ) + service.service_config = PlaybookGenerationServiceConfig( + request_id=_REQUEST_ID, + agent_version="v1", + user_id=_USER_ID, + source="api", + auto_run=True, + ) + extracted_playbooks = [ + UserPlaybook( + user_playbook_id=playbook_id, + agent_version="v1", + request_id=_REQUEST_ID, + content=f"content-{playbook_id}", + trigger=f"trigger-{playbook_id}", + ) + for playbook_id in (1, 2) + ] + write_plan = PlaybookWritePlan( + request_id=_REQUEST_ID, + output_pending_status=False, + skip_aggregation=True, + new_playbooks=extracted_playbooks[:1], + superseded_ids=[], + merge_groups=[], + ) + + events: list[UsageEvent] = [] + configure_usage_event_recorder(events.append) + try: + with ( + patch.object(service, "_prepare_generation_run", return_value=_prepared()), + patch.object( + service, "_execute_extractor", return_value=extracted_playbooks + ), + patch.object(service, "_resolve_write_plan", return_value=write_plan), + patch.object(service, "_finalize_extraction_runs"), + patch.object(service, "_persist_write_plan"), + patch.object(service, "_extraction_input_text", return_value=""), + patch.object(service, "_dispatch_playbook_schedulers"), + ): + plan = service.compute_generation(MagicMock()) + assert plan is not None + service.persist_generation(plan) + service.emit_generation_side_effects(plan) + finally: + configure_usage_event_recorder(None) + + billed = [event for event in events if event.event_name == "learnings_generated"] + assert [event.count_value for event in billed] == [1] + + def test_should_run_skip_emits_no_learning_billing(tmp_path, monkeypatch): """A should_run-gated skip emits NO extraction_tokens / learnings_generated. @@ -376,6 +504,25 @@ def test_non_learning_service_emits_no_learning_billing_events(): ) +def test_base_finalization_requires_receipt_aware_override_for_run_id(): + """A receipt-less service cannot finalize a resumable extraction run.""" + service = _StubService( + llm_client=LiteLLMClient(LiteLLMConfig(model="gpt-4o-mini")), + request_context=_make_minimal_request_context(), + ) + + with patch.object(service, "_process_results") as process_results: + assert service._finalize_extracted_items(["legacy-item"]) is None + process_results.assert_called_once_with([["legacy-item"]]) + + with pytest.raises(NotImplementedError, match=r"(?i)receipt-aware"): + service._finalize_extracted_items( + ["resumable-item"], finalization_run_id="run-1" + ) + + process_results.assert_called_once() + + # --------------------------------------------------------------------------- # Dedup: billing reuses the should-run gate's already-fetched window instead of # re-querying storage purely to recompute billing_input_tokens. diff --git a/tests/server/services/test_non_extraction_learning_metering.py b/tests/server/services/test_non_extraction_learning_metering.py index 3536b39cc..2b2c76ad6 100644 --- a/tests/server/services/test_non_extraction_learning_metering.py +++ b/tests/server/services/test_non_extraction_learning_metering.py @@ -1,18 +1,29 @@ from __future__ import annotations -from unittest.mock import MagicMock +from datetime import UTC, datetime +from unittest.mock import MagicMock, patch import pytest +from reflexio.models.api_schema.service_schemas import AgentPlaybook, UserPlaybook +from reflexio.models.config_schema import PlaybookAggregatorConfig, PlaybookConfig from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.services.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, AgentRunStatus, ) -from reflexio.server.usage_metrics import UsageEvent, configure_usage_event_recorder +from reflexio.server.usage_metrics import ( + UsageEvent, + UsageEventDeliveryStatus, + configure_usage_event_recorder, +) @pytest.fixture(autouse=True) @@ -31,6 +42,14 @@ def _request_context() -> RequestContext: return ctx +def _capture_events(events: list[UsageEvent]) -> None: + def recorder(event: UsageEvent) -> UsageEventDeliveryStatus: + events.append(event) + return UsageEventDeliveryStatus.APPENDED + + configure_usage_event_recorder(recorder) + + def _agent_run(*, extractor_kind: str) -> AgentRunRecord: return AgentRunRecord( id="run-1", @@ -44,77 +63,47 @@ def _agent_run(*, extractor_kind: str) -> AgentRunRecord: ), status=AgentRunStatus.FINALIZING, generation_request_snapshot={}, + created_at=datetime(2026, 8, 1, tzinfo=UTC), ) -def test_resumable_finalization_falls_back_when_items_lack_ids() -> None: - """Items with no durable id (e.g. plain objects) fall back to the - count-based aggregate event -- Task A3's documented fallback, since there - is no safe per-record id to key a dedup event on. - """ +def test_resumable_profile_bills_only_ids_returned_by_finalization() -> None: + """A preassigned candidate ID is not billed when finalization drops it.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), ) run = _agent_run(extractor_kind="profile") - 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:") - + 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]) -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(), + finalizer.assert_called_once_with( + [dropped_candidate], model_provenance=None, finalization_run_id=run.id ) - 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", - ] + assert events == [] def test_resumable_finalization_emits_one_event_per_profile_id() -> None: - """When every item carries a durable ``profile_id`` (the common case -- - profile ids are assigned by the extractor before finalize runs), emit one - entity-backed event per profile instead of the count-only fallback. - """ + """Finalization survivor IDs emit one entity-backed event per profile.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), ) run = _agent_run(extractor_kind="profile") - class _FakeProfile: - def __init__(self, profile_id: str) -> None: - self.profile_id = profile_id - worker._record_finalized_learnings( run, - [_FakeProfile("prof-1"), _FakeProfile("prof-2")], + ["prof-1", "prof-2"], entity_type="profile", ) @@ -134,22 +123,18 @@ def __init__(self, profile_id: str) -> None: def test_resumable_finalization_emits_one_event_per_playbook_id() -> None: - """Same as above for the playbook (``user_playbook_id``) kind.""" + """Finalization survivor IDs emit one entity-backed event per playbook.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), ) run = _agent_run(extractor_kind="playbook") - class _FakePlaybook: - def __init__(self, user_playbook_id: int) -> None: - self.user_playbook_id = user_playbook_id - worker._record_finalized_learnings( run, - [_FakePlaybook(11), _FakePlaybook(12), _FakePlaybook(13)], + ["11", "12", "13"], entity_type="user_playbook", ) @@ -166,106 +151,104 @@ def __init__(self, user_playbook_id: int) -> None: assert event.event_key == f"learn:{event.entity_type}:{event.entity_id}" -def test_resumable_finalization_falls_back_when_a_playbook_id_is_unset() -> None: - """A ``user_playbook_id=0`` (default, unset) mixed in with real ids means - dedup dropped that item before persist -- fall back to the count-based - aggregate rather than emit a colliding ``learn:0`` key or fabricate an id. - """ +def test_resumable_playbook_bills_consolidation_replacement_id() -> None: + """Billing follows the persisted replacement, not its input candidate.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) worker = ExtractionResumeWorker( request_context=_request_context(), llm_client=MagicMock(), ) run = _agent_run(extractor_kind="playbook") - class _FakePlaybook: - def __init__(self, user_playbook_id: int) -> None: - self.user_playbook_id = user_playbook_id + 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]) - worker._record_finalized_learnings( - run, - [_FakePlaybook(21), _FakePlaybook(0)], - entity_type="user_playbook", + finalizer.assert_called_once_with( + [original_candidate], + model_provenance=None, + extraction_run=run, + finalization_run_id=run.id, ) - assert len(events) == 1 - assert events[0].count_value == 2 # total unchanged vs old count=2 - assert events[0].event_key is not None and events[0].event_key.startswith( - "learn-batch:" - ) - + assert events[0].count_value == 1 + assert events[0].event_key == "learn:user_playbook:88" + assert events[0].entity_id == "88" -def test_aggregation_records_attributed_learnings_generated() -> None: - """Aggregation emits one entity-backed event per generated playbook. - ``saved_playbook_list`` entries always carry a real ``agent_playbook_id`` - (``save_agent_playbooks`` raises rather than - returning a partial row) -- aggregator.py is the one caller with a clean, - always-populated per-record id list, so it uses the entity-backed path - (Task A3) rather than the count-only fallback. - """ +def test_aggregation_emits_no_learnings_generated( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A completed aggregation remains observable but adds no billable learning.""" events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) + _capture_events(events) + request_context = _request_context() + storage = MagicMock() + configurator = MagicMock() + request_context.storage = storage + request_context.configurator = configurator aggregator = PlaybookAggregator( llm_client=MagicMock(), - request_context=_request_context(), + request_context=request_context, agent_version="v1", ) - - aggregator._record_learnings_generated( - learning_ids=["101", "102", "103"], - playbook_name="agent_rules", - request_id="agg-run-1", - metadata={"playbooks_generated": 3}, + config = MagicMock() + configurator.get_config.return_value = config + config.user_playbook_extractor_config = PlaybookConfig( + extractor_name="billing-boundary", + extraction_definition_prompt="Extract user playbooks.", + aggregation_config=PlaybookAggregatorConfig( + min_cluster_size=2, + reaggregation_trigger_count=2, + ), ) - - assert len(events) == 3 - assert sum(e.count_value for e in events) == 3 # total unchanged vs old count=3 - assert {e.event_key for e in events} == { - "learn:agent_playbook:101", - "learn:agent_playbook:102", - "learn:agent_playbook:103", - } - assert {e.entity_id for e in events} == {"101", "102", "103"} - for event in events: - assert event.event_name == "learnings_generated" - assert event.count_value == 1 - assert event.pipeline == "playbook" - assert event.source == "aggregation" - assert event.entity_type == "agent_playbook" - assert event.agent_version == "v1" - assert event.playbook_name == "agent_rules" - # tie key<->entity together so a swapped association can't pass on sets alone - assert event.event_key == f"learn:{event.entity_type}:{event.entity_id}" - - -def test_aggregation_falls_back_when_an_agent_playbook_id_is_falsy() -> None: - """Whole-branch-review finding (2): a falsy/0 ``agent_playbook_id`` mixed - into the run must not mint a colliding ``learn:agent_playbook:0`` key -- - fall back to the count-based aggregate event instead, matching - ``ExtractionResumeWorker``'s guard for the same failure mode. - """ - events: list[UsageEvent] = [] - configure_usage_event_recorder(events.append) - aggregator = PlaybookAggregator( - llm_client=MagicMock(), - request_context=_request_context(), + user_playbooks = [ + UserPlaybook( + user_playbook_id=1, + agent_version="v1", + request_id="request-1", + playbook_name="user_playbook", + content="Document deployment decisions.", + ), + UserPlaybook( + user_playbook_id=2, + agent_version="v1", + request_id="request-2", + playbook_name="user_playbook", + content="Verify deployment outcomes.", + ), + ] + generated = AgentPlaybook( + agent_playbook_id=101, + playbook_name="user_playbook", agent_version="v1", + content="Deploy changes and verify results.", ) - - aggregator._record_learnings_generated( - learning_ids=["201"], # one id missing relative to total_count=2 - playbook_name="agent_rules", - request_id="agg-run-2", - metadata={"playbooks_generated": 2}, - total_count=2, + storage.count_user_playbooks.return_value = len(user_playbooks) + storage.get_agent_playbooks.return_value = [] + storage.get_user_playbooks.return_value = user_playbooks + storage.save_agent_playbooks.return_value = [generated] + monkeypatch.setattr( + aggregator, + "get_clusters", + lambda *_args: {0: user_playbooks}, ) - - assert len(events) == 1 - assert events[0].count_value == 2 # total unchanged vs old count=2 - assert events[0].event_key is not None and events[0].event_key.startswith( - "learn-batch:" + monkeypatch.setattr( + aggregator, + "_generate_playbooks_with_source_clusters", + lambda *_args, **_kwargs: [(generated, user_playbooks, None)], ) - assert events[0].entity_type == "agent_playbook" - assert events[0].source == "aggregation" + monkeypatch.setattr( + aggregator, "_enqueue_playbook_optimization", lambda _items: None + ) + + stats = aggregator.run(PlaybookAggregatorRequest(agent_version="v1", rerun=True)) + + assert stats["playbooks_generated"] == 1 + assert any(event.event_name == "aggregation_succeeded" for event in events) + assert not any(event.event_name == "learnings_generated" for event in events) diff --git a/tests/server/services/test_search_exposure.py b/tests/server/services/test_search_exposure.py new file mode 100644 index 000000000..546f00b53 --- /dev/null +++ b/tests/server/services/test_search_exposure.py @@ -0,0 +1,303 @@ +"""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, + governance_subject_ref="subject:user-1", + retired_at=1_700_000_050, + ) + + +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) + ) + + +def test_internal_fields_are_excluded_from_user_playbook_serialization() -> None: + serialized = _playbook().model_dump(mode="json") + + assert "governance_subject_ref" not in serialized + assert "retired_at" not in serialized + + +# Every current UserPlaybook field is persisted except its derived embedding vector. +_PERSISTED_FIELD_CHANGES = [ + ("user_playbook_id", 102), + ("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), + ("governance_subject_ref", "subject:user-2"), + ("retired_at", 1_700_000_051), +] + + +@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.py b/tests/server/test_billing_meter.py index 740d8d477..416bbbc6e 100644 --- a/tests/server/test_billing_meter.py +++ b/tests/server/test_billing_meter.py @@ -1,10 +1,18 @@ -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +import pytest from reflexio.server.billing_meter import ( + ReceiptBillingDeliveryError, + emit_learnings_generated_records_strict, record_applied_learnings, record_extraction_tokens, record_learnings_generated, ) +from reflexio.server.usage_metrics import ( + UsageEventDeliveryError, + UsageEventDeliveryStatus, +) HOOK = "reflexio.server.billing_meter.record_usage_event" @@ -110,3 +118,29 @@ def test_record_applied_learnings_noop_for_empty_result(): platform_storage=None, ) hook.assert_not_called() + + +@pytest.mark.parametrize( + "status", + [UsageEventDeliveryStatus.FAILED, UsageEventDeliveryStatus.REJECTED], +) +def test_strict_receipt_billing_preserves_delivery_status(status): + configurator = MagicMock() + configurator.get_config.return_value = None + + with ( + patch( + "reflexio.server.billing_meter.record_usage_event_strict", + side_effect=UsageEventDeliveryError(status), + ), + pytest.raises(ReceiptBillingDeliveryError) as exc_info, + ): + emit_learnings_generated_records_strict( + org_id="org1", + configurator=configurator, + learning_ids=["profile-1"], + source="resumable_extraction", + entity_type="profile", + ) + + assert exc_info.value.status is status diff --git a/tests/server/test_billing_meter_events.py b/tests/server/test_billing_meter_events.py index b9ab61401..0012b753d 100644 --- a/tests/server/test_billing_meter_events.py +++ b/tests/server/test_billing_meter_events.py @@ -13,12 +13,11 @@ mint the same ``event_key`` and collapse into one event downstream. The existing count-based ``record_learnings_generated`` / ``emit_learnings_generated`` -remain as the documented FALLBACK for callers that genuinely lack a per-record -id list (e.g. dedup/consolidation can reduce the persisted count below the raw -extracted count, so there is no safe 1:1 id per unit of ``count``). The -fallback path now also carries a synthesized ``event_key=f"learn-batch:{uuid4()}"`` -so every ``learnings_generated`` event -- record-backed or batch -- has a -dedup key. +remain for online extraction callers that have a known billable count but do +not retain per-record ids. Resumable finalization uses only the record-backed +path and skips items without durable ids. The count-based helper carries a +synthesized ``event_key=f"learn-batch:{uuid4()}"`` so every +``learnings_generated`` event -- record-backed or batch -- has a dedup key. Totals are preserved in both paths: the sum of ``count_value`` across the per-record events equals ``len(learning_ids)``; the fallback emits exactly diff --git a/tests/server/test_create_app_capabilities.py b/tests/server/test_create_app_capabilities.py index 45a1e23da..95ae4c542 100644 --- a/tests/server/test_create_app_capabilities.py +++ b/tests/server/test_create_app_capabilities.py @@ -15,8 +15,13 @@ from fastapi import APIRouter from fastapi.testclient import TestClient +from reflexio.server import usage_metrics from reflexio.server.api import create_app from reflexio.server.extensions import AppContext, Capability, CapabilityRegistry +from reflexio.server.usage_metrics import ( + UsageEventDeliveryError, + UsageEventDeliveryStatus, +) class RouterCap(Capability): @@ -92,6 +97,30 @@ def test_capabilities_none_is_unchanged() -> None: assert app is not None +def test_oss_startup_failure_clears_exempt_usage_recorder(monkeypatch) -> None: + usage_metrics.configure_usage_event_recorder(None) + + def fail_startup_guard() -> None: + raise RuntimeError("startup guard failed") + + monkeypatch.setattr( + "reflexio.server.llm.model_defaults.validate_llm_availability", + fail_startup_guard, + ) + + app = create_app(capabilities=None, mount_data_plane=True) + with pytest.raises(RuntimeError, match="startup guard failed"), TestClient(app): + pass + + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + ) + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + + def test_partial_cleanup_invariant() -> None: """Already-started caps must be shut down; caps that never started must not be.""" LifecycleCap.started = LifecycleCap.stopped = False diff --git a/tests/server/test_usage_metrics.py b/tests/server/test_usage_metrics.py index 5bf359332..823503404 100644 --- a/tests/server/test_usage_metrics.py +++ b/tests/server/test_usage_metrics.py @@ -1,5 +1,13 @@ +import logging + +import pytest + from reflexio.server import usage_metrics -from reflexio.server.usage_metrics import UsageEvent +from reflexio.server.usage_metrics import ( + UsageEvent, + UsageEventDeliveryError, + UsageEventDeliveryStatus, +) def test_usage_event_carries_event_key(): @@ -18,5 +26,127 @@ def test_usage_event_carries_event_key(): assert captured and captured[0].event_key == "search:abc" +def test_ordinary_delivery_is_silent_without_recorder(caplog): + usage_metrics.configure_usage_event_recorder(None) + + with caplog.at_level(logging.WARNING, logger=usage_metrics.__name__): + usage_metrics.record_usage_event( + org_id="7", + event_name="search_request", + event_category="application", + ) + + assert caplog.records == [] + + +def test_ordinary_delivery_is_silent_for_legacy_recorder(caplog): + captured = [] + usage_metrics.configure_usage_event_recorder(captured.append) + try: + with caplog.at_level(logging.WARNING, logger=usage_metrics.__name__): + usage_metrics.record_usage_event( + org_id="7", + event_name="search_request", + event_category="application", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert len(captured) == 1 + assert caplog.records == [] + + def test_event_key_defaults_none(): assert UsageEvent(org_id="7", event_name="x", event_category="y").event_key is None + + +def test_strict_delivery_rejects_legacy_none_recorder_as_unknown(): + captured = [] + usage_metrics.configure_usage_event_recorder(captured.append) + try: + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + assert [event.event_key for event in captured] == ["learn:profile:1"] + + +def test_strict_delivery_rejects_missing_recorder_as_unknown(): + usage_metrics.configure_usage_event_recorder(None) + + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + + +def test_strict_delivery_accepts_explicit_deployment_exemption(): + usage_metrics.configure_usage_event_recorder( + usage_metrics.exempt_usage_event_recorder + ) + try: + outcome = usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert outcome is UsageEventDeliveryStatus.EXEMPT + + +@pytest.mark.parametrize( + "outcome", + [UsageEventDeliveryStatus.FAILED, UsageEventDeliveryStatus.REJECTED], +) +def test_strict_delivery_raises_for_unaccepted_outcome(outcome): + usage_metrics.configure_usage_event_recorder(lambda _event: outcome) + try: + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None) + + assert exc_info.value.status is outcome + + +def test_strict_delivery_propagates_recorder_exception_while_ordinary_stays_fail_open(): + def fail(_event): + raise RuntimeError("sink unavailable") + + usage_metrics.configure_usage_event_recorder(fail) + try: + usage_metrics.record_usage_event( + org_id="7", + event_name="search_request", + event_category="application", + event_key="search:1", + ) + with pytest.raises(RuntimeError, match="sink unavailable"): + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + event_key="learn:profile:1", + ) + finally: + usage_metrics.configure_usage_event_recorder(None)