diff --git a/AI_AGENT_INTEGRATION.md b/AI_AGENT_INTEGRATION.md index 4807eb2a..e9538ec4 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,13 +234,17 @@ 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` | | `force_extraction` | No | `False` for normal background publishing. `True` for manual learn-now, tests, or final flushes where you intentionally want extraction to run immediately. | `false` | | `wait_for_response` | SDK/query option | `False` on interactive paths. `True` only when the caller is prepared to wait for extraction results. | `false` | +The source contract applies to new publish and filter inputs. Read responses may +return historical source values created before this contract; Reflexio preserves +those values exactly because source participates in session-outcome identity. + Each interaction row should resemble Reflexio's `InteractionData` shape: ```json @@ -315,7 +319,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 035ba2b1..31d820b0 100644 --- a/README.md +++ b/README.md @@ -105,8 +105,15 @@ Publish conversations from your agent, and Reflexio closes the self-improvement | --- | --- | | [Python](https://www.python.org/) >= 3.12 | Required for PyPI and source installs | | [uv](https://docs.astral.sh/uv/getting-started/installation/) | Required when running from source | +| Python's linked SQLite runtime | >= 3.35.0 (required for local SQLite storage; not the standalone `sqlite3` CLI) | | [Node.js](https://nodejs.org/) >= 18 | Required only for the local docs site in a source checkout | +Check the SQLite runtime linked to Python with: + +```shell +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 4efbe928..9b7be37e 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 76e2004b..b6347072 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 36333c5f..d7d88f0a 100644 --- a/docs/lib/methods/user-playbooks.ts +++ b/docs/lib/methods/user-playbooks.ts @@ -70,6 +70,18 @@ export const userPlaybookMethods: MethodDef[] = [ required: false, description: "Filter by user (via request_id linkage to requests table)", }, + { + name: "request_id", + type: "string", + required: false, + description: "Caller correlation ID, at most 255 characters", + }, + { + name: "session_id", + type: "string", + required: false, + description: "Caller session ID, at most 255 characters", + }, { name: "agent_version", type: "string", @@ -105,7 +117,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 c0680c50..eb420b8b 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/cli/utils.py b/reflexio/cli/utils.py index 27d32ca0..97ef3b9f 100644 --- a/reflexio/cli/utils.py +++ b/reflexio/cli/utils.py @@ -504,8 +504,24 @@ def run_services( for path in supervisor.stop_request_paths.values(): remove_pidfile(path) + service_names = {svc.name for svc in services} + gate_local_embedding = {"embedding", "backend"}.issubset(service_names) + try: + if gate_local_embedding: + embedding = next(svc for svc in services if svc.name == "embedding") + supervisor.start_service(embedding) + if not _wait_for_all_ready( + {"embedding": supervisor.ready_events["embedding"]}, + {"embedding": supervisor.processes["embedding"]}, + ): + raise RuntimeError( + "embedding service did not become ready before backend startup" + ) + for svc in services: + if gate_local_embedding and svc.name == "embedding": + continue supervisor.start_service(svc) supervisor.write_current_pidfile() except (OSError, RuntimeError): diff --git a/reflexio/client/client.py b/reflexio/client/client.py index 71a2d869..b16acb1b 100644 --- a/reflexio/client/client.py +++ b/reflexio/client/client.py @@ -528,7 +528,9 @@ def publish_interaction( Args: user_id: The user ID. interactions: List of interaction data. - source: The source of the interaction. + source: Non-sensitive producer/workflow label. A non-empty value + must match ``^[a-z0-9][a-z0-9._:-]{0,127}$`` and must not + contain user identifiers or PII. agent_version: The agent version. session_id: Required non-empty session ID for grouping requests. wait_for_response: If True, the **server** waits for @@ -893,6 +895,8 @@ def search_user_playbooks( threshold: float | None = None, enable_reformulation: bool | None = None, search_mode: SearchMode | None = None, + request_id: str | None = None, + session_id: str | None = None, ) -> SearchUserPlaybooksViewResponse: """Search for user playbooks with semantic/text search and filtering. @@ -906,10 +910,14 @@ 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) + request_id (Optional[str]): Caller correlation ID for the search turn, + at most 255 characters. + session_id (Optional[str]): Caller session ID for the search turn, + at most 255 characters. Returns: SearchUserPlaybooksViewResponse: Response containing matching user playbooks @@ -929,6 +937,8 @@ def search_user_playbooks( threshold=threshold, enable_reformulation=enable_reformulation, search_mode=search_mode, + request_id=request_id, + session_id=session_id, ) response = self._make_request( "POST", "/api/search_user_playbooks", json=req.model_dump(mode="json") @@ -1210,13 +1220,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, @@ -2822,12 +2839,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". @@ -2843,9 +2862,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 d49cbcea..7a39ba51 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 295f0ac2..9e61513f 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -26,6 +26,8 @@ from ..validators import ( EmbeddingVector, NonEmptyStr, + PersistedSessionOutcomeSource, + SessionOutcomeSource, TimeRangeValidatorMixin, _validate_image_url, ) @@ -278,7 +280,8 @@ class Request(BaseModel): user_id (str): Owner of the request. created_at (int): Unix epoch seconds at request creation. Defaults to the current UTC time. - source (str): Free-form origin tag (integration name, etc.). + source (str): Producer/workflow label. Persisted reads preserve legacy + values verbatim; new publish inputs use the strict source contract. agent_version (str): The agent version that handled this request. session_id (str): Non-empty session this request belongs to. evaluation_only (bool): Whether this request is stored for @@ -293,7 +296,7 @@ class Request(BaseModel): request_id: str user_id: str created_at: int = Field(default_factory=lambda: int(datetime.now(UTC).timestamp())) - source: str = "" + source: PersistedSessionOutcomeSource = "" agent_version: str = "" session_id: NonEmptyStr evaluation_only: bool = False @@ -364,6 +367,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 +404,7 @@ class AgentPlaybook(BaseModel): OptimizerKind = Literal[ "gepa", "offline_tuner_replay", + "offline_tuner_open_world", "offline_tuner_legacy", "optimizer_legacy_unknown", ] @@ -439,6 +445,7 @@ class AgentPlaybook(BaseModel): "replay_manifest", "candidate", "candidate_search_projection", + "open_world_evidence_bundle", ] Sha256Digest = str @@ -825,16 +832,47 @@ class DeleteSessionResponse(BaseModel): class SessionOutcomeRecord(BaseModel): + """Persisted outcome row, including its immutable historical source.""" + + outcome_id: NonEmptyStr | None = None + outcome_revision: int | None = Field(default=None, ge=1) user_id: str session_id: NonEmptyStr outcome: SessionOutcomeKind occurred_at: int = Field(ge=0) - source: str + source: PersistedSessionOutcomeSource label: str | None = Field(default=None, max_length=128) value: float | None = Field(default=None, allow_inf_nan=False) metadata: dict[str, Any] | None = None + 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 +917,42 @@ class SetSessionOutcomeResponse(BaseModel): reason: SessionOutcomeFailureReason | None = None message: str = "" user_id: str | None = None - source: str | None = None + source: PersistedSessionOutcomeSource | None = None + outcome_id: NonEmptyStr | None = None + outcome_revision: int | None = Field(default=None, ge=1) + outcome_contract_digest: Sha256Digest | None = None + 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 +1023,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 +1249,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 +1737,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 +1764,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 +1784,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 +1810,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 41c99be3..6ea0d9fa 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 eab73202..b34be97a 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 50bf9d9d..008acc91 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, TypeAdapter # Embedding vector dimensions — must match config_schema.EMBEDDING_DIMENSIONS. # Duplicated here to avoid circular imports (config_schema imports from this module). @@ -104,6 +104,28 @@ 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.""" + +PersistedSessionOutcomeSource = str +"""Historical outcome source returned exactly as stored, without new-input validation.""" + +_SESSION_OUTCOME_SOURCE_ADAPTER = TypeAdapter(SessionOutcomeSource) + + +def validate_session_outcome_source(value: str) -> SessionOutcomeSource: + """Validate an outcome source before writing a new request.""" + return _SESSION_OUTCOME_SOURCE_ADAPTER.validate_python(value) + + # ============================================================================= # Security Validators — SSRF Prevention # ============================================================================= diff --git a/reflexio/server/README.md b/reflexio/server/README.md index 35707f65..b4b8d957 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 5117b2aa..ed860fe1 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 4e0380d5..4a5cdf6a 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/middleware.py b/reflexio/server/middleware.py index 86331792..40bbec92 100644 --- a/reflexio/server/middleware.py +++ b/reflexio/server/middleware.py @@ -23,9 +23,11 @@ # Bot protection configuration REQUEST_TIMEOUT_SECONDS = 60 SYNC_REQUEST_TIMEOUT_SECONDS = ( - 600 # Longer timeout for synchronous processing (wait_for_response=true) + 600 # Longer timeout for synchronous long-running processing. +) +SYNC_REQUEST_PATHS = frozenset( + {"/api/review_user_playbooks", "/api/run_playbook_aggregation"} ) -SYNC_REQUEST_PATHS = frozenset({"/api/review_user_playbooks"}) SUSPICIOUS_USER_AGENTS = ["bot", "crawler", "spider", "scraper", "curl", "wget"] ALLOWED_EMPTY_UA_PATHS = ["/health", "/"] # Paths that allow empty user agents DEFAULT_MAX_BODY_BYTES = 10 * 1024 * 1024 diff --git a/reflexio/server/routes/search.py b/reflexio/server/routes/search.py index cb83bf6b..7253a757 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 caller_type == "production_agent" and 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,17 @@ def run_search() -> Any: rehydrated_text=response.rehydrated_text, experiment=assignment, ) + if caller_type == "production_agent": + record_search_exposures( + SearchExposureBatch( + org_id=org_id, + request_id=payload.request_id, + session_id=payload.session_id, + interaction_id=payload.interaction_id, + user_id=payload.user_id, + user_playbooks=tuple(response.user_playbooks), + ) + ) enqueue_search_metering( org_id=org_id, caller_type=caller_type, diff --git a/reflexio/server/services/README.md b/reflexio/server/services/README.md index 2e2062aa..cce70010 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 1b1304f9..5aee8a7d 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 2334d374..89a46cbd 100644 --- a/reflexio/server/services/base_generation_service.py +++ b/reflexio/server/services/base_generation_service.py @@ -29,6 +29,7 @@ ) from reflexio.server.services.deferred_learning_plan import ( ExtractorBookmarkAdvance, + FinalizationResult, GenerationComputePlan, ) from reflexio.server.services.extractor_config_utils import ( @@ -210,9 +211,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 +343,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: @@ -578,7 +588,11 @@ def _run_generation(self, request: TRequest) -> None: plan = self.compute_generation(request) if plan is None: return - self.persist_generation(plan) + try: + self.persist_generation(plan) + except Exception as exc: + self._mark_extraction_runs_finalization_failed(exc) + raise self.emit_generation_side_effects(plan) except Exception as e: self._record_generation_event( @@ -599,11 +613,11 @@ def _run_generation(self, request: TRequest) -> None: def compute_generation(self, request: TRequest) -> GenerationComputePlan | None: """Compute half of one generation run — NO learning DB write, NO fence. - Runs the prepare gate, the extractor (thread-pool LLM tool-loop), dedup - + embedding resolution (``_resolve_write_plan``), and drives the - ``agent_run`` rows to their terminal state (``_finalize_extraction_runs`` - — agent_run only, §4.3). Snapshots the billing inputs onto the returned - plan so ``emit_generation_side_effects`` never depends on the reused + Runs the prepare gate, the extractor (thread-pool LLM tool-loop), and + dedup + embedding resolution (``_resolve_write_plan``). Agent runs stay + non-terminal until receipt-aware persistence commits and + ``emit_generation_side_effects`` finalizes them. Snapshots the billing + inputs onto the returned plan so emit does not depend on the reused instance's mutable ``_last_*``. Returns a resolved ``GenerationComputePlan`` for persist/emit, or @@ -634,18 +648,24 @@ 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 - self._finalize_extraction_runs() except Exception as exc: 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, @@ -667,6 +687,31 @@ def persist_generation(self, plan: GenerationComputePlan) -> None: atomically with the row writes it corresponds to (F1) — and it is applied in exactly one place, so it is never double-applied. """ + if self.EMITS_LEARNING_BILLING and plan.extraction_run_ids: + if len(plan.extraction_run_ids) != 1: + raise RuntimeError( + "Receipt-aware generation requires exactly one extraction run" + ) + finalization_run_id = plan.extraction_run_ids[0] + extraction_run = ( + self.storage.get_agent_run(finalization_run_id) + if self.storage is not None + else None + ) + if extraction_run is None: + raise RuntimeError( + "Receipt-aware generation extraction run could not be resolved" + ) + # A pending-tool run has not produced its final learning set yet. + # Its resume worker must remain free to persist that set and create + # the immutable receipt after the tool dependency resolves. + if not extraction_run.pending_tool_call_ids: + plan.finalization_result = self._finalize_write_plan_with_outcome( + plan.write_plan, + finalization_run_id=finalization_run_id, + bookmark_advance=plan.bookmark_advance, + ) + return if plan.write_plan is not None: self._persist_write_plan(plan.write_plan) self._apply_bookmark_advance(plan.bookmark_advance) @@ -676,7 +721,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`` @@ -693,6 +739,16 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: is exactly as compute left it when emit runs. The plan still snapshots the billing inputs so a future durable reuse can re-plumb locally. """ + try: + self._finalize_extraction_runs() + except Exception as exc: + self._mark_extraction_runs_finalization_failed(exc) + raise + if ( + plan.finalization_result is not None + and not plan.finalization_result.won_receipt + ): + return self._record_generation_event( event_name="generation_succeeded", outcome="success", @@ -708,7 +764,19 @@ 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 + ) + + def _finalize_write_plan_with_outcome( + self, + write_plan: Any, + *, + finalization_run_id: str, + bookmark_advance: ExtractorBookmarkAdvance | None, + ) -> FinalizationResult: + """Persist one precomputed inline plan through a receipt-aware finalizer.""" + raise NotImplementedError( + "Receipt-aware plan finalization must be implemented by learning services" ) @abstractmethod diff --git a/reflexio/server/services/deferred_learning_plan.py b/reflexio/server/services/deferred_learning_plan.py index eb0a87cd..1896ff11 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). @@ -127,23 +147,27 @@ class GenerationComputePlan: """Resolved compute output of one ``BaseGenerationService`` run (gate b). ``compute_generation`` runs the prepare gate + extractor + dedup/embedding - resolution (``_resolve_write_plan``) and drives the ``agent_run`` rows to - their terminal state (``_finalize_extraction_runs`` — agent_run only, §4.3), - issuing **no** learning DB write. ``persist_generation`` applies - ``write_plan`` + the extractor bookmark advance inside the fence; - ``emit_generation_side_effects`` fires the post-commit telemetry + billing. + resolution (``_resolve_write_plan``), issuing **no** learning DB write and + leaving receipt-aware ``agent_run`` rows non-terminal. ``persist_generation`` + applies ``write_plan`` + the extractor bookmark advance + immutable receipt + inside the fence; ``emit_generation_side_effects`` terminalizes the run and + fires post-commit telemetry + billing. The billing inputs (``extraction_run_ids`` / ``token_totals`` / - ``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. @@ -155,15 +179,19 @@ class GenerationComputePlan: extraction_run_ids: Snapshot of the run's ``agent_run`` ids. token_totals: Snapshot of the run's LLM token totals (billing cost facet), or ``None`` when the extractor reported none. + finalization_result: Receipt ownership recorded by persistence. ``None`` + for services that do not use receipt-aware inline finalization. """ prepared: PreparedGenerationRun[Any] generated_count: int + billable_count: int write_plan: Any bookmark_advance: ExtractorBookmarkAdvance | None generation_start: float extraction_run_ids: list[str] token_totals: RunTokenTotals | None + finalization_result: FinalizationResult | None = None @dataclass diff --git a/reflexio/server/services/durable_learning/worker.py b/reflexio/server/services/durable_learning/worker.py index fed22341..9f783cfe 100644 --- a/reflexio/server/services/durable_learning/worker.py +++ b/reflexio/server/services/durable_learning/worker.py @@ -24,7 +24,9 @@ from reflexio.server.api_endpoints.request_context import RequestContext from reflexio.server.cache.reflexio_cache import get_reflexio +from reflexio.server.services.deferred_learning_plan import DeferredLearningPlan from reflexio.server.services.generation_service import GenerationService +from reflexio.server.services.storage.storage_base import AgentRunStatus, BaseStorage from reflexio.server.services.storage.storage_base._learning_jobs import LearningJob logger = logging.getLogger(__name__) @@ -151,6 +153,7 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: return False gen: GenerationService | None = None + plan: DeferredLearningPlan | None = None try: reflexio = get_reflexio( org_id=ctx.org_id, storage_base_dir=ctx.storage_base_dir @@ -202,7 +205,16 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: # POST-COMMIT — billing / telemetry / tagging / off-thread schedulers # + the per-user lock release, only for the winning worker. - gen.emit_deferred_learning_side_effects(plan) + try: + gen.emit_deferred_learning_side_effects(plan) + except Exception: + logger.exception( + "event=learning_job_side_effects_failed " + "job_id=%s org_id=%s user_id=%s", + job.job_id, + job.org_id, + job.user_id, + ) logger.info( "event=learning_job_done job_id=%s org_id=%s user_id=%s", job.job_id, @@ -210,23 +222,25 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: job.user_id, ) return True - except _SupersededError: + except _SupersededError as exc: logger.info( "event=learning_job_superseded job_id=%s org_id=%s", job.job_id, job.org_id, ) + self._abandon_computed_agent_runs(storage, plan, exc) # The persist rolled back and emit never ran, so the per-user lock is # still held by this compute — release it so the reclaim isn't blocked. self._release_user_lock(gen, job) return False - except Exception: + except Exception as exc: logger.exception( "event=learning_job_failed job_id=%s org_id=%s user_id=%s", job.job_id, job.org_id, job.user_id, ) + self._abandon_computed_agent_runs(storage, plan, exc) # emit (which releases the lock) never ran on this path — release the # per-user lock so a failed job doesn't strand it. self._release_user_lock(gen, job) @@ -237,6 +251,28 @@ def _process_job(self, ctx: RequestContext, job: LearningJob) -> bool: ) return False + @staticmethod + def _abandon_computed_agent_runs( + storage: BaseStorage, + plan: DeferredLearningPlan | None, + exc: Exception, + ) -> None: + """Make runs from a rolled-back durable attempt non-resumable.""" + if plan is None: + return + run_ids = { + run_id + for pair in (plan.profile, plan.playbook) + if pair is not None + for run_id in pair[1].extraction_run_ids + } + for run_id in run_ids: + storage.update_agent_run_status( + run_id, + AgentRunStatus.FAILED, + last_error=f"durable learning attempt abandoned: {exc}", + ) + def _release_user_lock( self, gen: GenerationService | None, job: LearningJob ) -> None: diff --git a/reflexio/server/services/extraction/README.md b/reflexio/server/services/extraction/README.md index 66845520..a0dfe2b5 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 ff173202..fc68e8de 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 fca64cdb..5b60a9be 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 37d6e9dc..8e33701a 100644 --- a/reflexio/server/services/lineage/gc_scheduler.py +++ b/reflexio/server/services/lineage/gc_scheduler.py @@ -118,7 +118,8 @@ def register_global_sweep(fn: Callable[[int], int]) -> None: fn (Callable[[int], int]): Called with the current unix epoch; returns the number of rows it deleted. """ - _global_sweep_hooks.append(fn) + if fn not in _global_sweep_hooks: + _global_sweep_hooks.append(fn) def clear_global_sweeps() -> None: @@ -126,6 +127,31 @@ 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. + """ + if fn not in _always_global_sweep_hooks: + _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 @@ -148,7 +174,8 @@ def register_per_org_sweep(fn: Callable[[str, int], int]) -> None: fn (Callable[[str, int], int]): Called with ``(org_id, now)`` where ``now`` is the current unix epoch; returns the number of rows deleted. """ - _per_org_sweep_hooks.append(fn) + if fn not in _per_org_sweep_hooks: + _per_org_sweep_hooks.append(fn) def clear_per_org_sweeps() -> None: @@ -447,7 +474,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 +578,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 b0887d0e..0b53055c 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 ce12ec49..4ff6dd7c 100644 --- a/reflexio/server/services/playbook/service.py +++ b/reflexio/server/services/playbook/service.py @@ -9,7 +9,10 @@ if TYPE_CHECKING: from reflexio.server.api_endpoints.request_context import RequestContext from reflexio.server.llm.litellm_client import LiteLLMClient - from reflexio.server.services.deferred_learning_plan import GenerationComputePlan + from reflexio.server.services.deferred_learning_plan import ( + ExtractorBookmarkAdvance, + GenerationComputePlan, + ) from reflexio.server.services.storage.storage_base import ( AgentRunRecord, BaseStorage, @@ -34,7 +37,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 +635,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) @@ -659,6 +666,11 @@ def emit_generation_side_effects(self, plan: GenerationComputePlan) -> None: superseded one) — the phantom-billing gate. """ super().emit_generation_side_effects(plan) + if ( + plan.finalization_result is not None + and not plan.finalization_result.won_receipt + ): + return write_plan = plan.write_plan if write_plan is not None: self._dispatch_playbook_schedulers(write_plan) @@ -669,16 +681,42 @@ 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.""" + 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="user_playbook", + ) + 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 +729,84 @@ 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 + 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, + ) + + result = self._finalize_write_plan_with_outcome( + plan, + finalization_run_id=finalization_run_id, + bookmark_advance=None, + ) + if plan is not None and result.won_receipt: + self._dispatch_playbook_schedulers(plan) + return result + + def _finalize_write_plan_with_outcome( + self, + write_plan: PlaybookWritePlan | None, + *, + finalization_run_id: str, + bookmark_advance: ExtractorBookmarkAdvance | None, + ) -> FinalizationResult: + """Commit a resolved playbook plan, bookmark, and receipt atomically.""" + entity_type = "user_playbook" + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + learning_ids: list[str] = [] + + try: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + if write_plan is not None: + self._persist_write_plan(write_plan) + learning_ids = [ + str(playbook.user_playbook_id) + for playbook in write_plan.new_playbooks + if playbook.user_playbook_id + ] + self._apply_bookmark_advance(bookmark_advance) + inserted = self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + 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 _apply_consolidation_lineage( self, diff --git a/reflexio/server/services/profile/service.py b/reflexio/server/services/profile/service.py index 4b2bd105..b75671d9 100644 --- a/reflexio/server/services/profile/service.py +++ b/reflexio/server/services/profile/service.py @@ -30,7 +30,12 @@ BaseGenerationService, StatusChangeOperation, ) -from reflexio.server.services.deferred_learning_plan import ProfileWritePlan +from reflexio.server.services.deferred_learning_plan import ( + ExtractorBookmarkAdvance, + 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 +342,116 @@ 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.""" + 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="profile", + ) + 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) + + return self._finalize_write_plan_with_outcome( + plan, + finalization_run_id=finalization_run_id, + bookmark_advance=None, + ) + + def _finalize_write_plan_with_outcome( + self, + write_plan: ProfileWritePlan | None, + *, + finalization_run_id: str, + bookmark_advance: ExtractorBookmarkAdvance | None, + ) -> FinalizationResult: + """Commit a resolved profile plan, bookmark, and receipt atomically.""" + entity_type = "profile" + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + learning_ids = ( + [ + str(profile.profile_id) + for profile in write_plan.new_profiles + if profile.profile_id + ] + if write_plan is not None + else [] + ) + + try: + with self.storage.commit_scope(): # type: ignore[reportOptionalMemberAccess] + receipt = self.storage.get_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + ) + if receipt is not None: + return FinalizationResult(receipt, won_receipt=False) + if write_plan is not None: + self._persist_write_plan(write_plan) + self._apply_bookmark_advance(bookmark_advance) + inserted = self.storage.save_agent_run_finalization_receipt( # type: ignore[reportOptionalMemberAccess] + run_id=finalization_run_id, + entity_type=entity_type, + 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 00000000..e2ab47a4 --- /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/search_metering_worker.py b/reflexio/server/services/search_metering_worker.py index 8b879af5..4dc19d6f 100644 --- a/reflexio/server/services/search_metering_worker.py +++ b/reflexio/server/services/search_metering_worker.py @@ -65,7 +65,7 @@ def start(self) -> bool: return False self._stop_event.clear() self._abort_event.clear() - self._threads = [ + threads = [ threading.Thread( target=self._worker_loop, name=f"search-metering-worker-{index}", @@ -73,8 +73,19 @@ def start(self) -> bool: ) for index in range(self.worker_count) ] - for thread in self._threads: - thread.start() + started_threads: list[threading.Thread] = [] + try: + for thread in threads: + thread.start() + started_threads.append(thread) + except BaseException: + self._stop_event.set() + for thread in started_threads: + thread.join() + self._stop_event.clear() + self._threads = [] + raise + self._threads = threads self._started = True logger.info( "event=search_metering_worker_started workers=%d queue_capacity=%d", diff --git a/reflexio/server/services/storage/error.py b/reflexio/server/services/storage/error.py index cf696ed7..fe14734b 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 00000000..147d3386 --- /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 a89fa1de..3358604e 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 d4803568..1296545d 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,8 @@ class RetentionTarget: order_column: str id_columns: tuple[str, ...] priority_statuses: tuple[str, ...] = () + minimum_age_seconds: int = 0 + fixed_row_limit: int | None = None @dataclass(frozen=True, slots=True) @@ -130,6 +133,14 @@ 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, + fixed_row_limit=DEFAULT_ROW_RETENTION_LIMIT, + ), RetentionTarget("skills", "skills", "created_at", ("skill_id",)), ) @@ -186,12 +197,16 @@ class CascadeRef: def get_row_retention_limits() -> dict[str, int]: """Return per-target row limits from env with code defaults. - ``REFLEXIO_ROW_LIMIT_`` takes precedence for every target. + ``REFLEXIO_ROW_LIMIT_`` takes precedence for targets without a + ``fixed_row_limit``. Fixed targets explicitly reject that override path. ``INTERACTION_CLEANUP_THRESHOLD`` remains the legacy override for interactions when the new variable is not present. """ limits: dict[str, int] = {} for target in RETENTION_TARGETS: + if target.fixed_row_limit is not None: + limits[target.name] = target.fixed_row_limit + continue env_name = f"REFLEXIO_ROW_LIMIT_{target.name.upper()}" default = DEFAULT_ROW_RETENTION_LIMIT if target.name == "interactions": diff --git a/reflexio/server/services/storage/retention_mixin.py b/reflexio/server/services/storage/retention_mixin.py index 969cab1e..f7462b75 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``. @@ -207,6 +233,8 @@ def _retention_select_oldest_keys( statuses (tuple[str, ...] | None): When not None, restrict the select to rows whose ``status`` is one of these values. An empty tuple matches nothing (never every row). + older_than_epoch (int | None): When set, restrict the select to + rows whose target ordering column is strictly older. """ raise NotImplementedError diff --git a/reflexio/server/services/storage/session_outcome_identity.py b/reflexio/server/services/storage/session_outcome_identity.py new file mode 100644 index 00000000..ff0820d2 --- /dev/null +++ b/reflexio/server/services/storage/session_outcome_identity.py @@ -0,0 +1,347 @@ +"""Canonical identities for immutable session outcomes.""" + +import json +from collections.abc import Collection, Mapping, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from hashlib import sha256 +from typing import TypedDict + +from reflexio.server.services.playbook.publication import canonical_json_bytes + +__all__ = [ + "CanonicalSessionTrajectory", + "CanonicalTrajectoryDigestResult", + "CanonicalTrajectoryDigestAccumulator", + "OUTCOME_ALLOWED_VALUES", + "OUTCOME_FINALIZATION_RULE", + "OUTCOME_SCHEMA_VERSION", + "canonical_json_bytes", + "canonical_session_trajectory", + "canonical_trajectory_bytes", + "outcome_contract_digest", + "trajectory_digest", +] + +MAX_CANONICAL_TRAJECTORY_JSON_DEPTH = 100 +OUTCOME_SCHEMA_VERSION = 1 +OUTCOME_FINALIZATION_RULE = "first_write" +OUTCOME_ALLOWED_VALUES = ("success", "failure", "unknown") + + +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] + + +@dataclass(frozen=True) +class CanonicalTrajectoryDigestResult: + """Digest and bounded context derived from one ordered trajectory stream.""" + + digest: str + first_request: dict[str, object] | None + request_count: int + + +def _canonical_timestamp(value: object) -> str: + if isinstance(value, datetime): + if value.tzinfo is None: + 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_request(row: Mapping[str, object]) -> CanonicalRequest: + return { + "request_id": str(row["request_id"]), + "user_id": str(row["user_id"]), + "created_at": _canonical_timestamp(row["created_at"]), + "source": str(row["source"] or ""), + "agent_version": str(row["agent_version"] or ""), + "session_id": str(row["session_id"]), + "evaluation_only": bool(row["evaluation_only"]), + "retrieval_experiment_id": ( + str(row["retrieval_experiment_id"]) + if row["retrieval_experiment_id"] is not None + else None + ), + "retrieval_experiment_arm": ( + str(row["retrieval_experiment_arm"]) + if row["retrieval_experiment_arm"] is not None + else None + ), + } + + +def canonical_session_trajectory( + session_id: str, + request_rows: Sequence[Mapping[str, object]], + 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 = _canonical_request(row) + 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 _canonical_trajectory_bytes_at_depth(value: object, *, depth: int) -> bytes: + return _canonical_trajectory_json(value, depth=depth).encode("utf-8") + + +class CanonicalTrajectoryDigestAccumulator: + """Hash one complete canonical trajectory while retaining only its current row.""" + + def __init__(self, session_id: str) -> None: + self._digest = sha256() + self._byte_count = 0 + self._update(b'{"requests":[') + self._session_id = session_id + self._request: CanonicalRequest | None = None + self._request_count = 0 + self._interaction_count = 0 + self._hexdigest: str | None = None + self._poisoned = False + + def _raise_if_poisoned(self) -> None: + if self._poisoned: + raise RuntimeError("canonical trajectory digest accumulator is invalid") + + def _update(self, encoded: bytes) -> None: + self._digest.update(encoded) + self._byte_count += len(encoded) + + def start_request(self, row: Mapping[str, object]) -> None: + self._raise_if_poisoned() + try: + if self._hexdigest is not None: + raise RuntimeError("canonical trajectory digest is already finalized") + if self._request is not None: + raise RuntimeError("previous canonical request is not finished") + request = _canonical_request(row) + if self._request_count: + self._update(b",") + self._update(b'{"interactions":[') + self._request = request + self._interaction_count = 0 + except Exception: + self._poisoned = True + raise + + def add_interaction(self, row: Mapping[str, object]) -> None: + self._raise_if_poisoned() + try: + if self._request is None: + raise RuntimeError("canonical interaction has no active request") + if str(row["request_id"]) != self._request["request_id"]: + raise ValueError( + "canonical interaction does not belong to active request" + ) + encoded = _canonical_trajectory_bytes_at_depth( + _canonical_interaction(row), depth=4 + ) + if self._interaction_count: + self._update(b",") + self._update(encoded) + self._interaction_count += 1 + except Exception: + self._poisoned = True + raise + + def finish_request(self) -> None: + self._raise_if_poisoned() + try: + if self._request is None: + raise RuntimeError("canonical trajectory has no active request") + encoded = _canonical_trajectory_bytes_at_depth(self._request, depth=3) + self._update(b'],"request":') + self._update(encoded) + self._update(b"}") + self._request = None + self._request_count += 1 + except Exception: + self._poisoned = True + raise + + def byte_count_if_finalized(self) -> int: + """Return exact canonical bytes if the current stream ended now.""" + self._raise_if_poisoned() + try: + byte_count = self._byte_count + if self._hexdigest is not None: + return byte_count + if self._request is not None: + encoded_request = _canonical_trajectory_bytes_at_depth( + self._request, depth=3 + ) + byte_count += len(b'],"request":') + len(encoded_request) + 1 + encoded_session_id = _canonical_trajectory_bytes_at_depth( + self._session_id, depth=1 + ) + return byte_count + len(b'],"session_id":') + len(encoded_session_id) + 1 + except Exception: + self._poisoned = True + raise + + def hexdigest(self) -> str: + self._raise_if_poisoned() + try: + if self._request is not None: + raise RuntimeError("canonical trajectory has an unfinished request") + if self._hexdigest is None: + encoded_session_id = _canonical_trajectory_bytes_at_depth( + self._session_id, depth=1 + ) + self._update(b'],"session_id":') + self._update(encoded_session_id) + self._update(b"}") + self._hexdigest = self._digest.hexdigest() + return self._hexdigest + except Exception: + self._poisoned = True + raise + + +def trajectory_digest(trajectory: object) -> str: + """Hash the canonical finalized session trajectory.""" + return sha256(canonical_trajectory_bytes(trajectory)).hexdigest() diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index faba6627..c1b25ae4 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -18,6 +18,7 @@ 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 @@ -58,7 +59,15 @@ 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 ( + OUTCOME_ALLOWED_VALUES, + OUTCOME_FINALIZATION_RULE, + CanonicalTrajectoryDigestAccumulator, + CanonicalTrajectoryDigestResult, + canonical_json_bytes, + outcome_contract_digest, +) from reflexio.server.services.storage.storage_base import BaseStorage from reflexio.server.site_var.site_var_manager import SiteVarManager @@ -67,6 +76,8 @@ logger = logging.getLogger(__name__) +_TRAJECTORY_FETCH_SIZE = 256 +_SESSION_OUTCOME_MIGRATION_BATCH_SIZE = 256 _MINIMUM_SQLITE_VERSION = (3, 35, 0) _SQLITE_INITIALIZATION_LOCK_STRIPES = 64 _sqlite_initialization_locks = tuple( @@ -101,6 +112,164 @@ def _json_loads(text: str | None) -> Any: return json.loads(text) +def _canonical_session_trajectory_snapshot( + conn: sqlite3.Connection, session_id: str +) -> CanonicalTrajectoryDigestResult: + """Hash and describe a complete bounded, ordered row stream.""" + cursor = conn.execute( + """SELECT requests.request_id, requests.user_id, requests.created_at, + requests.source, requests.agent_version, requests.session_id, + requests.evaluation_only, requests.retrieval_experiment_id, + requests.retrieval_experiment_arm, + requests.governance_subject_ref, + interactions.interaction_id, + interactions.user_id AS interaction_user_id, + interactions.request_id AS interaction_request_id, + interactions.created_at AS interaction_created_at, + interactions.content, interactions.role, interactions.token_count, + interactions.user_action, interactions.user_action_description, + interactions.interacted_image_url, interactions.image_encoding, + interactions.shadow_content, interactions.expert_content, + interactions.tools_used, interactions.citations, + interactions.retrieved_learnings + FROM requests + LEFT JOIN interactions + ON interactions.request_id = requests.request_id + WHERE requests.session_id = ? + ORDER BY requests.created_at ASC, requests.request_id ASC, + interactions.created_at ASC, interactions.interaction_id ASC""", + (session_id,), + ) + accumulator = CanonicalTrajectoryDigestAccumulator(session_id) + active_request_id: str | None = None + first_request: dict[str, object] | None = None + request_count = 0 + while rows := cursor.fetchmany(_TRAJECTORY_FETCH_SIZE): + for row in rows: + request_id = str(row["request_id"]) + if request_id != active_request_id: + if active_request_id is not None: + accumulator.finish_request() + accumulator.start_request(row) + active_request_id = request_id + request_count += 1 + if first_request is None: + first_request = { + key: row[key] + for key in ( + "request_id", + "user_id", + "created_at", + "source", + "agent_version", + "session_id", + "evaluation_only", + "retrieval_experiment_id", + "retrieval_experiment_arm", + "governance_subject_ref", + ) + } + if row["interaction_id"] is None: + continue + interaction = dict(row) + interaction["user_id"] = row["interaction_user_id"] + interaction["request_id"] = row["interaction_request_id"] + interaction["created_at"] = row["interaction_created_at"] + accumulator.add_interaction(interaction) + if active_request_id is not None: + accumulator.finish_request() + return CanonicalTrajectoryDigestResult( + digest=accumulator.hexdigest(), + first_request=first_request, + request_count=request_count, + ) + + +def _canonical_session_trajectory_digest( + conn: sqlite3.Connection, session_id: str +) -> str: + return _canonical_session_trajectory_snapshot(conn, session_id).digest + + +def _prefetch_canonical_session_trajectory_digests( + conn: sqlite3.Connection, session_ids: Sequence[str] +) -> dict[str, str]: + """Derive trajectory digests from ordered joined rows in bounded chunks.""" + digests: dict[str, str] = {} + for session_id_chunk in chunked(list(dict.fromkeys(session_ids))): + placeholders = ",".join("?" for _ in session_id_chunk) + cursor = conn.execute( + f"""SELECT requests.request_id, requests.user_id, requests.created_at, + requests.source, requests.agent_version, requests.session_id, + requests.evaluation_only, requests.retrieval_experiment_id, + requests.retrieval_experiment_arm, + requests.governance_subject_ref, + interactions.interaction_id, + interactions.user_id AS interaction_user_id, + interactions.request_id AS interaction_request_id, + interactions.created_at AS interaction_created_at, + interactions.content, interactions.role, + interactions.token_count, interactions.user_action, + interactions.user_action_description, + interactions.interacted_image_url, + interactions.image_encoding, interactions.shadow_content, + interactions.expert_content, interactions.tools_used, + interactions.citations, interactions.retrieved_learnings + FROM requests + LEFT JOIN interactions + ON interactions.request_id = requests.request_id + WHERE requests.session_id IN ({placeholders}) + ORDER BY requests.session_id ASC, requests.created_at ASC, + requests.request_id ASC, + interactions.created_at ASC, + interactions.interaction_id ASC""", # noqa: S608 + session_id_chunk, + ) + active_session_id: str | None = None + active_request_id: str | None = None + accumulator: CanonicalTrajectoryDigestAccumulator | None = None + while rows := cursor.fetchmany(_TRAJECTORY_FETCH_SIZE): + for row in rows: + session_id = str(row["session_id"]) + request_id = str(row["request_id"]) + if accumulator is None or session_id != active_session_id: + if accumulator is not None: + if active_request_id is not None: + accumulator.finish_request() + assert active_session_id is not None # noqa: S101 + digests[active_session_id] = accumulator.hexdigest() + accumulator = CanonicalTrajectoryDigestAccumulator(session_id) + active_session_id = session_id + active_request_id = None + if request_id != active_request_id: + if active_request_id is not None: + accumulator.finish_request() + accumulator.start_request(row) + active_request_id = request_id + if row["interaction_id"] is not None: + interaction = dict(row) + interaction["user_id"] = row["interaction_user_id"] + interaction["request_id"] = row["interaction_request_id"] + interaction["created_at"] = row["interaction_created_at"] + accumulator.add_interaction(interaction) + if accumulator is not None: + if active_request_id is not None: + accumulator.finish_request() + assert active_session_id is not None # noqa: S101 + digests[active_session_id] = accumulator.hexdigest() + for session_id in session_id_chunk: + digests.setdefault( + session_id, + CanonicalTrajectoryDigestAccumulator(session_id).hexdigest(), + ) + 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 +1084,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 +1131,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 +1204,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 +1307,220 @@ 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 + + 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) + )""" + ) + last_user_id: str | None = None + last_session_id: str | None = None + while True: + if last_user_id is None: + legacy_rows = self.conn.execute( + """SELECT * FROM session_outcomes_legacy + ORDER BY user_id, session_id + LIMIT ?""", + (_SESSION_OUTCOME_MIGRATION_BATCH_SIZE,), + ).fetchall() + else: + legacy_rows = self.conn.execute( + """SELECT * FROM session_outcomes_legacy + WHERE user_id > ? + OR (user_id = ? AND session_id > ?) + ORDER BY user_id, session_id + LIMIT ?""", + ( + last_user_id, + last_user_id, + last_session_id, + _SESSION_OUTCOME_MIGRATION_BATCH_SIZE, + ), + ).fetchall() + if not legacy_rows: + break + trajectory_digests = _prefetch_canonical_session_trajectory_digests( + self.conn, + [ + str(row["session_id"]) + for row in legacy_rows + if "finalized_trajectory_digest" not in columns + or not row["finalized_trajectory_digest"] + ], + ) + for row in legacy_rows: + source = str(row["source"]) + subject_ref = ( + row["governance_subject_ref"] + if "governance_subject_ref" in columns + else None + ) + if subject_ref is None or not str(subject_ref).strip(): + subject_ref = self._subject_ref_for_user_id( + str(row["user_id"]) + ) + existing_outcome_id = ( + row["outcome_id"] if "outcome_id" in columns else None + ) + existing_contract_digest = ( + row["outcome_contract_digest"] + if "outcome_contract_digest" in columns + else None + ) + existing_trajectory_digest = ( + row["finalized_trajectory_digest"] + if "finalized_trajectory_digest" in columns + else None + ) + self.conn.execute( + """INSERT INTO session_outcomes ( + outcome_id, outcome_revision, user_id, session_id, + outcome, occurred_at, source, label, value, metadata, + outcome_contract_digest, finalized_trajectory_digest, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + existing_outcome_id + or _legacy_session_outcome_id( + str(row["user_id"]), str(row["session_id"]) + ), + ( + row["outcome_revision"] + if "outcome_revision" in columns + and row["outcome_revision"] is not None + else 1 + ), + row["user_id"], + row["session_id"], + row["outcome"], + row["occurred_at"], + source, + row["label"] if "label" in columns else None, + row["value"] if "value" in columns else None, + row["metadata"] if "metadata" in columns else None, + existing_contract_digest + or outcome_contract_digest( + source=source, + schema_version=1, + allowed_values=OUTCOME_ALLOWED_VALUES, + finalization_rule=OUTCOME_FINALIZATION_RULE, + ), + existing_trajectory_digest + or trajectory_digests[str(row["session_id"])], + subject_ref, + row["created_at"], + ), + ) + last_user_id = str(legacy_rows[-1]["user_id"]) + last_session_id = str(legacy_rows[-1]["session_id"]) + if len(legacy_rows) < _SESSION_OUTCOME_MIGRATION_BATCH_SIZE: + break + self.conn.execute("DROP TABLE session_outcomes_legacy") + self.conn.execute( + "CREATE INDEX idx_session_outcomes_occurred_at " + "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 +2173,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 +2185,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 +2280,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 +2320,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 +2687,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 +2704,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 +2717,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 +2733,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 +2743,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 +3039,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 +3089,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 +3219,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 +3464,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 +3611,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 5e4abad0..922691fc 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/_requests.py b/reflexio/server/services/storage/sqlite_storage/_requests.py index 46b557ee..9f9143b0 100644 --- a/reflexio/server/services/storage/sqlite_storage/_requests.py +++ b/reflexio/server/services/storage/sqlite_storage/_requests.py @@ -11,6 +11,7 @@ from reflexio.models.api_schema.service_schemas import ( Request, ) +from reflexio.models.api_schema.validators import validate_session_outcome_source from ._base import ( SQLiteStorageBase, @@ -40,6 +41,7 @@ class RequestMixin: @SQLiteStorageBase.handle_exceptions def add_request(self, request: Request) -> None: + source = validate_session_outcome_source(request.source) created_at_iso = _epoch_to_iso(request.created_at) subject_ref = self._subject_ref_for_user_id(request.user_id) with self._lock: @@ -58,7 +60,7 @@ def add_request(self, request: Request) -> None: request.request_id, request.user_id, created_at_iso, - request.source, + source, request.agent_version, request.session_id, 1 if request.evaluation_only else 0, diff --git a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py index aaec5a1b..aaec427a 100644 --- a/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py +++ b/reflexio/server/services/storage/sqlite_storage/_session_outcomes.py @@ -2,7 +2,8 @@ import json import sqlite3 -from typing import Any +from typing import Any, cast +from uuid import uuid4 from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, @@ -11,12 +12,49 @@ SetSessionOutcomeRequest, ) from reflexio.server.services.storage.error import SubjectWriteBarrierError +from reflexio.server.services.storage.session_outcome_identity import ( + OUTCOME_ALLOWED_VALUES, + OUTCOME_FINALIZATION_RULE, + OUTCOME_SCHEMA_VERSION, + outcome_contract_digest, +) from reflexio.server.services.storage.storage_base._session_outcomes import ( SessionOutcomeContext, SessionOutcomeWriteResult, ) -from ._base import SQLiteStorageBase, _iso_to_epoch +from ._base import ( + SQLiteStorageBase, + _canonical_session_trajectory_snapshot, + _iso_to_epoch, +) + + +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: @@ -38,21 +76,26 @@ def get_session_outcome_context(self, session_id: str) -> SessionOutcomeContext: source=str(existing["source"]), existing=True, ) - rows = self.conn.execute( + first = self.conn.execute( """SELECT user_id, source, created_at, request_id FROM requests WHERE session_id = ? - ORDER BY created_at ASC, request_id ASC""", + ORDER BY created_at ASC, request_id ASC LIMIT 1""", (session_id,), - ).fetchall() - if not rows: + ).fetchone() + if first is None: return SessionOutcomeContext() - first = rows[0] + counts = self.conn.execute( + """SELECT COUNT(DISTINCT user_id) AS user_count, + COUNT(DISTINCT COALESCE(source, '')) AS source_count + FROM requests WHERE session_id = ?""", + (session_id,), + ).fetchone() return SessionOutcomeContext( user_id=str(first["user_id"]), - source=str(first["source"]), + source=str(first["source"] or ""), first_request_at=_iso_to_epoch(first["created_at"]), - user_contract_violation=len({str(row["user_id"]) for row in rows}) > 1, - source_contract_violation=len({str(row["source"]) for row in rows}) > 1, + user_contract_violation=int(counts["user_count"]) > 1, + source_contract_violation=int(counts["source_count"]) > 1, ) @SQLiteStorageBase.handle_exceptions @@ -67,31 +110,139 @@ 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: + early_first = self.conn.execute( + """SELECT user_id, source, governance_subject_ref + FROM requests WHERE session_id = ? + ORDER BY created_at ASC, request_id ASC LIMIT 1""", + (request.session_id,), + ).fetchone() + snapshot = ( + _canonical_session_trajectory_snapshot( + self.conn, request.session_id + ) + if early_first is not None + else None + ) + first = snapshot.first_request if snapshot is not None else None + source = ( + str(first["source"] or "") + 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 = ( + snapshot.digest if snapshot 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"] or "") + 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 + barrier_first = self.conn.execute( + """SELECT user_id, source, created_at, request_id, + governance_subject_ref FROM requests WHERE session_id = ? ORDER BY created_at ASC, request_id ASC LIMIT 1""", (request.session_id,), ).fetchone() - if first is None: + if barrier_first is None: + self.conn.rollback() + return SessionOutcomeWriteResult( + recorded=False, + reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, + ) + barrier_user_id = str(barrier_first["user_id"]) + subject_ref = str( + barrier_first["governance_subject_ref"] + or self._subject_ref_for_user_id(barrier_user_id) + ) + try: + self._assert_subject_writable_locked(subject_ref) + except SubjectWriteBarrierError: + self.conn.rollback() + return SessionOutcomeWriteResult( + recorded=False, + user_id=barrier_user_id, + reason=SessionOutcomeFailureReason.SUBJECT_NOT_WRITABLE, + ) + snapshot = _canonical_session_trajectory_snapshot( + self.conn, request.session_id + ) + if snapshot.request_count == 0 or snapshot.first_request is None: self.conn.rollback() return SessionOutcomeWriteResult( recorded=False, reason=SessionOutcomeFailureReason.UNKNOWN_SESSION, ) + first = snapshot.first_request user_id = str(first["user_id"]) - source = str(first["source"]) - first_request_at = _iso_to_epoch(first["created_at"]) + source = str(first["source"] or "") + first_request_at = _iso_to_epoch(cast(str, first["created_at"])) if ( expected_context.user_id != user_id or expected_context.source != source @@ -112,41 +263,46 @@ def record_session_outcome( source=source, reason=SessionOutcomeFailureReason.OCCURRED_BEFORE_SESSION, ) - subject_ref = self._subject_ref_for_user_id(user_id) - try: - self._assert_subject_writable_locked(subject_ref) - except SubjectWriteBarrierError: - self.conn.rollback() - return SessionOutcomeWriteResult( - recorded=False, - user_id=user_id, - reason=SessionOutcomeFailureReason.SUBJECT_NOT_WRITABLE, - ) + subject_ref = str( + first["governance_subject_ref"] + or self._subject_ref_for_user_id(user_id) + ) + contract_digest = self._outcome_contract_digest(source) + snapshot_digest = snapshot.digest + outcome_id = uuid4().hex self.conn.execute( """INSERT INTO session_outcomes - (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 +335,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 +354,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 +363,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 3858b29b..4ecc0fc0 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 a7227d22..3f339417 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 954eae6c..4333632a 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 1f2a7f0e..2087ee2c 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 eea1397b..ead650ff 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 3390533f..da3506a9 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 d97fd6f2..41b938d9 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 4a790ee8..2ac8a0f6 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 d417e606..14695d9a 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 b8bba037..d5f2eec6 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 cc4016f7..4650f726 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 ef63fdc2..f3f1b462 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 6020245a..66f905a9 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 7752ec5c..2de2f5a2 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 de7f6f1b..d505f091 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/cli/test_utils.py b/tests/cli/test_utils.py index 10414c67..98b0f1f5 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -240,6 +240,7 @@ def fake_popen(*_args, **_kwargs) -> _FakeProcess: return proc monkeypatch.setattr(utils.subprocess, "Popen", fake_popen) + monkeypatch.setattr(utils, "_wait_for_all_ready", lambda *_args: True) with pytest.raises(OSError, match="simulated initial spawn failure"): utils.run_services( @@ -293,6 +294,71 @@ def fake_popen(*_args, **_kwargs) -> _FakeProcess: assert ready_calls == [{"backend": 8071}] +@pytest.mark.unit +def test_run_services_waits_for_local_embedding_before_starting_backend( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Do not start backend lifespan work before local inference is ready.""" + _patch_run_services_environment(monkeypatch, tmp_path) + lifecycle: list[str] = [] + + def fake_popen(command, **_kwargs) -> _FakeProcess: + lifecycle.append(f"start:{command[0]}") + return _FakeProcess(pid=1000 + len(lifecycle), polls_to_exit=1) + + def fake_wait_for_ready(ready_events, processes) -> bool: + assert set(ready_events) == {"embedding"} + assert set(processes) == {"embedding"} + lifecycle.append("ready:embedding") + return True + + monkeypatch.setattr(utils.subprocess, "Popen", fake_popen) + monkeypatch.setattr(utils, "_wait_for_all_ready", fake_wait_for_ready) + + utils.run_services( + [ + utils.ServiceConfig(name="embedding", command=["embedding"]), + utils.ServiceConfig(name="backend", command=["backend"]), + ], + {"embedding": 8072, "backend": 8071}, + ) + + assert lifecycle[:3] == [ + "start:embedding", + "ready:embedding", + "start:backend", + ] + + +@pytest.mark.unit +def test_run_services_aborts_when_local_embedding_is_not_ready( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Clean up local inference and never start backend when readiness fails.""" + _patch_run_services_environment(monkeypatch, tmp_path) + started: list[tuple[str, _FakeProcess]] = [] + + def fake_popen(command, **_kwargs) -> _FakeProcess: + proc = _FakeProcess(pid=1000 + len(started), polls_to_exit=1) + started.append((command[0], proc)) + return proc + + monkeypatch.setattr(utils.subprocess, "Popen", fake_popen) + monkeypatch.setattr(utils, "_wait_for_all_ready", lambda *_args: False) + + with pytest.raises(RuntimeError, match=r"embedding.*ready"): + utils.run_services( + [ + utils.ServiceConfig(name="embedding", command=["embedding"]), + utils.ServiceConfig(name="backend", command=["backend"]), + ], + {"embedding": 8072, "backend": 8071}, + ) + + assert [name for name, _proc in started] == ["embedding"] + assert started[0][1].terminated is True + + @pytest.mark.unit def test_run_services_does_not_respawn_an_explicitly_stopped_service( monkeypatch: pytest.MonkeyPatch, tmp_path: Path diff --git a/tests/client/test_search.py b/tests/client/test_search.py index 285abc8a..ee9dc43c 100644 --- a/tests/client/test_search.py +++ b/tests/client/test_search.py @@ -1,8 +1,25 @@ +import inspect +from pathlib import Path from typing import Any import pytest 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: @@ -42,6 +59,24 @@ def fake_make_request(method: str, path: str, **kwargs: Any) -> dict[str, Any]: assert captured["json"]["user_id"] == "user-1" +def test_user_playbook_search_serializes_correlation_ids(monkeypatch) -> None: + client = ReflexioClient(api_key="test-key", url_endpoint="http://localhost:8000") + captured: dict[str, Any] = {} + + def fake_make_request(method: str, path: str, **kwargs: Any) -> dict[str, Any]: + captured.update(method=method, path=path, **kwargs) + return {"success": True, "user_playbooks": []} + + monkeypatch.setattr(client, "_make_request", fake_make_request) + + client.search_user_playbooks( + query="billing", request_id="request-1", session_id="session-1" + ) + + assert captured["json"]["request_id"] == "request-1" + assert captured["json"]["session_id"] == "session-1" + + @pytest.mark.asyncio async def test_unified_search_async_uses_native_transport(monkeypatch) -> None: client = ReflexioClient(api_key="test-key", url_endpoint="http://localhost:8000") @@ -69,3 +104,40 @@ async def fake_make_async_request( assert captured["json"]["user_id"] == "u1" assert captured["json"]["agent_version"] == "a1" assert captured["json"]["top_k"] == 4 + + +def test_unified_search_docs_track_schema_contract() -> None: + top_k_schema = _non_null_schema("top_k") + identifier_limit = _non_null_schema("request_id")["maxLength"] + interaction_minimum = _non_null_schema("interaction_id")["exclusiveMinimum"] + 1 + client_docs = inspect.getdoc(ReflexioClient.search) or "" + registry_docs = ( + Path(__file__).parents[2] / "docs/lib/methods/unified-search.ts" + ).read_text(encoding="utf-8") + + assert f"1 to {top_k_schema['maximum']}" in client_docs + assert client_docs.count(f"at most {identifier_limit} characters") >= 3 + assert f"positive integer (minimum {interaction_minimum})" in client_docs + + assert f"1 to {top_k_schema['maximum']}" in registry_docs + for field_name in ("user_id", "request_id", "session_id"): + assert f'name: "{field_name}"' in registry_docs + assert registry_docs.count(f"at most {identifier_limit} characters") >= 3 + assert 'name: "interaction_id"' in registry_docs + assert f"positive integer (minimum {interaction_minimum})" in registry_docs + + +def test_user_playbook_search_docs_track_top_k_schema_contract() -> None: + top_k_schema = _non_null_schema("top_k", request_model=SearchUserPlaybookRequest) + identifier_limit = _non_null_schema( + "request_id", request_model=SearchUserPlaybookRequest + )["maxLength"] + client_docs = inspect.getdoc(ReflexioClient.search_user_playbooks) or "" + registry_docs = ( + Path(__file__).parents[2] / "docs/lib/methods/user-playbooks.ts" + ).read_text(encoding="utf-8") + + assert f"1 to {top_k_schema['maximum']}" in client_docs + assert client_docs.count(f"at most {identifier_limit} characters") >= 2 + assert f"1 to {top_k_schema['maximum']}" in registry_docs + assert registry_docs.count(f"at most {identifier_limit} characters") >= 2 diff --git a/tests/client/test_session_outcomes_client.py b/tests/client/test_session_outcomes_client.py index 662a7436..9b72405e 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 4085fbc3..af76368a 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 91129474..ee35fe46 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 00000000..09c58d9d --- /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 00000000..614ea6a8 --- /dev/null +++ b/tests/models/test_session_outcome_identity.py @@ -0,0 +1,885 @@ +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 ( + CanonicalTrajectoryDigestAccumulator, + 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: ... + + +_STRICT_OUTCOME_SOURCE_INPUT_FACTORIES: tuple[Callable[[str], _HasSource], ...] = ( + lambda source: PublishUserInteractionRequest( + user_id="user-1", + session_id="session-1", + interaction_data_list=[InteractionData(content="hello")], + source=source, + ), + lambda source: 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), +) +_STRICT_OUTCOME_SOURCE_INPUT_IDS = ( + "publish", + "get-outcomes", + "get-requests", + "search-profiles", + "get-profiles", + "rerun-profiles", + "manual-profiles", + "manual-playbooks", + "rerun-playbooks", +) + + +@pytest.mark.parametrize( + "factory", + _STRICT_OUTCOME_SOURCE_INPUT_FACTORIES, + ids=_STRICT_OUTCOME_SOURCE_INPUT_IDS, +) +def test_outcome_source_inputs_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", + _STRICT_OUTCOME_SOURCE_INPUT_FACTORIES, + ids=_STRICT_OUTCOME_SOURCE_INPUT_IDS, +) +@pytest.mark.parametrize( + "source", + [ + "Legacy Source", + "alice@example.com", + "https://example.com/hook", + "support agent", + " support-agent", + "support/agent", + "Support-Agent", + "support-agént", + "a" * 129, + ], +) +def test_outcome_source_inputs_reject_sensitive_or_free_form_values( + factory: Callable[[str], _HasSource], + source: str, +) -> None: + with pytest.raises(ValidationError): + factory(source) + + +@pytest.mark.parametrize( + "factory", + _STRICT_OUTCOME_SOURCE_INPUT_FACTORIES, + ids=_STRICT_OUTCOME_SOURCE_INPUT_IDS, +) +def test_outcome_source_inputs_preserve_empty_source( + factory: Callable[[str], _HasSource], +) -> None: + assert factory("").source == "" + + +_PERSISTED_OUTCOME_SOURCE_MODEL_FACTORIES: tuple[Callable[[str], _HasSource], ...] = ( + lambda source: Request( + request_id="request-1", + user_id="user-1", + session_id="session-1", + source=source, + ), + lambda source: SessionOutcomeRecord( + user_id="user-1", + session_id="session-1", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=1, + source=source, + created_at=2, + ), + lambda source: SetSessionOutcomeResponse(success=True, source=source), +) + + +@pytest.mark.parametrize( + "factory", + _PERSISTED_OUTCOME_SOURCE_MODEL_FACTORIES, + ids=("request", "outcome", "finalization-response"), +) +@pytest.mark.parametrize( + "source", + ["Legacy Source", "legacy/source", "legacy-sourcé", "x" * 256], +) +def test_persisted_outcome_source_models_preserve_legacy_values( + factory: Callable[[str], _HasSource], + source: str, +) -> None: + assert factory(source).source == source + + +def test_optional_outcome_source_models_preserve_absence() -> None: + assert SetSessionOutcomeResponse(success=False).source is None + assert GetSessionOutcomesRequest().source is None + 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() + + +def test_streaming_trajectory_digest_matches_materialized_multi_request_json() -> None: + request_rows = [ + { + "request_id": "request-b", + "user_id": "user-1", + "created_at": datetime(2026, 8, 10, 9, 0, tzinfo=UTC), + "source": "workflow:v2", + "agent_version": "agent-2", + "session_id": "session-1", + "evaluation_only": False, + "retrieval_experiment_id": "experiment-1", + "retrieval_experiment_arm": "treatment", + }, + { + "request_id": "request-a", + "user_id": "user-1", + "created_at": "2026-08-10T09:01:00+00:00", + "source": "workflow:v2", + "agent_version": "agent-2", + "session_id": "session-1", + "evaluation_only": True, + "retrieval_experiment_id": None, + "retrieval_experiment_arm": None, + }, + ] + interactions_by_request = { + "request-b": [ + { + "interaction_id": 11, + "user_id": "user-1", + "request_id": "request-b", + "created_at": datetime(2026, 8, 10, 9, 0, 1, tzinfo=UTC), + "content": "nested payload", + "role": "User", + "token_count": 3, + "user_action": "none", + "user_action_description": None, + "interacted_image_url": "", + "image_encoding": None, + "shadow_content": "", + "expert_content": None, + "tools_used": '[{"input":{"z":1,"a":[true,null,{"x":"y"}]} }]', + "citations": [{"metadata": {"beta": 2, "alpha": 1}}], + "retrieved_learnings": {"items": [[1, 2], {"nested": [3.5]}]}, + }, + { + "interaction_id": "12", + "user_id": "user-1", + "request_id": "request-b", + "created_at": "2026-08-10T09:00:02+00:00", + "content": "second", + "role": "Assistant", + "token_count": None, + "user_action": "none", + "user_action_description": "", + "interacted_image_url": None, + "image_encoding": "", + "shadow_content": None, + "expert_content": "", + "tools_used": [], + "citations": "[]", + "retrieved_learnings": None, + }, + ], + "request-a": [], + } + materialized = canonical_session_trajectory( + "session-1", request_rows, interactions_by_request + ) + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + + for request_row in request_rows: + accumulator.start_request(request_row) + for interaction_row in interactions_by_request[request_row["request_id"]]: + accumulator.add_interaction(interaction_row) + accumulator.finish_request() + + assert accumulator.hexdigest() == trajectory_digest(materialized) + + +def test_streaming_trajectory_byte_count_matches_each_materialized_prefix() -> None: + request_row = _streaming_request_row() + interaction_row = _streaming_interaction_row() + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(canonical_session_trajectory("session-1", [], {})) + ) + + accumulator.start_request(request_row) + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes( + canonical_session_trajectory("session-1", [request_row], {}) + ) + ) + + accumulator.add_interaction(interaction_row) + materialized = canonical_session_trajectory( + "session-1", [request_row], {"request-1": [interaction_row]} + ) + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(materialized) + ) + + accumulator.finish_request() + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(materialized) + ) + assert accumulator.hexdigest() == trajectory_digest(materialized) + + +def test_streaming_trajectory_byte_count_handles_one_huge_interaction() -> None: + request_row = _streaming_request_row() + interaction_row = _streaming_interaction_row() + interaction_row["content"] = "x" * 100_000 + materialized = canonical_session_trajectory( + "session-1", [request_row], {"request-1": [interaction_row]} + ) + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + + accumulator.start_request(request_row) + accumulator.add_interaction(interaction_row) + + assert accumulator.byte_count_if_finalized() == len( + canonical_trajectory_bytes(materialized) + ) + + +def _streaming_request_row() -> dict[str, object]: + return { + "request_id": "request-1", + "user_id": "user-1", + "created_at": "2026-08-10T09:00:00+00:00", + "source": "workflow:v2", + "agent_version": "agent-2", + "session_id": "session-1", + "evaluation_only": False, + "retrieval_experiment_id": None, + "retrieval_experiment_arm": None, + } + + +def _streaming_interaction_row(*, tools_used: object = ()) -> dict[str, object]: + return { + "interaction_id": 1, + "user_id": "user-1", + "request_id": "request-1", + "created_at": "2026-08-10T09:00:01+00:00", + "content": "nested payload", + "role": "User", + "token_count": 3, + "user_action": "none", + "user_action_description": "", + "interacted_image_url": "", + "image_encoding": "", + "shadow_content": "", + "expert_content": "", + "tools_used": tools_used, + "citations": [], + "retrieved_learnings": [], + } + + +@pytest.mark.parametrize( + ("nested_container_count", "raises"), + [(95, False), (96, True)], + ids=["exact-boundary", "over-boundary"], +) +def test_streaming_trajectory_digest_has_materialized_depth_parity( + nested_container_count: int, + raises: bool, +) -> None: + tools_used: object = "leaf" + for _ in range(nested_container_count): + tools_used = [tools_used] + request_row = _streaming_request_row() + interaction_row = _streaming_interaction_row(tools_used=tools_used) + materialized = canonical_session_trajectory( + "session-1", [request_row], {"request-1": [interaction_row]} + ) + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + accumulator.start_request(request_row) + + if raises: + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + trajectory_digest(materialized) + with pytest.raises( + ValueError, match="canonical trajectory JSON exceeds maximum depth" + ): + accumulator.add_interaction(interaction_row) + else: + accumulator.add_interaction(interaction_row) + accumulator.finish_request() + assert accumulator.hexdigest() == trajectory_digest(materialized) + + +@pytest.mark.parametrize( + "first_error", + ["bad-request", "no-request", "wrong-request", "bad-json", "unfinished"], +) +def test_streaming_trajectory_digest_is_permanently_poisoned_after_error( + first_error: str, +) -> None: + accumulator = CanonicalTrajectoryDigestAccumulator("session-1") + request_row = _streaming_request_row() + + with pytest.raises((KeyError, RuntimeError, ValueError)): + if first_error == "bad-request": + accumulator.start_request({}) + elif first_error == "no-request": + accumulator.add_interaction(_streaming_interaction_row()) + else: + accumulator.start_request(request_row) + if first_error == "wrong-request": + accumulator.add_interaction( + { + **_streaming_interaction_row(), + "request_id": "request-2", + } + ) + elif first_error == "bad-json": + accumulator.add_interaction(_streaming_interaction_row(tools_used="{")) + else: + accumulator.hexdigest() + + for operation in ( + lambda: accumulator.start_request(request_row), + lambda: accumulator.add_interaction(_streaming_interaction_row()), + accumulator.finish_request, + accumulator.hexdigest, + ): + with pytest.raises( + RuntimeError, + match=r"canonical trajectory digest accumulator is invalid$", + ): + operation() diff --git a/tests/server/api_endpoints/test_session_outcomes_integration.py b/tests/server/api_endpoints/test_session_outcomes_integration.py index 5471dd74..645c0ba2 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 00000000..bb0dc661 --- /dev/null +++ b/tests/server/routes/test_search_exposure_boundary.py @@ -0,0 +1,392 @@ +"""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(caller_type: str = "production_agent") -> TestClient: + return TestClient( + create_app( + get_org_id=lambda: "org-1", + get_caller_type=lambda: caller_type, + ), + 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 + + +@pytest.mark.parametrize("path", ["/api/search", "/api/search_user_playbooks"]) +def test_only_production_searches_record_exposures(path: str) -> None: + recorder = _Recorder() + register_service(SEARCH_EXPOSURE_RECORDER, recorder) + + with _search_results([_playbook(11, "First")]): + dashboard_response = _client("dashboard").post( + path, + json={ + "query": "inspect", + "user_id": "user-1", + "request_id": "dashboard-request", + }, + ) + production_response = _client("production_agent").post( + path, + json={ + "query": "answer", + "user_id": "user-1", + "request_id": "production-request", + }, + ) + + assert dashboard_response.status_code == 200, dashboard_response.text + assert production_response.status_code == 200, production_response.text + assert len(recorder.batches) == 1 + assert recorder.batches[0].request_id == "production-request" + + +def test_unified_search_records_the_final_user_playbook_set_before_return() -> None: + playbooks = [_playbook(11, "First"), _playbook(12, "Second")] + recorder = _Recorder() + 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/durable_learning/test_worker.py b/tests/server/services/durable_learning/test_worker.py index fdc766a8..4e606fa0 100644 --- a/tests/server/services/durable_learning/test_worker.py +++ b/tests/server/services/durable_learning/test_worker.py @@ -28,15 +28,25 @@ import tempfile import threading import time +from datetime import UTC, datetime, timedelta from unittest import mock import pytest from reflexio.models.api_schema.domain.entities import Interaction, Request from reflexio.server.api_endpoints.request_context import RequestContext -from reflexio.server.services.deferred_learning_plan import DeferredLearningPlan +from reflexio.server.services.base_generation_service import PreparedGenerationRun +from reflexio.server.services.deferred_learning_plan import ( + DeferredLearningPlan, + GenerationComputePlan, +) from reflexio.server.services.durable_learning.worker import DurableLearningWorker from reflexio.server.services.generation_service import GenerationService +from reflexio.server.services.storage.storage_base import ( + AgentBinding, + AgentRunRecord, + AgentRunStatus, +) # --------------------------------------------------------------------------- # Module-level fixture: disable the local ONNX embedder for all tests. @@ -124,6 +134,78 @@ def _setup_job( ) +def _seed_completed_agent_run( + storage, + *, + org_id: str, + request_id: str, + run_id: str, + extractor_kind: str, +) -> None: + storage.create_agent_run( + AgentRunRecord( + id=run_id, + binding=AgentBinding( + org_id=org_id, + extractor_kind=extractor_kind, + user_id="test_user", + request_id=request_id, + agent_version="v1", + source="test_src", + ), + status=AgentRunStatus.AGENT_COMPLETED, + generation_request_snapshot={"request_id": request_id}, + committed_output={f"{extractor_kind}s": []}, + ) + ) + + +def _deferred_plan_with_runs( + *, + request_id: str, + profile_run_id: str, + playbook_run_id: str, +) -> DeferredLearningPlan: + def generation_plan(run_id: str, extractor_name: str) -> GenerationComputePlan: + return GenerationComputePlan( + prepared=PreparedGenerationRun( + extractor_config=mock.Mock(), + extractor_name=extractor_name, + identifier=f"{extractor_name}_generation", + ), + generated_count=0, + billable_count=0, + write_plan=None, + bookmark_advance=None, + generation_start=0.0, + extraction_run_ids=[run_id], + token_totals=None, + ) + + return DeferredLearningPlan( + request_id=request_id, + user_id="test_user", + agent_version="v1", + lock_acquired=True, + profile=(mock.Mock(), generation_plan(profile_run_id, "profile")), + playbook=(mock.Mock(), generation_plan(playbook_run_id, "playbook")), + ) + + +def _assert_runs_abandoned(storage, *run_ids: str) -> None: + for run_id in run_ids: + run = storage.get_agent_run(run_id) + assert run is not None and run.status == AgentRunStatus.FAILED + assert ( + storage.claim_finalization_failed_agent_run( + org_id=storage.org_id, + worker_id="resume-worker", + now=datetime.now(UTC) + timedelta(days=1), + ) + is None + ) + + # --------------------------------------------------------------------------- # Test 1: Exactly-once under a claim race (the headline) # --------------------------------------------------------------------------- @@ -556,6 +638,233 @@ def test_superseded_job_does_not_commit_outputs(): ) +def test_persist_failure_abandons_computed_agent_runs(): + """A rolled-back durable persist cannot leave resumable agent runs behind.""" + with tempfile.TemporaryDirectory() as tmp_dir: + factory = _factory(tmp_dir) + ctx = factory("org_persist_failure") + assert ctx.storage is not None + _setup_job( + ctx.storage, + org_id="org_persist_failure", + user_id="test_user", + request_id="req_persist_failure", + ) + run_ids = ("profile-persist-failure", "playbook-persist-failure") + _seed_completed_agent_run( + ctx.storage, + org_id="org_persist_failure", + request_id="req_persist_failure", + run_id=run_ids[0], + extractor_kind="profile", + ) + _seed_completed_agent_run( + ctx.storage, + org_id="org_persist_failure", + request_id="req_persist_failure", + run_id=run_ids[1], + extractor_kind="playbook", + ) + [job] = ctx.storage.claim_learning_jobs( + claimed_by="persist-failure-worker", limit=1, lease_seconds=300 + ) + plan = _deferred_plan_with_runs( + request_id="req_persist_failure", + profile_run_id=run_ids[0], + playbook_run_id=run_ids[1], + ) + + with ( + mock.patch.object( + GenerationService, "compute_deferred_learning", return_value=plan + ), + mock.patch.object( + GenerationService, + "persist_deferred_learning", + side_effect=RuntimeError("persist failed"), + ), + ): + processed = DurableLearningWorker(factory)._process_job( + factory("org_persist_failure"), job + ) + + assert processed is False + _assert_runs_abandoned(ctx.storage, *run_ids) + + +def test_superseded_job_abandons_computed_agent_runs(): + """Fence loss abandons every run computed by the superseded attempt.""" + with tempfile.TemporaryDirectory() as tmp_dir: + factory = _factory(tmp_dir) + ctx = factory("org_run_supersede") + assert ctx.storage is not None + _setup_job( + ctx.storage, + org_id="org_run_supersede", + user_id="test_user", + request_id="req_run_supersede", + ) + [stale_job] = ctx.storage.claim_learning_jobs( + claimed_by="stale-worker", limit=1, lease_seconds=-1 + ) + [live_job] = ctx.storage.claim_learning_jobs( + claimed_by="live-worker", limit=1, lease_seconds=300 + ) + assert stale_job.claim_token != live_job.claim_token + + run_ids = ("profile-run-supersede", "playbook-run-supersede") + _seed_completed_agent_run( + ctx.storage, + org_id="org_run_supersede", + request_id="req_run_supersede", + run_id=run_ids[0], + extractor_kind="profile", + ) + _seed_completed_agent_run( + ctx.storage, + org_id="org_run_supersede", + request_id="req_run_supersede", + run_id=run_ids[1], + extractor_kind="playbook", + ) + plan = _deferred_plan_with_runs( + request_id="req_run_supersede", + profile_run_id=run_ids[0], + playbook_run_id=run_ids[1], + ) + + with mock.patch.object( + GenerationService, "compute_deferred_learning", return_value=plan + ): + processed = DurableLearningWorker(factory)._process_job( + factory("org_run_supersede"), stale_job + ) + + assert processed is False + _assert_runs_abandoned(ctx.storage, *run_ids) + + +def test_post_commit_emit_failure_preserves_completed_job_and_agent_runs(): + """A post-commit side-effect failure cannot undo the durable winner.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserProfile + from reflexio.server.services.deferred_learning_plan import ( + ProfileWritePlan, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + with tempfile.TemporaryDirectory() as tmp_dir: + factory = _factory(tmp_dir) + ctx = factory("org_post_commit_emit_failure") + assert ctx.storage is not None + _setup_job( + ctx.storage, + org_id="org_post_commit_emit_failure", + user_id="test_user", + request_id="req_post_commit_emit_failure", + force_extraction=True, + skip_aggregation=True, + ) + [job] = ctx.storage.claim_learning_jobs( + claimed_by="post-commit-worker", limit=1, lease_seconds=300 + ) + + computed_plans: list[DeferredLearningPlan] = [] + run_id = "profile-post-commit-emit-failure" + _seed_completed_agent_run( + ctx.storage, + org_id="org_post_commit_emit_failure", + request_id="req_post_commit_emit_failure", + run_id=run_id, + extractor_kind="profile", + ) + + def build_plan(self, **_kwargs): + profile = UserProfile( + profile_id="profile-post-commit", + user_id="test_user", + content="This durable profile survived a post-commit emit failure.", + last_modified_timestamp=1_000, + generated_from_request_id="req_post_commit_emit_failure", + ) + generation_plan = GenerationComputePlan( + prepared=PreparedGenerationRun( + extractor_config=mock.Mock(), + identifier="profile_generation", + extractor_name="profile", + ), + generated_count=1, + billable_count=1, + write_plan=ProfileWritePlan( + user_id="test_user", + request_id="req_post_commit_emit_failure", + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ), + bookmark_advance=None, + generation_start=0.0, + extraction_run_ids=[run_id], + token_totals=None, + ) + plan = DeferredLearningPlan( + request_id="req_post_commit_emit_failure", + user_id="test_user", + agent_version="v1", + lock_acquired=True, + profile=( + ProfileGenerationService(self.client, self.request_context), + generation_plan, + ), + playbook=None, + ) + computed_plans.append(plan) + return plan + + def emit_then_raise(self, plan): + assert plan.profile is not None + for extraction_run_id in plan.profile[1].extraction_run_ids: + self.request_context.storage.update_agent_run_status( + extraction_run_id, + AgentRunStatus.FINALIZED, + pending_tool_call_ids=[], + ) + raise RuntimeError("post-commit emit failed") + + with ( + mock.patch.object( + GenerationService, + "compute_deferred_learning", + build_plan, + ), + mock.patch.object( + GenerationService, + "emit_deferred_learning_side_effects", + emit_then_raise, + ), + ): + processed = DurableLearningWorker(factory)._process_job( + factory("org_post_commit_emit_failure"), job + ) + + assert len(computed_plans) == 1 + plan = computed_plans[0] + assert plan.profile is not None + generation_plan = plan.profile[1] + assert generation_plan.finalization_result is not None + assert generation_plan.finalization_result.won_receipt is True + run = ctx.storage.get_agent_run(run_id) + assert run is not None and run.status == AgentRunStatus.FINALIZED + assert ctx.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) == ["profile-post-commit"] + + assert processed is True + assert ctx.storage.count_learning_jobs_by_status("done") == 1 + assert ctx.storage.count_all_profiles() == 1 + + # --------------------------------------------------------------------------- # Test 6: Failed job is re-claimable without manual status reset # --------------------------------------------------------------------------- diff --git a/tests/server/services/extraction/test_resume_worker.py b/tests/server/services/extraction/test_resume_worker.py index f26a88d0..40cead5f 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 d8ee06b0..e717a807 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 c649bcaa..c8254711 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 642d9a12..62ed4cfb 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/playbook/test_playbook_generation_service_integration.py b/tests/server/services/playbook/test_playbook_generation_service_integration.py index ffd8e77f..599db9e7 100644 --- a/tests/server/services/playbook/test_playbook_generation_service_integration.py +++ b/tests/server/services/playbook/test_playbook_generation_service_integration.py @@ -80,6 +80,7 @@ def mock_request_context(): context = MagicMock(spec=RequestContext) context.org_id = "test_org_123" context.storage = MagicMock() + context.storage.get_agent_run_finalization_receipt.return_value = None agent_runs = {} def create_agent_run(record): diff --git a/tests/server/services/playbook/test_playbook_reviewer.py b/tests/server/services/playbook/test_playbook_reviewer.py index cde7141f..e0d8a193 100644 --- a/tests/server/services/playbook/test_playbook_reviewer.py +++ b/tests/server/services/playbook/test_playbook_reviewer.py @@ -465,7 +465,9 @@ def test_reviewer_prompt_preserves_grounded_procedures_and_forbids_substitutes() # fixes is the reviewer going down the evidence axis and never asking. "First name, in your own words, what the entry is ABOUT", "before asking\nwhether any of its clauses are supported".replace("\n", " "), - "is not a core to preserve, it is\nthe same subject in gentler words".replace("\n", " "), + "is not a core to preserve, it is\nthe same subject in gentler words".replace( + "\n", " " + ), ) for invariant in required_invariants: assert invariant in normalized 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 78300ee3..cb369c7d 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 4ff0df37..17d1c385 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 880cbad0..8ed1b9f8 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 fb40107b..2a651283 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 54a8d5d1..4e484c81 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 00000000..70d41161 --- /dev/null +++ b/tests/server/services/storage/sqlite_storage/test_session_outcome_migration.py @@ -0,0 +1,794 @@ +"""Regression coverage for the SQLite session-outcome identity migration.""" + +import json +import sqlite3 +from hashlib import sha256 +from math import ceil +from typing import Any, cast + +import pytest + +from reflexio.models.api_schema.domain import ( + GetSessionOutcomesRequest, + Request, + SessionOutcomeKind, + SetSessionOutcomeRequest, + SetSessionOutcomeResponse, +) +from reflexio.server.services.storage.retention_mixin import RETENTION_DELETE_CHUNK +from reflexio.server.services.storage.session_outcome_identity import ( + outcome_contract_digest, +) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._base import ( + _SESSION_OUTCOME_MIGRATION_BATCH_SIZE, + _TRAJECTORY_FETCH_SIZE, + _canonical_session_trajectory_digest, + _epoch_to_iso, + _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) +); +""" + + +class _NoTrajectoryFetchAllCursor: + def __init__(self, cursor: Any, fetch_sizes: list[int]) -> None: + self._cursor = cursor + self._fetch_sizes = fetch_sizes + + def fetchall(self) -> Any: + raise AssertionError("trajectory migration must not call fetchall") + + def __iter__(self) -> Any: + raise AssertionError( + "trajectory migration must not iterate the cursor directly" + ) + + def fetchmany(self, size: int) -> Any: + self._fetch_sizes.append(size) + return self._cursor.fetchmany(size) + + def __getattr__(self, name: str) -> Any: + return getattr(self._cursor, name) + + +class _NoTrajectoryFetchAllConnection: + def __init__(self, connection: sqlite3.Connection) -> None: + self._connection = connection + self.fetch_sizes: list[int] = [] + + def execute(self, *args: Any, **kwargs: Any) -> Any: + cursor = self._connection.execute(*args, **kwargs) + statement = str(args[0]) if args else str(kwargs.get("sql", "")) + if "FROM requests" in statement and "LEFT JOIN interactions" in statement: + return _NoTrajectoryFetchAllCursor(cursor, self.fetch_sizes) + return cursor + + def __getattr__(self, name: str) -> Any: + return getattr(self._connection, name) + + +def _replace_session_outcomes_with_unconstrained_revision( + storage: SQLiteStorage, revision: int | None +) -> None: + 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", + "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, + }, + ] + storage.conn.executemany( + """INSERT INTO requests ( + request_id, user_id, created_at, source, session_id, + governance_subject_ref + ) VALUES (?, ?, ?, ?, ?, ?)""", + [ + ( + f"request-{row['session_id']}", + row["user_id"], + _epoch_to_iso(row["occurred_at"] - 1), + row["source"], + row["session_id"], + row["governance_subject_ref"], + ) + for row in legacy_rows + ], + ) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + for row in legacy_rows: + 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 + persisted_request = migrated.get_request("request-c") + assert persisted_request is not None + assert persisted_request.source == "Legacy Source" + + for row in legacy_rows: + record = records_by_session[row["session_id"]] + 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 == ( + _canonical_session_trajectory_digest(migrated.conn, row["session_id"]) + ) + + legacy_identity_before = dict(rows_by_session["c"]) + retry = migrated.record_session_outcome( + SetSessionOutcomeRequest( + session_id="c", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + label="resolved:one", + value=2.5, + metadata={"nested": {"a": 1}, "legacy": True}, + ), + created_at=103, + expected_context=migrated.get_session_outcome_context("c"), + ) + response = SetSessionOutcomeResponse( + success=True, + recorded=retry.recorded, + user_id=retry.user_id, + source=retry.source, + outcome_id=retry.outcome_id, + outcome_revision=retry.outcome_revision, + outcome_contract_digest=retry.outcome_contract_digest, + finalized_trajectory_digest=retry.finalized_trajectory_digest, + ) + legacy_identity_after = dict( + migrated.conn.execute( + "SELECT * FROM session_outcomes WHERE session_id = ?", ("c",) + ).fetchone() + ) + + assert retry.recorded is False + assert retry.reason is None + assert response.source == "Legacy Source" + assert legacy_identity_after == legacy_identity_before + assert legacy_identity_after["outcome_contract_digest"] == ( + "a3faaa8073272084ffdcb3d1b12c410676829d8376495062b9600e55bfcc4dd4" + ) + + +def test_migration_defaults_null_legacy_outcome_revision_to_one(tmp_path) -> None: + db_path = str(tmp_path / "legacy-null-revision.db") + 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_streams_legacy_outcomes_and_trajectories_in_bounded_batches( + tmp_path, +) -> None: + db_path = str(tmp_path / "legacy-session-outcome-query-scaling.db") + storage = SQLiteStorage(org_id="legacy-query-scaling", db_path=db_path) + row_count = 501 + 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) + + legacy_outcome_queries = [ + statement + for statement in statements + if statement.lstrip().upper().startswith("SELECT") + and "FROM session_outcomes_legacy" in statement + ] + migration_batch_count = ceil(row_count / _SESSION_OUTCOME_MIGRATION_BATCH_SIZE) + assert len(legacy_outcome_queries) == migration_batch_count + assert all( + f"LIMIT {_SESSION_OUTCOME_MIGRATION_BATCH_SIZE}" in statement + for statement in legacy_outcome_queries + ) + assert not any( + statement.strip() == "SELECT * FROM session_outcomes" + for statement in statements + ) + trajectory_input_queries = [ + statement + for statement in statements + if statement.lstrip().upper().startswith("SELECT") + and "FROM requests" in statement + ] + expected_trajectory_query_count = sum( + ceil( + min(_SESSION_OUTCOME_MIGRATION_BATCH_SIZE, row_count - batch_start) + / RETENTION_DELETE_CHUNK + ) + for batch_start in range(0, row_count, _SESSION_OUTCOME_MIGRATION_BATCH_SIZE) + ) + assert len(trajectory_input_queries) == expected_trajectory_query_count + assert all("LEFT JOIN interactions" in query for query in trajectory_input_queries) + assert ( + storage.conn.execute("SELECT COUNT(*) FROM session_outcomes").fetchone()[0] + == row_count + ) + + +def test_migration_streams_complete_digest_for_trajectory_over_fetch_window( + tmp_path, +) -> None: + storage = SQLiteStorage( + org_id="legacy-large-trajectory", + db_path=str(tmp_path / "legacy-large-trajectory.db"), + ) + session_id = "large-trajectory-session" + request_id = "large-trajectory-request" + user_id = "large-trajectory-user" + source = "legacy-source" + storage.add_request( + Request( + request_id=request_id, + user_id=user_id, + session_id=session_id, + source=source, + created_at=100, + ) + ) + interaction_count = _TRAJECTORY_FETCH_SIZE + 1 + storage.conn.executemany( + """INSERT INTO interactions ( + interaction_id, user_id, request_id, created_at, content + ) VALUES (?, ?, ?, ?, ?)""", + [ + ( + index + 1, + user_id, + request_id, + _epoch_to_iso(101 + index), + f"trajectory-row-{index}", + ) + for index in range(interaction_count) + ], + ) + expected_digest = _canonical_session_trajectory_digest(storage.conn, session_id) + storage.conn.execute("DROP TABLE session_outcomes") + storage.conn.executescript(_LEGACY_SESSION_OUTCOMES_DDL) + storage.conn.execute( + """INSERT INTO session_outcomes ( + user_id, session_id, outcome, occurred_at, source, + governance_subject_ref, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?)""", + (user_id, session_id, "success", 500, source, "subject-large", 501), + ) + storage.conn.commit() + + raw_connection = storage.conn + guarded_connection = _NoTrajectoryFetchAllConnection(raw_connection) + cast(Any, storage).conn = guarded_connection + try: + storage._migrate_session_outcomes_schema() + finally: + cast(Any, storage).conn = raw_connection + + migrated = raw_connection.execute( + """SELECT finalized_trajectory_digest FROM session_outcomes + WHERE session_id = ?""", + (session_id,), + ).fetchone() + assert migrated is not None + assert migrated["finalized_trajectory_digest"] == expected_digest + assert len(guarded_connection.fetch_sizes) == ( + ceil(interaction_count / _TRAJECTORY_FETCH_SIZE) + 1 + ) + assert set(guarded_connection.fetch_sizes) == {_TRAJECTORY_FETCH_SIZE} + + +def test_migration_prefetch_retains_only_trajectory_digests(tmp_path) -> None: + storage = SQLiteStorage( + org_id="legacy-digest-retention", + 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": _canonical_session_trajectory_digest( + 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 55732634..0dd748a1 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 d6439ec3..0a6c64f1 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 81584e44..e187523c 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 cf550ae2..1ef97753 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_requests.py b/tests/server/services/storage/test_storage_contract_requests.py index 8bf182bf..bdfc7747 100644 --- a/tests/server/services/storage/test_storage_contract_requests.py +++ b/tests/server/services/storage/test_storage_contract_requests.py @@ -9,6 +9,7 @@ Request, UserActionType, ) +from reflexio.server.services.storage.error import StorageError from reflexio.server.services.storage.storage_base import BaseStorage pytestmark = pytest.mark.integration @@ -31,6 +32,14 @@ def _make_request( class TestRequestCRUD: + def test_add_request_rejects_legacy_source(self, storage: BaseStorage) -> None: + request = _make_request("legacy-source-write", "u1", source="Legacy Source") + + with pytest.raises(StorageError): + storage.add_request(request) + + assert storage.get_request(request.request_id) is None + def test_add_and_get_request(self, storage: BaseStorage) -> None: req = _make_request("r1", "u1") storage.add_request(req) diff --git a/tests/server/services/storage/test_storage_contract_retention.py b/tests/server/services/storage/test_storage_contract_retention.py index 6cdd1708..2612dc1c 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 b07a2223..b1d1a746 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,30 @@ """Session outcome storage contract.""" +from threading import RLock +from typing import Any, cast +from unittest.mock import MagicMock + +import pytest + from reflexio.models.api_schema.domain import ( GetSessionOutcomesRequest, + Interaction, Request, + SessionOutcomeFailureReason, SessionOutcomeKind, SetSessionOutcomeRequest, ) +from reflexio.server.services.storage.session_outcome_identity import ( + canonical_session_trajectory, + trajectory_digest, +) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._base import ( + _TRAJECTORY_FETCH_SIZE, +) +from reflexio.server.services.storage.sqlite_storage._session_outcomes import ( + SessionOutcomeStoreMixin, +) from reflexio.server.services.storage.storage_base import BaseStorage @@ -39,14 +58,545 @@ 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_generic_retention_cannot_delete_finalized_session_outcomes( + storage: BaseStorage, +) -> None: + storage.add_request( + Request( + request_id="retention-r1", + user_id="u1", + session_id="retention-session", + source="published", + created_at=100, + ) + ) + storage.record_session_outcome( + SetSessionOutcomeRequest( + session_id="retention-session", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ), + created_at=102, + expected_context=storage.get_session_outcome_context("retention-session"), + ) + + with pytest.raises(ValueError, match="Unknown retention target: session_outcomes"): + storage.delete_oldest_retention_target_rows("session_outcomes", 1) # type: ignore[attr-defined] + + [record] = storage.get_session_outcomes( + GetSessionOutcomesRequest(session_ids=["retention-session"]) + ) + assert record.outcome == SessionOutcomeKind.SUCCESS + + +def test_exact_finalization_retry_is_idempotent(storage: BaseStorage) -> None: + storage.add_request( + Request( + 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_sqlite_legacy_null_source_finalization_retry_is_idempotent( + storage: BaseStorage, +) -> None: + sqlite_storage = cast(SQLiteStorage, storage) + storage.add_request( + Request( + request_id="legacy-null-source-r1", + user_id="u1", + session_id="legacy-null-source", + source="", + created_at=100, + ) + ) + schema_sql = sqlite_storage.conn.execute( + "SELECT sql FROM sqlite_schema WHERE type = 'table' AND name = 'requests'" + ).fetchone()["sql"] + nullable_schema_sql = schema_sql.replace( + "source TEXT NOT NULL DEFAULT ''", "source TEXT DEFAULT ''", 1 + ) + sqlite_storage.conn.execute("PRAGMA writable_schema = ON") + sqlite_storage.conn.execute( + "UPDATE sqlite_schema SET sql = ? WHERE type = 'table' AND name = 'requests'", + (nullable_schema_sql,), + ) + schema_version = sqlite_storage.conn.execute("PRAGMA schema_version").fetchone()[0] + sqlite_storage.conn.execute(f"PRAGMA schema_version = {schema_version + 1}") + sqlite_storage.conn.execute("PRAGMA writable_schema = OFF") + sqlite_storage.conn.execute( + "UPDATE requests SET source = NULL WHERE request_id = ?", + ("legacy-null-source-r1",), + ) + sqlite_storage.conn.commit() + + request = SetSessionOutcomeRequest( + session_id="legacy-null-source", + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=101, + ) + context = storage.get_session_outcome_context("legacy-null-source") + first = storage.record_session_outcome( + request, created_at=102, expected_context=context + ) + retry = storage.record_session_outcome( + request, + created_at=103, + expected_context=storage.get_session_outcome_context("legacy-null-source"), + ) + + assert context.source == "" + assert first.recorded is True + assert first.source == "" + assert retry.recorded is False + assert retry.reason is None + assert retry.source == "" + + +def test_legacy_all_null_identity_exact_retry_uses_available_context( + storage: BaseStorage, +) -> None: + 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 + + +class _NoFetchAllCursor: + def __init__(self, cursor: Any, fetch_sizes: list[int]) -> None: + self._cursor = cursor + self._fetch_sizes = fetch_sizes + + def fetchall(self) -> Any: + raise AssertionError("session outcome finalization must not call fetchall") + + def fetchmany(self, size: int) -> Any: + self._fetch_sizes.append(size) + return self._cursor.fetchmany(size) + + def __getattr__(self, name: str) -> Any: + return getattr(self._cursor, name) + + +class _NoFetchAllConnection: + def __init__(self, connection: Any) -> None: + self._connection = connection + self.fetch_sizes: list[int] = [] + + def execute(self, *args: Any, **kwargs: Any) -> _NoFetchAllCursor: + return _NoFetchAllCursor( + self._connection.execute(*args, **kwargs), self.fetch_sizes + ) + + def __getattr__(self, name: str) -> Any: + return getattr(self._connection, name) + + +def test_sqlite_large_finalization_streams_complete_digest_and_preserves_retry_contract( + storage: BaseStorage, +) -> None: + sqlite_storage = cast(SQLiteStorage, storage) + session_id = "large-streamed-session" + request_id = "large-streamed-request" + storage.add_request( + Request( + request_id=request_id, + user_id="stream-user", + session_id=session_id, + source="published", + created_at=100, + ) + ) + interactions = [ + Interaction( + interaction_id=index + 1, + user_id="stream-user", + request_id=request_id, + created_at=101 + index, + content=f"complete-row-{index}", + token_count=index, + ) + for index in range(300) + ] + sqlite_storage.add_user_interactions_bulk( + "stream-user", interactions, embeddings_prepared=True + ) + request_rows = sqlite_storage.conn.execute( + """SELECT request_id, user_id, created_at, source, agent_version, session_id, + evaluation_only, retrieval_experiment_id, retrieval_experiment_arm + FROM requests WHERE session_id = ? + ORDER BY created_at ASC, request_id ASC""", + (session_id,), + ).fetchall() + interaction_rows = sqlite_storage.conn.execute( + """SELECT interaction_id, user_id, request_id, created_at, content, role, + token_count, user_action, user_action_description, + interacted_image_url, image_encoding, shadow_content, + expert_content, tools_used, citations, retrieved_learnings + FROM interactions WHERE request_id = ? + ORDER BY created_at ASC, interaction_id ASC""", + (request_id,), + ).fetchall() + expected_digest = trajectory_digest( + canonical_session_trajectory( + session_id, + [dict(row) for row in request_rows], + {request_id: [dict(row) for row in interaction_rows]}, + ) + ) + outcome = SetSessionOutcomeRequest( + session_id=session_id, + outcome=SessionOutcomeKind.SUCCESS, + occurred_at=500, + metadata={"complete": True}, + ) + raw_connection = sqlite_storage.conn + initial_retry_connection = _NoFetchAllConnection(raw_connection) + cast(Any, sqlite_storage).conn = initial_retry_connection + try: + context = storage.get_session_outcome_context(session_id) + first = storage.record_session_outcome( + outcome, created_at=501, expected_context=context + ) + retry = storage.record_session_outcome( + outcome, + created_at=502, + expected_context=storage.get_session_outcome_context(session_id), + ) + finally: + cast(Any, sqlite_storage).conn = raw_connection + + sqlite_storage.add_user_interactions_bulk( + "stream-user", + [ + Interaction( + interaction_id=301, + user_id="stream-user", + request_id=request_id, + created_at=401, + content="post-finalization-row", + ) + ], + embeddings_prepared=True, + ) + conflict_connection = _NoFetchAllConnection(raw_connection) + cast(Any, sqlite_storage).conn = conflict_connection + try: + conflict = storage.record_session_outcome( + outcome, + created_at=503, + expected_context=storage.get_session_outcome_context(session_id), + ) + finally: + cast(Any, sqlite_storage).conn = raw_connection + + assert first.recorded is True + assert first.finalized_trajectory_digest == expected_digest + assert retry.recorded is False + assert retry.reason is None + assert retry.finalized_trajectory_digest == expected_digest + assert conflict.recorded is False + assert conflict.reason == SessionOutcomeFailureReason.CONFLICTING_FINALIZATION + assert initial_retry_connection.fetch_sizes + assert conflict_connection.fetch_sizes + assert all( + size == _TRAJECTORY_FETCH_SIZE for size in initial_retry_connection.fetch_sizes + ) + assert all( + size == _TRAJECTORY_FETCH_SIZE for size in conflict_connection.fetch_sizes + ) + + +def test_changed_contract_identity_is_conflicting_finalization( + 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( @@ -127,6 +677,41 @@ def test_empty_request_source_is_preserved(storage: BaseStorage) -> None: assert record.source == "" +def test_sqlite_context_normalizes_nullable_request_source() -> None: + existing_cursor = MagicMock() + existing_cursor.fetchone.return_value = None + first_cursor = MagicMock() + first_cursor.fetchone.return_value = { + "user_id": "u1", + "source": None, + "created_at": "1970-01-01T00:01:40+00:00", + } + counts_cursor = MagicMock() + counts_cursor.fetchone.return_value = {"user_count": 1, "source_count": 1} + connection = MagicMock() + + def execute(statement: str, _parameters: tuple[str]) -> MagicMock: + if "FROM session_outcomes" in statement: + return existing_cursor + if "ORDER BY created_at ASC" in statement: + return first_cursor + if "COUNT(DISTINCT user_id)" in statement: + return counts_cursor + raise AssertionError( + f"unexpected SQL in nullable-source context test: {statement}" + ) + + connection.execute.side_effect = execute + reader = SessionOutcomeStoreMixin() + cast(Any, reader).conn = connection + cast(Any, reader)._lock = RLock() + + context = reader.get_session_outcome_context("nullable-source") + + assert context.source == "" + assert context.first_request_at == 100 + + def test_clear_outcomes_survives_governance_secret_rotation( storage: BaseStorage, monkeypatch ) -> None: diff --git a/tests/server/services/test_base_generation_service.py b/tests/server/services/test_base_generation_service.py index 2359f6af..aac56574 100644 --- a/tests/server/services/test_base_generation_service.py +++ b/tests/server/services/test_base_generation_service.py @@ -30,7 +30,11 @@ _weighted_content_length, ) from reflexio.server.services.extraction.outcome import ExtractionOutcome -from reflexio.server.services.storage.storage_base import AgentRunStatus +from reflexio.server.services.storage.storage_base import ( + AgentBinding, + AgentRunRecord, + AgentRunStatus, +) # =============================== # Test Data Classes @@ -2052,6 +2056,427 @@ def _create_extractor(self, extractor_config, service_config): pending_tool_call_ids=[], ) + @staticmethod + def _seed_inline_agent_run( + request_context, + *, + run_id: str, + request_id: str, + extractor_kind: str, + pending_tool_call_ids: list[str] | None = None, + ) -> None: + request_context.storage.create_agent_run( + AgentRunRecord( + id=run_id, + binding=AgentBinding( + org_id=request_context.org_id, + extractor_kind=extractor_kind, + user_id="test_user", + request_id=request_id, + agent_version="v1", + source="api", + ), + status=AgentRunStatus.AGENT_COMPLETED, + generation_request_snapshot={"request_id": request_id}, + committed_output={f"{extractor_kind}s": []}, + pending_tool_call_ids=pending_tool_call_ids or [], + ) + ) + + @staticmethod + def _build_receipt_required_plan( + entity_kind: str, + llm_client, + request_context, + *, + run_ids: list[str], + ): + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserPlaybook, UserProfile + from reflexio.server.services.deferred_learning_plan import ( + GenerationComputePlan, + PlaybookWritePlan, + ProfileWritePlan, + ) + from reflexio.server.services.playbook.service import PlaybookGenerationService + from reflexio.server.services.profile.service import ProfileGenerationService + + request_id = f"request-{entity_kind}-receipt-boundary" + if entity_kind == "profile": + profile = UserProfile( + profile_id="profile-receipt-boundary", + user_id="test_user", + content="This row requires a finalization receipt.", + last_modified_timestamp=1_000, + generated_from_request_id=request_id, + ) + service = ProfileGenerationService(llm_client, request_context) + write_plan = ProfileWritePlan( + user_id="test_user", + request_id=request_id, + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + else: + playbook = UserPlaybook( + user_id="test_user", + agent_version="v1", + request_id=request_id, + content="This row requires a finalization receipt.", + trigger="when validating persistence", + ) + service = PlaybookGenerationService( + llm_client, request_context, skip_aggregation=True + ) + write_plan = PlaybookWritePlan( + request_id=request_id, + output_pending_status=False, + skip_aggregation=True, + new_playbooks=[playbook], + superseded_ids=[], + merge_groups=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + return service, GenerationComputePlan( + prepared=SimpleNamespace(), + generated_count=1, + billable_count=1, + write_plan=write_plan, + bookmark_advance=None, + generation_start=0.0, + extraction_run_ids=run_ids, + token_totals=None, + ) + + @staticmethod + def _stored_inline_learning_count(request_context, entity_kind: str) -> int: + if entity_kind == "profile": + return len(request_context.storage.get_user_profile("test_user")) + return len( + request_context.storage.get_user_playbooks( + user_id="test_user", agent_version="v1" + ) + ) + + @pytest.mark.parametrize("entity_kind", ["profile", "playbook"]) + def test_receipt_required_plan_rejects_missing_run( + self, entity_kind, llm_client, request_context + ): + service, plan = self._build_receipt_required_plan( + entity_kind, + llm_client, + request_context, + run_ids=[f"missing-{entity_kind}-run"], + ) + + with pytest.raises(RuntimeError): + service.persist_generation(plan) + + assert self._stored_inline_learning_count(request_context, entity_kind) == 0 + + @pytest.mark.parametrize("entity_kind", ["profile", "playbook"]) + def test_receipt_required_plan_rejects_multiple_runs( + self, entity_kind, llm_client, request_context + ): + run_ids = [f"{entity_kind}-run-one", f"{entity_kind}-run-two"] + for run_id in run_ids: + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=f"request-{entity_kind}-receipt-boundary", + extractor_kind=entity_kind, + ) + service, plan = self._build_receipt_required_plan( + entity_kind, + llm_client, + request_context, + run_ids=run_ids, + ) + + with pytest.raises(RuntimeError): + service.persist_generation(plan) + + assert self._stored_inline_learning_count(request_context, entity_kind) == 0 + + def test_inline_profile_outcome_persists_receipt_before_terminal_status( + self, llm_client, request_context + ): + """A no-tool inline profile run retains an immutable receipt.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserProfile + from reflexio.server.services.deferred_learning_plan import ProfileWritePlan + from reflexio.server.services.profile.profile_generation_service_utils import ( + ProfileGenerationRequest, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + run_id = "run_inline_profile" + request_id = "request_inline_profile" + profile = UserProfile( + profile_id="profile-inline", + user_id="test_user", + content="The user deploys services to AWS ECS.", + last_modified_timestamp=1_000, + generated_from_request_id=request_id, + ) + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="profile", + ) + + class InlineProfileService(ProfileGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="profile") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([profile], run_id=run_id) + ) + + def _resolve_write_plan(self, results): + return ProfileWritePlan( + user_id="test_user", + request_id=request_id, + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + + service = InlineProfileService(llm_client, request_context) + service.run( + ProfileGenerationRequest( + user_id="test_user", + request_id=request_id, + source="api", + auto_run=False, + ) + ) + + stored = request_context.storage.get_user_profile("test_user") + run = request_context.storage.get_agent_run(run_id) + receipt = request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) + assert [item.profile_id for item in stored] == ["profile-inline"] + assert run is not None and run.status == AgentRunStatus.FINALIZED + assert receipt == ["profile-inline"] + + def test_inline_playbook_outcome_persists_receipt_before_terminal_status( + self, llm_client, request_context + ): + """A no-tool inline playbook run retains its database-assigned ids.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserPlaybook + from reflexio.server.services.deferred_learning_plan import PlaybookWritePlan + from reflexio.server.services.playbook.playbook_service_utils import ( + PlaybookGenerationRequest, + ) + from reflexio.server.services.playbook.service import PlaybookGenerationService + + run_id = "run_inline_playbook" + request_id = "request_inline_playbook" + playbook = UserPlaybook( + user_id="test_user", + agent_version="v1", + request_id=request_id, + content="Prefer AWS ECS for production deployments.", + trigger="when selecting a deployment target", + ) + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="playbook", + ) + + class InlinePlaybookService(PlaybookGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="playbook") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([playbook], run_id=run_id) + ) + + def _resolve_write_plan(self, results): + return PlaybookWritePlan( + request_id=request_id, + output_pending_status=False, + skip_aggregation=True, + new_playbooks=[playbook], + superseded_ids=[], + merge_groups=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + + def _dispatch_playbook_schedulers(self, plan): + return None + + service = InlinePlaybookService( + llm_client, + request_context, + skip_aggregation=True, + ) + service.run( + PlaybookGenerationRequest( + request_id=request_id, + agent_version="v1", + user_id="test_user", + source="api", + auto_run=False, + ) + ) + + stored = request_context.storage.get_user_playbooks( + user_id="test_user", + agent_version="v1", + ) + run = request_context.storage.get_agent_run(run_id) + receipt = request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="user_playbook", + ) + assert len(stored) == 1 + assert run is not None and run.status == AgentRunStatus.FINALIZED + assert receipt == [str(stored[0].user_playbook_id)] + + def test_inline_pending_tool_outcome_does_not_create_finalization_receipt( + self, llm_client, request_context + ): + """A pending-tool run remains available for resume finalization.""" + from reflexio.server.services.profile.profile_generation_service_utils import ( + ProfileGenerationRequest, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + run_id = "run_inline_pending_tool" + request_id = "request_inline_pending_tool" + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="profile", + pending_tool_call_ids=["pending-tool-call"], + ) + + class InlinePendingToolService(ProfileGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="profile") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([], run_id=run_id) + ) + + service = InlinePendingToolService(llm_client, request_context) + service.run( + ProfileGenerationRequest( + user_id="test_user", + request_id=request_id, + source="api", + auto_run=False, + ) + ) + + run = request_context.storage.get_agent_run(run_id) + assert run is not None + assert run.status == AgentRunStatus.FINALIZED_PENDING_TOOL + assert ( + request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) + is None + ) + + def test_inline_profile_persistence_failure_leaves_run_retryable( + self, llm_client, request_context + ): + """A failed inline receipt transaction never exposes terminal state.""" + from reflexio.models.api_schema.domain.entities import LineageContext + from reflexio.models.api_schema.service_schemas import UserProfile + from reflexio.server.services.deferred_learning_plan import ProfileWritePlan + from reflexio.server.services.profile.profile_generation_service_utils import ( + ProfileGenerationRequest, + ) + from reflexio.server.services.profile.service import ProfileGenerationService + + run_id = "run_inline_profile_failure" + request_id = "request_inline_profile_failure" + profile = UserProfile( + profile_id="profile-inline-failure", + user_id="test_user", + content="This profile must roll back.", + last_modified_timestamp=1_000, + generated_from_request_id=request_id, + ) + self._seed_inline_agent_run( + request_context, + run_id=run_id, + request_id=request_id, + extractor_kind="profile", + ) + + class FailingInlineProfileService(ProfileGenerationService): + def _load_extractor_config(self): + return MockExtractorConfig(extractor_name="profile") + + def _should_run_before_extraction(self, extractor_config): + return True + + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([profile], run_id=run_id) + ) + + def _resolve_write_plan(self, results): + return ProfileWritePlan( + user_id="test_user", + request_id=request_id, + new_profiles=[profile], + superseded_ids=[], + lineage_contexts=[LineageContext(op_kind="create")], + ) + + def _persist_write_plan(self, plan): + super()._persist_write_plan(plan) + raise RuntimeError("receipt persistence failed") + + service = FailingInlineProfileService(llm_client, request_context) + service.run( + ProfileGenerationRequest( + user_id="test_user", + request_id=request_id, + source="api", + auto_run=False, + ) + ) + + run = request_context.storage.get_agent_run(run_id) + assert request_context.storage.get_user_profile("test_user") == [] + assert ( + request_context.storage.get_agent_run_finalization_receipt( + run_id=run_id, + entity_type="profile", + ) + is None + ) + assert run is not None and run.status == AgentRunStatus.FINALIZATION_FAILED + assert run.finalization_attempts == 1 + def test_extraction_outcome_finalization_failure_marks_run_retryable( self, llm_client, request_context ): @@ -2086,6 +2511,41 @@ def _process_results(self, results): assert kwargs["last_error"] == "persist failed" assert kwargs["increment_finalization_attempts"] is True + def test_extraction_outcome_terminal_status_failure_marks_run_retryable( + self, llm_client, request_context + ): + """A failed terminal status write remains eligible for finalization retry.""" + + class FailingTerminalStatusService(ConcreteGenerationService): + def _create_extractor(self, extractor_config, service_config): + return MockExtractor( + result=ExtractionOutcome.completed([{"name": "x"}], run_id="run_1") + ) + + service = FailingTerminalStatusService( + llm_client, + request_context, + extractor_configs=[MockExtractorConfig(extractor_name="ext1")], + ) + service.storage = MagicMock() + service.storage.get_agent_run.return_value = SimpleNamespace( + pending_tool_call_ids=[], + committed_output={"items": []}, + finalization_attempts=0, + ) + service.storage.update_agent_run_status.side_effect = [ + RuntimeError("terminal status failed"), + None, + ] + + service.run(MockServiceConfig(user_id="test_user", request_id="test_request")) + + _, status = service.storage.update_agent_run_status.call_args.args[:2] + kwargs = service.storage.update_agent_run_status.call_args.kwargs + assert status == AgentRunStatus.FINALIZATION_FAILED + assert kwargs["last_error"] == "terminal status failed" + assert kwargs["increment_finalization_attempts"] is True + def test_configured_extractor_timeout_fails_generation( self, llm_client, request_context, monkeypatch ): diff --git a/tests/server/services/test_generation_billing_emission.py b/tests/server/services/test_generation_billing_emission.py index 6138f19f..ec622bc1 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 3536b39c..2b2c76ad 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 00000000..546f00b5 --- /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/services/test_search_metering_worker.py b/tests/server/services/test_search_metering_worker.py index 6c4bf910..37d793cf 100644 --- a/tests/server/services/test_search_metering_worker.py +++ b/tests/server/services/test_search_metering_worker.py @@ -125,6 +125,43 @@ def blocked_process(_job: SearchMeteringJob) -> None: assert worker.stop() == 0 +def test_partial_start_failure_stops_started_threads_and_allows_retry( + monkeypatch, +) -> None: + worker = SearchMeteringWorker(worker_count=2) + original_start = threading.Thread.start + start_calls = 0 + started_threads: list[threading.Thread] = [] + + def fail_second_start(thread: threading.Thread) -> None: + nonlocal start_calls + start_calls += 1 + if start_calls == 2: + raise RuntimeError("thread start failed") + original_start(thread) + started_threads.append(thread) + + try: + with monkeypatch.context() as patch: + patch.setattr(threading.Thread, "start", fail_second_start) + with pytest.raises(RuntimeError, match="thread start failed"): + worker.start() + + assert len(started_threads) == 1 + assert not started_threads[0].is_alive() + assert worker._threads == [] + assert worker._started is False + + assert worker.start() is True + assert worker.stop() == 0 + assert worker.start() is False + finally: + worker._stop_event.set() + for thread in worker._threads: + if thread.ident is not None: + thread.join(timeout=1.0) + + def test_full_queue_drops_without_backpressure_and_reports(monkeypatch) -> None: started = threading.Event() release = threading.Event() diff --git a/tests/server/test_api_security_middleware.py b/tests/server/test_api_security_middleware.py index 900ba9fc..e8ca7a2b 100644 --- a/tests/server/test_api_security_middleware.py +++ b/tests/server/test_api_security_middleware.py @@ -174,6 +174,40 @@ async def call_next(_request): assert observed["timeout"] == SYNC_REQUEST_TIMEOUT_SECONDS +def test_playbook_aggregation_post_uses_synchronous_request_timeout_without_wait_query( + monkeypatch, +): + observed: dict[str, float | None] = {} + + async def fake_wait_for(awaitable, *, timeout=None): + observed["timeout"] = timeout + return await awaitable + + async def call_next(_request): + from starlette.responses import Response + + return Response() + + monkeypatch.setattr(asyncio, "wait_for", fake_wait_for) + request = Request( + { + "type": "http", + "method": "POST", + "scheme": "http", + "path": "/api/run_playbook_aggregation", + "raw_path": b"/api/run_playbook_aggregation", + "query_string": b"", + "headers": [], + "client": ("testclient", 50000), + "server": ("testserver", 80), + } + ) + + asyncio.run(TimeoutMiddleware(FastAPI()).dispatch(request, call_next)) + + assert observed["timeout"] == SYNC_REQUEST_TIMEOUT_SECONDS + + def test_security_headers_are_added(monkeypatch): monkeypatch.delenv("REFLEXIO_ALLOWED_ORIGINS", raising=False) diff --git a/tests/server/test_billing_meter.py b/tests/server/test_billing_meter.py index 740d8d47..416bbbc6 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 b9ab6140..0012b753 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 45a1e23d..eb9c2f2b 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,56 @@ 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_oss_metering_startup_failure_clears_exempt_usage_recorder( + monkeypatch, +) -> None: + usage_metrics.configure_usage_event_recorder(None) + + def fail_start_search_metering_worker() -> None: + raise RuntimeError("startup guard failed") + + monkeypatch.setattr( + "reflexio.server.services.search_metering_worker.start_search_metering_worker", + fail_start_search_metering_worker, + ) + + app = create_app(capabilities=None, mount_data_plane=False) + with pytest.raises(RuntimeError, match="startup guard failed"), TestClient(app): + pass + + with pytest.raises(UsageEventDeliveryError) as exc_info: + usage_metrics.record_usage_event_strict( + org_id="7", + event_name="learnings_generated", + event_category="learning", + ) + assert exc_info.value.status is UsageEventDeliveryStatus.UNKNOWN + + def test_partial_cleanup_invariant() -> None: """Already-started caps must be shut down; caps that never started must not be.""" LifecycleCap.started = LifecycleCap.stopped = False diff --git a/tests/server/test_usage_metrics.py b/tests/server/test_usage_metrics.py index 5bf35933..82350340 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)